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

docs: Phase 14 design doc + benchmark, fix Phase 11/12 honesty

Browse Source 
Phase 14 (Flash Attention):
- Design doc: FA2 algorithm, SM120 hardware constraints (FA4 incompatible),
  kernel config (BR=BC=64, 32KB smem), GQA mapping, causal tile-skip,
  known limitations and optimization roadmap
- Benchmark doc: correctness (9/10 top-1 match, identical to pre-FA baseline),
  performance tracking (6.9→10.3→12.9 tok/s across phases), memory savings
  analysis, remaining bottleneck breakdown

Phase 11 doc: title corrected from "Paged Attention" to "GPU-Resident KV Cache"
with explicit note that paged allocation was not implemented.

Phase 12 doc: "当前状态" updated from "未实现" to reflect actual state —
iteration-level scheduling implemented + verified (6.0x concurrent speedup),
batched GPU forward explicitly marked as not yet implemented.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-22 18:51:29 +08:00
parent d67dda404e
commit 6cc1c9332d
4 changed files with 285 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
docs/11-paged-attention.md
View File

@@ -1,8 +1,10 @@
# Phase 11: Paged Attention + KV Cache Manager — Design Document # Phase 11: GPU-Resident KV Cache — Design Document
> **注意**: 原计划为 "Paged Attention + KV Cache Manager"，实际实现为 GPU 连续预分配 KV cache（非 paged）。Paged allocation 留待后续优化。
## Goal ## Goal
将 KV cache 从 CPU Vec 迁移到 GPU，使用 block-based paging 管理显存。消除每步 decode 的 CPU round-trip（当前 KV cache 最大性能瓶颈之一）。 将 KV cache 从 CPU Vec 迁移到 GPU，消除每步 decode 的 CPU round-trip（当前 KV cache 最大性能瓶颈之一）。
## 当前问题 ## 当前问题

6
docs/12-continuous-batching.md
View File

@@ -150,4 +150,8 @@ HTTP Handler Engine Thread
## 当前状态 ## 当前状态
**未实现**。当前是 FIFO 串行，一次只处理一个请求。本文档是实现的设计规格。 **已实现: iteration-level scheduling**。多请求可以并发进入 batch (max_batch_size)，新请求在 mid-generation 动态加入。Prefill 和 decode 阶段在每轮迭代内分离处理。
**未实现: batched GPU forward**。每个 seq 的 model forward 仍是串行调用 (per-seq forward_gpu_cache)。真正的 batched decode (多 seq 的 token 合并为一次 GPU forward) 需要 Flash Attention 的 variable-length attention 支持。Phase 14 实现了 FA2 kernel，为后续 batched forward 提供了基础。
**验证**: 8 个并发请求 (max_batch=4) 总 wall clock 22.5s，各请求延迟之和 135.0s，调度加速 6.0x。Server log 确认 `decode batch_size=4`。

167
docs/14-flash-attention.md Normal file
View File

@@ -0,0 +1,167 @@
# Phase 14: Flash Attention 2 for SM120 — Design Document
## Goal
用自写的 Flash Attention 2 CUDA kernel 替换 naive attention (Phase 5)。消除 O(S²) 显存分配，支持 GQA kernel 内部索引（消除 repeat_kv 开销）。
## 硬件约束: FA4 不适用于 RTX 5090
Flash Attention 已发展到第 4 代 (FA4, arxiv 2603.05451)，但各版本有明确硬件依赖：
| 版本 | 目标架构 | 关键硬件特性 | RTX 5090 (SM120) |
|------|---------|------------|-----------------|
| FA2 | 通用 CUDA (SM75+) | shared memory + HMMA | **兼容** |
| FA3 | Hopper SM90 (H100) | TMA + WGMMA + warp specialization | 不兼容 |
| FA4 | Blackwell SM100 (B200/B300) | TMEM + async MMA + 2-CTA mode | 不兼容 |
RTX 5090 使用消费级 Blackwell (GB202, SM120)，与数据中心 Blackwell (B200, SM100) 是不同硅片。SM120 **没有 TMEM (Tensor Memory)**，这是 FA4 kernel 设计的核心硬件依赖。这不是软件限制，是硬件级差异。
因此本项目实现 **FA2 算法**，使用标准 CUDA (shared memory + 标准 HMMA)。
## Naive Attention 的问题
Phase 5 的 naive attention 流程:
```
k_t = K.transpose(2,3).contiguous() ← 分配 K^T 显存
scores = batched_matmul(Q, k_t) ← 分配 [B,H,S,S] score 矩阵 (O(S²) 显存)
scores = scale(scores, 1/sqrt(d)) ← 逐元素 kernel
causal_mask(scores) ← 逐元素 kernel
weights = softmax(scores) ← 分配 [B,H,S,S] weight 矩阵
output = batched_matmul(weights, V) ← 最终结果
```
问题:
1. **显存 O(S²)**: score 和 weight 矩阵各需 `B × H × S × S × dtype_size`。S=2048, H=32, BF16 → 256 MB。S=8192 → 4 GB。
2. **GQA 预处理**: 在调用 attention 前需要 `repeat_kv_gpu` 将 K/V 从 8 heads 扩展到 32 heads，每层额外分配和拷贝。
3. **多次 kernel launch**: scale, mask, softmax 各一次 kernel launch + global memory round-trip。
4. **K^T materialization**: `K.transpose().contiguous()` 需要分配和拷贝。
## FA2 算法
核心思想: **不 materialize S×S 矩阵**。将 Q, K, V 分成 tiles，在 shared memory (SRAM) 中计算，使用 **online softmax trick** 边算边更新 running max 和 sum。
FA2 (Dao 2023) 相比 FA1 的改进: 外层循环遍历 Q tiles (而非 K/V)，减少 HBM 读写次数，提高并行性。
```
scale = 1 / sqrt(head_dim)
for each Q tile (q_start..q_start + BR): // 外层: Q tiles
load Q_tile [BR, D] to shared memory (一次加载，内层复用)
init per-row: O[D] = 0, m = -inf, l = 0
for each K/V tile j (kv_start..kv_start + BC): // 内层: K/V tiles
// Causal tile-skip: 如果整个 K tile 在 Q tile "未来"，跳过
if causal && kv_start > max_q_pos + kv_offset: skip
load K_tile [BC, D] to shared memory
S = Q_tile @ K_tile^T * scale // [BR, BC], in registers
if causal: mask S[r][c] = -inf where kv_pos > q_pos
// Online softmax update
m_new = max(m, rowmax(S))
P = exp(S - m_new)
l_new = exp(m - m_new) * l + rowsum(P)
O = exp(m - m_new) * O // rescale accumulator
load V_tile [BC, D] to shared memory (复用 K 的空间)
O += P @ V_tile // accumulate
m = m_new, l = l_new
O = O / l // final normalize
write O[BR, D] to HBM (convert FP32 → BF16)
```
## 实现细节
### Kernel 配置
| 参数 | 值 | 说明 |
|------|---|------|
| BR (Q tile rows) | 64 | Q tile 大小 |
| BC (K/V tile rows) | 64 | K/V tile 大小 |
| head_dim | 运行时参数 (≤128) | 支持 64 (GPT-2) 和 128 (Qwen3) |
| Block size | 128 threads | 64 线程各 own 一行 Q，其余协助加载 |
| Grid | (q_tiles, batch × num_q_heads) | 每个 block 处理一个 Q tile + 一个 head |
### Shared Memory (BF16 存储)
```
smem_q [BR × head_dim] BF16 = 64 × 128 × 2 = 16 KB (加载一次，内层复用)
smem_kv[BC × head_dim] BF16 = 64 × 128 × 2 = 16 KB (K 和 V 交替使用)
────────────────────────────────────────────
Total: 32 KB (SM120 默认 48 KB，余量充足)
```
### 线程映射
- Thread 0..63: 各 own Q_tile 的一行。负责该行的全部计算：dot products、softmax、PV 累加。
- Thread 64..127: 协助 shared memory 加载 (K/V tile)，不参与计算。
- 加载模式: 每个 thread 加载 `(BR × head_dim) / 128 = 64` 个 BF16 元素。
### Per-Thread Register 使用
```
O_acc[128] FP32 = 512 bytes (128 regs) — 输出累加器
P[64] FP32 = 256 bytes (64 regs) — 当前 tile 的 softmax 后权重
m, l FP32 = 8 bytes (2 regs) — online softmax running state
循环变量 + 临时 ≈ 16 regs
────────────────────────────────────────────
Total: ~210 regs/thread (max 255，在限制内)
```
### GQA 支持
每个 thread block 处理一个 Q head，通过 `kv_head = q_head / (num_q_heads / num_kv_heads)` 映射到对应的 KV head。K/V 的数据指针直接指向 KV head 的存储，无需 repeat_kv。
```
// 32 Q heads, 8 KV heads → heads_per_group = 4
// Q head 0,1,2,3 → KV head 0
// Q head 4,5,6,7 → KV head 1
// ...
kv_head = q_head / heads_per_group;
K_ptr = K + (batch * num_kv_heads + kv_head) * kv_len * head_dim;
```
### Causal Mask
两级优化:
1. **Tile-level skip**: 如果 `kv_tile_start > max_q_pos + kv_offset`，整个 K/V tile 都在未来，跳过（减少 ~50% 计算）。
2. **Element-level mask**: 在 tile 内部，`if kv_pos > q_pos + kv_offset: S = -inf`。
`kv_offset = kv_len - q_len` 处理 decode 时 KV cache 长于 Q 的情况。
## 与 Naive Attention 的对比
| 特性 | Naive (Phase 5) | FA2 (Phase 14) |
|------|----------------|----------------|
| 显存 | O(B × H × S²) | O(B × H × S × D) |
| GQA | 需要 repeat_kv (分配+拷贝) | Kernel 内部索引 (零开销) |
| K^T | 需要 transpose+contiguous | Kernel 内部计算 |
| Kernel launches | 6 (matmul, scale, mask, softmax, matmul, ...) | 1 (单个 fused kernel) |
| S=8192 可行性 | OOM (~4 GB score matrix) | 可行 (32 KB shared memory) |
## 源码结构
```
csrc/attention/flash_attention.cu — FA2 kernel (BF16 in, FP32 accumulate, BF16 out)
crates/xserv-kernels/src/attention.rs — flash_attention() Rust wrapper + 原 attention() 保留
crates/xserv-model/src/qwen3.rs — forward_gpu_cache 调用 flash_attention
```
## 已知局限与后续优化方向
1. **Decode (Q_len=1) 效率低**: BR=64 线程中只有 1 个 active（owns_row）。应写专用 decode attention kernel，沿 KV 维度 parallel reduction。
2. **无向量化加载**: 当前逐元素 bf16→f32 转换，应改用 `float4` 或 `__nv_bfloat162` 批量加载。
3. **Register tiling**: 每个 thread 目前串行计算 dot product (128 MADs per K column)。可改为多线程协作。
4. **K/V double buffering**: 可在计算当前 tile 时预加载下一个 tile 到另一半 shared memory。
5. **Tile size 调优**: 更大的 tile (BR=128) 可能在长 sequence 时更优，需要 opt-in shared memory。
## Test Plan
- [x] 正确性: logits 与 HF transformers 对比 (top-1 match 9/10, top-5 overlap 4.0/5)
- [x] 生成质量: 52/52 prompt 生成连贯文本，中英文均可
- [x] SSE streaming 正常工作
- [x] 性能: 12.9 tok/s (vs naive 10.3 tok/s, +25%)
- [ ] 长 sequence (S=4096, S=8192): 验证 naive OOM 而 FA2 正常
- [ ] ncu profile: compute utilization, memory throughput

109
docs/benchmarks/phase14-flash-attention.md Normal file
View File

@@ -0,0 +1,109 @@
# Phase 14 Benchmark: Flash Attention 2
**Date**: 2026-05-22
**Hardware**: RTX 5090 (32GB GDDR7, SM120 CC 12.0, 170 SMs)
**Model**: Qwen3-8B (BF16, 36 layers, 4096 hidden, 32 Q / 8 KV GQA heads, head_dim=128)
**Config**: greedy decoding (temperature=0), max_tokens=64, single-request serial
## Correctness
Logits comparison with HuggingFace transformers (10 prompts, raw text without ChatML):
| Metric | Result |
|--------|--------|
| Prefill Top-1 match vs HF | **9/10 (90%)** |
| Avg Top-5 overlap vs HF | **4.0/5** |
| Result vs pre-FA2 naive attention | **Identical** (same 9/10 top-1, same 4.0/5 overlap) |
The single top-1 mismatch ("Explain quantum computing.") has logits differing by 0.125
(22.000 vs 21.875) — within BF16 precision. The top-5 sets are identical (5/5 overlap).
FA2 introduces no precision degradation compared to the naive attention path.
## API Generation
52 diverse prompts (English, Chinese, code) via `/v1/chat/completions`:
| Metric | Result |
|--------|--------|
| Success rate | **52/52 (100%)** |
| SSE streaming | **Working** (role chunk, content chunks, finish_reason, [DONE]) |
| Usage stats | Correct (prompt_tokens + completion_tokens = total_tokens) |
## Performance
### xserv vs HuggingFace transformers
8 prompts (short/medium/long) × max_tokens=64, greedy:
| Category | Prompt Tokens | xserv (tok/s) | HF (tok/s) | Ratio |
|----------|--------------|---------------|------------|-------|
| Short (~12 tok) | 12-14 | 12.5 | 38.5 | 0.32x |
| Medium (~28 tok) | 27-28 | 13.6 | 44.1 | 0.31x |
| Long (~60 tok) | 58-64 | 13.0 | 36.0 | 0.36x |
| **Overall** | — | **12.9** | **36.6** | **0.35x** |
### Phase-over-Phase Improvement
| Phase | Attention | repeat_kv | tok/s | vs HF |
|-------|-----------|-----------|-------|-------|
| 10 | Naive (O(S²), cuBLAS batched) | CPU round-trip | 6.9 | 15% |
| 11 | Naive + GPU KV cache | GPU repeat_kv | 10.3 | 30% |
| **14** | **FA2 (O(1), fused kernel)** | **None (GQA in kernel)** | **12.9** | **35%** |
Phase 14 vs Phase 11: **+25% throughput** (10.3 → 12.9 tok/s).
### Improvement Breakdown (estimated)
| Factor | Contribution |
|--------|-------------|
| Eliminating repeat_kv GPU alloc + copy (per layer) | ~10% |
| Eliminating K^T transpose + contiguous | ~5% |
| Eliminating S×S score matrix alloc | ~5% |
| Fused kernel (1 launch vs 6) | ~5% |
### Concurrent Requests
8 concurrent requests, max_batch=4:
| Metric | Result |
|--------|--------|
| Wall clock | 22.5s |
| Sum of individual latencies | 135.0s |
| Scheduling speedup | **6.0x** |
| Throughput | 11.4 tok/s |
Continuous batching scheduling confirmed working (decode batch_size=4 in logs).
## Remaining Performance Gap
35% of HF throughput. Main bottlenecks:
| Bottleneck | Impact | Fix |
|-----------|--------|-----|
| **Decode Q_len=1 inefficiency** | FA2 kernel: 64 threads, only 1 active (owns_row=true for single query) | Specialized decode attention kernel (vector-dot against KV, parallel reduction along S) |
| **No kernel fusion** | RMSNorm+residual, SiLU*up: separate kernels, redundant HBM reads/writes | Fused kernels (Phase 15) |
| **No CUDA Graphs** | ~100+ kernel launches per decode step, each has host-side overhead | Capture decode iteration as CUDA Graph (Phase 15) |
| **Per-seq forward (no batched decode)** | With batch=4, 4 serial forward passes per iteration | Batched projections + per-seq attention (Phase 15, depends on FA2 decode kernel) |
| **No vectorized loads in FA2** | Scalar bf16→f32 conversion in dot product loop | float4 / bfloat162 vectorized loads |
## Memory Usage
| Component | Naive (Phase 11) | FA2 (Phase 14) |
|-----------|-----------------|----------------|
| Score matrix [1, 32, S, S] | S² × 32 × 2B | **0** |
| repeat_kv K/V [1, 32, S, 128] | 2 × S × 32 × 128 × 2B per layer | **0** |
| K^T contiguous copy | S × 32 × 128 × 2B per layer | **0** |
For S=256 (current max): savings ~6 MB per layer × 36 layers ≈ 216 MB.
For S=2048: savings ~384 MB per layer × 36 layers ≈ 13.5 GB (naive would OOM).
## Tracking
| Phase | Attention | tok/s | vs HF | Correctness |
|-------|-----------|-------|-------|-------------|
| 8 | Naive (no cache) | 2.5 | 5% | 50/50 vs HF |
| 9 | Naive + CPU KV cache | 44.3 (GPT-2) | — | 50/50 self |
| 10 | Naive + CPU KV cache | 6.9 (Qwen3-8B) | 15% | 100% top-5 |
| 11 | Naive + GPU KV cache | 10.3 | 30% | 9/10 top-1 |
| **14** | **FA2 + GQA in kernel** | **12.9** | **35%** | **9/10 top-1** |