diff --git a/Cargo.toml b/Cargo.toml index 1dc707c..cb498f0 100644 --- a/Cargo.toml +++ b/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" } \ No newline at end of file +mimalloc = { version = "0.1.39" } + diff --git a/HANDOFF-2026-04-21.md b/HANDOFF-2026-04-21.md new file mode 100644 index 0000000..7f0fda3 --- /dev/null +++ b/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`) 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` 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`, `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 diff --git a/src/lib.rs b/src/lib.rs index 934cfec..9ce3a51 100644 --- a/src/lib.rs +++ b/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 { #[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, } #[cfg(feature = "gpu")] @@ -87,11 +90,11 @@ impl TokenizerEngine { #[new] fn new(model_path: String) -> PyResult { 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) -> PyResult> { + fn encode_batch(&self, py: Python, texts: Vec) -> PyResult> { self.inner.batch_encode_to_gpu(py, texts) } } diff --git a/src/llm.rs b/src/llm.rs index 402e93d..9dbb17c 100644 --- a/src/llm.rs +++ b/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::(); - 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, + shape: Box<[i64]>, +} + pub struct PinnedBatcher { tokenizer: Tokenizer, device: Arc, - gpu_buffer: CudaSlice, - host_buffer: PinnedHostBuffer, - max_elements: usize, } impl PinnedBatcher { pub fn new(model_path: &str) -> PyResult { 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::(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) -> PyResult> { + pub fn batch_encode_to_gpu(&self, py: Python, texts: Vec) -> PyResult> { 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::(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::(); - 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); + } } -} \ No newline at end of file +}