Performance report性能报告
Performance Report B1: A Fused Indexer Score性能分析报告 B1：融合 indexer 打分
A single fused kernel replaces the chunked einsum for indexer scoring: 12.2× faster per layer, +31% real 4-GPU prefill throughput, 3.1 GiB less peak memory.用一个融合 kernel 取代分块 einsum 做 indexer 打分：单层快 12.2 倍，真实四卡 prefill 吞吐 +31%，峰值显存降低 3.1 GiB。
2026-10-22 · code 1868b01 (before the change: f759f32, i.e. stage A)
Conclusion
Real V4.1 4-GPU prefill throughput rises from 2,572 to 3,373 tok/s/GPU (+31%), and the step time falls from 1,593 ms to 1,214 ms. The gap to the official engine shrinks from 14.1× to 10.7×; mini reaches 9.3% of the official engine. The indexer drops from 461 ms/step to 76.5 ms/step (6.0×), and peak memory falls by 3.1 GiB. MoE (701 ms) is now 56% of each step and the only big item left.
Setup
Same as report A: real DeepSeek-V4.1-Flash (dba1be0), 1 node 4 × B200, EP4 / DP4, the shared workload (131,072 context, 126,976 cached, 4,096 new tokens per step per rank), timing over 10 steady-state steps.
- Final version: measured with
tools/bench_rev.py run perf/fused-index-scores --suite real4(job 9342). - Profile: taken from the first version (per-sequence launches, job 9334). Its end-to-end throughput differs from the final version by 0.1% (3,375.9 vs 3,373.1).
End-to-end results
| Metric | Before (A) | After (B1) | Official | Change |
|---|---|---|---|---|
| Input throughput (tok/s/GPU) | 2,571.7 | 3,373.1 | 36,170.9 | +31.2% |
| Mean step (ms) | 1,592.7 | 1,214.3 | 113.2 | −23.8% |
| Share of official | 7.1% | 9.3% | 100% | |
| Gap to official | 14.1× | 10.7× | — | |
| Peak memory (GiB) | 147.6 | 144.5 | — | −3.1 (the [512, 32, n] intermediate is no longer materialized) |
Per-stage time (profile, ms/step)
| Stage | Start | After A | After B1 | Official | Current gap |
|---|---|---|---|---|---|
| Sparse attention core | 2,254.9 | 98.8 | 98.9 | 14.9 | 6.6× |
| Indexer | 460.7 | 460.8 | 76.5 | 7.9 | 9.7× |
| of which scoring | ~400 | ~400 | 28.0 | 5.1 (mqa_logits + sparse_mqa_logits) | 5.5× |
| of which top-k, mask, candidate blocks | ~60 | ~60 | 48.5 | 2.8 | 17× |
| Other attention | 253 | 254 | 254 | ~19 | 13× |
| MoE (gate + experts + communication) | 746 | 726 | 701 | 52.8 | 13× |
| HC, normalization and other | 123 | 122 | 122 | ~10 | 12× |
| Total | 3,838 | 1,662 | 1,252 | 110.6 | 11.3× |
Kernel analysis
| Shape | Old: chunked einsum (ms) | New: fused (ms) | Speedup | TFLOP/s | Max error vs. fp32 reference (new / old) | Top-512 overlap (new / old) |
|---|---|---|---|---|---|---|
| T=4,096, n=65,536 (ratio-2 layer) | 30.81 | 2.52 | 12.2× | 874 | 1.49e-2 / 1.73e-2 | 0.9965 / 0.9957 |
| T=4,096, n=131,072 (ratio-1 layer) | 61.76 | 5.08 | 12.2× | 866 | 1.55e-2 / 2.10e-2 | 0.9934 / 0.9935 |
- Why the old one was slow: scoring itself is one bf16 GEMM, but the old path wrote the [512, 32, n] result back to memory, then ran relu, the weight multiply and the per-head sum as separate passes. In one ratio-1 layer those three elementwise passes move about 4 × 34 GB, about 80% of the indexer’s time.
- What the new kernel does: one [64 query × 128 entry] tile loads its keys once; the 32 heads run
tl.dot, relu and the weight multiply in turn, accumulating in registers, and a single [T, n] bf16 score is written at the end. The kernel reaches about 870 TFLOP/s, 39% of B200’s dense bf16 peak, already close to the compute limit. - Why a single launch: the first version launched one kernel per sequence. A tiny-model prefill batch has dozens of short sequences, and Triton’s CPU-side launch overhead cut throughput by 9.5%. The final version launches the whole batch once using the device-side
token_seq, computing row blocks that span sequences sequence by sequence. A row block that belongs to one sequence takes a fast path; on real shapes it is only 4% slower than the per-sequence version (5.08 vs 4.87 ms/layer) and introduces no host synchronization. - The remaining 5.5× vs. official: the official kernels use FP8 Q × FP4 K (
fp8_fp4_mqa_logits) and, in layers 24–36, score only the candidate blocks (sparse_mqa_logits). The second point is B2.
Accuracy
- The new kernel accumulates relu and the weighted sum in fp32 and rounds to bf16 once at the end; the old path had a bf16 intermediate at every step. So the new kernel’s error against the fp32 reference is actually smaller (1.55e-2 vs 2.10e-2, with a maximum score around 4.8), and top-512 overlap is on par.
- All end-to-end reference tests pass (golden tokens, per-position logits along the reference trajectory, chunked prefill, preemption). Full suite: 174 passed.
Tiny-model goodput (1 GPU, tools/bench_rev.py --suite tiny, 3 interleaved runs in one job, median)
| Version | Throughput (tok/s) | TTFT p50 (ms) | TPOT p50/p99 (ms) | goodput (req/s) |
|---|---|---|---|---|
A (f759f32) | 7,604 | 240 | 4.3 / 13.7 | 55.3 |
B1 (1868b01) | 7,469 (−1.8%) | 237 | 4.3 / 13.1 | 54.3 |
The tiny model’s context is at most 1,280 tokens, so scoring was never its bottleneck. The −1.8% is within the run-to-run spread of 3 runs: version A itself measured 7,420 and 7,604 in two jobs. The per-sequence-launch first version measured 6,718 on tiny (−9.5%); that regression is fixed by the single launch.
Tests
tests/test_indexer.py:
- Runs the fused path and the chunked fallback path.
- New “short-40”: 40 sequences of 2 tokens, so a 64-row block spans up to 33 sequences.
- New real head shape (32 × 128): row counts that are not multiples of 64, entry counts of 1, 129 and 1,000.
Next
Ordered by remaining time: MoE 701 ms, other attention 254 ms, sparse attention 99 ms, indexer 76.5 ms.
- C: run MoE grouped GEMM directly on FP4 weights. Today dequantization takes 128 ms, bf16 conversion and copies about 300 ms, grouped GEMM 67 ms and communication 70 ms. The target is around 120 ms, for about 6,000 tok/s/GPU (17%).
- B2: scoring inside candidate blocks, a faster top-k and mask, saving about 50 ms more.
Raw data: profile-2026-10-22/stage-b1/.
2026-10-22 · 代码 1868b01（改动前 f759f32，即阶段 A）
结论
真实 V4.1 四卡 prefill 吞吐从 2,572 提高到 3,373 tok/s/GPU（+31%），每步从 1,593 ms 降到 1,214 ms。 到官方的差距从 14.1 倍缩小到 10.7 倍，mini 达到官方的 9.3%。indexer 从 461 ms/步降到 76.5 ms/步（6.0 倍）， 峰值显存降低 3.1 GiB。现在 MoE（701 ms）占每步的 56%，是唯一的大头。
实验设置
与报告 A 相同：真实 DeepSeek-V4.1-Flash（dba1be0），单机 4 × B200，EP4 / DP4，共享 workload（131,072 上下文，
126,976 已缓存，每步每 rank 4,096 新 token），计时取 10 个稳态步。
- 最终版本：用
tools/bench_rev.py run perf/fused-index-scores --suite real4测量（作业 9342）。 - profile：取自第一版（逐序列启动，作业 9334）。它和最终版的端到端吞吐相差 0.1%（3,375.9 vs 3,373.1）。
端到端结果
| 指标 | 改动前（A） | 改动后（B1） | 官方 | 变化 |
|---|---|---|---|---|
| 输入吞吐 (tok/s/GPU) | 2,571.7 | 3,373.1 | 36,170.9 | +31.2% |
| 平均每步 (ms) | 1,592.7 | 1,214.3 | 113.2 | −23.8% |
| 占官方 | 7.1% | 9.3% | 100% | |
| 与官方差距 | 14.1× | 10.7× | — | |
| 峰值显存 (GiB) | 147.6 | 144.5 | — | −3.1（[512, 32, n] 中间结果不再落地） |
分环节耗时（profile，ms/步）
| 环节 | 起点 | A 之后 | B1 之后 | 官方 | 当前差距 |
|---|---|---|---|---|---|
| 稀疏注意力核心 | 2,254.9 | 98.8 | 98.9 | 14.9 | 6.6× |
| Indexer | 460.7 | 460.8 | 76.5 | 7.9 | 9.7× |
| 其中打分 | 约 400 | 约 400 | 28.0 | 5.1（mqa_logits + sparse_mqa_logits） | 5.5× |
| 其中 top-k、mask、候选块 | 约 60 | 约 60 | 48.5 | 2.8 | 17× |
| 注意力其余 | 253 | 254 | 254 | 约 19 | 13× |
| MoE（gate + 专家 + 通信） | 746 | 726 | 701 | 52.8 | 13× |
| HC、归一化及其他 | 123 | 122 | 122 | 约 10 | 12× |
| 合计 | 3,838 | 1,662 | 1,252 | 110.6 | 11.3× |
Kernel 分析
| 形状 | 旧：分块 einsum (ms) | 新：融合 (ms) | 加速 | TFLOP/s | 与 fp32 参考的最大误差（新 / 旧） | top-512 重合率（新 / 旧） |
|---|---|---|---|---|---|---|
| T=4,096, n=65,536（ratio-2 层） | 30.81 | 2.52 | 12.2× | 874 | 1.49e-2 / 1.73e-2 | 0.9965 / 0.9957 |
| T=4,096, n=131,072（ratio-1 层） | 61.76 | 5.08 | 12.2× | 866 | 1.55e-2 / 2.10e-2 | 0.9934 / 0.9935 |
- 旧实现为什么慢： 打分本身是一次 bf16 GEMM，但它把 [512, 32, n] 的结果写回显存，再分别做 relu、乘权重、 按 head 求和。这三遍逐元素运算在一个 ratio-1 层要搬运约 4 × 34 GB，占了 indexer 耗时的约 80%。
- 新 kernel 怎么做： 一个 [64 query × 128 条目] 的 tile 加载一次 key，32 个 head 依次
tl.dot、relu、 乘权重，结果累加在寄存器里，最后只写一次 [T, n] 的 bf16 分数。这个 kernel 达到约 870 TFLOP/s，是 B200 稠密 bf16 峰值的 39%，已经接近计算上限。 - 为什么是单次启动： 第一版按序列逐个启动 kernel。tiny 模型一个 prefill batch 有几十条短序列，Triton 在 CPU
侧的启动开销让吞吐下降了 9.5%。最终版用设备端的
token_seq一次启动整个 batch，跨序列的行块逐序列计算。 一个行块只属于一条序列时走快速路径，真实形状下只比逐序列版本慢 4%（5.08 vs 4.87 ms/层），而且没有引入任何 主机同步。 - 和官方还差 5.5 倍： 官方用 FP8 Q × FP4 K（
fp8_fp4_mqa_logits），而且第 24–36 层只对候选块打分 （sparse_mqa_logits）。第二点就是 B2。
精度
- 新 kernel 在 fp32 中累加 relu 和加权求和，最后只舍入一次 bf16。旧实现每一步都有 bf16 中间结果。因此新 kernel 与 fp32 参考的误差反而更小（1.55e-2 vs 2.10e-2，分数最大值约 4.8），top-512 选择的重合率持平。
- 端到端参考测试全部通过（golden token、参考轨迹逐位置 logits、chunked prefill、抢占）。全量 174 passed。
tiny 模型 goodput（单卡，tools/bench_rev.py --suite tiny，同一作业中交替运行 3 次，取中位数）
| 版本 | 吞吐 (tok/s) | TTFT p50 (ms) | TPOT p50/p99 (ms) | goodput (req/s) |
|---|---|---|---|---|
A（f759f32） | 7,604 | 240 | 4.3 / 13.7 | 55.3 |
B1（1868b01） | 7,469（−1.8%） | 237 | 4.3 / 13.1 | 54.3 |
tiny 模型的上下文不超过 1,280 token，打分本来就不是瓶颈。−1.8% 在 3 次运行的波动范围内：同一版本 A 在两次作业中 测得 7,420 和 7,604。逐序列启动的第一版在 tiny 上是 6,718（−9.5%），这个退化已经由单次启动修复。
测试
tests/test_indexer.py：
- 融合路径和分块回退路径各跑一遍。
- 新增 “short-40”：40 条 2-token 序列，64 行的块会跨越最多 33 条序列。
- 新增真实 head 形状（32 × 128）：行数不是 64 的倍数，条目数取 1、129、1,000。
下一步
按剩余耗时排序：MoE 701 ms，注意力其余 254 ms，稀疏注意力 99 ms，indexer 76.5 ms。
- C： MoE 改为直接用 FP4 权重做 grouped GEMM。目前反量化 128 ms、bf16 转换和拷贝约 300 ms、grouped GEMM 67 ms、 通信 70 ms。目标 120 ms 左右，预计吞吐约 6,000 tok/s/GPU（17%）。
- B2： 候选块内打分、更快的 top-k 和 mask，预计再省约 50 ms。