从 1,084.8 到 4,315.9 tok/s/GPU：mini-DeepInfer 三轮性能优化

为什么一个约 2000 行的 PyTorch 推理引擎，能在 4 张 B200 上把 prefill 从 1,085 做到 4,316 tok/s/GPU？我做 mini-DeepInfer，是想从最小实现开始，一步步追问 DeepInfer 的每个设计为什么出现、到底解决了什么问题。接下来把实现、性能、失败和新瓶颈全部摊开。项目代码在 mini-DeepInfer，官方实现参考 DeepSeek-Inc/DeepInfer；本文的三轮数据来自 Issue #1，剩余差距见 Roadmap。

项目代码： mini-DeepInfer

先固定基准

所有项目级数据来自同一组 prefill 配置：4× NVIDIA B200、EP4×DP4、CP1、131072 token context、126976 cached token、每个 rank 每步 4096 个 uncached token，预热后统计 10 个 steady-state steps。官方对比使用同类 4×B200 配置，带单节点 NVLink adapter。

当前 mini-DeepInfer 达到官方吞吐的 11.9%，与官方仍差 8.4×。这个数字只描述同一 prefill 口径下的一次配对运行，decode、speculative decoding 和完整 CUDA Graph 没有混入比较。

初始 profile 的 3,838 ms 主要花在三处：

普通 GEMM 只占很小一部分，继续调它不会改变主路径。三轮优化分别击中了这三个热点。

A：把 sparse attention 改成真正的矩阵计算

旧 kernel 对每个 query head 逐元素计算：

score = tl.sum(kv * q[None, :], 1) * SCALE acc += p[:, None] * kv

这个写法让编译器看到的是逐元素乘加，KV 读取也无法按 Tensor Core 需要的形状组织。A 版把 32 个 query head 放进 M 维 tile，使用 tl.dot 完成 score 和加权累加：

s = tl.dot(q, tl.trans(kv)) * SCALE p = tl.where(mask, tl.exp2(s - m), 0.0) acc = tl.dot(p.to(tl.bfloat16), kv, acc)

M=64 会带来寄存器压力，最终选择 M=32；这个取值和 V4.1 的真实 head 布局匹配。源码入口在 attention.pyattention.pyattention.py。

单层真实压缩层的 kernel 测试：

端到端 sparse attention core 从 2,254.9 ms 降到 98.8 ms，约 22.8×；完整 step 从 3,775.9 ms 降到 1,592.7 ms，吞吐从 1,084.8 提升到 2,571.7 tok/s/GPU。fp32 CPU reference 的最大误差不超过 4.9e-4。kernel 加速没有直接等于整机加速，key concatenation、投影和同步随后变成新的可见成本。

B1：把 Indexer scoring 留在一个 launch 里

原始 PyTorch 路径先生成完整 score，再做 ReLU、weight 乘法和求和：

s = torch.einsum("shd,nd->shn", q[rows], keys[b]) score[rows] = (s.relu_() * weights[rows].unsqueeze(-1)).sum(dim=1)

长上下文下，显存往返比乘法本身更贵。B1 使用 64 个 query × 128 个 entry 的 tile，一个 batch 只发一个 kernel；token_seq 在 device 侧读取，ReLU、权重和 reduction 保留在寄存器中，最后只写出 [T,n]。

Indexer 阶段从 460.7 ms 降到 76.5 ms，约 6.0×；完整 step 从 1,592.7 ms 降到 1,214.3 ms，吞吐提升 31.2%，达到 3,373.1 tok/s/GPU。峰值显存从 147.6 GiB 降到 144.5 GiB。fused 结果与旧路径的 top-512 overlap 为 0.9934，对照值 0.9935；误差分别为 1.55e-2 和 2.10e-2。

关键点在布局复用：Indexer 直接产出 sparse attention 能消费的 candidate block。后续 attention 如果再次复制 index，B1 的局部收益会被写回显存的代价抵消。

C1：MoE 只组合本卡真正需要的 rows

MoE 的旧路径先把完整的 [T×k, 5120] 中间块搬出来，再做 dtype 转换、where、index_copy 和求和。以 98,304 rows 为例，本卡实际只负责约四分之一，剩余 rows 的搬运都在制造无效流量。

C1 将 dequant 直接写到本卡的 output slice，只让 combine kernel 读取 local rows：

for s in tl.static_range(K): y = tl.load(Y + row * dim + d, ok & local, other=0.0) acc += y.to(tl.float32) tl.store(O + row * dim + d, acc.to(tl.bfloat16), mask=ok)

单层、1×B200、16,384 tokens、top-6、96 个 local experts 的测试：

单层延迟从 12.93 ms 降到 6.24 ms，checksum 都是 5.480333e+04，峰值显存从 8.7 GiB 降到 7.7 GiB。放回整机后，MoE 阶段从 701 ms 降到 410 ms，完整吞吐从 3,373.1 提升到 4,315.9 tok/s/GPU，提升 27.8%，平均 step 降到 949 ms。

Issue 中还加了严格回归：把非本地 rows 填成 NaN，覆盖 T=1/5/300、k=1/6/4、dim=1030；256 字节 FP4 block 和 256 个 E8M0 exponent 都按 bitwise 比较，确保 kernel 没有偷偷读取 remote rows。

三轮之后，时间花到哪里

当前 profile 的剩余大项如下：

Gate 的 CUDA torch.mm(x, weight.t(), out_dtype=torch.float32) 已经在单独 issue 中把 55.5 ms 从 step 中拿掉；indexed keys 的独立实验再减少约 130 ms，达到 5,004.4 tok/s/GPU。两项来自不同 job，不能相加写成一个结果。

下一步先处理 key concatenation、MoE dispatch/GEMM/combine 融合和 Engram peer read，再验证 SWA bounded replay。后续方向包括 DSpark 静态验证、Breakable CUDA Graph、FP8/FP4 KV cache、prefix cache/fork 和 decode grouped expert GEMM。每次提交按三层证据检查：kernel 时间、stage 时间、固定基准下的端到端吞吐，并同时跑 reference 对齐。

完整实现可以从 attention.py、indexer.py、moe.py 和 Issue #1 对照阅读。