xserv/docs/15-performance.md

Phase 15: Performance Optimization — Design Document (Milestone ④)

Goal

系统性 profiling + 优化，从 12.9 tok/s (Phase 14 结束) 逼近 RTX 5090 的理论带宽上限 (112 tok/s)。

硬件 Roofline

RTX 5090 (SM120, CC 12.0) 的 decode 理论极限：

模型权重:          16 GB (Qwen3-8B BF16)
内存带宽:          1.79 TB/s (GDDR7)
理论最优 decode:   16 GB / 1.79 TB/s = 8.9 ms/step = 112 tok/s (batch=1)

Decode 阶段 100% memory-bound：每步读取全部 16 GB 权重（252 个 GEMV），计算量可忽略。

瓶颈分析

Phase 14 结束时性能 12.9 tok/s = 77.5 ms/step，roofline 利用率仅 12%。

量化瓶颈分解

来源	估计耗时	占比
cuBLAS M=1 GEMV (252 calls, 带宽利用 ~8%)	~60 ms	77%
非 matmul 内核 (attention, norm, activation, reshape)	~8 ms	10%
Tensor 分配 + cudaMemset (1440+ allocs/step)	~5 ms	7%
Kernel launch overhead (200+ launches × 5μs)	~1 ms	1%
其他 (sampling CPU round-trip, etc.)	~3.5 ms	5%

核心发现: cuBLAS 对 M=1 GEMM (GEMV) 的带宽利用率极低（~8%），是 9x gap 的根本原因。

cuBLAS 设计用于大 M 的 GEMM，对 M=1 场景存在：

• Kernel launch dispatch overhead 无法被大量计算掩盖
• TensorCore tile (16×16) 无法被 M=1 充分利用
• 内部 heuristic 选择了次优算法

优化实施

Opt 1: Decode Attention Kernel

目标: 替换 FA2 在 Q_len=1 时的低效路径（64 线程仅 1 个 active）。

实现 (csrc/attention/flash_attention.cu):

• 专用 decode_attention_bf16_kernel: 256 线程并行沿 KV 序列维度
• 每个 thread 加载完整 Q vector (128 dim) 到寄存器
• 处理其分配的 KV 位置块: dot product → online softmax
• Block-level warp-shuffle + shared memory reduction 合并结果
• GQA 支持: kv_head = q_head / heads_per_group

效果: 在当前短序列 (kv_len ≤ 79) 下效果微小——attention 不是瓶颈。在长序列时会显著受益。

Opt 2: Fused SiLU×Mul

目标: silu(gate) * up 两个 element-wise op 合并为一个 kernel。

实现 (csrc/activation/activations.cu):

Before: read gate → silu → write temp → read temp + up → mul → write out
After:  read gate + up → silu(gate) * up → write out
Saved:  1 HBM read + 1 HBM write per element

效果: 每层省 1 次 HBM round-trip，36 层总计可观但在 GEMV 瓶颈下被掩盖。

Opt 3: Fused Add+RMSNorm

目标: x = residual + attn_proj; normed = rmsnorm(x) 合并为一个 kernel。

实现 (csrc/normalization/rmsnorm.cu):

Before: read residual + x → add → write sum → read sum + gamma → norm → write out
After:  read residual + x + gamma → add + norm → write sum + normed
Saved:  1 full HBM round-trip per attention block

Opt 4: Batched Decode Forward ⭐

目标: 多序列 decode token 合并为 M=batch_size 的 GEMM，提升 cuBLAS 效率。

实现 (crates/xserv-model/src/qwen3.rs + crates/xserv-server/src/engine.rs):

• 新增 Qwen3::forward_decode_batch(tokens, positions, caches)
• Batched ops: embedding, norm, projections, FFN — [B, hidden] × [hidden, X]
• Per-seq ops: RoPE, KV cache, attention（各序列位置/长度不同）
• Row extraction (row_view) + concatenation (concat_rows) 在 batched/per-seq 间切换
• Engine Step 4b: batch≥2 时自动使用 batched decode

效果: batch=4 时 cuBLAS 从 1008× M=1 → 252× M=4，吞吐 35.1 tok/s (vs serial 13.2)。

Opt 5: Custom GEMV Kernel ⭐⭐⭐ (决定性优化)

目标: 替换 cuBLAS 的 M=1 GEMV，手写带宽最优化 kernel。

实现 (csrc/gemm/gemv.cu):

设计: K-split tiled GEMV
- TILE_N = 128 (output columns per block, one thread per column)
- TILE_K = 256 (K-dimension slice per block)
- BLOCK_SIZE = 128 threads
- Grid: (ceil(N/128), ceil(K/256)) — 对 K=N=4096 得到 512 blocks
  512 blocks / 170 SMs ≈ 3 blocks/SM (良好 occupancy)

内存访问:
- 相邻线程读 W 矩阵的相邻列 → 完美 coalesced
- x vector 加载到 shared memory (每 K-chunk 仅加载一次)
- FP32 accumulation via atomicAdd (K-split partial sums)
- 独立 kernel 做 FP32→BF16 转换

调度:
- matmul() 中检测 M==1 && dtype==BF16 → 自动使用 custom GEMV
- M>1 保持 cuBLAS

效果: 13.2 → 46.6 tok/s (+253%)。带宽利用率从 ~8% 提升到 ~42%。

Opt 6: Tensor::empty() (消除无用 cudaMemset)

目标: kernel 输出 tensor 全量覆写时，跳过分配后的 cudaMemset 清零。

实现:

• Storage::empty() + Tensor::empty(): 分配不清零
• 21 个 kernel wrapper (activation, attention, embedding, gemm, norm, softmax, transpose) 从 zeros 改为 empty
• GEMV FP32 accumulator buffer 保持 cudaMemsetAsync（atomicAdd 需要零初始化）

效果: 46.6 → 50.3 tok/s (+8%)。消除 ~756 个 cudaMemset/step。

Infra: CUDA Graph 基础设施

• FFI bindings: cudaStreamBeginCapture, cudaGraphInstantiate, cudaGraphLaunch
• RAII wrapper: CudaGraph (capture/instantiate/launch lifecycle)
• 当前未在 forward path 使用（variable kv_len 限制），为后续优化预留

Ablation 结果

dash5, RTX 5090, Qwen3-8B BF16, greedy decode, max_tokens=64:

优化叠加	tok/s	增量	vs HF	Roofline
Phase 14 baseline (FA2)	12.9	—	36%	12%
+ Decode attention	12.9	+0%	36%	12%
+ Fused SiLU×Mul	13.0	+1%	36%	12%
+ Fused Add+RMSNorm	13.2	+2%	37%	12%
+ Batched decode (batch=4)	35.1	—	97%	—
+ Custom GEMV (M=1)	46.6	+253%	130%	42%
+ Tensor::empty	50.3	+8%	140%	45%

对比:

系统	tok/s	Roofline
HF transformers	36.0	32%
xserv (Phase 15)	50.3	45%
理论极限 (1.79 TB/s)	112.0	100%

剩余 55% Roofline Gap 分析

来源	估计占比	优化方向
GEMV kernel 非满带宽 (atomicAdd contention, K-split overhead)	25%	无 K-split GEMV (更大 block), 向量化加载
Non-matmul kernels (attention, norm, RoPE, reshape)	15%	Fused layer kernel, 更高效的 decode attention
Kernel launch overhead (200+ launches/step)	5%	CUDA Graphs (需解决 variable kv_len)
Memory allocator overhead (Arc, SmallVec per tensor)	5%	Pre-allocated decode workspace
Sampling D2H copy (pipeline stall)	3%	GPU-side argmax kernel
其他 (host-side logic, channel overhead)	2%	—

下一步

Phase 15 的 Milestone ④ 目标 (50% of HF) 已远超 — 达到 140% of HF, 45% of roofline。

后续优化路径（按 ROI 排序）：

1. 无 K-split GEMV: 消除 atomicAdd，减少 kernel launches → 预期 +15-20%
2. 向量化 GEMV loads: float4 加载 W 矩阵 → 预期 +10%
3. Pre-allocated workspace: 消除 Tensor 对象分配开销 → 预期 +5%
4. CUDA Graphs: 需要 fixed-shape decode path → 预期 +5%
5. GPU-side sampling: 消除 logits D2H pipeline stall → 预期 +3%