212 lines
7.0 KiB
Rust
212 lines
7.0 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,
|
|
xserv_cuda::current_stream_raw(),
|
|
);
|
|
}
|
|
|
|
Self {
|
|
cos,
|
|
sin,
|
|
max_seq_len,
|
|
half_dim,
|
|
}
|
|
}
|
|
|
|
/// YaRN (Yet another RoPE extensioN) RoPE cache. Applies frequency-dependent
|
|
/// interpolation so the model can extrapolate beyond its training context.
|
|
pub fn new_yarn(
|
|
max_seq_len: usize,
|
|
head_dim: usize,
|
|
theta: f64,
|
|
factor: f64,
|
|
original_max_pos: usize,
|
|
beta_fast: f64,
|
|
beta_slow: f64,
|
|
) -> Self {
|
|
let half_dim = head_dim / 2;
|
|
let dim = head_dim as f64;
|
|
|
|
// find_correction_dim: inverse formula to find dimension from number of rotations
|
|
let find_correction_dim = |num_rotations: f64| -> f64 {
|
|
dim * (original_max_pos as f64 / (num_rotations * 2.0 * std::f64::consts::PI)).ln()
|
|
/ (2.0 * theta.ln())
|
|
};
|
|
|
|
let low_raw = find_correction_dim(beta_fast);
|
|
let high_raw = find_correction_dim(beta_slow);
|
|
// config has truncate=false, so use raw values (no floor/ceil)
|
|
let low = low_raw.max(0.0);
|
|
let high = high_raw.min((half_dim - 1) as f64);
|
|
|
|
// Compute inv_freq with YaRN interpolation
|
|
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
|
|
|
|
// Linear ramp: 0 where we keep original, 1 where we interpolate
|
|
let ramp = if (high - low).abs() < 0.001 {
|
|
0.5
|
|
} else {
|
|
((i as f64 - low) / (high - low)).clamp(0.0, 1.0)
|
|
};
|
|
let extrapolation_factor = 1.0 - ramp;
|
|
|
|
inv_freq[i] = inv_freq_interpolation * (1.0 - extrapolation_factor)
|
|
+ inv_freq_extrapolation * extrapolation_factor;
|
|
}
|
|
|
|
// Attention scaling factor for YaRN: 0.1 * ln(factor) + 1.0
|
|
let attn_factor = 0.1 * factor.ln() + 1.0;
|
|
|
|
// Build cos/sin cache on CPU then upload
|
|
let total = max_seq_len * half_dim;
|
|
let mut cos_host = vec![0.0f32; total];
|
|
let mut sin_host = vec![0.0f32; total];
|
|
for pos in 0..max_seq_len {
|
|
for i in 0..half_dim {
|
|
let angle = pos as f64 * inv_freq[i];
|
|
cos_host[pos * half_dim + i] = (angle.cos() * attn_factor) as f32;
|
|
sin_host[pos * half_dim + i] = (angle.sin() * attn_factor) as f32;
|
|
}
|
|
}
|
|
|
|
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) };
|
|
cos.copy_from_host(cos_bytes).unwrap();
|
|
sin.copy_from_host(sin_bytes).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 =
|
|
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);
|
|
}
|
|
|
|
/// RoPE in-place with positions already on the GPU (u32, [num_tokens]).
|
|
/// Used by the CUDA-graph decode path, where the position lives in a
|
|
/// persistent device buffer updated outside the captured region.
|
|
pub fn rope_inplace_device_pos(x: &Tensor, cache: &RopeCache, pos_gpu: *const c_void) {
|
|
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);
|
|
|
|
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,
|
|
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 _,
|
|
pos_gpu,
|
|
num_tokens as i32,
|
|
num_heads as i32,
|
|
head_dim as i32,
|
|
xserv_cuda::current_stream_raw(),
|
|
),
|
|
_ => panic!("unsupported dtype for rope"),
|
|
}
|
|
}
|
|
}
|