bill/unchecked-io
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

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

Browse Source 
- 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>
This commit is contained in:
Bill Holcombe
2026-04-21 14:10:31 -06:00
parent 6f8b6440aa
commit dc8494bb2a
4 changed files with 149 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
Cargo.toml
View File

@@ -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"]
@@ -53,3 +53,4 @@ tracing = {version ="0.1.41", optional = true }
[target.'cfg(not(target_env = "msvc"))'.dependencies]
mimalloc = { version = "0.1.39" }

84
HANDOFF-2026-04-21.md Normal file
View File

@@ -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

9
src/lib.rs
View File

@@ -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)
}
}

112
src/llm.rs
View File

@@ -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)))?;
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,
})
Ok(PinnedBatcher { tokenizer, device })
}
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 {
return Err(PyValueError::new_err(format!(
"Batch too large: {} tokens (Max: {})",
total_elements, self.max_elements
)));
}
// 1.4: pad shorter sequences with the tokenizer's pad token (fallback: 0)
let pad_id = self.tokenizer.get_padding()
.map(|p| p.pad_id as i64)
.unwrap_or(0);
// Parallel Write into Pinned Memory
// Explicit types added to closure to satisfy type inference
self.host_buffer.data[..total_elements].par_chunks_mut(seq_len)
// 1.1: allocate pinned host buffer sized to the actual batch — no hardcoded max
let mut host_buffer = PinnedHostBuffer::new(total_elements)?;
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 dst_ptr = *self.gpu_buffer.device_ptr();
let src_ptr = host_buffer.data.as_ptr() as *const std::ffi::c_void;
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
let ptr_address = *self.gpu_buffer.device_ptr() as *mut std::ffi::c_void;
let shape = Box::into_raw(vec![batch_size as i64, seq_len as i64].into_boxed_slice()) as *mut i64;
// 1.2: transfer gpu_buffer ownership into DLPackContext so each returned tensor owns its
// own device allocation. The deleter frees it when PyTorch releases the tensor.
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);
if !managed.dl_tensor.shape.is_null() {
let _ = Box::from_raw(std::slice::from_raw_parts_mut(
managed.dl_tensor.shape,
managed.dl_tensor.ndim as usize
));
unsafe {
let managed = Box::from_raw(managed_ptr);
if !managed.manager_ctx.is_null() {
// Drops DLPackContext: frees CudaSlice (device memory) and shape array.
let _ = Box::from_raw(managed.manager_ctx as *mut DLPackContext);
}
}
}