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

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(