gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
74 Commits 1 Branch 10 Tags
3d6bb1918e4d0ff6bac640fbc4c1ace08c7a7552
Commit Graph

9 Commits

Author SHA1 Message Date
Gahow Wang
4368e79695 model: fused GPU MoE kernel — eliminate CPU roundtrip 
Replace the per-token CPU-routed MoE forward with an all-GPU path:

  1. moe_topk_softmax: GPU top-k + softmax (was CPU sort + softmax)
  2. moe_replicate: broadcast input to all local experts
  3. cublasGemmStridedBatchedEx: batched expert matmul (was per-expert cuBLAS)
  4. moe_weighted_sum: FP32-accumulated weighted sum on GPU (was GPU→CPU→F32→BF16→GPU)

Expert weights stored as contiguous 3D tensors for strided batched GEMM.
Zero CPU↔GPU transfers per MoE layer (was ~40 per token per layer).

Also: configurable geglu_alpha, LayerNorm bias auto-detect, unused-weight
diagnostic at load time.

GSM8K 30-problem: 11/30 → 23/30 (76.7%) vs llama.cpp 30/30 (100%).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
 2026-05-31 13:22:59 +08:00
Gahow Wang
13ae3de69e kernels: reshape_and_cache, GPU argmax, single-launch GEMV 
Three new CUDA kernels and one rewrite:

- reshape_and_cache: scatter K/V into paged pool in a single kernel per
  layer, replacing the Rust-side per-token per-head cudaMemcpy loop.
  Includes both single-sequence (prefill) and batched (decode) variants.

- argmax: GPU-side BF16 argmax with warp-shuffle reduction. Greedy
  decode now only D2H-transfers B×4 bytes (token ids) instead of the
  full [B, vocab] logits tensor.

- GEMV rewrite: fused zero-init inside the K-split kernel eliminates
  the cudaMemsetAsync call, reducing launches from 3 to 2 per GEMV.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
 2026-05-30 12:50:17 +08:00
Gahow Wang
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
Gahow Wang
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
Gahow Wang
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
Gahow Wang
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
Gahow Wang
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
Gahow Wang
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
Gahow Wang
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