Gahow Wang
c8e8153702
CUDA kernels (csrc/): - common.cuh: shared warp_reduce_sum/max, block_reduce_sum/max - normalization/rmsnorm.cu: RMSNorm (F32 + BF16) - normalization/layernorm.cu: LayerNorm with Welford (F32 + BF16) - activation/activations.cu: GELU tanh-approx + SiLU (F32 + BF16) - reduce/softmax.cu: safe softmax, 3-pass (F32 + BF16) - embedding/embedding.cu: gather lookup (F32 + BF16) - embedding/rope.cu: RoPE in-place + precomputed cos/sin cache (F32 + BF16) Rust wrappers (xserv-kernels/src/): - rmsnorm.rs, layernorm.rs, activation.rs, softmax.rs, embedding.rs, rope.rs - RopeCache struct with GPU-side precomputation Tests: 12 new tests (ops_test.rs), all passing with good precision: - F32: max_err 1e-6 ~ 1e-9 - BF16: max_err 2e-3 ~ 7e-3 Total: 29 kernel tests + 27 prior = 56 tests passing Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
86 lines
3.2 KiB
Rust
86 lines
3.2 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(),
|
|
);
|
|
}
|
|
xserv_cuda::device::synchronize().unwrap();
|
|
|
|
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 = GpuBuffer::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"),
|
|
}
|
|
}
|
|
xserv_cuda::device::synchronize().unwrap();
|
|
}
|