Performance report性能报告
Performance Report C1: MoE Combine Over Local Rows, In-Place Dequantization性能分析报告 C1:MoE 只合并本地行,原地反量化
Combining only the local expert rows and dequantizing in place cuts routed experts from 526 to 252 ms/step and lifts real 4-GPU prefill to 4,316 tok/s/GPU.只合并本地专家的行并原地反量化,路由专家从 526 降到 252 ms/步,真实四卡 prefill 提高到 4,316 tok/s/GPU。
2026-10-22 · code 5f5fb73 (before the change: 32e2b5b, i.e. B1)
Conclusion
Real V4.1 4-GPU prefill throughput rises from 3,377 to 4,316 tok/s/GPU (+27.8%), and the step time falls from 1,213 ms to 949 ms, matching the roughly 950 ms predicted from the micro-benchmark before the change. The gap to the official engine shrinks from 10.7× to 8.4×; mini reaches 11.9% of the official engine. MoE drops from 701 ms/step to 410 ms/step, and the expert part from 526 ms to 252 ms. Attention (432 ms) is again the largest item, 182 ms of which is glue that concatenates the keys.
Setup
Same as reports A and B1: real V4.1-Flash (dba1be0), 4 × B200, EP4 / DP4, the shared workload, 10 steady-state steps.
- Throughput:
tools/bench_rev.py run perf/fused-index-scores perf/moe-combine --suite real4, the two versions interleaved in the same job 9349. - Profile: job 9351. Within the profile job, the same version measured 4,320.4 tok/s/GPU, 0.1% from job 9349.
End-to-end results
| Metric | Before (B1) | After (C1) | Official | Change |
|---|---|---|---|---|
| Input throughput (tok/s/GPU) | 3,376.7 | 4,315.9 | 36,170.9 | +27.8% |
| Mean step (ms) | 1,213.0 | 949.0 | 113.2 | −21.8% |
| Share of official | 9.3% | 11.9% | 100% | |
| Gap to official | 10.7× | 8.4× | — |
Per-stage time (profile, ms/step)
| Stage | Start | After B1 | After C1 | Official | Current gap |
|---|---|---|---|---|---|
| Sparse attention core | 2,254.9 | 98.9 | 99.0 | 14.9 | 6.6× |
| Attention key concatenation (gather / cat / where) | — | ~180 | 182 | 0 (FlashMLA reads by index) | — |
| Indexer | 460.7 | 76.5 | 79.3 | 7.9 | 10× |
| Other attention (projections, compression, fake quant) | — | ~74 | 72 | ~19 | 3.8× |
| MoE total | 746 | 701 | 410 | 52.8 | 7.8× |
| routed experts | 535 | 526 | 252 | 50.1 (mega-MoE) | 5.0× |
| gate | 76 | 76 | 76 | 1.2 | 63× |
| communication (reduce-scatter + all-gather) | ~113 | ~70 | 53 | 1.5 | 35× |
| HC, normalization and other | 123 | 122 | 122 | ~10 | 12× |
| Total (step under profile) | 3,838 | 1,252 | 965 | 110.6 | 8.7× |
The “other attention 254 ms” of the earlier reports is split this time into 182 ms of key concatenation and 72 ms of projections and the rest.
Inside the routed experts (ms/step)
| Kernel | Before | After | Note |
|---|---|---|---|
| FP4 dequantization | 128.1 | 132.3 | Reads the E8M0 scale directly; the compute is unchanged |
| Whole-block copy of the dequantized result (Memcpy DtoD) | 85.7 | 0.1 | Dequantization writes straight into the bf16 weight stack |
bf16 conversion (.float() etc.) | 66.0 | 0 | The GEMM output stays bf16 |
index_copy_ to restore order | 59.0 | 0 | Replaced by the combine kernel |
where to zero non-local rows | 55.8 | 0 | combine never reads non-local rows |
| Sum over slots | 16.5 | 0 | Folded into combine |
| combine kernel | — | 4.0 | |
| Grouped GEMM | 67.1 | 66.0 | |
| SwiGLU, gather, sort | ~45 | ~47 | |
| Total | 526 | 252 | −52% |
- Why the old version was slow:
GroupedExpertslays its output out in sorted [T·k, 5120] rows with T·k = 98,304, and only about 1/4 of them come from this rank’s experts. The old code converted the whole block to fp32 (2 GB per layer), zeroed non-local rows withwhere, restored order withindex_copy_and summed over slots; most of what it moved was zeros. - What the combine kernel does: for each token, in slot order, it reads only the rows of local experts, accumulates in fp32 and writes [T, 5120] once. The summation order is unchanged, so the result is deterministic. The GEMM never writes non-local rows in the first place, so they are never read.
- Dequantization: it writes straight into the matching slice of the stack, saving a 6.8 GB bf16 copy per layer. The E8M0 scale is decoded in the kernel, with the two boundary values 0 and 255 handled separately (an fp32 subnormal 2^-127 and NaN, respectively).
Single-layer micro-benchmark (1 × B200, 16,384 tokens, top-6, 96 experts on this rank): 12.93 → 6.24 ms, identical output checksum (5.480333e+04), peak memory 8.7 → 7.7 GiB.
Accuracy
- The combine’s summation order matches the old implementation, and the unit test requires an error ≤ 1e-6 against it.
- Dequantization is bit-exact against the CPU reference:
test_fp4_all_codes_and_scaleswalks all 256 byte values and all 256 E8M0 exponents, including signed zero, overflow, underflow and NaN. - All end-to-end reference tests pass. Full suite: 178 passed.
Tiny-model goodput (1 GPU, 3 interleaved runs in one job, median)
| Version | Throughput (tok/s) | TTFT p50 (ms) | TPOT p50/p99 (ms) | goodput (req/s) |
|---|---|---|---|---|
B1 (32e2b5b) | 7,871 | 221 | 4.2 / 12.8 | 57.2 |
C1 (5f5fb73) | 8,036 (+2.1%) | 238 | 4.1 / 12.5 | 58.4 (+2.1%) |
Tests
test_combine_reads_only_local_rows_in_slot_order:- fills non-local rows with NaN and requires the output to be fully finite and equal to the old implementation;
- covers T = 1 / 5 / 300, k = 1 / 6 / 4, and a dimension of 1,030 that is not a multiple of 1,024.
test_chunked_dequant_writes_the_stack_in_place: the in-place weight stack is bit-equal to per-expert dequantization.
Next
Ordered by remaining time (965 ms in total):
| Item | Time (ms) |
|---|---|
| Attention key concatenation | 182 |
| FP4 dequantization | 132 |
| HC, normalization and other | 122 |
| Sparse attention | 99 |
| Indexer | 79 |
| gate | 76 |
| Attention projections etc. | 72 |
| Grouped GEMM | 66 |
| Communication | 53 |
- A2: let the sparse-attention kernel read the KV cache directly by index, removing the key concatenation; about −170 ms expected.
- C2: dequantize two values per byte and write them vectorized; about −80 ms expected.
- C3: run the gate in tf32x3, close to fp32 accuracy; about −40 ms expected.
Raw data: profile-2026-10-22/stage-c1/.
2026-10-22 · 代码 5f5fb73(改动前 32e2b5b,即 B1)
结论
真实 V4.1 四卡 prefill 吞吐从 3,377 提高到 4,316 tok/s/GPU(+27.8%),每步从 1,213 ms 降到 949 ms,和改动前 根据微基准做的预测(约 950 ms)一致。到官方的差距从 10.7 倍缩小到 8.4 倍,mini 达到官方的 11.9%。MoE 从 701 ms/步降到 410 ms/步,专家部分从 526 ms 降到 252 ms。现在注意力(432 ms)重新成为最大的一项,其中 182 ms 是 拼接 key 的胶水操作。
实验设置
与报告 A、B1 相同:真实 V4.1-Flash(dba1be0),4 × B200,EP4 / DP4,共享 workload,10 个稳态步。
- 吞吐:
tools/bench_rev.py run perf/fused-index-scores perf/moe-combine --suite real4,两个版本在同一个作业 9349 里交替运行。 - profile:作业 9351。profile 作业里同一版本测得 4,320.4 tok/s/GPU,和 9349 相差 0.1%。
端到端结果
| 指标 | 改动前(B1) | 改动后(C1) | 官方 | 变化 |
|---|---|---|---|---|
| 输入吞吐 (tok/s/GPU) | 3,376.7 | 4,315.9 | 36,170.9 | +27.8% |
| 平均每步 (ms) | 1,213.0 | 949.0 | 113.2 | −21.8% |
| 占官方 | 9.3% | 11.9% | 100% | |
| 与官方差距 | 10.7× | 8.4× | — |
分环节耗时(profile,ms/步)
| 环节 | 起点 | B1 之后 | C1 之后 | 官方 | 当前差距 |
|---|---|---|---|---|---|
| 稀疏注意力核心 | 2,254.9 | 98.9 | 99.0 | 14.9 | 6.6× |
| 注意力 key 拼接(gather / cat / where) | — | 约 180 | 182 | 0(FlashMLA 按索引读) | — |
| Indexer | 460.7 | 76.5 | 79.3 | 7.9 | 10× |
| 注意力其余(投影、压缩、伪量化) | — | 约 74 | 72 | 约 19 | 3.8× |
| MoE 合计 | 746 | 701 | 410 | 52.8 | 7.8× |
| 路由专家 | 535 | 526 | 252 | 50.1(mega-MoE) | 5.0× |
| gate | 76 | 76 | 76 | 1.2 | 63× |
| 通信(reduce-scatter + all-gather) | 约 113 | 约 70 | 53 | 1.5 | 35× |
| HC、归一化及其他 | 123 | 122 | 122 | 约 10 | 12× |
| 合计(profile 下的步长) | 3,838 | 1,252 | 965 | 110.6 | 8.7× |
之前报告里的”注意力其余 254 ms”,这次拆成了 key 拼接 182 ms 和投影等 72 ms 两项。
路由专家内部(ms/步)
| kernel | 改动前 | 改动后 | 说明 |
|---|---|---|---|
| FP4 反量化 | 128.1 | 132.3 | 直接读 E8M0 scale,计算量不变 |
| 反量化结果整块拷贝(Memcpy DtoD) | 85.7 | 0.1 | 反量化直接写进 bf16 权重栈 |
bf16 转换(.float() 等) | 66.0 | 0 | GEMM 输出保持 bf16 |
index_copy_ 还原排序 | 59.0 | 0 | 换成 combine kernel |
where 置零非本地行 | 55.8 | 0 | combine 不读非本地行 |
| 按 slot 求和 | 16.5 | 0 | 并入 combine |
| combine kernel | — | 4.0 | |
| grouped GEMM | 67.1 | 66.0 | |
| SwiGLU、gather、排序 | 约 45 | 约 47 | |
| 合计 | 526 | 252 | −52% |
- 为什么旧实现慢:
GroupedExperts的输出按排序后的 [T·k, 5120] 行排列,T·k = 98,304,其中只有约 1/4 来自 本 rank 的专家。旧代码把整块转成 fp32(2 GB/层),用where把非本地行置零,再index_copy_还原顺序、按 slot 求和,大部分搬运的是零。 - combine kernel 怎么做: 对每个 token 按 slot 顺序,只读本地专家的那几行,在 fp32 里累加,一次写出 [T, 5120]。 求和顺序和原来一样,结果是确定的。非本地行 GEMM 本来就不写,也不会被读到。
- 反量化: 直接写进栈里对应的切片,省掉了每层 6.8 GB 的 bf16 拷贝。E8M0 scale 在 kernel 里解码,0 和 255 这 两个边界值单独处理(分别对应 fp32 次正规数 2^-127 和 NaN)。
单层微基准(1 × B200,16,384 个 token,top-6,本 rank 96 个专家):12.93 → 6.24 ms,输出校验和相同 (5.480333e+04),峰值显存 8.7 → 7.7 GiB。
精度
- combine 的求和顺序和旧实现一致,单测要求和旧实现的误差 ≤ 1e-6。
- 反量化与 CPU 参考逐位相等:
test_fp4_all_codes_and_scales遍历了全部 256 种字节和 256 种 E8M0 指数,包括有符号 零、溢出、下溢和 NaN。 - 端到端参考测试全部通过。全量 178 passed。
tiny 模型 goodput(单卡,同一作业中交替运行 3 次,取中位数)
| 版本 | 吞吐 (tok/s) | TTFT p50 (ms) | TPOT p50/p99 (ms) | goodput (req/s) |
|---|---|---|---|---|
B1(32e2b5b) | 7,871 | 221 | 4.2 / 12.8 | 57.2 |
C1(5f5fb73) | 8,036(+2.1%) | 238 | 4.1 / 12.5 | 58.4(+2.1%) |
测试
test_combine_reads_only_local_rows_in_slot_order:- 把非本地行填成 NaN,要求输出全部有限,且和旧实现一致;
- 覆盖 T = 1 / 5 / 300、k = 1 / 6 / 4,以及不是 1,024 整数倍的维度 1,030。
test_chunked_dequant_writes_the_stack_in_place:原地写入的权重栈和逐专家反量化结果逐位相等。
下一步
按剩余耗时排序(共 965 ms):
| 项 | 耗时 (ms) |
|---|---|
| 注意力 key 拼接 | 182 |
| FP4 反量化 | 132 |
| HC、归一化及其他 | 122 |
| 稀疏注意力 | 99 |
| Indexer | 79 |
| gate | 76 |
| 注意力投影等 | 72 |
| grouped GEMM | 66 |
| 通信 | 53 |
- A2: 稀疏注意力 kernel 按索引直接读 KV cache,省掉 key 拼接,预计 −170 ms。
- C2: 反量化按字节一次解两个值并向量化写出,预计 −80 ms。
- C3: gate 改用 tf32x3,精度接近 fp32,预计 −40 ms。