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

338
crates/xserv-model/src/decode_graph.rs
View File

@@ -18,19 +18,19 @@ use crate::kv_cache::GpuKVCache;
/// All buffers have stable GPU addresses for CUDA Graph replay.
struct DecodeBuffers {
// Hidden-size buffers: [1, hidden]
x: GpuBuffer, // running hidden state
normed: GpuBuffer, // rmsnorm output
attn_out: GpuBuffer, // attention output [1, num_heads, 1, head_dim]
attn_merged: GpuBuffer, // merge_heads output [1, hidden]
o_proj: GpuBuffer, // O projection output [1, hidden]
normed2: GpuBuffer, // post-attn norm output [1, hidden]
sum_out: GpuBuffer, // add_rmsnorm sum output [1, hidden]
down: GpuBuffer, // down projection output [1, hidden]
x: GpuBuffer, // running hidden state
normed: GpuBuffer, // rmsnorm output
attn_out: GpuBuffer, // attention output [1, num_heads, 1, head_dim]
attn_merged: GpuBuffer, // merge_heads output [1, hidden]
o_proj: GpuBuffer, // O projection output [1, hidden]
normed2: GpuBuffer, // post-attn norm output [1, hidden]
sum_out: GpuBuffer, // add_rmsnorm sum output [1, hidden]
down: GpuBuffer, // down projection output [1, hidden]
// QKV projection outputs
q_proj: GpuBuffer, // [1, num_heads * head_dim]
k_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
v_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
q_proj: GpuBuffer, // [1, num_heads * head_dim]
k_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
v_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
// Reshaped: [1, H, 1, D]
q_reshaped: GpuBuffer,
@@ -50,23 +50,23 @@ struct DecodeBuffers {
k_final: GpuBuffer,
// FFN intermediates
gate: GpuBuffer, // [1, intermediate]
up: GpuBuffer, // [1, intermediate]
silu_out: GpuBuffer, // [1, intermediate]
gate: GpuBuffer, // [1, intermediate]
up: GpuBuffer, // [1, intermediate]
silu_out: GpuBuffer, // [1, intermediate]
// GEMV fp32 accumulators (separate per output dimension)
fp32_hidden: GpuBuffer, // for hidden-sized GEMV outputs
fp32_q: GpuBuffer, // for Q projection
fp32_kv: GpuBuffer, // for K/V projection
fp32_intermediate: GpuBuffer,// for gate/up projections
fp32_vocab: GpuBuffer, // for lm_head
fp32_hidden: GpuBuffer, // for hidden-sized GEMV outputs
fp32_q: GpuBuffer, // for Q projection
fp32_kv: GpuBuffer, // for K/V projection
fp32_intermediate: GpuBuffer, // for gate/up projections
fp32_vocab: GpuBuffer, // for lm_head
// Token ID and position (GPU-resident, updated before replay)
token_id_gpu: GpuBuffer, // 4 bytes (u32)
position_gpu: GpuBuffer, // 4 bytes (u32)
token_id_gpu: GpuBuffer, // 4 bytes (u32)
position_gpu: GpuBuffer, // 4 bytes (u32)
// Final output
logits: GpuBuffer, // [1, vocab_size]
logits: GpuBuffer, // [1, vocab_size]
}
pub struct DecodeGraphState {
@@ -199,127 +199,296 @@ impl DecodeGraphState {
let cublas = cublas_handle();
// Set cuBLAS to use our stream
unsafe { dispatch::set_cublas_stream(cublas, s); }
unsafe {
dispatch::set_cublas_stream(cublas, s);
}
for (l, lw) in layers.iter().enumerate() {
// === Pre-attention graph ===
self.pre_attn_graphs[l].begin_capture(&self.stream).expect("begin pre-attn capture");
self.pre_attn_graphs[l]
.begin_capture(&self.stream)
.expect("begin pre-attn capture");
unsafe {
// RMSNorm
dispatch::rmsnorm_bf16(
self.buffers.x.as_ptr() as _, lw.input_norm, self.buffers.normed.as_mut_ptr() as _,
1, h, eps, s,
self.buffers.x.as_ptr() as _,
lw.input_norm,
self.buffers.normed.as_mut_ptr() as _,
1,
h,
eps,
s,
);
// Q projection (GEMV)
dispatch::gemv_bf16(
self.buffers.normed.as_ptr() as _, lw.q_proj_wt, self.buffers.q_proj.as_mut_ptr() as _,
self.buffers.normed.as_ptr() as _,
lw.q_proj_wt,
self.buffers.q_proj.as_mut_ptr() as _,
self.buffers.fp32_q.as_mut_ptr() as _,
h, nh * hd, s,
h,
nh * hd,
s,
);
// K projection (GEMV)
dispatch::gemv_bf16(
self.buffers.normed.as_ptr() as _, lw.k_proj_wt, self.buffers.k_proj.as_mut_ptr() as _,
self.buffers.normed.as_ptr() as _,
lw.k_proj_wt,
self.buffers.k_proj.as_mut_ptr() as _,
self.buffers.fp32_kv.as_mut_ptr() as _,
h, nkv * hd, s,
h,
nkv * hd,
s,
);
// V projection (GEMV)
dispatch::gemv_bf16(
self.buffers.normed.as_ptr() as _, lw.v_proj_wt, self.buffers.v_proj.as_mut_ptr() as _,
self.buffers.normed.as_ptr() as _,
lw.v_proj_wt,
self.buffers.v_proj.as_mut_ptr() as _,
self.buffers.fp32_kv.as_mut_ptr() as _,
h, nkv * hd, s,
h,
nkv * hd,
s,
);
// Reshape heads: [1, H*D] -> [1, H, 1, D]
dispatch::reshape_heads_bf16(self.buffers.q_proj.as_ptr() as _, self.buffers.q_reshaped.as_mut_ptr() as _, 1, nh, hd, s);
dispatch::reshape_heads_bf16(self.buffers.k_proj.as_ptr() as _, self.buffers.k_reshaped.as_mut_ptr() as _, 1, nkv, hd, s);
dispatch::reshape_heads_bf16(self.buffers.v_proj.as_ptr() as _, self.buffers.v_reshaped.as_mut_ptr() as _, 1, nkv, hd, s);
dispatch::reshape_heads_bf16(
self.buffers.q_proj.as_ptr() as _,
self.buffers.q_reshaped.as_mut_ptr() as _,
1,
nh,
hd,
s,
);
dispatch::reshape_heads_bf16(
self.buffers.k_proj.as_ptr() as _,
self.buffers.k_reshaped.as_mut_ptr() as _,
1,
nkv,
hd,
s,
);
dispatch::reshape_heads_bf16(
self.buffers.v_proj.as_ptr() as _,
self.buffers.v_reshaped.as_mut_ptr() as _,
1,
nkv,
hd,
s,
);
// QK norm (head-level rmsnorm: treat [1,H,1,D] as [H, D])
dispatch::rmsnorm_bf16(self.buffers.q_reshaped.as_ptr() as _, lw.q_norm, self.buffers.q_normed.as_mut_ptr() as _, nh, hd, eps, s);
dispatch::rmsnorm_bf16(self.buffers.k_reshaped.as_ptr() as _, lw.k_norm, self.buffers.k_normed.as_mut_ptr() as _, nkv, hd, eps, s);
dispatch::rmsnorm_bf16(
self.buffers.q_reshaped.as_ptr() as _,
lw.q_norm,
self.buffers.q_normed.as_mut_ptr() as _,
nh,
hd,
eps,
s,
);
dispatch::rmsnorm_bf16(
self.buffers.k_reshaped.as_ptr() as _,
lw.k_norm,
self.buffers.k_normed.as_mut_ptr() as _,
nkv,
hd,
eps,
s,
);
// Transpose for RoPE: [1,H,1,D] -> [1,H,D]
dispatch::transpose_hsd_to_shd_bf16(self.buffers.q_normed.as_ptr() as _, self.buffers.q_rope.as_mut_ptr() as _, 1, nh, hd, s);
dispatch::transpose_hsd_to_shd_bf16(self.buffers.k_normed.as_ptr() as _, self.buffers.k_rope.as_mut_ptr() as _, 1, nkv, hd, s);
dispatch::transpose_hsd_to_shd_bf16(
self.buffers.q_normed.as_ptr() as _,
self.buffers.q_rope.as_mut_ptr() as _,
1,
nh,
hd,
s,
);
dispatch::transpose_hsd_to_shd_bf16(
self.buffers.k_normed.as_ptr() as _,
self.buffers.k_rope.as_mut_ptr() as _,
1,
nkv,
hd,
s,
);
// RoPE (in-place, reads position_gpu)
dispatch::rope_bf16(self.buffers.q_rope.as_mut_ptr() as _, rope_cos, rope_sin, self.buffers.position_gpu.as_ptr() as _, 1, nh, hd, s);
dispatch::rope_bf16(self.buffers.k_rope.as_mut_ptr() as _, rope_cos, rope_sin, self.buffers.position_gpu.as_ptr() as _, 1, nkv, hd, s);
dispatch::rope_bf16(
self.buffers.q_rope.as_mut_ptr() as _,
rope_cos,
rope_sin,
self.buffers.position_gpu.as_ptr() as _,
1,
nh,
hd,
s,
);
dispatch::rope_bf16(
self.buffers.k_rope.as_mut_ptr() as _,
rope_cos,
rope_sin,
self.buffers.position_gpu.as_ptr() as _,
1,
nkv,
hd,
s,
);
// Transpose back: [1,H,D] -> [1,H,1,D]
dispatch::transpose_shd_to_hsd_bf16(self.buffers.q_rope.as_ptr() as _, self.buffers.q_final.as_mut_ptr() as _, 1, nh, hd, s);
dispatch::transpose_shd_to_hsd_bf16(self.buffers.k_rope.as_ptr() as _, self.buffers.k_final.as_mut_ptr() as _, 1, nkv, hd, s);
dispatch::transpose_shd_to_hsd_bf16(
self.buffers.q_rope.as_ptr() as _,
self.buffers.q_final.as_mut_ptr() as _,
1,
nh,
hd,
s,
);
dispatch::transpose_shd_to_hsd_bf16(
self.buffers.k_rope.as_ptr() as _,
self.buffers.k_final.as_mut_ptr() as _,
1,
nkv,
hd,
s,
);
}
self.pre_attn_graphs[l].end_capture(&self.stream).expect("end pre-attn capture");
self.pre_attn_graphs[l]
.end_capture(&self.stream)
.expect("end pre-attn capture");
// === Post-attention graph ===
self.post_attn_graphs[l].begin_capture(&self.stream).expect("begin post-attn capture");
self.post_attn_graphs[l]
.begin_capture(&self.stream)
.expect("begin post-attn capture");
unsafe {
// Merge heads: [1,H,1,D] -> [1, hidden]
// attn_out is written by ungraphed attention
dispatch::merge_heads_bf16(self.buffers.attn_out.as_ptr() as _, self.buffers.attn_merged.as_mut_ptr() as _, 1, nh, hd, s);
dispatch::merge_heads_bf16(
self.buffers.attn_out.as_ptr() as _,
self.buffers.attn_merged.as_mut_ptr() as _,
1,
nh,
hd,
s,
);
// O projection
dispatch::gemv_bf16(
self.buffers.attn_merged.as_ptr() as _, lw.o_proj_wt, self.buffers.o_proj.as_mut_ptr() as _,
self.buffers.attn_merged.as_ptr() as _,
lw.o_proj_wt,
self.buffers.o_proj.as_mut_ptr() as _,
self.buffers.fp32_hidden.as_mut_ptr() as _,
nh * hd, h, s,
nh * hd,
h,
s,
);
// Fused Add+RMSNorm: normed2 = rmsnorm(o_proj + x), sum_out = o_proj + x
dispatch::add_rmsnorm_bf16(
self.buffers.o_proj.as_ptr() as _, self.buffers.x.as_ptr() as _, lw.post_norm,
self.buffers.normed2.as_mut_ptr() as _, self.buffers.sum_out.as_mut_ptr() as _,
1, h, eps, s,
self.buffers.o_proj.as_ptr() as _,
self.buffers.x.as_ptr() as _,
lw.post_norm,
self.buffers.normed2.as_mut_ptr() as _,
self.buffers.sum_out.as_mut_ptr() as _,
1,
h,
eps,
s,
);
// Gate projection
dispatch::gemv_bf16(
self.buffers.normed2.as_ptr() as _, lw.gate_proj_wt, self.buffers.gate.as_mut_ptr() as _,
self.buffers.normed2.as_ptr() as _,
lw.gate_proj_wt,
self.buffers.gate.as_mut_ptr() as _,
self.buffers.fp32_intermediate.as_mut_ptr() as _,
h, inter, s,
h,
inter,
s,
);
// Up projection
dispatch::gemv_bf16(
self.buffers.normed2.as_ptr() as _, lw.up_proj_wt, self.buffers.up.as_mut_ptr() as _,
self.buffers.normed2.as_ptr() as _,
lw.up_proj_wt,
self.buffers.up.as_mut_ptr() as _,
self.buffers.fp32_intermediate.as_mut_ptr() as _,
h, inter, s,
h,
inter,
s,
);
// Fused SiLU x Mul
dispatch::silu_mul_bf16(self.buffers.gate.as_ptr() as _, self.buffers.up.as_ptr() as _, self.buffers.silu_out.as_mut_ptr() as _, inter, s);
dispatch::silu_mul_bf16(
self.buffers.gate.as_ptr() as _,
self.buffers.up.as_ptr() as _,
self.buffers.silu_out.as_mut_ptr() as _,
inter,
s,
);
// Down projection
dispatch::gemv_bf16(
self.buffers.silu_out.as_ptr() as _, lw.down_proj_wt, self.buffers.down.as_mut_ptr() as _,
self.buffers.silu_out.as_ptr() as _,
lw.down_proj_wt,
self.buffers.down.as_mut_ptr() as _,
self.buffers.fp32_hidden.as_mut_ptr() as _,
inter, h, s,
inter,
h,
s,
);
// x = sum_out + down (residual connection for next layer)
dispatch::add_bf16(self.buffers.sum_out.as_ptr() as _, self.buffers.down.as_ptr() as _, self.buffers.x.as_mut_ptr() as _, h, s);
dispatch::add_bf16(
self.buffers.sum_out.as_ptr() as _,
self.buffers.down.as_ptr() as _,
self.buffers.x.as_mut_ptr() as _,
h,
s,
);
}
self.post_attn_graphs[l].end_capture(&self.stream).expect("end post-attn capture");
self.post_attn_graphs[l]
.end_capture(&self.stream)
.expect("end post-attn capture");
}
// === Final graph: norm + lm_head ===
self.final_graph.begin_capture(&self.stream).expect("begin final capture");
self.final_graph
.begin_capture(&self.stream)
.expect("begin final capture");
unsafe {
dispatch::rmsnorm_bf16(self.buffers.x.as_ptr() as _, norm_weight, self.buffers.normed.as_mut_ptr() as _, 1, h, eps, s);
dispatch::rmsnorm_bf16(
self.buffers.x.as_ptr() as _,
norm_weight,
self.buffers.normed.as_mut_ptr() as _,
1,
h,
eps,
s,
);
dispatch::gemv_bf16(
self.buffers.normed.as_ptr() as _, lm_head_wt, self.buffers.logits.as_mut_ptr() as _,
self.buffers.normed.as_ptr() as _,
lm_head_wt,
self.buffers.logits.as_mut_ptr() as _,
self.buffers.fp32_vocab.as_mut_ptr() as _,
h, vocab, s,
h,
vocab,
s,
);
}
self.final_graph.end_capture(&self.stream).expect("end final capture");
self.final_graph
.end_capture(&self.stream)
.expect("end final capture");
// Reset cuBLAS back to null stream
unsafe { dispatch::set_cublas_stream(cublas, std::ptr::null_mut()); }
unsafe {
dispatch::set_cublas_stream(cublas, std::ptr::null_mut());
}
self.captured = true;
}
@@ -343,8 +512,14 @@ impl DecodeGraphState {
let es = 2usize; // BF16
// Upload token ID and position to fixed GPU buffers
self.buffers.token_id_gpu.copy_from_host(&token_id.to_le_bytes()).unwrap();
self.buffers.position_gpu.copy_from_host(&position.to_le_bytes()).unwrap();
self.buffers
.token_id_gpu
.copy_from_host(&token_id.to_le_bytes())
.unwrap();
self.buffers
.position_gpu
.copy_from_host(&position.to_le_bytes())
.unwrap();
// Embedding (outside graph since token_id changes each step)
unsafe {
@@ -352,13 +527,18 @@ impl DecodeGraphState {
embed_table,
self.buffers.token_id_gpu.as_ptr() as _,
self.buffers.x.as_mut_ptr() as _,
1, hidden_size, vocab_size, s,
1,
hidden_size,
vocab_size,
s,
);
}
for l in 0..self.num_layers {
// Pre-attention graph (norm + QKV + reshape + QK-norm + RoPE)
self.pre_attn_graphs[l].launch(&self.stream).expect("launch pre-attn graph");
self.pre_attn_graphs[l]
.launch(&self.stream)
.expect("launch pre-attn graph");
// Ungraphed: KV cache append
// k_final shape: [1, num_kv_heads, 1, head_dim] (after RoPE pipeline)
@@ -402,9 +582,13 @@ impl DecodeGraphState {
k_full.data_ptr() as _,
v_full.data_ptr() as _,
self.buffers.attn_out.as_mut_ptr() as _,
1, nh as i32, nkv as i32,
kv_len, hd as i32,
scale, s,
1,
nh as i32,
nkv as i32,
kv_len,
hd as i32,
scale,
s,
);
}
@@ -412,11 +596,15 @@ impl DecodeGraphState {
self.stream.synchronize().expect("sync before post-attn");
// Post-attention graph (merge + O-proj + add_rmsnorm + FFN + residual)
self.post_attn_graphs[l].launch(&self.stream).expect("launch post-attn graph");
self.post_attn_graphs[l]
.launch(&self.stream)
.expect("launch post-attn graph");
}
// Final graph (norm + lm_head)
self.final_graph.launch(&self.stream).expect("launch final graph");
self.final_graph
.launch(&self.stream)
.expect("launch final graph");
// Sync to ensure logits are ready
self.stream.synchronize().expect("sync after decode");