为什么一个约 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。阶段吞吐tok/s/GPU平均 step相对初始版baseline1,084.83,775.9 ms1.00×A稀疏注意力2,571.71,592.7 ms2.37×B1Indexer scoring3,373.11,214.3 ms3.11×C1MoE combine4,315.9949.0 ms3.98×官方 DeepInfer36,170.9113.2 ms33.34×当前 mini-DeepInfer 达到官方吞吐的 11.9%与官方仍差 8.4×。这个数字只描述同一 prefill 口径下的一次配对运行decode、speculative decoding 和完整 CUDA Graph 没有混入比较。初始 profile 的 3,838 ms 主要花在三处阶段时间占比sparse attention2,254.9 ms59%MoE746.0 ms19%indexer460.7 ms12%其余 attention253.0 ms7%HC、norm 和其他123.0 ms3%普通 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)M64 会带来寄存器压力最终选择 M32这个取值和 V4.1 的真实 head 布局匹配。源码入口在 attention.pyattention.pyattention.py。单层真实压缩层的 kernel 测试workload旧实现A 版加速compressed layer3.672 ms0.166 ms22.1×plain layer0.686 ms0.069 ms9.9×tiny layer0.284 ms0.105 ms2.7×端到端 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(dim1)长上下文下显存往返比乘法本身更贵。B1 使用 64 个 query × 128 个 entry 的 tile一个 batch 只发一个 kerneltoken_seq 在 device 侧读取ReLU、权重和 reduction 保留在寄存器中最后只写出 [T,n]。n旧实现B1 fused kernel加速实测吞吐65,53630.81 ms2.52 ms12.2×874 TFLOP/s131,07261.76 ms5.08 ms12.2×866 TFLOP/sIndexer 阶段从 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 如果再次复制 indexB1 的局部收益会被写回显存的代价抵消。C1MoE 只组合本卡真正需要的 rowsMoE 的旧路径先把完整的 [T×k, 5120] 中间块搬出来再做 dtype 转换、where、index_copy 和求和。以 98,304 rows 为例本卡实际只负责约四分之一剩余 rows 的搬运都在制造无效流量。C1 将 dequant 直接写到本卡的 output slice只让 combine kernel 读取 local rowsfor s in tl.static_range(K): y tl.load(Y row * dim d, ok local, other0.0) acc y.to(tl.float32) tl.store(O row * dim d, acc.to(tl.bfloat16), maskok)单层、1×B200、16,384 tokens、top-6、96 个 local experts 的测试项目旧路径C1变化FP4 dequant128.1 ms132.3 ms3.3%whole-block copy85.7 ms0.1 ms基本消失bf16→fp3266.0 ms0 ms消失index_copy59.0 ms0 ms消失where55.8 ms0 ms消失sum16.5 ms0 ms消失新 combine0 ms4.0 ms新增grouped GEMM67.1 ms66.0 ms基本不变routed experts 总计526 ms252 ms-52%单层延迟从 12.93 ms 降到 6.24 mschecksum 都是 5.480333e04峰值显存从 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覆盖 T1/5/300、k1/6/4、dim1030256 字节 FP4 block 和 256 个 E8M0 exponent 都按 bitwise 比较确保 kernel 没有偷偷读取 remote rows。三轮之后时间花到哪里当前 profile 的剩余大项如下阶段当前耗时官方对应耗时key concat / gather182 ms约 0使用 indexed KVFP4 dequant132 ms更深的融合路径HC、norm 和其他122 ms约 10 mssparse attention core99 ms14.9 msindexer79 ms7.9 msgate76 ms约 1.2 msattention projection72 ms已融合grouped GEMM66 ms约 40 ms 级别communication53 ms更大程度隐藏Gate 的 CUDA torch.mm(x, weight.t(), out_dtypetorch.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 对照阅读。 
长河君 · 2026/10/27 2:05:07 · 阅读 1.2k 从 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。阶段吞吐tok/s/GPU平均 step相对初始版baseline1,084.83,775.9 ms1.00×A稀疏注意力2,571.71,592.7 ms2.37×B1Indexer scoring3,373.11,214.3 ms3.11×C1MoE combine4,315.9949.0 ms3.98×官方 DeepInfer36,170.9113.2 ms33.34×当前 mini-DeepInfer 达到官方吞吐的 11.9%与官方仍差 8.4×。这个数字只描述同一 prefill 口径下的一次配对运行decode、speculative decoding 和完整 CUDA Graph 没有混入比较。初始 profile 的 3,838 ms 主要花在三处阶段时间占比sparse attention2,254.9 ms59%MoE746.0 ms19%indexer460.7 ms12%其余 attention253.0 ms7%HC、norm 和其他123.0 ms3%普通 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)M64 会带来寄存器压力最终选择 M32这个取值和 V4.1 的真实 head 布局匹配。源码入口在 attention.pyattention.pyattention.py。单层真实压缩层的 kernel 测试workload旧实现A 版加速compressed layer3.672 ms0.166 ms22.1×plain layer0.686 ms0.069 ms9.9×tiny layer0.284 ms0.105 ms2.7×端到端 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(dim1)长上下文下显存往返比乘法本身更贵。B1 使用 64 个 query × 128 个 entry 的 tile一个 batch 只发一个 kerneltoken_seq 在 device 侧读取ReLU、权重和 reduction 保留在寄存器中最后只写出 [T,n]。n旧实现B1 fused kernel加速实测吞吐65,53630.81 ms2.52 ms12.2×874 TFLOP/s131,07261.76 ms5.08 ms12.2×866 TFLOP/sIndexer 阶段从 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 如果再次复制 indexB1 的局部收益会被写回显存的代价抵消。C1MoE 只组合本卡真正需要的 rowsMoE 的旧路径先把完整的 [T×k, 5120] 中间块搬出来再做 dtype 转换、where、index_copy 和求和。以 98,304 rows 为例本卡实际只负责约四分之一剩余 rows 的搬运都在制造无效流量。C1 将 dequant 直接写到本卡的 output slice只让 combine kernel 读取 local rowsfor s in tl.static_range(K): y tl.load(Y row * dim d, ok local, other0.0) acc y.to(tl.float32) tl.store(O row * dim d, acc.to(tl.bfloat16), maskok)单层、1×B200、16,384 tokens、top-6、96 个 local experts 的测试项目旧路径C1变化FP4 dequant128.1 ms132.3 ms3.3%whole-block copy85.7 ms0.1 ms基本消失bf16→fp3266.0 ms0 ms消失index_copy59.0 ms0 ms消失where55.8 ms0 ms消失sum16.5 ms0 ms消失新 combine0 ms4.0 ms新增grouped GEMM67.1 ms66.0 ms基本不变routed experts 总计526 ms252 ms-52%单层延迟从 12.93 ms 降到 6.24 mschecksum 都是 5.480333e04峰值显存从 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覆盖 T1/5/300、k1/6/4、dim1030256 字节 FP4 block 和 256 个 E8M0 exponent 都按 bitwise 比较确保 kernel 没有偷偷读取 remote rows。三轮之后时间花到哪里当前 profile 的剩余大项如下阶段当前耗时官方对应耗时key concat / gather182 ms约 0使用 indexed KVFP4 dequant132 ms更深的融合路径HC、norm 和其他122 ms约 10 mssparse attention core99 ms14.9 msindexer79 ms7.9 msgate76 ms约 1.2 msattention projection72 ms已融合grouped GEMM66 ms约 40 ms 级别communication53 ms更大程度隐藏Gate 的 CUDA torch.mm(x, weight.t(), out_dtypetorch.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 对照阅读。