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

189
crates/xserv-kernels/src/activation.rs
View File

@@ -6,78 +6,189 @@ unsafe extern "C" {
fn launch_gelu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_silu_f32(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_silu_bf16(x: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_scale_f32(x: *const c_void, out: *mut c_void, scale: f32, n: i32, stream: *mut c_void);
fn launch_scale_bf16(x: *const c_void, out: *mut c_void, scale: f32, n: i32, stream: *mut c_void);
fn launch_add_f32(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_add_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_mul_f32(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_mul_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_silu_mul_bf16(gate: *const c_void, up: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_gpt_oss_glu_bf16(gate_up: *const c_void, out: *mut c_void, n_elements: i32,
alpha: f32, limit: f32, stream: *mut c_void);
fn launch_bias_add_2d_bf16(x: *const c_void, bias: *const c_void, out: *mut c_void,
rows: i32, cols: i32, stream: *mut c_void);
fn launch_scale_f32(
x: *const c_void,
out: *mut c_void,
scale: f32,
n: i32,
stream: *mut c_void,
);
fn launch_scale_bf16(
x: *const c_void,
out: *mut c_void,
scale: f32,
n: i32,
stream: *mut c_void,
);
fn launch_add_f32(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_add_bf16(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_mul_f32(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_mul_bf16(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_silu_mul_bf16(
gate: *const c_void,
up: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_gpt_oss_glu_bf16(
gate_up: *const c_void,
out: *mut c_void,
n_elements: i32,
alpha: f32,
limit: f32,
stream: *mut c_void,
);
fn launch_bias_add_2d_bf16(
x: *const c_void,
bias: *const c_void,
out: *mut c_void,
rows: i32,
cols: i32,
stream: *mut c_void,
);
}
fn dispatch_unary(x: &Tensor, f32_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
bf16_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void)) -> Tensor {
fn dispatch_unary(
x: &Tensor,
f32_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
bf16_fn: unsafe extern "C" fn(*const c_void, *mut c_void, i32, *mut c_void),
) -> Tensor {
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
let out = Tensor::empty(x.shape(), x.dtype(), x.device());
let n = x.numel();
assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
assert!(
n <= i32::MAX as usize,
"tensor too large for i32 kernel param ({n} elements)"
);
let n = n as i32;
unsafe {
match x.dtype() {
DType::F32 => f32_fn(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, xserv_cuda::current_stream_raw()),
DType::BF16 => bf16_fn(x.data_ptr() as _, out.data_ptr() as *mut c_void, n, xserv_cuda::current_stream_raw()),
DType::F32 => f32_fn(
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
n,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => bf16_fn(
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
n,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype"),
}
}
out
}
fn dispatch_binary(a: &Tensor, b: &Tensor,
f32_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void),
bf16_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void)) -> Tensor {
fn dispatch_binary(
a: &Tensor,
b: &Tensor,
f32_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void),
bf16_fn: unsafe extern "C" fn(*const c_void, *const c_void, *mut c_void, i32, *mut c_void),
) -> Tensor {
assert_eq!(a.shape(), b.shape());
assert!(a.is_contiguous() && b.is_contiguous());
assert!(matches!(a.device(), Device::Cuda(_)));
assert_eq!(a.dtype(), b.dtype());
let out = Tensor::empty(a.shape(), a.dtype(), a.device());
let n = a.numel();
assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
assert!(
n <= i32::MAX as usize,
"tensor too large for i32 kernel param ({n} elements)"
);
let n = n as i32;
unsafe {
match a.dtype() {
DType::F32 => f32_fn(a.data_ptr() as _, b.data_ptr() as _, out.data_ptr() as *mut c_void, n, xserv_cuda::current_stream_raw()),
DType::BF16 => bf16_fn(a.data_ptr() as _, b.data_ptr() as _, out.data_ptr() as *mut c_void, n, xserv_cuda::current_stream_raw()),
DType::F32 => f32_fn(
a.data_ptr() as _,
b.data_ptr() as _,
out.data_ptr() as *mut c_void,
n,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => bf16_fn(
a.data_ptr() as _,
b.data_ptr() as _,
out.data_ptr() as *mut c_void,
n,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype"),
}
}
out
}
pub fn gelu(x: &Tensor) -> Tensor { dispatch_unary(x, launch_gelu_f32, launch_gelu_bf16) }
pub fn silu(x: &Tensor) -> Tensor { dispatch_unary(x, launch_silu_f32, launch_silu_bf16) }
pub fn gelu(x: &Tensor) -> Tensor {
dispatch_unary(x, launch_gelu_f32, launch_gelu_bf16)
}
pub fn silu(x: &Tensor) -> Tensor {
dispatch_unary(x, launch_silu_f32, launch_silu_bf16)
}
pub fn scale(x: &Tensor, scale_val: f32) -> Tensor {
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
let out = Tensor::empty(x.shape(), x.dtype(), x.device());
let n = x.numel();
assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
assert!(
n <= i32::MAX as usize,
"tensor too large for i32 kernel param ({n} elements)"
);
let n = n as i32;
unsafe {
match x.dtype() {
DType::F32 => launch_scale_f32(x.data_ptr() as _, out.data_ptr() as *mut c_void, scale_val, n, xserv_cuda::current_stream_raw()),
DType::BF16 => launch_scale_bf16(x.data_ptr() as _, out.data_ptr() as *mut c_void, scale_val, n, xserv_cuda::current_stream_raw()),
DType::F32 => launch_scale_f32(
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
scale_val,
n,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => launch_scale_bf16(
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
scale_val,
n,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype for scale"),
}
}
out
}
pub fn add(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_add_f32, launch_add_bf16) }
pub fn mul(a: &Tensor, b: &Tensor) -> Tensor { dispatch_binary(a, b, launch_mul_f32, launch_mul_bf16) }
pub fn add(a: &Tensor, b: &Tensor) -> Tensor {
dispatch_binary(a, b, launch_add_f32, launch_add_bf16)
}
pub fn mul(a: &Tensor, b: &Tensor) -> Tensor {
dispatch_binary(a, b, launch_mul_f32, launch_mul_bf16)
}
/// Row-broadcast bias add: out[r, c] = x[r, c] + bias[c] (BF16 only).
pub fn bias_add_2d(x: &Tensor, bias: &Tensor) -> Tensor {
@@ -89,13 +200,22 @@ pub fn bias_add_2d(x: &Tensor, bias: &Tensor) -> Tensor {
assert!(matches!(x.device(), Device::Cuda(_)));
let rows = x.shape()[0];
let cols = x.shape()[1];
assert_eq!(bias.shape()[0], cols, "bias size {} != cols {cols}", bias.shape()[0]);
assert_eq!(
bias.shape()[0],
cols,
"bias size {} != cols {cols}",
bias.shape()[0]
);
assert!(rows * cols <= i32::MAX as usize);
let out = Tensor::empty(&[rows, cols], DType::BF16, x.device());
unsafe {
launch_bias_add_2d_bf16(
x.data_ptr() as _, bias.data_ptr() as _, out.data_ptr() as *mut c_void,
rows as i32, cols as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
bias.data_ptr() as _,
out.data_ptr() as *mut c_void,
rows as i32,
cols as i32,
xserv_cuda::current_stream_raw(),
);
}
out
@@ -110,7 +230,10 @@ pub fn silu_mul(gate: &Tensor, up: &Tensor) -> Tensor {
assert_eq!(gate.dtype(), DType::BF16, "silu_mul requires BF16");
let out = Tensor::empty(gate.shape(), gate.dtype(), gate.device());
let n = gate.numel();
assert!(n <= i32::MAX as usize, "tensor too large for i32 kernel param ({n} elements)");
assert!(
n <= i32::MAX as usize,
"tensor too large for i32 kernel param ({n} elements)"
);
let n = n as i32;
unsafe {
launch_silu_mul_bf16(

23
crates/xserv-kernels/src/argmax.rs
View File

@@ -2,8 +2,13 @@ use std::ffi::c_void;
use xserv_tensor::{DType, Device, Tensor};
unsafe extern "C" {
fn launch_argmax_bf16(logits: *const c_void, out_idx: *mut c_void,
rows: i32, cols: i32, stream: *mut c_void);
fn launch_argmax_bf16(
logits: *const c_void,
out_idx: *mut c_void,
rows: i32,
cols: i32,
stream: *mut c_void,
);
}
/// GPU argmax over the last dim of a [rows, cols] BF16 tensor.
@@ -19,7 +24,10 @@ pub fn argmax_bf16_to_host(logits: &Tensor) -> Vec<u32> {
assert_eq!(logits.ndim(), 2, "argmax expects a 2D [rows, cols] tensor");
assert_eq!(logits.dtype(), DType::BF16, "argmax kernel is BF16-only");
assert!(logits.is_contiguous(), "argmax requires contiguous input");
assert!(matches!(logits.device(), Device::Cuda(_)), "argmax requires GPU input");
assert!(
matches!(logits.device(), Device::Cuda(_)),
"argmax requires GPU input"
);
let rows = logits.shape()[0];
let cols = logits.shape()[1];
@@ -35,7 +43,8 @@ pub fn argmax_bf16_to_host(logits: &Tensor) -> Vec<u32> {
launch_argmax_bf16(
logits.data_ptr() as *const c_void,
out.as_mut_ptr() as *mut c_void,
rows as i32, cols as i32,
rows as i32,
cols as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -44,9 +53,8 @@ pub fn argmax_bf16_to_host(logits: &Tensor) -> Vec<u32> {
out.copy_to_host(&mut host_bytes).expect("argmax D2H");
drop(out); // returned to pool
let host_i32: &[i32] = unsafe {
std::slice::from_raw_parts(host_bytes.as_ptr() as *const i32, rows)
};
let host_i32: &[i32] =
unsafe { std::slice::from_raw_parts(host_bytes.as_ptr() as *const i32, rows) };
host_i32.iter().map(|&v| v as u32).collect()
}
@@ -62,4 +70,3 @@ pub fn argmax_bf16_single(logits: &Tensor) -> u32 {
};
argmax_bf16_to_host(&view)[0]
}

237
crates/xserv-kernels/src/attention.rs
View File

@@ -6,28 +6,67 @@ use crate::gemm::batched_matmul;
use crate::softmax::softmax;
unsafe extern "C" {
fn launch_causal_mask_f32(scores: *mut c_void, batch: i32, rows: i32, cols: i32,
offset: i32, stream: *mut c_void);
fn launch_causal_mask_bf16(scores: *mut c_void, batch: i32, rows: i32, cols: i32,
offset: i32, stream: *mut c_void);
fn launch_causal_mask_f32(
scores: *mut c_void,
batch: i32,
rows: i32,
cols: i32,
offset: i32,
stream: *mut c_void,
);
fn launch_causal_mask_bf16(
scores: *mut c_void,
batch: i32,
rows: i32,
cols: i32,
offset: i32,
stream: *mut c_void,
);
fn launch_flash_attention_bf16(
q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
batch: i32, num_q_heads: i32, num_kv_heads: i32,
q_len: i32, kv_len: i32, head_dim: i32,
scale: f32, causal: i32, stream: *mut c_void,
q: *const c_void,
k: *const c_void,
v: *const c_void,
o: *mut c_void,
batch: i32,
num_q_heads: i32,
num_kv_heads: i32,
q_len: i32,
kv_len: i32,
head_dim: i32,
scale: f32,
causal: i32,
stream: *mut c_void,
);
fn launch_flash_attention_sinks_bf16(
q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
q: *const c_void,
k: *const c_void,
v: *const c_void,
o: *mut c_void,
sinks: *const c_void,
batch: i32, num_q_heads: i32, num_kv_heads: i32,
q_len: i32, kv_len: i32, head_dim: i32,
scale: f32, causal: i32, window_size: i32, stream: *mut c_void,
batch: i32,
num_q_heads: i32,
num_kv_heads: i32,
q_len: i32,
kv_len: i32,
head_dim: i32,
scale: f32,
causal: i32,
window_size: i32,
stream: *mut c_void,
);
fn launch_decode_attention_bf16(
q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
batch: i32, num_q_heads: i32, num_kv_heads: i32,
kv_len: i32, head_dim: i32,
scale: f32, causal: i32, stream: *mut c_void,
q: *const c_void,
k: *const c_void,
v: *const c_void,
o: *mut c_void,
batch: i32,
num_q_heads: i32,
num_kv_heads: i32,
kv_len: i32,
head_dim: i32,
scale: f32,
causal: i32,
stream: *mut c_void,
);
fn launch_paged_decode_attention_bf16(
q: *const c_void,
@@ -36,9 +75,13 @@ unsafe extern "C" {
o: *mut c_void,
block_tables: *const i32,
context_lens: *const i32,
batch: i32, num_q_heads: i32, num_kv_heads: i32,
head_dim: i32, max_blocks_per_seq: i32,
scale: f32, stream: *mut c_void,
batch: i32,
num_q_heads: i32,
num_kv_heads: i32,
head_dim: i32,
max_blocks_per_seq: i32,
scale: f32,
stream: *mut c_void,
);
fn launch_paged_decode_attention_sinks_bf16(
q: *const c_void,
@@ -48,24 +91,40 @@ unsafe extern "C" {
block_tables: *const i32,
context_lens: *const i32,
sinks: *const c_void,
batch: i32, num_q_heads: i32, num_kv_heads: i32,
head_dim: i32, max_blocks_per_seq: i32,
scale: f32, window_size: i32, stream: *mut c_void,
batch: i32,
num_q_heads: i32,
num_kv_heads: i32,
head_dim: i32,
max_blocks_per_seq: i32,
scale: f32,
window_size: i32,
stream: *mut c_void,
);
fn launch_reshape_and_cache_bf16(
k_src: *const c_void, v_src: *const c_void,
k_pool: *mut c_void, v_pool: *mut c_void,
k_src: *const c_void,
v_src: *const c_void,
k_pool: *mut c_void,
v_pool: *mut c_void,
block_ids: *const c_void,
num_tokens: i32, num_heads: i32,
head_dim: i32, start_pos: i32, block_size: i32,
num_tokens: i32,
num_heads: i32,
head_dim: i32,
start_pos: i32,
block_size: i32,
stream: *mut c_void,
);
fn launch_reshape_and_cache_batched_bf16(
k_src: *const c_void, v_src: *const c_void,
k_pool: *mut c_void, v_pool: *mut c_void,
block_tables: *const c_void, kv_lens: *const c_void,
batch: i32, num_heads: i32,
head_dim: i32, block_size: i32, max_blocks_per_seq: i32,
k_src: *const c_void,
v_src: *const c_void,
k_pool: *mut c_void,
v_pool: *mut c_void,
block_tables: *const c_void,
kv_lens: *const c_void,
batch: i32,
num_heads: i32,
head_dim: i32,
block_size: i32,
max_blocks_per_seq: i32,
stream: *mut c_void,
);
}
@@ -84,20 +143,30 @@ unsafe extern "C" {
/// `block_ids_gpu` must contain at least `(start_pos + num_tokens + block_size - 1) / block_size`
/// valid physical block ids.
pub unsafe fn reshape_and_cache_bf16(
k_src: *const c_void, v_src: *const c_void,
k_pool_ptr: *mut c_void, v_pool_ptr: *mut c_void,
k_src: *const c_void,
v_src: *const c_void,
k_pool_ptr: *mut c_void,
v_pool_ptr: *mut c_void,
block_ids_gpu: *const i32,
num_tokens: usize, num_heads: usize,
head_dim: usize, start_pos: usize, block_size: usize,
num_tokens: usize,
num_heads: usize,
head_dim: usize,
start_pos: usize,
block_size: usize,
stream: *mut c_void,
) {
unsafe {
launch_reshape_and_cache_bf16(
k_src, v_src,
k_pool_ptr, v_pool_ptr,
k_src,
v_src,
k_pool_ptr,
v_pool_ptr,
block_ids_gpu as *const c_void,
num_tokens as i32, num_heads as i32,
head_dim as i32, start_pos as i32, block_size as i32,
num_tokens as i32,
num_heads as i32,
head_dim as i32,
start_pos as i32,
block_size as i32,
stream,
);
}
@@ -113,21 +182,32 @@ pub unsafe fn reshape_and_cache_bf16(
/// All pointers must be on the same GPU. `block_tables` and `kv_lens` must
/// already be synced to the device for the active batch.
pub unsafe fn reshape_and_cache_batched_bf16(
k_src: *const c_void, v_src: *const c_void,
k_pool_ptr: *mut c_void, v_pool_ptr: *mut c_void,
block_tables_gpu: *const i32, kv_lens_gpu: *const i32,
batch: usize, num_heads: usize,
head_dim: usize, block_size: usize, max_blocks_per_seq: usize,
k_src: *const c_void,
v_src: *const c_void,
k_pool_ptr: *mut c_void,
v_pool_ptr: *mut c_void,
block_tables_gpu: *const i32,
kv_lens_gpu: *const i32,
batch: usize,
num_heads: usize,
head_dim: usize,
block_size: usize,
max_blocks_per_seq: usize,
stream: *mut c_void,
) {
unsafe {
launch_reshape_and_cache_batched_bf16(
k_src, v_src,
k_pool_ptr, v_pool_ptr,
k_src,
v_src,
k_pool_ptr,
v_pool_ptr,
block_tables_gpu as *const c_void,
kv_lens_gpu as *const c_void,
batch as i32, num_heads as i32,
head_dim as i32, block_size as i32, max_blocks_per_seq as i32,
batch as i32,
num_heads as i32,
head_dim as i32,
block_size as i32,
max_blocks_per_seq as i32,
stream,
);
}
@@ -143,12 +223,18 @@ fn apply_causal_mask(scores: &Tensor, offset: usize) {
match scores.dtype() {
DType::F32 => launch_causal_mask_f32(
scores.data_ptr() as *mut c_void,
batch as i32, rows as i32, cols as i32, offset as i32,
batch as i32,
rows as i32,
cols as i32,
offset as i32,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => launch_causal_mask_bf16(
scores.data_ptr() as *mut c_void,
batch as i32, rows as i32, cols as i32, offset as i32,
batch as i32,
rows as i32,
cols as i32,
offset as i32,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype for causal mask"),
@@ -214,11 +300,7 @@ pub fn decode_attention(q: &Tensor, k: &Tensor, v: &Tensor) -> Tensor {
let kv_len = k.shape()[2];
let scale = 1.0 / (head_dim as f32).sqrt();
let output = Tensor::empty(
&[batch, num_q_heads, 1, head_dim],
DType::BF16,
q.device(),
);
let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
unsafe {
launch_decode_attention_bf16(
@@ -266,8 +348,14 @@ pub fn flash_attention(q: &Tensor, k: &Tensor, v: &Tensor, causal: bool) -> Tens
assert_eq!(k.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
assert_eq!(v.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
assert!(num_q_heads % num_kv_heads == 0, "num_q_heads must be divisible by num_kv_heads");
assert!(head_dim <= 128, "flash_attention supports head_dim up to 128");
assert!(
num_q_heads % num_kv_heads == 0,
"num_q_heads must be divisible by num_kv_heads"
);
assert!(
head_dim <= 128,
"flash_attention supports head_dim up to 128"
);
// Dispatch to specialized decode kernel for single-token generation
if q_len == 1 {
@@ -333,10 +421,18 @@ pub fn flash_attention_sinks(
assert_eq!(v.shape(), &[batch, num_kv_heads, kv_len, head_dim]);
assert!(num_q_heads % num_kv_heads == 0);
assert!(head_dim <= 128);
assert_eq!(sinks.shape()[0], num_q_heads, "sinks must have num_q_heads entries");
assert_eq!(
sinks.shape()[0],
num_q_heads,
"sinks must have num_q_heads entries"
);
let scale = 1.0 / (head_dim as f32).sqrt();
let output = Tensor::empty(&[batch, num_q_heads, q_len, head_dim], DType::BF16, q.device());
let output = Tensor::empty(
&[batch, num_q_heads, q_len, head_dim],
DType::BF16,
q.device(),
);
unsafe {
launch_flash_attention_sinks_bf16(
@@ -383,17 +479,20 @@ pub fn paged_decode_attention(
max_blocks_per_seq: usize,
) -> Tensor {
assert_eq!(q.ndim(), 4);
assert_eq!(q.shape()[2], 1, "paged_decode_attention requires q_len == 1");
assert_eq!(
q.shape()[2],
1,
"paged_decode_attention requires q_len == 1"
);
assert_eq!(q.dtype(), DType::BF16);
assert!(num_q_heads % num_kv_heads == 0, "GQA: num_q_heads must be divisible by num_kv_heads");
assert!(
num_q_heads % num_kv_heads == 0,
"GQA: num_q_heads must be divisible by num_kv_heads"
);
assert!(head_dim <= 128);
let scale = 1.0 / (head_dim as f32).sqrt();
let output = Tensor::empty(
&[batch, num_q_heads, 1, head_dim],
DType::BF16,
q.device(),
);
let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
unsafe {
launch_paged_decode_attention_bf16(
@@ -442,11 +541,7 @@ pub fn paged_decode_attention_sinks(
assert!(head_dim <= 128);
let scale = 1.0 / (head_dim as f32).sqrt();
let output = Tensor::empty(
&[batch, num_q_heads, 1, head_dim],
DType::BF16,
q.device(),
);
let output = Tensor::empty(&[batch, num_q_heads, 1, head_dim], DType::BF16, q.device());
unsafe {
launch_paged_decode_attention_sinks_bf16(

302
crates/xserv-kernels/src/dispatch.rs
View File

@@ -5,104 +5,302 @@ use std::ffi::c_void;
// Re-declare the extern functions we need (same as in the individual modules)
unsafe extern "C" {
fn launch_rmsnorm_bf16(x: *const c_void, gamma: *const c_void, out: *mut c_void,
rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
fn launch_add_rmsnorm_bf16(x: *const c_void, residual: *const c_void, gamma: *const c_void,
normed_out: *mut c_void, sum_out: *mut c_void,
rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
fn launch_silu_mul_bf16(gate: *const c_void, up: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_add_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void);
fn launch_embedding_bf16(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void);
fn launch_reshape_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
fn launch_merge_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
fn launch_transpose_hsd_to_shd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
fn launch_transpose_shd_to_hsd_bf16(inp: *const c_void, out: *mut c_void, seq_len: 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_gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void, y_fp32_buf: *mut c_void,
k: i32, n: i32, stream: *mut c_void);
fn launch_rmsnorm_bf16(
x: *const c_void,
gamma: *const c_void,
out: *mut c_void,
rows: i32,
hidden_size: i32,
eps: f32,
stream: *mut c_void,
);
fn launch_add_rmsnorm_bf16(
x: *const c_void,
residual: *const c_void,
gamma: *const c_void,
normed_out: *mut c_void,
sum_out: *mut c_void,
rows: i32,
hidden_size: i32,
eps: f32,
stream: *mut c_void,
);
fn launch_silu_mul_bf16(
gate: *const c_void,
up: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_add_bf16(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
);
fn launch_embedding_bf16(
table: *const c_void,
token_ids: *const c_void,
out: *mut c_void,
num_tokens: i32,
hidden_size: i32,
vocab_size: i32,
stream: *mut c_void,
);
fn launch_reshape_heads_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_merge_heads_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_transpose_hsd_to_shd_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_transpose_shd_to_hsd_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: 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_gemv_bf16(
x: *const c_void,
w: *const c_void,
y_bf16: *mut c_void,
y_fp32_buf: *mut c_void,
k: i32,
n: i32,
stream: *mut c_void,
);
fn launch_decode_attention_bf16(
q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
batch: i32, num_q_heads: i32, num_kv_heads: i32,
kv_len: i32, head_dim: i32,
scale: f32, causal: i32, stream: *mut c_void,
q: *const c_void,
k: *const c_void,
v: *const c_void,
o: *mut c_void,
batch: i32,
num_q_heads: i32,
num_kv_heads: i32,
kv_len: i32,
head_dim: i32,
scale: f32,
causal: i32,
stream: *mut c_void,
);
}
/// Raw rmsnorm dispatch: writes to pre-allocated `out`.
pub unsafe fn rmsnorm_bf16(x: *const c_void, gamma: *const c_void, out: *mut c_void,
rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void) {
pub unsafe fn rmsnorm_bf16(
x: *const c_void,
gamma: *const c_void,
out: *mut c_void,
rows: i32,
hidden_size: i32,
eps: f32,
stream: *mut c_void,
) {
launch_rmsnorm_bf16(x, gamma, out, rows, hidden_size, eps, stream);
}
/// Raw add_rmsnorm dispatch.
pub unsafe fn add_rmsnorm_bf16(x: *const c_void, residual: *const c_void, gamma: *const c_void,
normed_out: *mut c_void, sum_out: *mut c_void,
rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void) {
launch_add_rmsnorm_bf16(x, residual, gamma, normed_out, sum_out, rows, hidden_size, eps, stream);
pub unsafe fn add_rmsnorm_bf16(
x: *const c_void,
residual: *const c_void,
gamma: *const c_void,
normed_out: *mut c_void,
sum_out: *mut c_void,
rows: i32,
hidden_size: i32,
eps: f32,
stream: *mut c_void,
) {
launch_add_rmsnorm_bf16(
x,
residual,
gamma,
normed_out,
sum_out,
rows,
hidden_size,
eps,
stream,
);
}
/// Raw silu_mul dispatch.
pub unsafe fn silu_mul_bf16(gate: *const c_void, up: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void) {
pub unsafe fn silu_mul_bf16(
gate: *const c_void,
up: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
) {
launch_silu_mul_bf16(gate, up, out, n, stream);
}
/// Raw add dispatch.
pub unsafe fn add_bf16(a: *const c_void, b: *const c_void, out: *mut c_void, n: i32, stream: *mut c_void) {
pub unsafe fn add_bf16(
a: *const c_void,
b: *const c_void,
out: *mut c_void,
n: i32,
stream: *mut c_void,
) {
launch_add_bf16(a, b, out, n, stream);
}
/// Raw embedding dispatch.
pub unsafe fn embedding_bf16(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void) {
launch_embedding_bf16(table, token_ids, out, num_tokens, hidden_size, vocab_size, stream);
pub unsafe fn embedding_bf16(
table: *const c_void,
token_ids: *const c_void,
out: *mut c_void,
num_tokens: i32,
hidden_size: i32,
vocab_size: i32,
stream: *mut c_void,
) {
launch_embedding_bf16(
table,
token_ids,
out,
num_tokens,
hidden_size,
vocab_size,
stream,
);
}
/// Raw reshape_heads dispatch.
pub unsafe fn reshape_heads_bf16(inp: *const c_void, out: *mut c_void,
seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void) {
pub unsafe fn reshape_heads_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
) {
launch_reshape_heads_bf16(inp, out, seq_len, num_heads, head_dim, stream);
}
/// Raw merge_heads dispatch.
pub unsafe fn merge_heads_bf16(inp: *const c_void, out: *mut c_void,
seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void) {
pub unsafe fn merge_heads_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
) {
launch_merge_heads_bf16(inp, out, seq_len, num_heads, head_dim, stream);
}
/// Raw transpose HSD->SHD dispatch.
pub unsafe fn transpose_hsd_to_shd_bf16(inp: *const c_void, out: *mut c_void,
seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void) {
pub unsafe fn transpose_hsd_to_shd_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
) {
launch_transpose_hsd_to_shd_bf16(inp, out, seq_len, num_heads, head_dim, stream);
}
/// Raw transpose SHD->HSD dispatch.
pub unsafe fn transpose_shd_to_hsd_bf16(inp: *const c_void, out: *mut c_void,
seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void) {
pub unsafe fn transpose_shd_to_hsd_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
) {
launch_transpose_shd_to_hsd_bf16(inp, out, seq_len, num_heads, head_dim, stream);
}
/// Raw RoPE dispatch (in-place).
pub unsafe fn 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) {
launch_rope_bf16(x, cos_cache, sin_cache, positions, num_tokens, num_heads, head_dim, stream);
pub unsafe fn 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,
) {
launch_rope_bf16(
x, cos_cache, sin_cache, positions, num_tokens, num_heads, head_dim, stream,
);
}
/// Raw GEMV dispatch (BF16, M=1). Caller must provide fp32 accumulator buffer.
pub unsafe fn gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void,
y_fp32_buf: *mut c_void, k: i32, n: i32, stream: *mut c_void) {
pub unsafe fn gemv_bf16(
x: *const c_void,
w: *const c_void,
y_bf16: *mut c_void,
y_fp32_buf: *mut c_void,
k: i32,
n: i32,
stream: *mut c_void,
) {
launch_gemv_bf16(x, w, y_bf16, y_fp32_buf, k, n, stream);
}
/// Raw decode attention dispatch.
pub unsafe fn decode_attention_bf16(q: *const c_void, k: *const c_void, v: *const c_void, o: *mut c_void,
batch: i32, num_q_heads: i32, num_kv_heads: i32,
kv_len: i32, head_dim: i32,
scale: f32, stream: *mut c_void) {
launch_decode_attention_bf16(q, k, v, o, batch, num_q_heads, num_kv_heads, kv_len, head_dim, scale, 1, stream);
pub unsafe fn decode_attention_bf16(
q: *const c_void,
k: *const c_void,
v: *const c_void,
o: *mut c_void,
batch: i32,
num_q_heads: i32,
num_kv_heads: i32,
kv_len: i32,
head_dim: i32,
scale: f32,
stream: *mut c_void,
) {
launch_decode_attention_bf16(
q,
k,
v,
o,
batch,
num_q_heads,
num_kv_heads,
kv_len,
head_dim,
scale,
1,
stream,
);
}
// cuBLAS FFI

56
crates/xserv-kernels/src/embedding.rs
View File

@@ -2,10 +2,24 @@ use std::ffi::c_void;
use xserv_tensor::{DType, Device, Tensor};
unsafe extern "C" {
fn launch_embedding_f32(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void);
fn launch_embedding_bf16(table: *const c_void, token_ids: *const c_void, out: *mut c_void,
num_tokens: i32, hidden_size: i32, vocab_size: i32, stream: *mut c_void);
fn launch_embedding_f32(
table: *const c_void,
token_ids: *const c_void,
out: *mut c_void,
num_tokens: i32,
hidden_size: i32,
vocab_size: i32,
stream: *mut c_void,
);
fn launch_embedding_bf16(
table: *const c_void,
token_ids: *const c_void,
out: *mut c_void,
num_tokens: i32,
hidden_size: i32,
vocab_size: i32,
stream: *mut c_void,
);
}
/// Embedding lookup: table[token_ids[i]] for each i.
@@ -18,8 +32,14 @@ pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
let hidden_size = table.shape()[1];
let num_tokens = token_ids.len();
let vocab_size = table.shape()[0];
assert!(num_tokens <= i32::MAX as usize, "too many tokens for i32 kernel param");
assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
assert!(
num_tokens <= i32::MAX as usize,
"too many tokens for i32 kernel param"
);
assert!(
hidden_size <= i32::MAX as usize,
"hidden_size too large for i32 kernel param"
);
// Upload token_ids to GPU
let ids_bytes = unsafe {
@@ -28,11 +48,15 @@ pub fn embedding(table: &Tensor, token_ids: &[u32]) -> Tensor {
num_tokens * std::mem::size_of::<u32>(),
)
};
let mut ids_gpu = xserv_cuda::allocator::cached_alloc(ids_bytes.len()).expect("alloc token_ids");
let mut ids_gpu =
xserv_cuda::allocator::cached_alloc(ids_bytes.len()).expect("alloc token_ids");
ids_gpu.copy_from_host(ids_bytes).unwrap();
for &tid in token_ids {
assert!((tid as usize) < vocab_size, "token_id {tid} out of bounds (vocab_size={vocab_size})");
assert!(
(tid as usize) < vocab_size,
"token_id {tid} out of bounds (vocab_size={vocab_size})"
);
}
embedding_device_ids(table, ids_gpu.as_ptr() as *const c_void, num_tokens)
@@ -53,14 +77,22 @@ pub fn embedding_device_ids(table: &Tensor, ids_gpu: *const c_void, num_tokens:
unsafe {
match table.dtype() {
DType::F32 => launch_embedding_f32(
table.data_ptr() as _, ids_gpu,
table.data_ptr() as _,
ids_gpu,
out.data_ptr() as *mut c_void,
num_tokens as i32, hidden_size as i32, vocab_size as i32, xserv_cuda::current_stream_raw(),
num_tokens as i32,
hidden_size as i32,
vocab_size as i32,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => launch_embedding_bf16(
table.data_ptr() as _, ids_gpu,
table.data_ptr() as _,
ids_gpu,
out.data_ptr() as *mut c_void,
num_tokens as i32, hidden_size as i32, vocab_size as i32, xserv_cuda::current_stream_raw(),
num_tokens as i32,
hidden_size as i32,
vocab_size as i32,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype for embedding"),
}

232
crates/xserv-kernels/src/gemm.rs
View File

@@ -1,14 +1,22 @@
use std::cell::RefCell;
use std::ffi::c_void;
use xserv_cuda::error::{self, Result};
use xserv_cuda::GpuBuffer;
use xserv_cuda::error::{self, Result};
use xserv_tensor::{DType, Device, Tensor};
const CUBLAS_WORKSPACE_BYTES: usize = 32 * 1024 * 1024;
// GEMV: single-kernel, no FP32 temp buffer needed
unsafe extern "C" {
fn launch_gemv_bf16(x: *const c_void, w: *const c_void, y_bf16: *mut c_void, y_fp32_buf: *mut c_void, k: i32, n: i32, stream: *mut c_void);
fn launch_gemv_bf16(
x: *const c_void,
w: *const c_void,
y_bf16: *mut c_void,
y_fp32_buf: *mut c_void,
k: i32,
n: i32,
stream: *mut c_void,
);
}
#[derive(Debug, Clone, Copy)]
@@ -20,10 +28,42 @@ pub enum GemmBackend {
// --- FFI: custom CUDA kernels ---
unsafe extern "C" {
fn launch_gemm_naive_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
fn launch_gemm_naive_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
fn launch_gemm_tiled_f32(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
fn launch_gemm_tiled_bf16(a: *const c_void, b: *const c_void, c: *mut c_void, m: i32, n: i32, k: i32, stream: *mut c_void);
fn launch_gemm_naive_f32(
a: *const c_void,
b: *const c_void,
c: *mut c_void,
m: i32,
n: i32,
k: i32,
stream: *mut c_void,
);
fn launch_gemm_naive_bf16(
a: *const c_void,
b: *const c_void,
c: *mut c_void,
m: i32,
n: i32,
k: i32,
stream: *mut c_void,
);
fn launch_gemm_tiled_f32(
a: *const c_void,
b: *const c_void,
c: *mut c_void,
m: i32,
n: i32,
k: i32,
stream: *mut c_void,
);
fn launch_gemm_tiled_bf16(
a: *const c_void,
b: *const c_void,
c: *mut c_void,
m: i32,
n: i32,
k: i32,
stream: *mut c_void,
);
}
// --- FFI: cuBLAS ---
@@ -46,25 +86,46 @@ unsafe extern "C" {
fn cublasSetWorkspace_v2(handle: CublasHandle, workspace: *mut c_void, size: usize) -> i32;
fn cublasGemmEx(
handle: CublasHandle,
transa: i32, transb: i32,
m: i32, n: i32, k: i32,
transa: i32,
transb: i32,
m: i32,
n: i32,
k: i32,
alpha: *const c_void,
a: *const c_void, a_type: i32, lda: i32,
b: *const c_void, b_type: i32, ldb: i32,
a: *const c_void,
a_type: i32,
lda: i32,
b: *const c_void,
b_type: i32,
ldb: i32,
beta: *const c_void,
c: *mut c_void, c_type: i32, ldc: i32,
c: *mut c_void,
c_type: i32,
ldc: i32,
compute_type: i32,
algo: i32,
) -> i32;
fn cublasGemmStridedBatchedEx(
handle: CublasHandle,
transa: i32, transb: i32,
m: i32, n: i32, k: i32,
transa: i32,
transb: i32,
m: i32,
n: i32,
k: i32,
alpha: *const c_void,
a: *const c_void, a_type: i32, lda: i32, stride_a: i64,
b: *const c_void, b_type: i32, ldb: i32, stride_b: i64,
a: *const c_void,
a_type: i32,
lda: i32,
stride_a: i64,
b: *const c_void,
b_type: i32,
ldb: i32,
stride_b: i64,
beta: *const c_void,
c: *mut c_void, c_type: i32, ldc: i32, stride_c: i64,
c: *mut c_void,
c_type: i32,
ldc: i32,
stride_c: i64,
batch_count: i32,
compute_type: i32,
algo: i32,
@@ -89,9 +150,16 @@ impl CublasContext {
// set, so we keep the GpuBuffer in this struct.
let mut workspace = GpuBuffer::alloc(CUBLAS_WORKSPACE_BYTES)?;
error::check(unsafe {
cublasSetWorkspace_v2(handle, workspace.as_mut_ptr() as *mut c_void, CUBLAS_WORKSPACE_BYTES)
cublasSetWorkspace_v2(
handle,
workspace.as_mut_ptr() as *mut c_void,
CUBLAS_WORKSPACE_BYTES,
)
})?;
Ok(Self { handle, _workspace: Some(workspace) })
Ok(Self {
handle,
_workspace: Some(workspace),
})
}
}
@@ -123,9 +191,7 @@ where
/// Get the thread-local cuBLAS handle for use with dispatch module.
pub fn cublas_handle() -> CublasHandle {
CUBLAS_CTX.with(|cell| {
cell.borrow().handle
})
CUBLAS_CTX.with(|cell| cell.borrow().handle)
}
/// Matrix multiplication: C = A @ B
@@ -136,8 +202,14 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
assert_eq!(b.ndim(), 2);
assert_eq!(a.shape()[1], b.shape()[0], "inner dimension mismatch");
assert_eq!(a.dtype(), b.dtype(), "dtype mismatch");
assert!(a.is_contiguous() && b.is_contiguous(), "matmul requires contiguous tensors");
assert!(matches!(a.device(), Device::Cuda(_)), "matmul requires GPU tensors");
assert!(
a.is_contiguous() && b.is_contiguous(),
"matmul requires contiguous tensors"
);
assert!(
matches!(a.device(), Device::Cuda(_)),
"matmul requires GPU tensors"
);
let m = a.shape()[0];
let k = a.shape()[1];
@@ -154,32 +226,63 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
let null_stream = xserv_cuda::current_stream_raw();
match backend {
GemmBackend::Naive => {
unsafe {
match dtype {
DType::F32 => launch_gemm_naive_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
DType::BF16 => launch_gemm_naive_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
_ => panic!("unsupported dtype for naive GEMM"),
}
GemmBackend::Naive => unsafe {
match dtype {
DType::F32 => launch_gemm_naive_f32(
a_ptr,
b_ptr,
c_ptr,
m as i32,
n as i32,
k as i32,
null_stream,
),
DType::BF16 => launch_gemm_naive_bf16(
a_ptr,
b_ptr,
c_ptr,
m as i32,
n as i32,
k as i32,
null_stream,
),
_ => panic!("unsupported dtype for naive GEMM"),
}
}
GemmBackend::Tiled => {
unsafe {
match dtype {
DType::F32 => launch_gemm_tiled_f32(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
DType::BF16 => launch_gemm_tiled_bf16(a_ptr, b_ptr, c_ptr, m as i32, n as i32, k as i32, null_stream),
_ => panic!("unsupported dtype for tiled GEMM"),
}
},
GemmBackend::Tiled => unsafe {
match dtype {
DType::F32 => launch_gemm_tiled_f32(
a_ptr,
b_ptr,
c_ptr,
m as i32,
n as i32,
k as i32,
null_stream,
),
DType::BF16 => launch_gemm_tiled_bf16(
a_ptr,
b_ptr,
c_ptr,
m as i32,
n as i32,
k as i32,
null_stream,
),
_ => panic!("unsupported dtype for tiled GEMM"),
}
}
},
GemmBackend::CuBlas => {
if m == 1 && dtype == DType::BF16 && n >= 256 {
let mut fp32_buf = xserv_cuda::allocator::cached_alloc(n * 4).unwrap();
unsafe {
launch_gemv_bf16(
a_ptr, b_ptr, c_ptr,
a_ptr,
b_ptr,
c_ptr,
fp32_buf.as_mut_ptr() as *mut c_void,
k as i32, n as i32,
k as i32,
n as i32,
null_stream,
);
}
@@ -197,16 +300,26 @@ pub fn matmul(a: &Tensor, b: &Tensor, backend: GemmBackend) -> Tensor {
cublasSetStream_v2(handle, null_stream);
error::check(cublasGemmEx(
handle,
CUBLAS_OP_N, CUBLAS_OP_N,
n as i32, m as i32, k as i32,
CUBLAS_OP_N,
CUBLAS_OP_N,
n as i32,
m as i32,
k as i32,
&alpha as *const f32 as *const c_void,
b_ptr, b_type, n as i32,
a_ptr, a_type, k as i32,
b_ptr,
b_type,
n as i32,
a_ptr,
a_type,
k as i32,
&beta as *const f32 as *const c_void,
c_ptr, c_type, n as i32,
c_ptr,
c_type,
n as i32,
CUBLAS_COMPUTE_32F,
-1,
)).expect("cuBLAS GEMM failed");
))
.expect("cuBLAS GEMM failed");
});
}
}
@@ -264,17 +377,30 @@ pub fn batched_matmul(a: &Tensor, b: &Tensor) -> Tensor {
// Row-major trick: C = A @ B ⟺ C^T = B^T @ A^T (col-major)
error::check(cublasGemmStridedBatchedEx(
handle,
CUBLAS_OP_N, CUBLAS_OP_N,
n as i32, m as i32, k as i32,
CUBLAS_OP_N,
CUBLAS_OP_N,
n as i32,
m as i32,
k as i32,
&alpha as *const f32 as *const c_void,
b.data_ptr() as _, b_type, n as i32, stride_b,
a.data_ptr() as _, a_type, k as i32, stride_a,
b.data_ptr() as _,
b_type,
n as i32,
stride_b,
a.data_ptr() as _,
a_type,
k as i32,
stride_a,
&beta as *const f32 as *const c_void,
c.data_ptr() as *mut c_void, c_type, n as i32, stride_c,
c.data_ptr() as *mut c_void,
c_type,
n as i32,
stride_c,
batch as i32,
CUBLAS_COMPUTE_32F,
-1,
)).expect("cuBLAS batched GEMM failed");
))
.expect("cuBLAS batched GEMM failed");
});
c
}

52
crates/xserv-kernels/src/layernorm.rs
View File

@@ -2,10 +2,26 @@ use std::ffi::c_void;
use xserv_tensor::{DType, Device, Tensor};
unsafe extern "C" {
fn launch_layernorm_f32(x: *const c_void, gamma: *const c_void, beta: *const c_void,
out: *mut c_void, rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
fn launch_layernorm_bf16(x: *const c_void, gamma: *const c_void, beta: *const c_void,
out: *mut c_void, rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
fn launch_layernorm_f32(
x: *const c_void,
gamma: *const c_void,
beta: *const c_void,
out: *mut c_void,
rows: i32,
hidden_size: i32,
eps: f32,
stream: *mut c_void,
);
fn launch_layernorm_bf16(
x: *const c_void,
gamma: *const c_void,
beta: *const c_void,
out: *mut c_void,
rows: i32,
hidden_size: i32,
eps: f32,
stream: *mut c_void,
);
}
pub fn layernorm(x: &Tensor, gamma: &Tensor, beta: &Tensor, eps: f32) -> Tensor {
@@ -17,21 +33,37 @@ pub fn layernorm(x: &Tensor, gamma: &Tensor, beta: &Tensor, eps: f32) -> Tensor
assert_eq!(beta.shape(), &[hidden_size]);
let rows = x.numel() / hidden_size;
assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
assert!(
rows <= i32::MAX as usize,
"too many rows for i32 kernel param"
);
assert!(
hidden_size <= i32::MAX as usize,
"hidden_size too large for i32 kernel param"
);
let out = Tensor::empty(x.shape(), x.dtype(), x.device());
unsafe {
match x.dtype() {
DType::F32 => launch_layernorm_f32(
x.data_ptr() as _, gamma.data_ptr() as _, beta.data_ptr() as _,
x.data_ptr() as _,
gamma.data_ptr() as _,
beta.data_ptr() as _,
out.data_ptr() as *mut c_void,
rows as i32, hidden_size as i32, eps, xserv_cuda::current_stream_raw(),
rows as i32,
hidden_size as i32,
eps,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => launch_layernorm_bf16(
x.data_ptr() as _, gamma.data_ptr() as _, beta.data_ptr() as _,
x.data_ptr() as _,
gamma.data_ptr() as _,
beta.data_ptr() as _,
out.data_ptr() as *mut c_void,
rows as i32, hidden_size as i32, eps, xserv_cuda::current_stream_raw(),
rows as i32,
hidden_size as i32,
eps,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype for layernorm"),
}

14
crates/xserv-kernels/src/lib.rs
View File

@@ -14,14 +14,20 @@ pub mod transpose;
pub use activation::{add, bias_add_2d, gelu, gpt_oss_glu, mul, scale, silu, silu_mul};
pub use argmax::{argmax_bf16_single, argmax_bf16_to_host};
pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
pub use attention::{attention, decode_attention, flash_attention, flash_attention_sinks, paged_decode_attention, paged_decode_attention_sinks, reshape_and_cache_bf16, reshape_and_cache_batched_bf16};
pub use attention::{
attention, decode_attention, flash_attention, flash_attention_sinks, paged_decode_attention,
paged_decode_attention_sinks, reshape_and_cache_batched_bf16, reshape_and_cache_bf16,
};
pub use embedding::{embedding, embedding_device_ids};
pub use gemm::{batched_matmul, matmul, GemmBackend};
pub use gemm::{GemmBackend, batched_matmul, matmul};
pub use layernorm::layernorm;
pub use rmsnorm::{add_rmsnorm, rmsnorm};
pub use rope::{rope_inplace, rope_inplace_device_pos, RopeCache};
pub use rope::{RopeCache, rope_inplace, rope_inplace_device_pos};
pub use softmax::softmax;
pub use transpose::{
merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, strided_to_contiguous_gpu,
transpose_for_rope_gpu, transpose_from_rope_gpu,
};
/// Register GPU kernels with the tensor crate. Call once at startup.
pub fn init() {

217
crates/xserv-kernels/src/moe.rs
View File

@@ -1,66 +1,113 @@
use std::ffi::c_void;
use xserv_tensor::{DType, Tensor};
use crate::gemm::{cublas_handle, CublasHandle};
use crate::gemm::{CublasHandle, cublas_handle};
unsafe extern "C" {
fn launch_moe_topk_softmax_bf16(
router_logits: *const c_void,
topk_ids: *mut c_void, topk_weights: *mut c_void,
num_tokens: i32, num_experts: i32, top_k: i32,
topk_ids: *mut c_void,
topk_weights: *mut c_void,
num_tokens: i32,
num_experts: i32,
top_k: i32,
stream: *mut c_void,
);
fn launch_moe_replicate_bf16(
x: *const c_void, x_rep: *mut c_void,
num_tokens: i32, hidden: i32, local_experts: i32,
x: *const c_void,
x_rep: *mut c_void,
num_tokens: i32,
hidden: i32,
local_experts: i32,
stream: *mut c_void,
);
fn launch_moe_bias_add_3d_bf16(
x: *mut c_void, bias: *const c_void,
batch: i32, num_tokens: i32, dim: i32,
x: *mut c_void,
bias: *const c_void,
batch: i32,
num_tokens: i32,
dim: i32,
stream: *mut c_void,
);
fn launch_moe_weighted_sum_bf16(
expert_out: *const c_void,
topk_ids: *const c_void, topk_weights: *const c_void,
topk_ids: *const c_void,
topk_weights: *const c_void,
out: *mut c_void,
num_tokens: i32, hidden: i32, top_k: i32,
expert_start: i32, local_experts: i32,
num_tokens: i32,
hidden: i32,
top_k: i32,
expert_start: i32,
local_experts: i32,
stream: *mut c_void,
);
fn launch_moe_sparse_gemv_fp8_bf16(
x: *const c_void, w: *const c_void, w_scales: *const c_void,
bias: *const c_void, topk_ids: *const c_void, y: *mut c_void,
num_tokens: i32, n: i32, k: i32, top_k: i32,
expert_start: i32, local_experts: i32, x_per_slot: i32,
x: *const c_void,
w: *const c_void,
w_scales: *const c_void,
bias: *const c_void,
topk_ids: *const c_void,
y: *mut c_void,
num_tokens: i32,
n: i32,
k: i32,
top_k: i32,
expert_start: i32,
local_experts: i32,
x_per_slot: i32,
stream: *mut c_void,
);
fn launch_moe_sparse_gemv_mxfp4_bf16(
x: *const c_void, w_packed: *const c_void, w_scales: *const c_void,
bias: *const c_void, topk_ids: *const c_void, y: *mut c_void,
num_tokens: i32, n: i32, k: i32, top_k: i32,
expert_start: i32, local_experts: i32, x_per_slot: i32,
x: *const c_void,
w_packed: *const c_void,
w_scales: *const c_void,
bias: *const c_void,
topk_ids: *const c_void,
y: *mut c_void,
num_tokens: i32,
n: i32,
k: i32,
top_k: i32,
expert_start: i32,
local_experts: i32,
x_per_slot: i32,
stream: *mut c_void,
);
fn launch_moe_weighted_sum_sparse_bf16(
down: *const c_void,
topk_ids: *const c_void, topk_weights: *const c_void,
topk_ids: *const c_void,
topk_weights: *const c_void,
out: *mut c_void,
num_tokens: i32, hidden: i32, top_k: i32,
expert_start: i32, local_experts: i32,
num_tokens: i32,
hidden: i32,
top_k: i32,
expert_start: i32,
local_experts: i32,
stream: *mut c_void,
);
fn cublasGemmStridedBatchedEx(
handle: CublasHandle,
transa: i32, transb: i32,
m: i32, n: i32, k: i32,
transa: i32,
transb: i32,
m: i32,
n: i32,
k: i32,
alpha: *const c_void,
a: *const c_void, a_type: i32, lda: i32, stride_a: i64,
b: *const c_void, b_type: i32, ldb: i32, stride_b: i64,
a: *const c_void,
a_type: i32,
lda: i32,
stride_a: i64,
b: *const c_void,
b_type: i32,
ldb: i32,
stride_b: i64,
beta: *const c_void,
c: *mut c_void, c_type: i32, ldc: i32, stride_c: i64,
c: *mut c_void,
c_type: i32,
ldc: i32,
stride_c: i64,
batch_count: i32,
compute_type: i32,
algo: i32,
@@ -99,7 +146,9 @@ pub fn moe_topk_softmax(
router_logits.data_ptr() as *const c_void,
topk_ids.data_ptr() as *mut c_void,
topk_weights.data_ptr() as *mut c_void,
num_tokens as i32, num_experts as i32, top_k as i32,
num_tokens as i32,
num_experts as i32,
top_k as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -114,13 +163,19 @@ pub fn moe_replicate(x: &Tensor, local_experts: usize) -> Tensor {
assert!(x.is_contiguous());
let num_tokens = x.shape()[0];
let hidden = x.shape()[1];
let out = Tensor::empty(&[local_experts, num_tokens, hidden], DType::BF16, x.device());
let out = Tensor::empty(
&[local_experts, num_tokens, hidden],
DType::BF16,
x.device(),
);
unsafe {
launch_moe_replicate_bf16(
x.data_ptr() as *const c_void,
out.data_ptr() as *mut c_void,
num_tokens as i32, hidden as i32, local_experts as i32,
num_tokens as i32,
hidden as i32,
local_experts as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -143,7 +198,9 @@ pub fn moe_bias_add_3d(x: &Tensor, bias: &Tensor) {
launch_moe_bias_add_3d_bf16(
x.data_ptr() as *mut c_void,
bias.data_ptr() as *const c_void,
batch as i32, num_tokens as i32, dim as i32,
batch as i32,
num_tokens as i32,
dim as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -175,8 +232,11 @@ pub fn moe_weighted_sum(
topk_ids.data_ptr() as *const c_void,
topk_weights.data_ptr() as *const c_void,
out.data_ptr() as *mut c_void,
num_tokens as i32, hidden as i32, top_k as i32,
expert_start as i32, local_experts as i32,
num_tokens as i32,
hidden as i32,
top_k as i32,
expert_start as i32,
local_experts as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -198,9 +258,16 @@ pub fn moe_weighted_sum(
/// consumer must skip them (see moe_weighted_sum_sparse).
#[allow(clippy::too_many_arguments)]
pub fn moe_sparse_gemv_fp8(
x: &Tensor, w_fp8_t: &Tensor, w_scales: &Tensor, bias: &Tensor,
topk_ids: &Tensor, num_tokens: usize, top_k: usize,
expert_start: usize, local_experts: usize, x_per_slot: bool,
x: &Tensor,
w_fp8_t: &Tensor,
w_scales: &Tensor,
bias: &Tensor,
topk_ids: &Tensor,
num_tokens: usize,
top_k: usize,
expert_start: usize,
local_experts: usize,
x_per_slot: bool,
) -> Tensor {
assert_eq!(x.dtype(), DType::BF16);
assert!(x.is_contiguous());
@@ -211,7 +278,14 @@ pub fn moe_sparse_gemv_fp8(
// silently skip a K%16 tail.
assert_eq!(k % 16, 0, "sparse FP8 GEMV requires K % 16 == 0, got {k}");
assert_eq!(x.shape()[x.ndim() - 1], k);
assert_eq!(x.shape()[0], if x_per_slot { num_tokens * top_k } else { num_tokens });
assert_eq!(
x.shape()[0],
if x_per_slot {
num_tokens * top_k
} else {
num_tokens
}
);
let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
unsafe {
@@ -222,8 +296,13 @@ pub fn moe_sparse_gemv_fp8(
bias.data_ptr() as *const c_void,
topk_ids.data_ptr() as *const c_void,
y.data_ptr() as *mut c_void,
num_tokens as i32, n as i32, k as i32, top_k as i32,
expert_start as i32, local_experts as i32, x_per_slot as i32,
num_tokens as i32,
n as i32,
k as i32,
top_k as i32,
expert_start as i32,
local_experts as i32,
x_per_slot as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -234,16 +313,32 @@ pub fn moe_sparse_gemv_fp8(
/// with packed 4-bit weights [E, N, K/2] + UE8M0 block scales [E, N, K/32].
#[allow(clippy::too_many_arguments)]
pub fn moe_sparse_gemv_mxfp4(
x: &Tensor, w_packed: &Tensor, w_scales: &Tensor, bias: &Tensor,
topk_ids: &Tensor, num_tokens: usize, top_k: usize, n: usize, k: usize,
expert_start: usize, local_experts: usize, x_per_slot: bool,
x: &Tensor,
w_packed: &Tensor,
w_scales: &Tensor,
bias: &Tensor,
topk_ids: &Tensor,
num_tokens: usize,
top_k: usize,
n: usize,
k: usize,
expert_start: usize,
local_experts: usize,
x_per_slot: bool,
) -> Tensor {
assert_eq!(x.dtype(), DType::BF16);
assert!(x.is_contiguous());
// 32-element MXFP4 blocks, read as uint4 (32 nibbles) per lane.
assert_eq!(k % 32, 0, "sparse MXFP4 GEMV requires K % 32 == 0, got {k}");
assert_eq!(x.shape()[x.ndim() - 1], k);
assert_eq!(x.shape()[0], if x_per_slot { num_tokens * top_k } else { num_tokens });
assert_eq!(
x.shape()[0],
if x_per_slot {
num_tokens * top_k
} else {
num_tokens
}
);
let y = Tensor::empty(&[num_tokens, top_k, n], DType::BF16, x.device());
unsafe {
@@ -254,8 +349,13 @@ pub fn moe_sparse_gemv_mxfp4(
bias.data_ptr() as *const c_void,
topk_ids.data_ptr() as *const c_void,
y.data_ptr() as *mut c_void,
num_tokens as i32, n as i32, k as i32, top_k as i32,
expert_start as i32, local_experts as i32, x_per_slot as i32,
num_tokens as i32,
n as i32,
k as i32,
top_k as i32,
expert_start as i32,
local_experts as i32,
x_per_slot as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -286,8 +386,11 @@ pub fn moe_weighted_sum_sparse(
topk_ids.data_ptr() as *const c_void,
topk_weights.data_ptr() as *const c_void,
out.data_ptr() as *mut c_void,
num_tokens as i32, hidden as i32, top_k as i32,
expert_start as i32, local_experts as i32,
num_tokens as i32,
hidden as i32,
top_k as i32,
expert_start as i32,
local_experts as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -341,13 +444,25 @@ pub fn batched_gemm_strided(a: &Tensor, b: &Tensor) -> Tensor {
cublasSetStream_v2(handle, xserv_cuda::current_stream_raw());
let status = cublasGemmStridedBatchedEx(
handle,
0, 0, // CUBLAS_OP_N, CUBLAS_OP_N
n as i32, m as i32, k as i32,
0,
0, // CUBLAS_OP_N, CUBLAS_OP_N
n as i32,
m as i32,
k as i32,
&alpha as *const f32 as *const c_void,
b.data_ptr() as *const c_void, CUDA_R_16BF, n as i32, stride_b,
a.data_ptr() as *const c_void, CUDA_R_16BF, k as i32, stride_a,
b.data_ptr() as *const c_void,
CUDA_R_16BF,
n as i32,
stride_b,
a.data_ptr() as *const c_void,
CUDA_R_16BF,
k as i32,
stride_a,
&beta as *const f32 as *const c_void,
c.data_ptr() as *mut c_void, CUDA_R_16BF, n as i32, stride_c,
c.data_ptr() as *mut c_void,
CUDA_R_16BF,
n as i32,
stride_c,
batch as i32,
CUBLAS_COMPUTE_32F,
CUBLAS_GEMM_DEFAULT,

220
crates/xserv-kernels/src/quantization.rs
View File

@@ -13,30 +13,46 @@ unsafe extern "C" {
src: *const c_void,
scales: *const c_void,
dst: *mut c_void,
num_experts: i32, rows: i32, cols: i32,
num_experts: i32,
rows: i32,
cols: i32,
stream: *mut c_void,
);
fn launch_quantize_bf16_to_fp8e4m3_rowwise(
src: *const c_void,
dst: *mut c_void,
scales: *mut c_void,
num_rows: i32, cols: i32,
num_rows: i32,
cols: i32,
stream: *mut c_void,
);
fn launch_rowwise_scale_moe_bf16(
data: *mut c_void,
a_scales: *const c_void,
b_scales: *const c_void,
num_rows: i32, cols: i32, tokens: i32,
num_rows: i32,
cols: i32,
tokens: i32,
stream: *mut c_void,
);
fn launch_batched_gemv_mxfp4_bf16(
x: *const c_void, w_packed: *const c_void, w_scales: *const c_void, y: *mut c_void,
e: i32, n: i32, k: i32, stream: *mut c_void,
x: *const c_void,
w_packed: *const c_void,
w_scales: *const c_void,
y: *mut c_void,
e: i32,
n: i32,
k: i32,
stream: *mut c_void,
);
fn launch_dequant_mxfp4_to_bf16_t(
w_packed: *const c_void, w_scales: *const c_void, out: *mut c_void,
e: i32, n: i32, k: i32, stream: *mut c_void,
w_packed: *const c_void,
w_scales: *const c_void,
out: *mut c_void,
e: i32,
n: i32,
k: i32,
stream: *mut c_void,
);
}
@@ -66,34 +82,68 @@ struct CublasLtMatmulHeuristicResult {
unsafe extern "C" {
fn cublasLtCreate(handle: *mut CublasLtHandle) -> i32;
fn cublasLtDestroy(handle: CublasLtHandle) -> i32;
fn cublasLtMatmulDescCreate(desc: *mut CublasLtMatmulDesc, compute_type: i32, scale_type: i32) -> i32;
fn cublasLtMatmulDescCreate(
desc: *mut CublasLtMatmulDesc,
compute_type: i32,
scale_type: i32,
) -> i32;
fn cublasLtMatmulDescDestroy(desc: CublasLtMatmulDesc) -> i32;
fn cublasLtMatmulDescSetAttribute(desc: CublasLtMatmulDesc, attr: i32, buf: *const c_void, size: usize) -> i32;
fn cublasLtMatrixLayoutCreate(layout: *mut CublasLtMatrixLayout, dtype: i32, rows: u64, cols: u64, ld: i64) -> i32;
fn cublasLtMatmulDescSetAttribute(
desc: CublasLtMatmulDesc,
attr: i32,
buf: *const c_void,
size: usize,
) -> i32;
fn cublasLtMatrixLayoutCreate(
layout: *mut CublasLtMatrixLayout,
dtype: i32,
rows: u64,
cols: u64,
ld: i64,
) -> i32;
fn cublasLtMatrixLayoutDestroy(layout: CublasLtMatrixLayout) -> i32;
fn cublasLtMatrixLayoutSetAttribute(layout: CublasLtMatrixLayout, attr: i32, buf: *const c_void, size: usize) -> i32;
fn cublasLtMatrixLayoutSetAttribute(
layout: CublasLtMatrixLayout,
attr: i32,
buf: *const c_void,
size: usize,
) -> i32;
fn cublasLtMatmulPreferenceCreate(pref: *mut CublasLtMatmulPreference) -> i32;
fn cublasLtMatmulPreferenceDestroy(pref: CublasLtMatmulPreference) -> i32;
fn cublasLtMatmulPreferenceSetAttribute(pref: CublasLtMatmulPreference, attr: i32, buf: *const c_void, size: usize) -> i32;
fn cublasLtMatmulPreferenceSetAttribute(
pref: CublasLtMatmulPreference,
attr: i32,
buf: *const c_void,
size: usize,
) -> i32;
fn cublasLtMatmulAlgoGetHeuristic(
handle: CublasLtHandle, desc: CublasLtMatmulDesc,
a_layout: CublasLtMatrixLayout, b_layout: CublasLtMatrixLayout,
c_layout: CublasLtMatrixLayout, d_layout: CublasLtMatrixLayout,
handle: CublasLtHandle,
desc: CublasLtMatmulDesc,
a_layout: CublasLtMatrixLayout,
b_layout: CublasLtMatrixLayout,
c_layout: CublasLtMatrixLayout,
d_layout: CublasLtMatrixLayout,
pref: CublasLtMatmulPreference,
requested: i32,
results: *mut CublasLtMatmulHeuristicResult,
found: *mut i32,
) -> i32;
fn cublasLtMatmul(
handle: CublasLtHandle, desc: CublasLtMatmulDesc,
handle: CublasLtHandle,
desc: CublasLtMatmulDesc,
alpha: *const c_void,
a: *const c_void, a_layout: CublasLtMatrixLayout,
b: *const c_void, b_layout: CublasLtMatrixLayout,
a: *const c_void,
a_layout: CublasLtMatrixLayout,
b: *const c_void,
b_layout: CublasLtMatrixLayout,
beta: *const c_void,
c: *const c_void, c_layout: CublasLtMatrixLayout,
d: *mut c_void, d_layout: CublasLtMatrixLayout,
c: *const c_void,
c_layout: CublasLtMatrixLayout,
d: *mut c_void,
d_layout: CublasLtMatrixLayout,
algo: *const CublasLtMatmulAlgo,
workspace: *mut c_void, workspace_size: usize,
workspace: *mut c_void,
workspace_size: usize,
stream: *mut c_void,
) -> i32;
}
@@ -153,8 +203,15 @@ impl CublasLtContext {
assert_eq!(status, 0, "cublasLtCreate failed: {status}");
let workspace = GpuBuffer::alloc(WORKSPACE_BYTES).expect("alloc cublasLt workspace");
let mut one_buf = GpuBuffer::alloc(4).expect("alloc cublasLt fp8 scale");
one_buf.copy_from_host(&1.0f32.to_le_bytes()).expect("init fp8 scale");
Self { handle, workspace, one_buf, plans: HashMap::new() }
one_buf
.copy_from_host(&1.0f32.to_le_bytes())
.expect("init fp8 scale");
Self {
handle,
workspace,
one_buf,
plans: HashMap::new(),
}
}
/// Get the cached strided-batched plan for (m, n, k, batch), building it on
@@ -210,10 +267,25 @@ unsafe fn build_fp8_plan(
// transA=T (required for FP8 on Blackwell)
let trans_a: i32 = 1;
cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSA, &trans_a as *const i32 as _, 4);
cublasLtMatmulDescSetAttribute(
desc,
CUBLASLT_MATMUL_DESC_TRANSA,
&trans_a as *const i32 as _,
4,
);
let ptr_sz = std::mem::size_of::<*const c_void>();
cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &one_ptr as *const _ as _, ptr_sz);
cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &one_ptr as *const _ as _, ptr_sz);
cublasLtMatmulDescSetAttribute(
desc,
CUBLASLT_MATMUL_DESC_A_SCALE_POINTER,
&one_ptr as *const _ as _,
ptr_sz,
);
cublasLtMatmulDescSetAttribute(
desc,
CUBLASLT_MATMUL_DESC_B_SCALE_POINTER,
&one_ptr as *const _ as _,
ptr_sz,
);
// Per-expert strides in ELEMENTS for the strided-batch layout.
let stride_a = (n * k) as i64; // weights [N, K]
@@ -221,10 +293,18 @@ unsafe fn build_fp8_plan(
let stride_c = (m * n) as i64; // output [M, N]
let bc = batch as i32;
let set_batch = |layout: CublasLtMatrixLayout, stride: i64| {
cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
&bc as *const i32 as _, 4);
cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
&stride as *const i64 as _, 8);
cublasLtMatrixLayoutSetAttribute(
layout,
CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT,
&bc as *const i32 as _,
4,
);
cublasLtMatrixLayoutSetAttribute(
layout,
CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET,
&stride as *const i64 as _,
8,
);
};
// "A" layout (weights, transposed): physical (K, N) col-major, ld=K
@@ -246,20 +326,39 @@ unsafe fn build_fp8_plan(
let mut pref: CublasLtMatmulPreference = std::ptr::null_mut();
cublasLtMatmulPreferenceCreate(&mut pref);
let ws_bytes = WORKSPACE_BYTES as u64;
cublasLtMatmulPreferenceSetAttribute(pref, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &ws_bytes as *const u64 as _, 8);
cublasLtMatmulPreferenceSetAttribute(
pref,
CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES,
&ws_bytes as *const u64 as _,
8,
);
let mut heuristic = std::mem::zeroed::<CublasLtMatmulHeuristicResult>();
let mut found: i32 = 0;
let status = cublasLtMatmulAlgoGetHeuristic(
handle, desc, a_layout, b_layout, c_layout, d_layout,
pref, 1, &mut heuristic, &mut found,
handle,
desc,
a_layout,
b_layout,
c_layout,
d_layout,
pref,
1,
&mut heuristic,
&mut found,
);
assert!(
status == 0 && found > 0,
"cublasLtMatmulAlgoGetHeuristic failed for batched FP8 GEMM (m={m}, n={n}, k={k}, batch={batch}): status={status}, found={found}"
);
assert!(status == 0 && found > 0,
"cublasLtMatmulAlgoGetHeuristic failed for batched FP8 GEMM (m={m}, n={n}, k={k}, batch={batch}): status={status}, found={found}");
cublasLtMatmulPreferenceDestroy(pref);
Fp8Plan {
desc, a_layout, b_layout, c_layout, d_layout,
desc,
a_layout,
b_layout,
c_layout,
d_layout,
algo: heuristic.algo,
workspace_size: heuristic.workspace_size,
}
@@ -299,7 +398,9 @@ pub fn dequant_fp8_to_bf16(src: &Tensor, scales: &Tensor) -> Tensor {
src.data_ptr() as *const c_void,
scales.data_ptr() as *const c_void,
out.data_ptr() as *mut c_void,
num_experts as i32, rows as i32, cols as i32,
num_experts as i32,
rows as i32,
cols as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -329,7 +430,8 @@ pub fn quantize_bf16_to_fp8_rowwise(src: &Tensor) -> (Tensor, Tensor) {
src.data_ptr() as *const c_void,
fp8_out.data_ptr() as *mut c_void,
scales.data_ptr() as *mut c_void,
num_rows as i32, cols as i32,
num_rows as i32,
cols as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -392,23 +494,27 @@ pub fn batched_gemm_fp8(
unsafe {
let status = cublasLtMatmul(
handle, plan.desc,
handle,
plan.desc,
&alpha as *const f32 as _,
b_fp8_t.data_ptr() as *const c_void, // cuBLASLt "A" = weights
plan.a_layout,
a_fp8.data_ptr() as *const c_void, // cuBLASLt "B" = activations
a_fp8.data_ptr() as *const c_void, // cuBLASLt "B" = activations
plan.b_layout,
&beta as *const f32 as _,
c.data_ptr() as *const c_void, // C (unused with beta=0)
c.data_ptr() as *const c_void, // C (unused with beta=0)
plan.c_layout,
c.data_ptr() as *mut c_void, // D = output
c.data_ptr() as *mut c_void, // D = output
plan.d_layout,
&plan.algo,
ws_ptr,
plan.workspace_size,
xserv_cuda::current_stream_raw(),
);
assert_eq!(status, 0, "batched cublasLtMatmul FP8 failed: status={status}");
assert_eq!(
status, 0,
"batched cublasLtMatmul FP8 failed: status={status}"
);
}
});
@@ -423,7 +529,9 @@ pub fn batched_gemm_fp8(
c.data_ptr() as *mut c_void,
a_scales.data_ptr() as *const c_void,
b_scales.data_ptr() as *const c_void,
total_rows, n as i32, m as i32,
total_rows,
n as i32,
m as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -442,7 +550,13 @@ pub fn batched_gemm_fp8(
/// w_scales: [E, N, K/32] byte tensor — UE8M0 scale per 32-element block
///
/// Returns: [E, N] BF16, where y[e,n] = sum_k x[e,k] * dequant(W[e,n,k]).
pub fn batched_gemv_mxfp4(x: &Tensor, w_packed: &Tensor, w_scales: &Tensor, n: usize, k: usize) -> Tensor {
pub fn batched_gemv_mxfp4(
x: &Tensor,
w_packed: &Tensor,
w_scales: &Tensor,
n: usize,
k: usize,
) -> Tensor {
assert_eq!(x.dtype(), DType::BF16);
assert!(x.is_contiguous());
let e = x.shape()[0];
@@ -455,7 +569,9 @@ pub fn batched_gemv_mxfp4(x: &Tensor, w_packed: &Tensor, w_scales: &Tensor, n: u
w_packed.data_ptr() as *const c_void,
w_scales.data_ptr() as *const c_void,
y.data_ptr() as *mut c_void,
e as i32, n as i32, k as i32,
e as i32,
n as i32,
k as i32,
xserv_cuda::current_stream_raw(),
);
}
@@ -464,14 +580,22 @@ pub fn batched_gemv_mxfp4(x: &Tensor, w_packed: &Tensor, w_scales: &Tensor, n: u
/// Dequantize MXFP4 weights [E, N, K] → BF16 [E, K, N] for the prefill GEMM path
/// (the BF16 batched GEMM expects weights as [E, K, N]).
pub fn dequant_mxfp4_to_bf16_t(w_packed: &Tensor, w_scales: &Tensor, e: usize, n: usize, k: usize) -> Tensor {
pub fn dequant_mxfp4_to_bf16_t(
w_packed: &Tensor,
w_scales: &Tensor,
e: usize,
n: usize,
k: usize,
) -> Tensor {
let out = Tensor::empty(&[e, k, n], DType::BF16, w_packed.device());
unsafe {
launch_dequant_mxfp4_to_bf16_t(
w_packed.data_ptr() as *const c_void,
w_scales.data_ptr() as *const c_void,
out.data_ptr() as *mut c_void,
e as i32, n as i32, k as i32,
e as i32,
n as i32,
k as i32,
xserv_cuda::current_stream_raw(),
);
}

74
crates/xserv-kernels/src/rmsnorm.rs
View File

@@ -2,13 +2,35 @@ use std::ffi::c_void;
use xserv_tensor::{DType, Device, Tensor};
unsafe extern "C" {
fn launch_rmsnorm_f32(x: *const c_void, gamma: *const c_void, out: *mut c_void,
rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
fn launch_rmsnorm_bf16(x: *const c_void, gamma: *const c_void, out: *mut c_void,
rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
fn launch_add_rmsnorm_bf16(x: *const c_void, residual: *const c_void, gamma: *const c_void,
normed_out: *mut c_void, sum_out: *mut c_void,
rows: i32, hidden_size: i32, eps: f32, stream: *mut c_void);
fn launch_rmsnorm_f32(
x: *const c_void,
gamma: *const c_void,
out: *mut c_void,
rows: i32,
hidden_size: i32,
eps: f32,
stream: *mut c_void,
);
fn launch_rmsnorm_bf16(
x: *const c_void,
gamma: *const c_void,
out: *mut c_void,
rows: i32,
hidden_size: i32,
eps: f32,
stream: *mut c_void,
);
fn launch_add_rmsnorm_bf16(
x: *const c_void,
residual: *const c_void,
gamma: *const c_void,
normed_out: *mut c_void,
sum_out: *mut c_void,
rows: i32,
hidden_size: i32,
eps: f32,
stream: *mut c_void,
);
}
pub fn rmsnorm(x: &Tensor, gamma: &Tensor, eps: f32) -> Tensor {
@@ -20,19 +42,35 @@ pub fn rmsnorm(x: &Tensor, gamma: &Tensor, eps: f32) -> Tensor {
assert_eq!(x.dtype(), gamma.dtype());
let rows = x.numel() / hidden_size;
assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
assert!(
rows <= i32::MAX as usize,
"too many rows for i32 kernel param"
);
assert!(
hidden_size <= i32::MAX as usize,
"hidden_size too large for i32 kernel param"
);
let out = Tensor::empty(x.shape(), x.dtype(), x.device());
unsafe {
match x.dtype() {
DType::F32 => launch_rmsnorm_f32(
x.data_ptr() as _, gamma.data_ptr() as _, out.data_ptr() as *mut c_void,
rows as i32, hidden_size as i32, eps, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
gamma.data_ptr() as _,
out.data_ptr() as *mut c_void,
rows as i32,
hidden_size as i32,
eps,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => launch_rmsnorm_bf16(
x.data_ptr() as _, gamma.data_ptr() as _, out.data_ptr() as *mut c_void,
rows as i32, hidden_size as i32, eps, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
gamma.data_ptr() as _,
out.data_ptr() as *mut c_void,
rows as i32,
hidden_size as i32,
eps,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype for rmsnorm"),
}
@@ -56,8 +94,14 @@ pub fn add_rmsnorm(x: &Tensor, residual: &Tensor, gamma: &Tensor, eps: f32) -> (
assert_eq!(gamma.shape(), &[hidden_size]);
let rows = x.numel() / hidden_size;
assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
assert!(hidden_size <= i32::MAX as usize, "hidden_size too large for i32 kernel param");
assert!(
rows <= i32::MAX as usize,
"too many rows for i32 kernel param"
);
assert!(
hidden_size <= i32::MAX as usize,
"hidden_size too large for i32 kernel param"
);
let normed_out = Tensor::empty(x.shape(), DType::BF16, x.device());
let sum_out = Tensor::empty(x.shape(), DType::BF16, x.device());

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"),

40
crates/xserv-kernels/src/softmax.rs
View File

@@ -2,8 +2,20 @@ use std::ffi::c_void;
use xserv_tensor::{DType, Device, Tensor};
unsafe extern "C" {
fn launch_softmax_f32(x: *const c_void, out: *mut c_void, rows: i32, cols: i32, stream: *mut c_void);
fn launch_softmax_bf16(x: *const c_void, out: *mut c_void, rows: i32, cols: i32, stream: *mut c_void);
fn launch_softmax_f32(
x: *const c_void,
out: *mut c_void,
rows: i32,
cols: i32,
stream: *mut c_void,
);
fn launch_softmax_bf16(
x: *const c_void,
out: *mut c_void,
rows: i32,
cols: i32,
stream: *mut c_void,
);
}
/// Softmax along the last dimension.
@@ -14,19 +26,31 @@ pub fn softmax(x: &Tensor) -> Tensor {
let cols = *x.shape().last().unwrap();
let rows = x.numel() / cols;
assert!(rows <= i32::MAX as usize, "too many rows for i32 kernel param");
assert!(cols <= i32::MAX as usize, "cols too large for i32 kernel param");
assert!(
rows <= i32::MAX as usize,
"too many rows for i32 kernel param"
);
assert!(
cols <= i32::MAX as usize,
"cols too large for i32 kernel param"
);
let out = Tensor::empty(x.shape(), x.dtype(), x.device());
unsafe {
match x.dtype() {
DType::F32 => launch_softmax_f32(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
rows as i32, cols as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
rows as i32,
cols as i32,
xserv_cuda::current_stream_raw(),
),
DType::BF16 => launch_softmax_bf16(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
rows as i32, cols as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
rows as i32,
cols as i32,
xserv_cuda::current_stream_raw(),
),
_ => panic!("unsupported dtype for softmax"),
}

188
crates/xserv-kernels/src/transpose.rs
View File

@@ -2,19 +2,79 @@ use std::ffi::c_void;
use xserv_tensor::{DType, Device, Tensor};
unsafe extern "C" {
fn launch_reshape_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
fn launch_merge_heads_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
fn launch_transpose_hsd_to_shd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
fn launch_transpose_shd_to_hsd_bf16(inp: *const c_void, out: *mut c_void, seq_len: i32, num_heads: i32, head_dim: i32, stream: *mut c_void);
fn launch_repeat_kv_bf16(inp: *const c_void, out: *mut c_void, kv_heads: i32, n_rep: i32, seq_len: i32, head_dim: i32, stream: *mut c_void);
fn launch_strided_copy_bf16(inp: *const c_void, out: *mut c_void, numel: i32, ndim: i32,
shape0: i32, shape1: i32, shape2: i32, shape3: i32,
in_stride0: i32, in_stride1: i32, in_stride2: i32, in_stride3: i32,
in_offset: i32, stream: *mut c_void);
fn launch_strided_copy_f32(inp: *const c_void, out: *mut c_void, numel: i32, ndim: i32,
shape0: i32, shape1: i32, shape2: i32, shape3: i32,
in_stride0: i32, in_stride1: i32, in_stride2: i32, in_stride3: i32,
in_offset: i32, stream: *mut c_void);
fn launch_reshape_heads_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_merge_heads_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_transpose_hsd_to_shd_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_transpose_shd_to_hsd_bf16(
inp: *const c_void,
out: *mut c_void,
seq_len: i32,
num_heads: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_repeat_kv_bf16(
inp: *const c_void,
out: *mut c_void,
kv_heads: i32,
n_rep: i32,
seq_len: i32,
head_dim: i32,
stream: *mut c_void,
);
fn launch_strided_copy_bf16(
inp: *const c_void,
out: *mut c_void,
numel: i32,
ndim: i32,
shape0: i32,
shape1: i32,
shape2: i32,
shape3: i32,
in_stride0: i32,
in_stride1: i32,
in_stride2: i32,
in_stride3: i32,
in_offset: i32,
stream: *mut c_void,
);
fn launch_strided_copy_f32(
inp: *const c_void,
out: *mut c_void,
numel: i32,
ndim: i32,
shape0: i32,
shape1: i32,
shape2: i32,
shape3: i32,
in_stride0: i32,
in_stride1: i32,
in_stride2: i32,
in_stride3: i32,
in_offset: i32,
stream: *mut c_void,
);
}
/// [S, H*D] → [1, H, S, D] on GPU (BF16)
@@ -24,8 +84,12 @@ pub fn reshape_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim:
let out = Tensor::empty(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
unsafe {
launch_reshape_heads_bf16(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
seq_len as i32, num_heads as i32, head_dim as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
seq_len as i32,
num_heads as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
);
}
out
@@ -39,36 +103,58 @@ pub fn merge_heads_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: u
let out = Tensor::empty(&[seq_len, hidden], DType::BF16, x.device());
unsafe {
launch_merge_heads_bf16(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
seq_len as i32, num_heads as i32, head_dim as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
seq_len as i32,
num_heads as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
);
}
out
}
/// [1, H, S, D] → [S, H, D] for RoPE on GPU (BF16)
pub fn transpose_for_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
pub fn transpose_for_rope_gpu(
x: &Tensor,
seq_len: usize,
num_heads: usize,
head_dim: usize,
) -> Tensor {
assert_eq!(x.dtype(), DType::BF16);
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
let out = Tensor::empty(&[seq_len, num_heads, head_dim], DType::BF16, x.device());
unsafe {
launch_transpose_hsd_to_shd_bf16(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
seq_len as i32, num_heads as i32, head_dim as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
seq_len as i32,
num_heads as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
);
}
out
}
/// [S, H, D] → [1, H, S, D] after RoPE on GPU (BF16)
pub fn transpose_from_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, head_dim: usize) -> Tensor {
pub fn transpose_from_rope_gpu(
x: &Tensor,
seq_len: usize,
num_heads: usize,
head_dim: usize,
) -> Tensor {
assert_eq!(x.dtype(), DType::BF16);
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
let out = Tensor::empty(&[1, num_heads, seq_len, head_dim], DType::BF16, x.device());
unsafe {
launch_transpose_shd_to_hsd_bf16(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
seq_len as i32, num_heads as i32, head_dim as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
seq_len as i32,
num_heads as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
);
}
out
@@ -76,7 +162,9 @@ pub fn transpose_from_rope_gpu(x: &Tensor, seq_len: usize, num_heads: usize, hea
/// [1, KV_H, S, D] → [1, KV_H*n_rep, S, D] on GPU (BF16)
pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
if n_rep == 1 { return x.clone(); }
if n_rep == 1 {
return x.clone();
}
assert_eq!(x.dtype(), DType::BF16);
assert!(x.is_contiguous() && matches!(x.device(), Device::Cuda(_)));
let kv_heads = x.shape()[1];
@@ -86,8 +174,13 @@ pub fn repeat_kv_gpu(x: &Tensor, n_rep: usize) -> Tensor {
let out = Tensor::empty(&[1, new_heads, seq_len, head_dim], DType::BF16, x.device());
unsafe {
launch_repeat_kv_bf16(
x.data_ptr() as _, out.data_ptr() as *mut c_void,
kv_heads as i32, n_rep as i32, seq_len as i32, head_dim as i32, xserv_cuda::current_stream_raw(),
x.data_ptr() as _,
out.data_ptr() as *mut c_void,
kv_heads as i32,
n_rep as i32,
seq_len as i32,
head_dim as i32,
xserv_cuda::current_stream_raw(),
);
}
out
@@ -122,20 +215,41 @@ pub fn strided_to_contiguous_gpu(x: &Tensor) -> Tensor {
unsafe {
match x.dtype() {
DType::BF16 => launch_strided_copy_bf16(
storage_ptr as _, out.data_ptr() as *mut c_void,
numel as i32, ndim as i32,
shape4[0], shape4[1], shape4[2], shape4[3],
strides4[0], strides4[1], strides4[2], strides4[3],
in_offset, xserv_cuda::current_stream_raw(),
storage_ptr as _,
out.data_ptr() as *mut c_void,
numel as i32,
ndim as i32,
shape4[0],
shape4[1],
shape4[2],
shape4[3],
strides4[0],
strides4[1],
strides4[2],
strides4[3],
in_offset,
xserv_cuda::current_stream_raw(),
),
DType::F32 => launch_strided_copy_f32(
storage_ptr as _, out.data_ptr() as *mut c_void,
numel as i32, ndim as i32,
shape4[0], shape4[1], shape4[2], shape4[3],
strides4[0], strides4[1], strides4[2], strides4[3],
in_offset, xserv_cuda::current_stream_raw(),
storage_ptr as _,
out.data_ptr() as *mut c_void,
numel as i32,
ndim as i32,
shape4[0],
shape4[1],
shape4[2],
shape4[3],
strides4[0],
strides4[1],
strides4[2],
strides4[3],
in_offset,
xserv_cuda::current_stream_raw(),
),
_ => panic!(
"strided_to_contiguous_gpu: unsupported dtype {:?}",
x.dtype()
),
_ => panic!("strided_to_contiguous_gpu: unsupported dtype {:?}", x.dtype()),
}
}
out