gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

style: format Rust workspace

Browse Source
This commit is contained in:
Gahow Wang
2026-06-18 18:11:58 +08:00
parent 013465fc06
commit 531cd3fe08
57 changed files with 4045 additions and 1204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

90
crates/xserv-kernels/src/rope.rs
View File

@@ -3,15 +3,34 @@ 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);
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 {
@@ -30,12 +49,21 @@ impl RopeCache {
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, xserv_cuda::current_stream_raw(),
cos.as_mut_ptr() as _,
sin.as_mut_ptr() as _,
max_seq_len as i32,
half_dim as i32,
theta,
xserv_cuda::current_stream_raw(),
);
}
Self { cos, sin, max_seq_len, half_dim }
Self {
cos,
sin,
max_seq_len,
half_dim,
}
}
/// YaRN (Yet another RoPE extensioN) RoPE cache. Applies frequency-dependent
@@ -68,8 +96,8 @@ impl RopeCache {
let mut inv_freq = vec![0.0f64; half_dim];
for i in 0..half_dim {
let pos_freq = theta.powf((2 * i) as f64 / dim);
let inv_freq_extrapolation = 1.0 / pos_freq; // original
let inv_freq_interpolation = 1.0 / (factor * pos_freq); // scaled
let inv_freq_extrapolation = 1.0 / pos_freq; // original
let inv_freq_interpolation = 1.0 / (factor * pos_freq); // scaled
// Linear ramp: 0 where we keep original, 1 where we interpolate
let ramp = if (high - low).abs() < 0.001 {
@@ -101,16 +129,19 @@ impl RopeCache {
let nbytes = total * std::mem::size_of::<f32>();
let mut cos = GpuBuffer::alloc(nbytes).expect("alloc yarn cos_cache");
let mut sin = GpuBuffer::alloc(nbytes).expect("alloc yarn sin_cache");
let cos_bytes = unsafe {
std::slice::from_raw_parts(cos_host.as_ptr() as *const u8, nbytes)
};
let sin_bytes = unsafe {
std::slice::from_raw_parts(sin_host.as_ptr() as *const u8, nbytes)
};
let cos_bytes =
unsafe { std::slice::from_raw_parts(cos_host.as_ptr() as *const u8, nbytes) };
let sin_bytes =
unsafe { std::slice::from_raw_parts(sin_host.as_ptr() as *const u8, nbytes) };
cos.copy_from_host(cos_bytes).unwrap();
sin.copy_from_host(sin_bytes).unwrap();
Self { cos, sin, max_seq_len, half_dim }
Self {
cos,
sin,
max_seq_len,
half_dim,
}
}
}
@@ -133,7 +164,8 @@ pub fn rope_inplace(x: &Tensor, cache: &RopeCache, positions: &[u32]) {
num_tokens * std::mem::size_of::<u32>(),
)
};
let mut pos_gpu = xserv_cuda::allocator::cached_alloc(pos_bytes.len()).expect("alloc positions");
let mut pos_gpu =
xserv_cuda::allocator::cached_alloc(pos_bytes.len()).expect("alloc positions");
pos_gpu.copy_from_host(pos_bytes).unwrap();
rope_inplace_device_pos(x, cache, pos_gpu.as_ptr() as *const c_void);
@@ -155,16 +187,22 @@ pub fn rope_inplace_device_pos(x: &Tensor, cache: &RopeCache, pos_gpu: *const c_
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 _,
cache.cos.as_ptr() as _,
cache.sin.as_ptr() as _,
pos_gpu,
num_tokens as i32, num_heads as i32, head_dim as i32,
num_tokens as i32,
num_heads as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => launch_rope_bf16(
x.data_ptr() as *mut c_void,
cache.cos.as_ptr() as _, cache.sin.as_ptr() as _,
cache.cos.as_ptr() as _,
cache.sin.as_ptr() as _,
pos_gpu,
num_tokens as i32, num_heads as i32, head_dim as i32,
num_tokens as i32,
num_heads as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype for rope"),