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>
84 lines
3.1 KiB
Rust
84 lines
3.1 KiB
Rust
use std::ffi::c_void;
|
|
use xserv_cuda::GpuBuffer;
|
|
use xserv_tensor::{DType, Device, Tensor};
|
|
|
|
unsafe extern "C" {
|
|
fn launch_rope_f32(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
|
|
positions: *const c_void, num_tokens: i32, num_heads: i32,
|
|
head_dim: i32, stream: *mut c_void);
|
|
fn launch_rope_bf16(x: *mut c_void, cos_cache: *const c_void, sin_cache: *const c_void,
|
|
positions: *const c_void, num_tokens: i32, num_heads: i32,
|
|
head_dim: i32, stream: *mut c_void);
|
|
fn launch_compute_rope_cache(cos_cache: *mut c_void, sin_cache: *mut c_void,
|
|
max_seq_len: i32, half_dim: i32, theta: f32,
|
|
stream: *mut c_void);
|
|
}
|
|
|
|
pub struct RopeCache {
|
|
pub cos: GpuBuffer,
|
|
pub sin: GpuBuffer,
|
|
pub max_seq_len: usize,
|
|
pub half_dim: usize,
|
|
}
|
|
|
|
impl RopeCache {
|
|
pub fn new(max_seq_len: usize, head_dim: usize, theta: f32) -> Self {
|
|
let half_dim = head_dim / 2;
|
|
let nbytes = max_seq_len * half_dim * std::mem::size_of::<f32>();
|
|
let mut cos = GpuBuffer::alloc(nbytes).expect("alloc cos_cache");
|
|
let mut sin = GpuBuffer::alloc(nbytes).expect("alloc sin_cache");
|
|
|
|
unsafe {
|
|
launch_compute_rope_cache(
|
|
cos.as_mut_ptr() as _, sin.as_mut_ptr() as _,
|
|
max_seq_len as i32, half_dim as i32, theta, std::ptr::null_mut(),
|
|
);
|
|
}
|
|
|
|
Self { cos, sin, max_seq_len, half_dim }
|
|
}
|
|
}
|
|
|
|
/// Apply RoPE in-place to x.
|
|
/// x: [num_tokens, num_heads, head_dim] on GPU
|
|
/// positions: [num_tokens] (u32 on CPU, will be uploaded)
|
|
pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) {
|
|
assert_eq!(x.ndim(), 3);
|
|
assert!(x.is_contiguous());
|
|
assert!(matches!(x.device(), Device::Cuda(_)));
|
|
let num_tokens = x.shape()[0];
|
|
let num_heads = x.shape()[1];
|
|
let head_dim = x.shape()[2];
|
|
assert_eq!(head_dim / 2, cache.half_dim);
|
|
assert_eq!(positions.len(), num_tokens);
|
|
|
|
let pos_bytes = unsafe {
|
|
std::slice::from_raw_parts(
|
|
positions.as_ptr() as *const u8,
|
|
num_tokens * std::mem::size_of::<u32>(),
|
|
)
|
|
};
|
|
let mut pos_gpu = xserv_cuda::allocator::cached_alloc(pos_bytes.len()).expect("alloc positions");
|
|
pos_gpu.copy_from_host(pos_bytes).unwrap();
|
|
|
|
unsafe {
|
|
match x.dtype() {
|
|
DType::F32 => launch_rope_f32(
|
|
x.data_ptr() as *mut c_void,
|
|
cache.cos.as_ptr() as _, cache.sin.as_ptr() as _,
|
|
pos_gpu.as_ptr() as _,
|
|
num_tokens as i32, num_heads as i32, head_dim as i32,
|
|
std::ptr::null_mut(),
|
|
),
|
|
DType::BF16 => launch_rope_bf16(
|
|
x.data_ptr() as *mut c_void,
|
|
cache.cos.as_ptr() as _, cache.sin.as_ptr() as _,
|
|
pos_gpu.as_ptr() as _,
|
|
num_tokens as i32, num_heads as i32, head_dim as i32,
|
|
std::ptr::null_mut(),
|
|
),
|
|
_ => panic!("unsupported dtype for rope"),
|
|
}
|
|
}
|
|
}
|