4c3f914459
kernels/cuda: paged-attention kernel, dispatch, pinned host memory
...
CUDA layer for the paged-KV + swap work:
- csrc: new paged_attention.cu plus updates across attention/gemm/norm/
activation/embedding/reduce kernels and common.cuh.
- xserv-kernels: new dispatch module and kernel-binding updates.
- xserv-cuda: cudaMallocHost/FreeHost bindings + PinnedBuffer (host swap
pool backing) and offset-aware D2H/H2D copies used to move KV blocks
between the GPU pool and pinned host memory.
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com >
2026-05-28 19:58:36 +08:00
e207523e21
phase 15: custom GEMV kernel — 46.6 tok/s serial (3.5x improvement, 130% of HF)
...
Custom bandwidth-optimized GEMV kernel for M=1 BF16 decode, replacing
cuBLAS which achieves only ~8% bandwidth utilization for tiny M=1 GEMMs.
Kernel design (csrc/gemm/gemv.cu):
- K-split tiled: TILE_N=128, TILE_K=256, Grid=(N/128, K/256)=512 blocks
- High occupancy: 512 blocks / 170 SMs = ~3 blocks/SM
- Coalesced memory access: adjacent threads read adjacent columns of W
- Shared memory for x vector (avoids redundant global reads)
- FP32 accumulation via atomicAdd (K-split partial sums)
- Separate fp32→bf16 conversion kernel
Integration:
- matmul() auto-dispatches to custom GEMV when M==1 && dtype==BF16
- Batched decode (M>1) continues to use cuBLAS
- Caching allocator provides FP32 temp buffer (pooled, no per-call malloc)
Ablation results (dash5, RTX 5090, Qwen3-8B BF16):
| Config | tok/s | vs HF (36) | vs roofline (112) |
|--------|-------|-----------|-------------------|
| Phase 14 (cuBLAS M=1) | 13.2 | 37% | 12% |
| + Custom GEMV (M=1) | 46.6 | 130% | 42% |
| Concurrent batch=4 | 28.2 | 78% | — |
Single-request throughput now EXCEEDS HuggingFace transformers by 30%.
The custom GEMV achieves ~42% of the theoretical roofline (vs 12% before).
Note: concurrent batch=4 (28.2 tok/s) is slower than serial (46.6 tok/s)
because the per-seq attention/reshape overhead in batched decode outweighs
the cuBLAS M=4 benefit when the custom GEMV already handles M=1 efficiently.
Engine should prefer serial decode when custom GEMV is available.
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com >
2026-05-22 22:22:31 +08:00
d67dda404e
phase 14: Flash Attention 2 for SM120 (RTX 5090)
...
Implement Flash Attention 2 forward kernel targeting SM120 (CC 12.0).
FA4 requires TMEM (only on data-center Blackwell SM100), so FA2 is the
correct target for consumer Blackwell GPUs like the RTX 5090.
CUDA kernel (csrc/attention/flash_attention.cu):
- Online softmax with tiled Q/K/V — O(1) extra memory, no S×S matrix
- Tile sizes: BR=BC=64, head_dim up to 128 (runtime parameter)
- BF16 input, FP32 accumulation, BF16 output
- Native GQA: kv_head = q_head / (num_q_heads / num_kv_heads)
- Causal mask with tile-level skip optimization
- Shared memory: 32 KB (Q_tile 16KB + KV_tile 16KB, fits in 48KB default)
- Grid: (q_tiles, batch × num_q_heads), Block: 128 threads
Integration:
- flash_attention() Rust wrapper in xserv-kernels with shape/dtype validation
- Qwen3 forward_gpu_cache uses flash_attention directly (no repeat_kv_gpu)
- Eliminates repeat_kv memory allocation + copy per layer per step
- Naive attention() preserved for testing/comparison
Validated on dash5 (RTX 5090, CUDA 12.9):
- Correctness: 9/10 top-1 match vs HF (identical to pre-FA baseline)
- Throughput: 12.9 tok/s (up from 10.3, +25% improvement)
- Now at 35% of HF transformers baseline (up from 30%)
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com >
2026-05-22 18:27:39 +08:00
2be27d6d94
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 >
2026-05-22 12:01:07 +08:00
6035ffdc0b
phase 5: naive multi-head attention
...
- Batched GEMM via cublasGemmStridedBatchedEx
- Causal mask CUDA kernel (F32 + BF16)
- Element-wise scale CUDA kernel (F32 + BF16)
- attention() composing: batched_matmul + scale + causal_mask + softmax
- Fixed to_device/contiguous infinite recursion (GPU contiguous via CPU round-trip)
- 5 attention tests passing (max_err < 3e-7 F32)
- Total: 61 tests passing across all crates
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com >
2026-05-21 21:17:23 +08:00
c8e8153702
phase 4: transformer core kernels
...
CUDA kernels (csrc/):
- common.cuh: shared warp_reduce_sum/max, block_reduce_sum/max
- normalization/rmsnorm.cu: RMSNorm (F32 + BF16)
- normalization/layernorm.cu: LayerNorm with Welford (F32 + BF16)
- activation/activations.cu: GELU tanh-approx + SiLU (F32 + BF16)
- reduce/softmax.cu: safe softmax, 3-pass (F32 + BF16)
- embedding/embedding.cu: gather lookup (F32 + BF16)
- embedding/rope.cu: RoPE in-place + precomputed cos/sin cache (F32 + BF16)
Rust wrappers (xserv-kernels/src/):
- rmsnorm.rs, layernorm.rs, activation.rs, softmax.rs, embedding.rs, rope.rs
- RopeCache struct with GPU-side precomputation
Tests: 12 new tests (ops_test.rs), all passing with good precision:
- F32: max_err 1e-6 ~ 1e-9
- BF16: max_err 2e-3 ~ 7e-3
Total: 29 kernel tests + 27 prior = 56 tests passing
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com >
2026-05-21 21:07:24 +08:00
d77f921a12
phase 3: GEMM kernels (naive, tiled, cuBLAS)
...
- Naive GEMM kernel: one thread per output element (F32 + BF16)
- Tiled GEMM kernel: 32x32 shared memory tiles (F32 + BF16)
- cuBLAS wrapper: cublasGemmEx with row-major trick
- GemmBackend enum for runtime backend selection
- CublasContext RAII handle
- Made error::check public for cross-crate use
- 17 GEMM tests: small/medium/rect sizes, all backends, F32+BF16
- Cross-backend consistency verified (naive vs tiled vs cuBLAS)
- All 44 tests pass across all crates
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com >
2026-05-21 19:48:05 +08:00
