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

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

114
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);
}
}
}
}