Fix all four Phase 1 correctness bugs (dynamic buffer, DLPack ownership, thread safety, padding)

- 1.1: Remove hardcoded 4096*512 buffer from PinnedBatcher; allocate host+device per-call
- 1.2: Wrap CudaSlice in DLPackContext owned by DLManagedTensor.manager_ctx; deleter frees
  device memory when PyTorch releases tensor — no more silent overwrite across calls
- 1.3: batch_encode_to_gpu now takes &self (no mutable shared state); TokenizerEngine wraps
  Arc<PinnedBatcher> and encode_batch takes &self — safe for concurrent Python threads
- 1.4: seq_len = max across all encodings; host buffer prefilled with pad_id before token write

Also: switch cudarc gpu feature from cuda-version-from-build-system to cuda-12050 to fix
build on machines with CUDA 13.x (cudarc 0.11.9 only knows up to 12.5)

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Cargo.toml

@@ -15,7 +15,7 @@ gpu = [
     "tokenizers",
     "cudarc",
     "pyo3-dlpack",
-    "cudarc/cuda-version-from-build-system",
+    "cudarc/cuda-12050",
     "dep:cuda-driver-sys"
 ]
 profiling = ["dep:tracing", "dep:tracing-subscriber", "dep:tracing-tracy"]
@@ -52,4 +52,5 @@ tracing-tracy = { version = "=0.11.2", optional = true }
 tracing = {version ="0.1.41", optional = true }
 
 [target.'cfg(not(target_env = "msvc"))'.dependencies]
-mimalloc = { version = "0.1.39" }
+mimalloc = { version = "0.1.39" }
+

HANDOFF-2026-04-21.md

@@ -0,0 +1,84 @@
+# Handoff — 2026-04-21
+
+## What Was Done This Session
+
+Fixed all four Phase 1 correctness bugs. Code compiles clean on temp-ubuntu with CUDA 13.1.
+
+---
+
+## Phase 1 Bugs — ALL FIXED
+
+### 1.1 Dynamic buffer allocation — DONE
+**Old behavior:** `max_elements = 4096 * 512` hardcoded in `PinnedBatcher::new`. Any batch exceeding ~2M tokens returned an error.
+**Fix:** Removed `gpu_buffer`, `host_buffer`, and `max_elements` from the `PinnedBatcher` struct entirely. Both buffers are now allocated per-call based on actual `batch_size * seq_len`.
+
+### 1.2 Buffer isolation (DLPack overwrite) — DONE
+**Old behavior:** Every call overwrote the same `CudaSlice`. Any tensor from call N got corrupted by call N+1.
+**Fix:** Each call allocates its own device buffer and wraps it in a `DLPackContext` struct stored in `DLManagedTensor.manager_ctx`. The deleter frees both the device memory (by dropping `CudaSlice<i64>`) and the shape array when PyTorch releases the tensor. Each returned tensor now owns its allocation.
+
+### 1.3 Thread safety — DONE
+**Old behavior:** `batch_encode_to_gpu` took `&mut self`. `TokenizerEngine.encode_batch` took `&mut self`. Not usable from multiple Python threads.
+**Fix:** Since `PinnedBatcher` no longer has mutable shared state (buffers are per-call), `batch_encode_to_gpu` now takes `&self`. `TokenizerEngine` wraps `Arc<PinnedBatcher>` instead of owning it directly. `encode_batch` takes `&self`. Concurrent calls from Python threads are safe.
+
+### 1.4 Variable sequence length + padding — DONE
+**Old behavior:** `seq_len = encodings[0].len()` — assumed uniform length. Variable-length batches produced wrong tensor shape.
+**Fix:** `seq_len = encodings.iter().map(|e| e.len()).max()`. Host buffer prefilled with `pad_id` (from `tokenizer.get_padding()`, fallback to 0). Rayon fill loop only writes actual tokens; shorter sequences retain the pad value.
+
+---
+
+## Build Issue Discovered and Resolved
+
+`cudarc 0.11.9` panics at build time on CUDA 13.1 (only knows 12.5 and below). This machine has CUDA 13.1.
+
+**Fix:** Changed the `gpu` feature in `Cargo.toml` from `cudarc/cuda-version-from-build-system` to `cudarc/cuda-12050`. This hard-pins to CUDA 12.5 headers and skips the `nvcc` version check. CUDA 13.x is backward-compatible with 12.5 driver API at runtime. If you ever upgrade cudarc to 0.12+, revert this change and re-test.
+
+---
+
+## Files Changed
+
+| File | What Changed |
+|------|-------------|
+| `src/llm.rs` | Full rewrite of `PinnedBatcher` — removed static buffers, per-call allocation, `DLPackContext` for ownership, `&self` throughout |
+| `src/lib.rs` | `TokenizerEngine.inner` changed to `Arc<PinnedBatcher>`, `encode_batch` takes `&self`, added `use std::sync::Arc` |
+| `Cargo.toml` | `gpu` feature: `cuda-version-from-build-system` → `cuda-12050` |
+
+---
+
+## This Week — What's Left (in order)
+
+### Today / Tomorrow (Tuesday 2026-04-22)
+- [ ] Rent H100 on Lambda Labs (~$2 for 30 min)
+- [ ] Run `benchmark_gpu.py` under `nsys profile --stats=true python benchmark_gpu.py`
+- [ ] Capture Nsight `.nsys-rep` file and download it
+- [ ] Note the speedup number — this is the headline for Thursday's talk
+
+**Still need:** H100 setup script (a single paste-into-terminal script for the cloud instance). Needs: Rust install, maturin build with `--features gpu`, nsys profile run, output capture. This was not written last session.
+
+### Wednesday (2026-04-23)
+- [ ] Update `docs/talk-10min.md` Slide 6 with the H100 result
+- [ ] Open Nsight Systems GUI locally, load the `.nsys-rep` file, screenshot the timeline showing both pipelines
+
+### Thursday (2026-04-24)
+- [ ] Talk to IT/security group (10 min)
+- [ ] If someone gives you a contact: do not show GPU source code without NDA signed first; first call is discovery
+
+---
+
+## Next Session Priorities (after Thursday)
+
+1. **Write H100 setup script** — single paste-into-terminal script (see above)
+2. **vLLM integration** — Phase 2.1, this is the demo that sells the concept
+3. **NM anonymous LLC** — Wyoming, single-member, before any contract is signed
+4. **C++ version scoping** — after Phase 1 solid, assess Whetstone transpilation of `llm.rs`
+
+---
+
+## Key Technical Context
+
+- Repo: `/home/bill/Documents/unchecked-io/`, on `gpu-zero` branch
+- Build: `maturin develop --release --features gpu` (maturin not installed on temp-ubuntu — install with `pip install maturin` first)
+- `cargo check --features gpu` confirms clean compile
+- The four Phase 1 bugs are fixed and verified to compile
+- `src/llm.rs` is the GPU core. `src/lib.rs` is the Python-facing wrapper.
+- `chrono-tz` pinned at `=0.10.4` — do not change
+- `gpu-zero` branch is private; `main` is open source Apache 2.0

src/lib.rs

@@ -5,6 +5,8 @@ mod parser;
 mod llm;
 
 use pyo3::prelude::*;
+#[cfg(feature = "gpu")]
+use std::sync::Arc;
 use pyo3::exceptions::PyValueError;
 use tokio;
 use pyo3::types::{PyModule, PyAny};
@@ -78,7 +80,8 @@ fn profile_data(config_path: String) -> PyResult<String> {
 #[cfg(feature = "gpu")]
 #[pyclass(name = "TokenizerEngine")]
 struct TokenizerEngine {
-    inner: llm::PinnedBatcher,
+    // 1.3: Arc allows cheap cloning; encode_batch takes &self so concurrent calls are safe
+    inner: Arc<llm::PinnedBatcher>,
 }
 
 #[cfg(feature = "gpu")]
@@ -87,11 +90,11 @@ impl TokenizerEngine {
     #[new]
     fn new(model_path: String) -> PyResult<Self> {
         Ok(TokenizerEngine {
-            inner: llm::PinnedBatcher::new(&model_path)?
+            inner: Arc::new(llm::PinnedBatcher::new(&model_path)?)
         })
     }
 
-    fn encode_batch(&mut self, py: Python, texts: Vec<String>) -> PyResult<Py<PyAny>> {
+    fn encode_batch(&self, py: Python, texts: Vec<String>) -> PyResult<Py<PyAny>> {
         self.inner.batch_encode_to_gpu(py, texts)
     }
 }

src/llm.rs

@@ -6,11 +6,9 @@ use tokenizers::Tokenizer;
 use cudarc::driver::{CudaDevice, DeviceSlice, CudaSlice, DevicePtr};
 use pyo3::prelude::*;
 use pyo3::exceptions::PyValueError;
-use pyo3::types::PyCapsule;
 use std::ffi::CStr;
 use rayon::prelude::*;
 
-// Import raw CUDA symbols. Note the _v2 suffixes which are required by the sys crate.
 use cuda_driver_sys::{
     cuMemHostRegister_v2,
     cuMemHostUnregister,
@@ -51,15 +49,12 @@ impl PinnedHostBuffer {
         let data = vec![0i64; capacity];
         let ptr = data.as_ptr() as *mut std::ffi::c_void;
         let bytes = capacity * std::mem::size_of::<i64>();
-
         unsafe {
-            // Use _v2 variant
             let result = cuMemHostRegister_v2(ptr, bytes, 0);
             if result != cudaError_enum::CUDA_SUCCESS {
-                return Err(PyValueError::new_err(format!("Failed to Pin Memory: {:?}", result)));
+                return Err(PyValueError::new_err(format!("Failed to pin memory: {:?}", result)));
             }
         }
-
         Ok(Self { data })
     }
 }
@@ -73,39 +68,27 @@ impl Drop for PinnedHostBuffer {
     }
 }
 
+// Owned by DLManagedTensor.manager_ctx — freed when PyTorch releases the tensor.
+struct DLPackContext {
+    _gpu_buffer: CudaSlice<i64>,
+    shape: Box<[i64]>,
+}
+
 pub struct PinnedBatcher {
     tokenizer: Tokenizer,
     device: Arc<CudaDevice>,
-    gpu_buffer: CudaSlice<i64>,
-    host_buffer: PinnedHostBuffer,
-    max_elements: usize,
 }
 
 impl PinnedBatcher {
     pub fn new(model_path: &str) -> PyResult<Self> {
         let tokenizer = Tokenizer::from_file(model_path)
             .map_err(|e| PyValueError::new_err(format!("Failed to load tokenizer: {}", e)))?;
-
         let device = CudaDevice::new(0)
             .map_err(|e| PyValueError::new_err(format!("No CUDA GPU found: {:?}", e)))?;
-
+        Ok(PinnedBatcher { tokenizer, device })
-        let max_elements = 4096 * 512;
-
-        let gpu_buffer = device.alloc_zeros::<i64>(max_elements)
-            .map_err(|e| PyValueError::new_err(format!("GPU Alloc Failed: {:?}", e)))?;
-
-        let host_buffer = PinnedHostBuffer::new(max_elements)?;
-
-        Ok(PinnedBatcher {
-            tokenizer,
-            device,
-            gpu_buffer,
-            host_buffer,
-            max_elements,
-        })
     }
 
-    pub fn batch_encode_to_gpu(&mut self, py: Python, texts: Vec<String>) -> PyResult<Py<PyAny>> {
+    pub fn batch_encode_to_gpu(&self, py: Python, texts: Vec<String>) -> PyResult<Py<PyAny>> {
         let encodings = self.tokenizer.encode_batch(texts, true)
             .map_err(|e| PyValueError::new_err(format!("Tokenization failed: {}", e)))?;
 
@@ -114,19 +97,25 @@ impl PinnedBatcher {
         }
 
         let batch_size = encodings.len();
-        let seq_len = encodings[0].len();
+
+        // 1.4: use max length across all encodings, not just encodings[0]
+        let seq_len = encodings.iter().map(|e| e.len()).max().unwrap_or(0);
+        if seq_len == 0 {
+            return Err(PyValueError::new_err("All sequences are empty"));
+        }
         let total_elements = batch_size * seq_len;
 
-        if total_elements > self.max_elements {
+        // 1.4: pad shorter sequences with the tokenizer's pad token (fallback: 0)
-            return Err(PyValueError::new_err(format!(
+        let pad_id = self.tokenizer.get_padding()
-                "Batch too large: {} tokens (Max: {})",
+            .map(|p| p.pad_id as i64)
-                total_elements, self.max_elements
+            .unwrap_or(0);
-            )));
-        }
 
-        // Parallel Write into Pinned Memory
+        // 1.1: allocate pinned host buffer sized to the actual batch — no hardcoded max
-        // Explicit types added to closure to satisfy type inference
+        let mut host_buffer = PinnedHostBuffer::new(total_elements)?;
-        self.host_buffer.data[..total_elements].par_chunks_mut(seq_len)
+        host_buffer.data.fill(pad_id);
+
+        // Parallel write into pinned memory; shorter sequences leave pad_id in remaining slots
+        host_buffer.data[..total_elements].par_chunks_mut(seq_len)
             .zip(encodings.par_iter())
             .for_each(|(dest, enc): (&mut [i64], &tokenizers::Encoding)| {
                 let ids = enc.get_ids();
@@ -136,36 +125,47 @@ impl PinnedBatcher {
                 }
             });
 
-        // Direct DMA Copy (Pinned Host -> Device)
+        // 1.1: allocate device buffer sized to this batch — no hardcoded max
+        let gpu_buffer = self.device.alloc_zeros::<i64>(total_elements)
+            .map_err(|e| PyValueError::new_err(format!("GPU alloc failed: {:?}", e)))?;
+
+        // Direct DMA copy (pinned host -> device)
         unsafe {
-            let src_ptr = self.host_buffer.data.as_ptr() as *const std::ffi::c_void;
+            let src_ptr = host_buffer.data.as_ptr() as *const std::ffi::c_void;
-            let dst_ptr = *self.gpu_buffer.device_ptr();
+            let dst_ptr = *gpu_buffer.device_ptr();
             let bytes = total_elements * std::mem::size_of::<i64>();
-
             let result = cuMemcpyHtoD_v2(dst_ptr, src_ptr, bytes);
-
             if result != cudaError_enum::CUDA_SUCCESS {
-                return Err(PyValueError::new_err(format!("Direct GPU Copy Failed: {:?}", result)));
+                return Err(PyValueError::new_err(format!("DMA copy failed: {:?}", result)));
             }
         }
+        // host_buffer drops here — unpins host memory
 
-        // Zero-Copy Handover
+        // 1.2: transfer gpu_buffer ownership into DLPackContext so each returned tensor owns its
-        let ptr_address = *self.gpu_buffer.device_ptr() as *mut std::ffi::c_void;
+        // own device allocation. The deleter frees it when PyTorch releases the tensor.
-        let shape = Box::into_raw(vec![batch_size as i64, seq_len as i64].into_boxed_slice()) as *mut i64;
+        let ctx = Box::new(DLPackContext {
+            _gpu_buffer: gpu_buffer,
+            shape: vec![batch_size as i64, seq_len as i64].into_boxed_slice(),
+        });
+
+        // Capture pointers before leaking the box
+        let shape_ptr = ctx.shape.as_ptr() as *mut i64;
+        let data_ptr = unsafe { *ctx._gpu_buffer.device_ptr() } as *mut std::ffi::c_void;
+        let manager_ctx = Box::into_raw(ctx) as *mut std::ffi::c_void;
 
         let dl_tensor = DLTensor {
-            data: ptr_address,
+            data: data_ptr,
             device: DLDevice { device_type: 2, device_id: 0 },
             ndim: 2,
             dtype: DLDataType { code: 0, bits: 64, lanes: 1 },
-            shape,
+            shape: shape_ptr,
             strides: std::ptr::null_mut(),
             byte_offset: 0,
         };
 
         let managed_tensor = Box::new(DLManagedTensor {
             dl_tensor,
-            manager_ctx: std::ptr::null_mut(),
+            manager_ctx,
             deleter: Some(dlpack_deleter),
         });
 
@@ -175,13 +175,11 @@ impl PinnedBatcher {
             let capsule_ptr = pyo3::ffi::PyCapsule_New(
                 managed_ptr as *mut _,
                 CAPSULE_NAME.as_ptr(),
-                None
+                None,
             );
-
             if capsule_ptr.is_null() {
                 return Err(PyValueError::new_err("Failed to create PyCapsule"));
             }
-
             Ok(Bound::from_owned_ptr(py, capsule_ptr).into_any().unbind())
         }
     }
@@ -189,11 +187,11 @@ impl PinnedBatcher {
 
 unsafe extern "C" fn dlpack_deleter(managed_ptr: *mut DLManagedTensor) {
     if managed_ptr.is_null() { return; }
-    let managed = Box::from_raw(managed_ptr);
+    unsafe {
-    if !managed.dl_tensor.shape.is_null() {
+        let managed = Box::from_raw(managed_ptr);
-        let _ = Box::from_raw(std::slice::from_raw_parts_mut(
+        if !managed.manager_ctx.is_null() {
-            managed.dl_tensor.shape,
+            // Drops DLPackContext: frees CudaSlice (device memory) and shape array.
-            managed.dl_tensor.ndim as usize
+            let _ = Box::from_raw(managed.manager_ctx as *mut DLPackContext);
-        ));
+        }
     }
-}
+}