docs: Phase 21 — decode CUDA graph + GPU argmax results
dash5, gpt-oss-20b FP8, warm-server vs llama.cpp MXFP4 (6 reps): TP=2 TPOT 5.76-5.89 vs 7.42-8.45 ms (xserv 1.26-1.47x), TTFT 2.4x ahead short/medium; TP=1 5.78-5.95 vs 2.80-3.22 ms (gap 2.5x -> 2.0x, TTFT now ahead short/medium). GSM8K-50 through the graph path: 94%. Lesson recorded: graphs bought ~0.6 ms (launches were already hidden by async execution), the GPU argmax ~1 ms — measure, don't guess. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
15
README.md
15
README.md
@@ -12,13 +12,14 @@ xserv 不依赖 PyTorch / vLLM / TensorRT 等现成框架，自己实现了张
|
||||
- **模型**：GPT-2（124M）、Qwen3-8B（BF16）、gpt-oss-20b（32 专家 top-4 MoE，harmony 格式）
|
||||
- **性能**（RTX 5090，贪心，单流）：
|
||||
- Qwen3-8B BF16 单卡：约 56 tok/s（HF transformers 的 1.4×）
|
||||
- gpt-oss-20b FP8 稀疏 MoE TP=2：**132 tok/s（TPOT 7.2ms），decode 快于
|
||||
llama.cpp 同配置**（7.5-8.4ms）；llama 单卡模式（2.9ms）仍领先，是下一阶段目标
|
||||
- gpt-oss-20b FP8 稀疏 MoE + CUDA Graph decode：**TPOT 5.8ms（~172 tok/s，
|
||||
TP=1/2 同速）**；同配置 TP=2 全面快于 llama.cpp（1.26-1.47×），llama
|
||||
单卡模式（2.8ms）仍领先，差距 2.0×
|
||||
- **精度**：GSM8K 全量与 llama.cpp 同权重持平（94.5% vs 94.4%）；FP8/MXFP4 量化无回归
|
||||
- **服务**：OpenAI 兼容 `/v1/chat/completions`，SSE 流式；gpt-oss 量化后可**单卡 32GB 服务**
|
||||
- **关键能力**：自写 GEMM / Flash-Attention 2(SM120，含 attention sinks + sliding window) /
|
||||
Paged-Attention kernel、分页 KV cache（含 **CPU 换出/换入**）、连续批处理、
|
||||
CUDA Graph 解码（Qwen3 单卡路径）、**Tensor/Pipeline 并行**（NCCL，TP=1/2/4、PP=2/4）、
|
||||
CUDA Graph 解码（Qwen3 单卡 + gpt-oss 全路径整图回放）、**Tensor/Pipeline 并行**（NCCL，TP=1/2/4、PP=2/4）、
|
||||
**FP8 W8A8 / MXFP4 W4A16 量化**、**稀疏 top-k MoE decode**（只算被路由的专家）
|
||||
|
||||
> 这是一个以学习为主的项目，逐 Phase 推进，每步都做数值/端到端验证。
|
||||
@@ -193,12 +194,14 @@ GSM8K 12 个格子全是 29/30，xserv 与 llama.cpp 完全一致；AIME 的 ±1
|
||||
|
||||
## 路线图（节选）
|
||||
|
||||
已完成 Phase 0–18：CUDA 基础设施 → Tensor → GEMM → Transformer kernels → Attention →
|
||||
已完成 Phase 0–21：CUDA 基础设施 → Tensor → GEMM → Transformer kernels → Attention →
|
||||
模型加载 → 分词器 → GPT-2 → KV cache → Qwen3-8B → Paged Attention → 连续批处理 →
|
||||
HTTP API → Flash Attention 2 → 性能优化 → **张量并行（TP）** → **流水线并行（PP）**；
|
||||
HTTP API → Flash Attention 2 → 性能优化 → **张量并行（TP）** → **流水线并行（PP）** →
|
||||
**gpt-oss MoE + FP8/MXFP4 量化** → **稀疏 top-k MoE decode** → **decode CUDA Graph 整图回放**；
|
||||
并加入了 **llama.cpp 对比基准** 与 **KV CPU 换出** 等基础设施。
|
||||
|
||||
后续方向：PP microbatch/1F1B 流水线重叠（吞吐收益）、2D TP×PP、投机解码、量化（FP8 / INT8）、多模态。
|
||||
后续方向：非专家权重量化（lm_head/qkv/o）、稀疏 prefill（grouped GEMM）、server 侧 harmony
|
||||
channel 分离、PP microbatch/1F1B、投机解码、多模态。详见 `docs/00-roadmap.md` 的实际进展记录。
|
||||
|
||||
## 许可
|
||||
|
||||
|
||||
@@ -1758,14 +1758,14 @@ Phase 0–17 按计划完成。Phase 18 起实际路线偏离了上面的原计
|
||||
| 18 | Pipeline Parallelism(PP=2/4) | `18-pipeline-parallelism.md`、`benchmarks/pp-sweep.md` |
|
||||
| 19 | **gpt-oss-20b MoE**:harmony 格式、attention sinks + sliding window、YaRN;两个 CUDA bug 实战(prefill sinks NaN、GEMV 未初始化 smem);GSM8K 94.5% 对齐 llama.cpp;FP8 W8A8 / MXFP4 W4A16 量化 | `19-gpt-oss-moe.md`、`benchmarks/{fp8-quantization,mxfp4-and-llama-decode}.md` |
|
||||
| 20 | **稀疏 top-k MoE decode**:只算被路由的专家,decode 13.9→7.0ms,TP=2 下 decode/TTFT 全面快于 llama.cpp 同配置;gpt-oss 单卡 serving | `20-sparse-moe.md`、`benchmarks/sparse-moe.md` |
|
||||
| 21 | **decode CUDA Graph + GPU argmax**:整个 decode step 录成一个图回放(thread-local launch stream、retained-warmup 分配策略、NCCL capture);greedy 采样换 GPU argmax。TPOT 7.5→5.9ms(TP=1)/ 5.8ms(TP=2);TP=2 全面领先 llama(1.26-1.47×),TP=1 差距 2.5×→2.0× | `21-cuda-graph-decode.md` |
|
||||
|
||||
**下一步候选(按预期收益排序):**
|
||||
|
||||
| 候选 Phase | 内容 | 预期 |
|
||||
|---|---|---|
|
||||
| 21 | **gpt-oss decode CUDA Graph**:把 Phase 15 的 split-graph 方案(`decode_graph.rs`,目前只用于 Qwen3 单卡)推广到 MoE/TP 路径,消除 ~200 launch/token | TPOT 7.0 → ~4-5ms,逼近 llama 单卡 2.9ms |
|
||||
| 22 | **非专家权重量化**:qkv/o + lm_head(1.16GB/token)仍是 BF16 | TPOT 再省 ~1-1.5ms |
|
||||
| 23 | **稀疏 prefill**(按专家 permute + grouped GEMM) | 长 prompt TTFT 79 → ~40ms |
|
||||
| 22 | **非专家权重量化**:qkv/o + lm_head(1.16GB/token)仍是 BF16 | TPOT 再省 ~1.5ms |
|
||||
| 23 | **稀疏 prefill**(按专家 permute + grouped GEMM) | 长 prompt TTFT 51-75 → ~30ms |
|
||||
| 24 | server 侧 harmony channel 分离(`reasoning_content` 流式输出,对齐 llama-server 行为) | API 易用性 |
|
||||
| — | Speculative Decoding、多模态(原 16/19) | 推迟 |
|
||||
|
||||
|
||||
111
docs/21-cuda-graph-decode.md
Normal file
111
docs/21-cuda-graph-decode.md
Normal file
@@ -0,0 +1,111 @@
|
||||
# Phase 21: gpt-oss decode CUDA Graph + GPU argmax
|
||||
|
||||
> 目标:消除 decode 的每 token 固定开销。Phase 20 之后 TPOT ~7ms,其中
|
||||
> GPU 实际计算只占一部分,剩下是 ~200 个 kernel launch 和 per-token 的
|
||||
> host 工作。本阶段把**整个 decode step 捕获成一个 CUDA graph**,每 token
|
||||
> 一次 `cudaGraphLaunch` 回放;顺带把 greedy 采样换成 GPU argmax。
|
||||
>
|
||||
> 实现:`crates/xserv-model/src/gpt_oss_graph.rs`(~150 行)+ 三块基础设施。
|
||||
|
||||
## 1. CUDA Graph 是什么,为什么有约束
|
||||
|
||||
`cudaStreamBeginCapture` 之后,发到该 stream 的 kernel 不执行而是被**录制**;
|
||||
`EndCapture + Instantiate` 得到可执行图;以后每步 `cudaGraphLaunch` 一次性
|
||||
重放全部 ~200 个 kernel,host 端开销从 ~200 次 launch 降到 1 次。
|
||||
|
||||
代价是三条硬约束,每条都对应一个工程问题:
|
||||
|
||||
1. **地址稳定**:录制时烤进图里的全部指针,回放时必须仍然有效且指向正确数据;
|
||||
2. **capture 期间禁止"不安全"调用**:`cudaMalloc`/同步 memcpy/`cudaDeviceSynchronize`
|
||||
都会让 capture 报错(error 900);
|
||||
3. **形状固定**:grid 尺寸被烤死,变 shape 就要重录。
|
||||
|
||||
## 2. 为什么 xserv 的 decode 本来就"差一点"就能整图捕获
|
||||
|
||||
逐项检查 decode step 的输入,发现绝大部分已经满足地址稳定:
|
||||
|
||||
| 每步会变的输入 | 地址 | 内容如何更新 |
|
||||
|---|---|---|
|
||||
| block table / context lens | PagedKVCache 的常驻 GPU 缓冲 ✓ | `decode_prepare` 在图外 H2D |
|
||||
| KV 写入位置 | scatter kernel **从 GPU 上的 context_lens 读** ✓ | 同上 |
|
||||
| attention 读取范围 | paged kernel 从同一缓冲读 ✓ | 同上 |
|
||||
| MoE 专家选择 | sparse GEMV 从图内刚写的 `topk_ids` 读 ✓ | 数据依赖,天然支持 |
|
||||
| token id / position | ✗ 原来是每步从 host slice 上传 | **本阶段改造点** |
|
||||
|
||||
也就是说,Phase 11(paged KV)和 Phase 20(sparse MoE)的"数据驱动"设计
|
||||
无意中已经为 graph 化铺平了路 —— 唯二需要动的是 embedding 的 token id 和
|
||||
RoPE 的 position:各加一个 device-buffer 变体(`embedding_device_ids` /
|
||||
`rope_inplace_device_pos`),id/pos 存进两个常驻 4 字节缓冲,每步图外更新。
|
||||
|
||||
重构后的结构:
|
||||
|
||||
```text
|
||||
forward_decode_paged = decode_prepare(host 簿记,图外)
|
||||
+ upload ids/pos(图外)
|
||||
+ decode_core(纯 GPU,可整段捕获)
|
||||
+ advance_seq_len(host 簿记,图外)
|
||||
```
|
||||
|
||||
## 3. 三个工程问题
|
||||
|
||||
### 3.1 null stream 不可捕获 → thread-local launch stream
|
||||
|
||||
全代码库的 kernel 都发射在 legacy null stream 上,而 capture 必须在显式
|
||||
stream 上。解法:`xserv_cuda::stream` 加一个 **thread-local 当前 stream**
|
||||
(默认 null,行为与从前逐字节一致),所有 kernel wrapper、cuBLAS 的
|
||||
`cublasSetStream`、NCCL 的 collective 全部改读它。capture 代码用 RAII guard
|
||||
(`push_stream`)把 capture stream 装进去,录完自动还原。
|
||||
顺序正确性:显式 stream 以默认(blocking)方式创建,legacy stream 与其
|
||||
双向隐式同步,所以图外的 H2D/采样 memcpy 与回放天然有序。
|
||||
|
||||
### 3.2 capture 期间禁止 cudaMalloc → "retained warmup" 二段式
|
||||
|
||||
中间张量来自 caching allocator;capture 中任何一次 pool miss 都会触发
|
||||
`cudaMalloc` → error 900。第一版实现就栽在这里:**隔离机制自己制造了
|
||||
pool miss**(capture 中释放的块被隔离,下一层同尺寸分配就找不到块了)。
|
||||
|
||||
解法是把同一个 step 先 eager 跑一遍、但**隔离打开**(`begin_retain`):
|
||||
释放的块全部扣下不回池 → 跑完后池外恰好积累了"这一步需要的每一块";
|
||||
把它们整批放回池,再开始 capture —— capture 重复完全相同的分配序列,
|
||||
每次分配都命中池,一次 cudaMalloc 都不会发生。
|
||||
(重复执行同一 step 是无害的:KV scatter 往同一个位置重写同样的值。)
|
||||
|
||||
### 3.3 回放引用的内存不能被别人拿走 → 隔离仓(quarantine)
|
||||
|
||||
capture 录下的中间缓冲在 host 侧早就 Drop 了,但图每次回放都会读写这些
|
||||
地址。若它们回到分配池、被后续 prefill 拿走长期持有,就是双写损坏。
|
||||
所以 capture 期间释放的块进入 `RetainedBlocks` 隔离仓,由 graph 对象持有,
|
||||
graph 销毁时才归还 —— 这些内存在 graph 存活期内被锁定为它专用。
|
||||
|
||||
### 3.4 两个顺手的点
|
||||
|
||||
- **THREAD_LOCAL capture mode**:GLOBAL 模式下,任何线程的 cudaMalloc 都会
|
||||
毒化 capture;TP 多 rank 线程并发 capture 必须用 THREAD_LOCAL。
|
||||
- **NCCL 可以被捕获**:rank 内 `ncclAllReduce` 发在 capture stream 上即可,
|
||||
TP=2 一次成功(各 rank 录各自的图,回放时 collective 自然配对)。
|
||||
|
||||
## 4. 意外的教训:launch 开销没有想象的大,argmax 才是大头
|
||||
|
||||
A/B 实测(in-process,FP8,96 tok):
|
||||
|
||||
| | TP=1 | TP=2 |
|
||||
|---|---|---|
|
||||
| eager + host argmax(Phase 20 末) | 7.5 ms | 7.6 ms |
|
||||
| graph + host argmax | 6.9 ms | 6.9 ms |
|
||||
| eager + **GPU argmax** | 6.5 ms | — |
|
||||
| **graph + GPU argmax** | **5.9 ms** | **5.8 ms** |
|
||||
|
||||
- **graph 只省了 ~0.6ms**:decode 循环本来就是全异步的,launch 大部分被
|
||||
GPU 执行掩盖,"~200 launch ≈ 4ms"的预估错了 —— 优化要测不要猜。
|
||||
- **GPU argmax 省了 ~1ms**:greedy 采样原来每 token 把 [1, 201088] 的
|
||||
logits(402KB)同步拷回 host、再扫描 201K 个 bf16。仓库里 Phase 15 就写好
|
||||
的 argmax kernel(kernel 内归约 + 4 字节 D2H)一直没接到 `sample()` 上。
|
||||
- 细节:GPU argmax 与 host `max_by` 对**完全相等**的 logits 平局取的索引
|
||||
不同,greedy 轨迹会在某个平局 token 处分叉 —— 输出同样合法(GSM8K 验证)。
|
||||
|
||||
## 5. 结果与剩余瓶颈
|
||||
|
||||
见 `docs/benchmarks/sparse-moe.md` 的 Phase 21 小节(warm-server 对打 llama
|
||||
的数字以那里为准)。剩余 TPOT 的构成:~3ms 是 HBM 字节(其中非专家权重
|
||||
仍是 BF16,含 1.16GB 的 lm_head —— **Phase 22 量化它们**),其余是 GEMV
|
||||
带宽效率与 attention。llama 单卡 2.9ms 的差距主要就在"全模型 4-bit"。
|
||||
@@ -78,19 +78,34 @@ reasons, both instructive:
|
||||
noise; W8A16 removes activation-quantization error so ≥ dense is expected).
|
||||
Avg 1.3 s/problem also reflects the decode speedup.
|
||||
|
||||
## Remaining gaps / next levers (to catch llama TP=1 at 2.9 ms)
|
||||
## Phase 21 update: decode CUDA graph + GPU argmax (docs/21-cuda-graph-decode.md)
|
||||
|
||||
Sparse MoE removed the dominant cost; the residual ~7 ms splits roughly into
|
||||
~3 ms HBM reads and ~4 ms fixed overhead. In impact order:
|
||||
The whole batch=1 decode step now replays as one CUDA graph, and greedy
|
||||
sampling uses the GPU argmax kernel (4-byte D2H instead of a 402 KB logits
|
||||
copy + 201k-element host scan). In-process A/B: graph −0.6 ms, GPU argmax
|
||||
−1.0 ms. Warm-server head-to-head (same harness/GPUs, 6 reps):
|
||||
|
||||
1. **CUDA graphs for decode** (~2–4 ms): with experts down to ~1–2 ms, the
|
||||
~200 un-graphed launches/token are now the single largest cost. (The old
|
||||
"graphs ≈ useless" conclusion was relative to a 13 ms dense TPOT — no
|
||||
longer true.)
|
||||
2. **Quantize non-expert weights** (~1–1.5 ms): attn qkv/o + the 1.16 GB BF16
|
||||
| | xserv FP8 (graph) | llama MXFP4 | |
|
||||
|---|---|---|---|
|
||||
| TP=2 TPOT | **5.76–5.89 ms** (170–174 tok/s) | 7.42–8.45 ms | **xserv 1.26–1.47×** |
|
||||
| TP=2 TTFT s/m/l | **25 / 28 / 51 ms** | 63 / 66 / 45 ms | xserv 2.4× s/m; long ~par |
|
||||
| TP=1 TPOT | 5.78–5.95 ms | **2.80–3.22 ms** | llama 2.0× (was 2.5×) |
|
||||
| TP=1 TTFT s/m | **32 / 35 ms** | 34 / 36 ms | xserv slightly ahead |
|
||||
|
||||
GSM8K-50 through the graph path: 47/50 = 94% (unchanged). Note: GPU argmax
|
||||
breaks exact-tie logits differently than the host scan, so greedy trajectories
|
||||
can legitimately diverge at a tie token.
|
||||
|
||||
## Remaining gaps / next levers (to catch llama TP=1 at 2.8 ms)
|
||||
|
||||
Per-token fixed overhead is now mostly gone; the residual ~5.8 ms is
|
||||
dominated by HBM bytes and kernel efficiency. In impact order:
|
||||
|
||||
1. **Quantize non-expert weights** (~1.5 ms): attn qkv/o + the 1.16 GB BF16
|
||||
lm_head read every token; FP8/MXFP4 them like llama quantizes everything.
|
||||
2. **GEMV/attention bandwidth tuning**: effective BW of the hand GEMVs is
|
||||
well under peak; llama's 2.8 ms implies ~85%+ efficiency on ~1.3 GB.
|
||||
3. **Sparse prefill** (permute tokens by expert + grouped GEMM): long-prompt
|
||||
TTFT 94–120 ms → llama's ~30 ms territory.
|
||||
TTFT 51–75 ms → llama's ~30 ms territory.
|
||||
4. **W4A4 FP4 tensor cores / bandwidth-tuned MXFP4 GEMV**: make 4-bit experts
|
||||
actually beat FP8 (today sparse MXFP4 is 8.4 ms vs FP8 7.6 ms — the 4-bit
|
||||
GEMV's lower effective bandwidth still cancels its byte advantage).
|
||||
actually beat FP8.
|
||||
|
||||
Reference in New Issue
Block a user