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

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:
2026-05-22 12:01:07 +08:00
parent 2d48f25e66
commit 2be27d6d94
5 changed files with 273 additions and 12 deletions

View File

@@ -6,8 +6,10 @@ pub mod layernorm;
pub mod rmsnorm;
pub mod rope;
pub mod softmax;
pub mod transpose;
pub use activation::{add, gelu, mul, scale, silu};
pub use transpose::{merge_heads_gpu, repeat_kv_gpu, reshape_heads_gpu, transpose_for_rope_gpu, transpose_from_rope_gpu};
pub use attention::attention;
pub use embedding::embedding;
pub use gemm::{batched_matmul, matmul, GemmBackend};