Gahow Wang
986a289616
P0 fixes (blocking usability): - FIX-01: thread-local cuBLAS handle (was creating/destroying per matmul) - FIX-16: EOS token no longer leaks into API responses - FIX-17: max_seq_len configurable via --max-seq-len (default 2048, was hardcoded 256) - FIX-18: max_tokens clamped to available seq space, prompt overflow returns 400 P1 fixes (bugs & performance): - FIX-07: CachingAllocator wired into all hot paths (to_device, embedding, rope, concat) - FIX-08: CudaDeviceProp buffer increased to 32KB for CUDA 12.9 safety - FIX-09: tokenizer byte_fallback graceful degradation (was panic) - FIX-19: causal mask uses -INFINITY instead of -1e9 (BF16 supports inf) - FIX-20: LayerNorm rewritten to numerically stable two-pass algorithm - FIX-21: min block size guard (32 threads) for LayerNorm/RMSNorm launches P2 fixes (improvements): - FIX-22: Option<GpuKVCache> + take() eliminates dummy KV cache allocations - FIX-23: RoPE cache no longer artificially capped at 8192 positions Verified on dash5 (RTX 5090): 51 tok/s batch=1, 74 tok/s 2-concurrent, 1.7-3.3x HF transformers. Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
94 lines
2.8 KiB
Rust
94 lines
2.8 KiB
Rust
use crate::error::{self, Result};
|
|
use crate::ffi;
|
|
use std::ffi::CStr;
|
|
use std::os::raw::c_char;
|
|
|
|
#[derive(Debug, Clone)]
|
|
pub struct DeviceInfo {
|
|
pub index: u32,
|
|
pub name: String,
|
|
pub total_memory: usize,
|
|
pub free_memory: usize,
|
|
pub compute_major: i32,
|
|
pub compute_minor: i32,
|
|
pub sm_count: i32,
|
|
pub shared_mem_per_block: usize,
|
|
pub warp_size: i32,
|
|
pub max_threads_per_block: i32,
|
|
}
|
|
|
|
unsafe extern "C" {
|
|
fn cudaDeviceGetAttribute(value: *mut i32, attr: i32, device: i32) -> i32;
|
|
fn cudaMemGetInfo(free: *mut usize, total: *mut usize) -> i32;
|
|
}
|
|
|
|
fn get_attr(attr: i32, device: u32) -> Result<i32> {
|
|
let mut value = 0;
|
|
error::check(unsafe { cudaDeviceGetAttribute(&mut value, attr, device as i32) })?;
|
|
Ok(value)
|
|
}
|
|
|
|
pub fn device_count() -> Result<i32> {
|
|
let mut count = 0;
|
|
error::check(unsafe { ffi::cudaGetDeviceCount(&mut count) })?;
|
|
Ok(count)
|
|
}
|
|
|
|
pub fn set_device(device: u32) -> Result<()> {
|
|
error::check(unsafe { ffi::cudaSetDevice(device as i32) })
|
|
}
|
|
|
|
pub fn current_device() -> Result<u32> {
|
|
let mut dev = 0;
|
|
error::check(unsafe { ffi::cudaGetDevice(&mut dev) })?;
|
|
Ok(dev as u32)
|
|
}
|
|
|
|
pub fn device_info(device: u32) -> Result<DeviceInfo> {
|
|
// Heap-allocate oversized buffer for cudaDeviceProp (layout varies by CUDA version).
|
|
// CUDA 12.x struct is ~5-6 KB; use 32 KB to guard against future growth.
|
|
let mut prop_buf = vec![0u8; 32768];
|
|
error::check(unsafe {
|
|
ffi::cudaGetDeviceProperties(prop_buf.as_mut_ptr(), device as i32)
|
|
})?;
|
|
// Name is always the first field: char[256].
|
|
let name = unsafe { CStr::from_ptr(prop_buf.as_ptr() as *const c_char) }
|
|
.to_string_lossy()
|
|
.into_owned();
|
|
|
|
// Get memory info via cudaMemGetInfo (layout-independent).
|
|
let prev = current_device()?;
|
|
set_device(device)?;
|
|
let mut free = 0usize;
|
|
let mut total = 0usize;
|
|
error::check(unsafe { cudaMemGetInfo(&mut free, &mut total) })?;
|
|
if prev != device {
|
|
set_device(prev)?;
|
|
}
|
|
|
|
// Attribute IDs from cuda_runtime_api.h
|
|
const SHARED_MEM_PER_BLOCK: i32 = 8;
|
|
const WARP_SIZE: i32 = 10;
|
|
const MAX_THREADS_PER_BLOCK: i32 = 1;
|
|
const MULTI_PROCESSOR_COUNT: i32 = 16;
|
|
const COMPUTE_MAJOR: i32 = 75;
|
|
const COMPUTE_MINOR: i32 = 76;
|
|
|
|
Ok(DeviceInfo {
|
|
index: device,
|
|
name,
|
|
total_memory: total,
|
|
free_memory: free,
|
|
compute_major: get_attr(COMPUTE_MAJOR, device)?,
|
|
compute_minor: get_attr(COMPUTE_MINOR, device)?,
|
|
sm_count: get_attr(MULTI_PROCESSOR_COUNT, device)?,
|
|
shared_mem_per_block: get_attr(SHARED_MEM_PER_BLOCK, device)? as usize,
|
|
warp_size: get_attr(WARP_SIZE, device)?,
|
|
max_threads_per_block: get_attr(MAX_THREADS_PER_BLOCK, device)?,
|
|
})
|
|
}
|
|
|
|
pub fn synchronize() -> Result<()> {
|
|
error::check(unsafe { ffi::cudaDeviceSynchronize() })
|
|
}
|