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

perf: GPU transpose/reshape/repeat_kv kernels (eliminate CPU round-trips)

Browse Source 
New CUDA kernels (csrc/embedding/transpose.cu):
- reshape_heads_bf16: [S, H*D] → [1, H, S, D]
- merge_heads_bf16: [1, H, S, D] → [S, H*D]
- transpose_hsd_to_shd_bf16: [1, H, S, D] → [S, H, D] (for RoPE)
- transpose_shd_to_hsd_bf16: [S, H, D] → [1, H, S, D] (from RoPE)
- repeat_kv_bf16: [1, KV_H, S, D] → [1, KV_H*n_rep, S, D]

Rust wrappers (xserv-kernels/src/transpose.rs):
- reshape_heads_gpu, merge_heads_gpu, transpose_for/from_rope_gpu, repeat_kv_gpu

Qwen3 forward_gpu_cache now uses all GPU kernels — zero CPU data round-trips.

Result: 50/50 self-consistent, 3-5% faster (TBT 142→137ms)
Remaining bottleneck: ~900 device::synchronize() calls + 252 cuBLAS handle
creations per token (Phase 15 targets)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-22 12:01:07 +08:00
parent 2d48f25e66
commit 2be27d6d94
5 changed files with 273 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
crates/xserv-model/src/qwen3.rs
View File

@@ -147,7 +147,7 @@ impl Qwen3 {
matmul_2d(&x, &self.lm_head_t)
}
/// Forward with GPU-resident KV cache (no CPU round-trips for KV).
/// Forward with GPU-resident KV cache and GPU transpose/reshape kernels.
pub fn forward_gpu_cache(&self, token_ids: &[u32], cache: &mut GpuKVCache) -> Tensor {
let new_tokens = token_ids.len();
let pos_offset = cache.seq_len();
@@ -168,30 +168,36 @@ impl Qwen3 {
let k = matmul_2d(&normed, &layer.k_proj_wt);
let v = matmul_2d(&normed, &layer.v_proj_wt);
let q = reshape_heads(&q, new_tokens, num_heads, head_dim);
let k = reshape_heads(&k, new_tokens, num_kv_heads, head_dim);
let v = reshape_heads(&v, new_tokens, num_kv_heads, head_dim);
// GPU reshape: [S, H*D] → [1, H, S, D]
let q = xserv_kernels::reshape_heads_gpu(&q, new_tokens, num_heads, head_dim);
let k = xserv_kernels::reshape_heads_gpu(&k, new_tokens, num_kv_heads, head_dim);
let v = xserv_kernels::reshape_heads_gpu(&v, new_tokens, num_kv_heads, head_dim);
// QK norm (reshape to [H*S, D], rmsnorm, reshape back — stays on GPU)
let q = head_rmsnorm(&q, &layer.q_norm, eps);
let k = head_rmsnorm(&k, &layer.k_norm, eps);
let q = transpose_for_rope(&q, new_tokens, num_heads, head_dim);
let k = transpose_for_rope(&k, new_tokens, num_kv_heads, head_dim);
// GPU transpose for RoPE: [1, H, S, D] → [S, H, D]
let q = xserv_kernels::transpose_for_rope_gpu(&q, new_tokens, num_heads, head_dim);
let k = xserv_kernels::transpose_for_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
rope_inplace(&q, &self.rope_cache, &positions);
rope_inplace(&k, &self.rope_cache, &positions);
let q = transpose_from_rope(&q, new_tokens, num_heads, head_dim);
let k = transpose_from_rope(&k, new_tokens, num_kv_heads, head_dim);
// GPU transpose back: [S, H, D] → [1, H, S, D]
let q = xserv_kernels::transpose_from_rope_gpu(&q, new_tokens, num_heads, head_dim);
let k = xserv_kernels::transpose_from_rope_gpu(&k, new_tokens, num_kv_heads, head_dim);
// GPU KV cache: D2D append, no CPU round-trip
// GPU KV cache
cache.append(layer_idx, &k, &v, new_tokens, pos_offset);
let (k_full, v_full) = cache.get_kv_len(layer_idx, pos_offset + new_tokens);
// GPU repeat KV for GQA
let n_rep = num_heads / num_kv_heads;
let k_full = repeat_kv(&k_full, n_rep);
let v_full = repeat_kv(&v_full, n_rep);
let k_full = xserv_kernels::repeat_kv_gpu(&k_full, n_rep);
let v_full = xserv_kernels::repeat_kv_gpu(&v_full, n_rep);
let attn_out = attention(&q, &k_full, &v_full, true);
let attn_merged = merge_heads_any(&attn_out, new_tokens, hidden);
// GPU merge_heads: [1, H, S, D] → [S, H*D]
let attn_merged = xserv_kernels::merge_heads_gpu(&attn_out, new_tokens, num_heads, head_dim);
let attn_proj = matmul_2d(&attn_merged, &layer.o_proj_wt);
x = add_any(&residual, &attn_proj);