[algebra-MCX]

English · click to expand

Motivation

The MoE gate takes about 76 ms per prefill step; the official DeepInfer takes 1.2 ms. The gate casts x and its weight to fp32 and runs an fp32 SIMT GEMM, which alone costs about 49 ms per step.

Change

On CUDA, Gate.forward uses torch.mm(x, weight.t(), out_dtype=torch.float32): bf16 inputs on tensor cores, fp32 output. Scoring, top-k and normalization stay fp32. The numerics match DeepInfer's gate (bf16 in, fp32 out). CPU does not support out_dtype, so it keeps the fp32 path.

Size: model.py +5/−1; one new test in tests/test_moe.py (3 scoring functions).

Validation

• pytest -q -m "not gpu" (CPU): 35 passed
• python -m compileall -q deepinfer_lab tests
• git diff --check
• GPU, 1 × B200:
  • tests/test_moe.py: 16 passed, including the new test. V4.1-Flash gate shape (5120 → 384, top-6), 300 tokens: bf16 and fp32 pick identical experts, weight error < 1e-4
  • tests/test_cuda_graph.py, tests/test_sync_free.py: 20 passed together with test_moe.py. The decode CUDA graph captures it; no host sync in the forward
  • tests/test_v41.py, tests/test_engine.py: 17 passed. Tiny model vs DeepSeek's reference implementation; logit gap still < 0.05
• Real V4.1-Flash 4 × B200 prefill throughput: +6.2%, −55.5 ms per step (about 50 ms predicted); see Measurement below

Measurement

Job 9816: 4 × B200, real V4.1-Flash, EP4 / DP4 / CP1, 131,072-token context (126,976 cached), 4,096 new tokens per rank per step, 10 steady-state steps per run, no CUDA graph. The three revisions alternate in one job: main → gate → attention → attention → gate → main.

Revision	Run 1	Run 2	Mean throughput (tok/s/GPU)	Median step	vs main
main (2818d6e)	4,322.6	4,321.6	4,322.1	947.7 ms	—
bf16 gate (048e6e3)	4,591.6	4,587.1	4,589.4	892.2 ms	throughput +6.2%, −55.5 ms per step
indexed keys (ba0c0d3)	4,997.1	5,011.7	5,004.4	817.7 ms	throughput +15.8%, −130.0 ms per step

• The two runs of each revision differ by < 0.3%; peak memory is 143.5 GiB for all three.
• The two changes are on separate branches; their combined effect is not measured yet.

Conflicts & rollback

base main (2818d6e), head perf/bf16-gate, no conflicts. Revert the single commit to roll back.

Related

Closes #4 · Roadmap: #3 (MoE)

简体中文 · 点击展开

动机

MoE gate 在 prefill 中每步约 76 ms，官方 DeepInfer 为 1.2 ms。原因是 gate 把 x 和权重转成 fp32，再做 fp32 SIMT 矩阵乘；仅这一个 kernel 每步就约 49 ms。

方案

Gate.forward 在 CUDA 上改用 torch.mm(x, weight.t(), out_dtype=torch.float32)：bf16 输入走 tensor core，输出 fp32。打分、top-k、归一化不变，仍为 fp32。数值上与 DeepInfer 的 gate 一致（bf16 输入、fp32 输出）。CPU 不支持 out_dtype，保留原 fp32 路径。

改动：model.py +5/−1，tests/test_moe.py 新增 1 个测试（3 种打分函数）。

验证

• pytest -q -m "not gpu"（CPU）：35 passed
• python -m compileall -q deepinfer_lab tests
• git diff --check
• GPU，1 × B200：
  • tests/test_moe.py：16 passed，含新测试。V4.1-Flash 的 gate 尺寸（5120 → 384，top-6），300 个 token：bf16 与 fp32 选出的专家完全相同，权重误差 < 1e-4
  • tests/test_cuda_graph.py、tests/test_sync_free.py：与 test_moe.py 合计 20 passed。decode CUDA graph 能捕获，forward 中没有 host 同步
  • tests/test_v41.py、tests/test_engine.py：17 passed。tiny 模型对照 DeepSeek 参考实现，logits 差距仍 < 0.05
• 真实 V4.1-Flash 4 × B200 prefill 测速：吞吐 +6.2%，每步 −55.5 ms（预测约 50 ms），见下方「实测」

实测

作业 9816：4 × B200，真实 V4.1-Flash，EP4 / DP4 / CP1，131,072-token 上下文（126,976 已缓存），每步每卡新增 4,096 token，每轮 10 个稳态步，未启用 CUDA graph。三个版本在同一作业内按 main → gate → 注意力 → 注意力 → gate → main 交替运行。

版本	第 1 轮	第 2 轮	平均吞吐（tok/s/GPU）	每步中位数	相对 main
main（2818d6e）	4,322.6	4,321.6	4,322.1	947.7 ms	—
bf16 gate（048e6e3）	4,591.6	4,587.1	4,589.4	892.2 ms	吞吐 +6.2%，每步 −55.5 ms
按索引读 key（ba0c0d3）	4,997.1	5,011.7	5,004.4	817.7 ms	吞吐 +15.8%，每步 −130.0 ms

• 同一版本两轮相差 < 0.3%；三个版本的峰值显存都是 143.5 GiB。
• 两项改动在不同分支上，合并后的效果尚未测量。

冲突与回滚

base main（2818d6e），head perf/bf16-gate，无冲突。若有问题，回退这一个提交即可。

关联

Closes #4 · 路线图：#3（MoE）