Performance report性能报告
Performance Report A: Sparse Attention on Tensor Cores性能分析报告 A：稀疏注意力改用 tensor core
Real V4.1 4-GPU prefill goes from 1,085 to 2,572 tok/s/GPU (2.37×) after the sparse-attention kernel moves from CUDA cores to tensor cores.稀疏注意力 kernel 从 CUDA core 改用 tensor core 后，真实 V4.1 四卡 prefill 吞吐从 1,085 提高到 2,572 tok/s/GPU（2.37 倍）。
2026-10-22 · PR #38 · code 804f4b0 (before the change: eac0347)
Conclusion
Real V4.1 4-GPU prefill throughput rises from 1,085 to 2,572 tok/s/GPU (2.37×), and the step time falls from 3,776 ms to 1,593 ms. The gap to the official engine shrinks from 33.3× to 14.1×; mini reaches 7.1% of the official engine. The sparse-attention kernel drops from 2,255 ms/step to 98.8 ms/step (22.8×) and is no longer the bottleneck. The two largest items now are MoE (726 ms) and the indexer (461 ms).
Setup
| Item | Configuration |
|---|---|
| Model | Real DeepSeek-V4.1-Flash, revision dba1be0 |
| Hardware / topology | 1 node, 4 × B200, EP4 / DP4 / CP1 |
| Workload | workload.json.gz (SHA-256 adf09ec9…), 131,072 context, 126,976 cached, 4,096 new tokens per step per rank |
| Timing | 10 complete steady-state steps on rank 0; plus a 2-step torch profile |
| Jobs | Before the change and official: 9297; after the change: 9329 |
End-to-end results
| Metric | Before | After | Official | Change |
|---|---|---|---|---|
| Input throughput (tok/s/GPU) | 1,084.8 | 2,571.7 | 36,170.9 | ×2.37 |
| Mean step (ms) | 3,775.9 | 1,592.7 | 113.2 | −57.8% |
| Step variation (ms) | 3,774.9–3,777.0 | 1,591.5–1,596.8 | 111.5–116.3 | Stable on both sides |
| Share of official | 3.0% | 7.1% | 100% | |
| Gap to official | 33.3× | 14.1× | — | |
| Peak memory (GiB) | 147.6 | 147.6 | — | Unchanged (scratch is 42 MB per query block) |
Per-stage time (profile, ms/step)
| Stage | Before | After | Official | Current gap |
|---|---|---|---|---|
| Sparse attention core | 2,254.9 | 98.8 | 14.9 | 6.6× |
| Indexer | 460.7 | 460.8 | 7.9 | 58× |
| Other attention (projections, compression, gather, fake quant) | 253 | 254 | ~19 | 13× |
| MoE (gate + experts + communication) | 746 | 726 | 52.8 | 14× |
| HC, normalization and other | 123 | 122 | ~10 | 12× |
| Total (step under profile) | 3,838 | 1,662 | 110.6 | 15× |
Only sparse attention changed; every other stage is within measurement error, which shows the change has no side effects. The step under the profiler is about 4% longer than the timing window; that is the profiler’s own overhead.
Kernel analysis
| Shape (T, H, K, D) | Old (ms) | New (ms) | Speedup | New TFLOP/s |
|---|---|---|---|---|
| 256, 64, 640, 512 (real compressed layer) | 3.672 | 0.166 | 22.1× | 129 |
| 256, 64, 128, 512 (real uncompressed layer) | 0.686 | 0.069 | 9.9× | 62 |
| 256, 4, 640, 512 (tiny model) | 0.284 | 0.105 | 2.7× | 13 |
- Why it is faster: the 64 heads of one query share the same 640 keys. The old kernel ran one program per (query, head) and summed with
tl.sumon CUDA cores, reading the same key set 64 × 2 times. The new kernel stacks 32 heads into the M dimension oftl.dot, runs on tensor cores, and reads each key set once for QK and once for PV per head block. - How far from peak: each step’s attention is about 13.2 TFLOP, so 98.8 ms is 134 TFLOP/s, roughly 6% of B200’s dense bf16 peak. The official 14.9 ms is about 0.89 PFLOP/s.
- Why a gap remains:
- M is only 32, so the tensor-core tile is small. With HEADS=64 the fp32 accumulator [64, 512] does not fit in registers and spills; the parameter sweep confirms it.
- The two-pass algorithm writes and reads the scores scratch once more.
- About 180 ms of gather/copy outside the kernel assembles the keys into [T, 640, 512].
- The tiny model gains little: it has only 4 heads, which cannot fill an M=16 tile, so sharing keys helps less.
Accuracy
- The maximum error against the fp32 CPU reference is the same order as the old kernel (≤ 4.9e-4). Tensor-core accumulation order differs, so results are not bit-identical to the old kernel.
- The bf16 probability-rounding regression fixture still passes; it fails if the bf16 rounding is removed or the sink is folded into the max.
- All end-to-end reference tests pass (golden tokens, per-position logits along the reference trajectory, chunked prefill, preemption). Full suite: 163 passed.
Tiny-model goodput (1 GPU, bench.py defaults, median of 3)
| Throughput (tok/s) | TTFT p50 (ms) | TPOT p50/p99 (ms) | goodput (req/s) | |
|---|---|---|---|---|
| Before | 6,791 | 294 | 4.8 / 16.1 | 49.36 |
| After | 7,202 (+6%) | 274 | 4.6 / 15.0 | 52.36 (+6%) |
Next
Ordered by remaining time: MoE 726 ms, indexer 461 ms, other attention 254 ms.
- B1: fuse indexer scoring. The micro-benchmark already drops from 61.8 ms to 4.9 ms per layer (904 TFLOP/s); this item should go from about 461 to 90 ms, for about 3,200 tok/s/GPU.
- B2: score only the candidate blocks in layers 24–36.
- C: run MoE grouped GEMM directly on FP4 weights.
Raw data: profile-2026-10-22/stage-a/; before the change: profile-2026-10-22/.
2026-10-22 · PR #38 · 代码 804f4b0（改动前 eac0347）
结论
真实 V4.1 四卡 prefill 吞吐从 1,085 提高到 2,572 tok/s/GPU（2.37 倍），每步从 3,776 ms 降到 1,593 ms。 到官方的差距从 33.3 倍缩小到 14.1 倍，mini 达到官方的 7.1%。稀疏注意力 kernel 从 2,255 ms/步降到 98.8 ms/步（22.8 倍），不再是瓶颈。现在最大的两项是 MoE（726 ms）和 indexer（461 ms）。
实验设置
| 项 | 配置 |
|---|---|
| 模型 | 真实 DeepSeek-V4.1-Flash，revision dba1be0 |
| 硬件 / 拓扑 | 单机 4 × B200，EP4 / DP4 / CP1 |
| 负载 | workload.json.gz（SHA-256 adf09ec9…），131,072 上下文，126,976 已缓存，每步每 rank 4,096 新 token |
| 计时 | 10 个完整稳态步，rank 0；另抓 2 步 torch profile |
| 作业 | 改动前与官方：9297；改动后：9329 |
端到端结果
| 指标 | 改动前 | 改动后 | 官方 | 变化 |
|---|---|---|---|---|
| 输入吞吐 (tok/s/GPU) | 1,084.8 | 2,571.7 | 36,170.9 | ×2.37 |
| 平均每步 (ms) | 3,775.9 | 1,592.7 | 113.2 | −57.8% |
| 每步波动 (ms) | 3,774.9–3,777.0 | 1,591.5–1,596.8 | 111.5–116.3 | 两边都稳定 |
| 占官方 | 3.0% | 7.1% | 100% | |
| 与官方差距 | 33.3× | 14.1× | — | |
| 峰值显存 (GiB) | 147.6 | 147.6 | — | 不变（scratch 为每块 query 42 MB） |
分环节耗时（profile，ms/步）
| 环节 | 改动前 | 改动后 | 官方 | 当前差距 |
|---|---|---|---|---|
| 稀疏注意力核心 | 2,254.9 | 98.8 | 14.9 | 6.6× |
| Indexer | 460.7 | 460.8 | 7.9 | 58× |
| 注意力其余（投影、压缩、gather、伪量化） | 253 | 254 | 约 19 | 13× |
| MoE（gate + 专家 + 通信） | 746 | 726 | 52.8 | 14× |
| HC、归一化及其他 | 123 | 122 | 约 10 | 12× |
| 合计（profile 下的步长） | 3,838 | 1,662 | 110.6 | 15× |
只有稀疏注意力一项发生了变化，其他环节都在测量误差以内，说明这次改动没有副作用。profile 下的步长比计时窗口长约 4%，这是 profiler 自身的开销。
Kernel 分析
| 形状 (T, H, K, D) | 旧 (ms) | 新 (ms) | 加速 | 新 TFLOP/s |
|---|---|---|---|---|
| 256, 64, 640, 512（真实压缩层） | 3.672 | 0.166 | 22.1× | 129 |
| 256, 64, 128, 512（真实无压缩层） | 0.686 | 0.069 | 9.9× | 62 |
| 256, 4, 640, 512（tiny 模型） | 0.284 | 0.105 | 2.7× | 13 |
- 为什么快： 一个 query 的 64 个 head 共享同一组 640 个 key。旧 kernel 每个 (query, head) 是一个 program，用
tl.sum在 CUDA core 上算，同一组 key 被读了 64 × 2 次。新 kernel 把 32 个 head 拼成tl.dot的 M 维，走 tensor core，每组 key 每个 head 块只读一遍 QK，再读一遍 PV。 - 离峰值还有多远： 每步注意力约 13.2 TFLOP，98.8 ms 相当于 134 TFLOP/s，约为 B200 稠密 bf16 峰值的 6%。官方 14.9 ms 约 0.89 PFLOP/s。
- 剩余差距的原因：
- M 只有 32，tensor core tile 偏小。HEADS=64 时 fp32 累加器 [64, 512] 放不进寄存器，会溢出，扫参也证实了这点。
- 两遍算法要多写读一次 scores scratch。
- kernel 外面还有约 180 ms 的 gather/拷贝，用来把 key 拼成 [T, 640, 512]。
- tiny 模型收益小： 它只有 4 个 head，凑不满 M=16 的 tile，共享 key 带来的收益有限。
精度
- 与 fp32 CPU 参考的最大误差和旧 kernel 同一量级（≤ 4.9e-4）。tensor core 的累加顺序不同，所以和旧 kernel 不是 逐位相等。
- bf16 概率舍入回归 fixture 仍然通过：去掉 bf16 舍入，或把 sink 算进最大值，这个 fixture 都会失败。
- 端到端参考测试全部通过（golden token、参考轨迹逐位置 logits、chunked prefill、抢占）。全量 163 passed。
tiny 模型 goodput（单卡，bench.py 默认参数，3 次取中位数）
| 吞吐 (tok/s) | TTFT p50 (ms) | TPOT p50/p99 (ms) | goodput (req/s) | |
|---|---|---|---|---|
| 改动前 | 6,791 | 294 | 4.8 / 16.1 | 49.36 |
| 改动后 | 7,202（+6%） | 274 | 4.6 / 15.0 | 52.36（+6%） |
下一步
按剩余耗时排序：MoE 726 ms，indexer 461 ms，注意力其余 254 ms。
- B1： 融合 indexer 打分。微基准已经从 61.8 ms 降到 4.9 ms/层（904 TFLOP/s），预计这一项约 461 → 90 ms， 吞吐约 3,200 tok/s/GPU。
- B2： 第 24–36 层只对候选块打分。
- C： MoE 直接用 FP4 权重做 grouped GEMM。
原始数据：profile-2026-10-22/stage-a/，改动前：
profile-2026-10-22/。