coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-31 00:36:20 +08:00
Code Issues Packages Projects Releases Wiki Activity
DeepInfer/H20-FLAGGEMS-PLAN.md
loopyt dcf00aca7c feat(h20): enable FlagGems end to end, from an allow-list of ops proven equal to PyTorch 
`DEEPINFER_ENABLE_FLAGGEMS=1` turns it on. Verified on all 8 DP ranks: exit=0, 8/8 ranks agreeing,
no kernel fallbacks, `hard failures: none`, correct answers. 19 ops actually dispatched, 152 call
sites, 43 registrations from 37 requested names.

The default is `flag_gems.only_enable(include=[...])`, not `enable()`, and that is the whole design:

* allow-list -- every intermediate state is safe (unproven ops stay on PyTorch, which is correct by
  construction); an incomplete list costs throughput, not correctness.
* deny-list -- I tried it first. Five full engine runs, one new defect each: `cat` (e8m0 dtype),
  `mul` (broadcast shape), `topk` (compile blowup), `rms_norm` (signature), `mm` (compile assert).
  The method requires being bitten before it can exclude, and every state along the way is unsafe.

That matches the instruction this was built to: use what works, leave the rest.

Where it is called is not incidental. `import flag_gems` initialises a CUDA context -- measured,
`torch.cuda.is_initialized()` flips on the import alone, before `enable()`. A `sitecustomize`
injector therefore put all 8 ranks on device 0 and NCCL refused with "Multiple Ranks are using the
same GPU". It now runs inside `create_llm`, after `setup_cuda_visible_devices`, which is the same
hazard already documented there for the profiler.

Defects found and documented rather than worked around, because each one is worth knowing:

* `mul` returns the *first operand's shape* when broadcasting expands it: (1,64,8)*(4,1,8) gives
  (1,64,8), and (64,1)*(1,128) gives (64,1). Silent, no error, fewer elements computed. It made
  `engram_hash_ref` disagree with PyTorch during a real run. add/sub/div/maximum/minimum all
  broadcast correctly, so it is mul's own bug.
* `cumsum` drifts by up to 1.0 over a 512-element bf16 row (tolerance ~0.17) -- bf16 accumulation.
* `topk`'s compile time explodes with k: 4.8 s at k=8, 19 s at k=64, >180 s at k=512, and the
  sampler asks for 4096. py-spy put all 8 ranks at 100% CPU inside `make_llir` with every GPU idle.
* `cat` raises `KeyError: 'float8_e8m0fnu'` on the model's UE8M0 scale tensors during weight load.
* `mm` passes every isolated probe yet fails inside the engine at `general_mm`; trigger unknown, so
  GEMMs stay on cuBLAS, which is faster anyway.

Two methodological errors in my own measurements, both fixed and both recorded:

1. After `enable()`, `torch.mul` *is* `flag_gems.mul`, so the first version of the mul A/B compared
   FlagGems with itself and reported a perfect score. All references are now computed before enable.
2. "ok" did not mean FlagGems ran. `torch.bitwise_xor` has no FlagGems registration at all, so that
   "match" was PyTorch against PyTorch. `screen_ops.py` now checks registration first and reports
   such cases as `skip` rather than pass.

Name mismatch is its own trap: FlagGems registers `add` bare but `eq.Scalar`, `where.self`,
`repeat_interleave.self_int`, `_softmax` (not `softmax`), `div.true_divide`. Asking for `softmax`
silently registers nothing. The bootstrap now expands names against `flag_gems._FULL_CONFIG` and
warns on anything unresolved -- which is how `softmax` was found to need `_softmax`.

Output is not token-identical to the no-FlagGems build (diverges at the first token), but two runs of
the same FlagGems config are bit-identical, so the difference is Triton-versus-eager accumulation
order, not nondeterminism. Answers stay correct and coherent.
2026-10-26 12:53:47 +08:00

11 KiB
Raw Blame History

FlagGems 全面接入：计划与已完成的 de-risk

分支 h20-sm90-flaggems，基于 h20-sm90-sparse-mla-triton（6d337a8）。 截止 2026-10-25 14:40。

起点：三个已经实测的事实

不是推测，脚本在 scripts/flaggems_port/，都可复跑。

1. import flag_gems 没有全局副作用 ✅

ab_import_side_effects.py：对 aten::rms_norm / silu / softmax / topk / lerp 的 torch._C._dispatch_dump_table 做内容哈希，import 前后全部相同。

这推翻了我上一轮的说法（当时写在 vendor/flashmla_sparse.py 头部和 README 里，说 「import 会进程级替换 ATen 算子」）。实际情况：__init__ 只创建 aten_lib 句柄并 import 各算子模块， 注册发生在 enable() / only_enable() 被显式调用时。已在两处文档里改正并注明是修正。

2. 可以按白名单注册 ✅

only_enable(include=["rms_norm"]) → 实测只有 aten::rms_norm 变化，其余不动。 这是「受控接入」的机制，也是整个计划能成立的前提。

3. 首要候选已数值验证：逐位相同 ✅

ab_quant.py：FlagGems per_token_group_quant_fp8(scale_ue8m0=True, column_major_scales=True) vs DeepInfer per_token_cast(fmt="e4m3", 128, round_sf=True, use_packed_ue8m0=False, use_tma_aligned_col_major_sf=True)，5 个形状：

形状 量化值 scale 反量化
512×4096 bitwise 相同 0 误差 0 误差
1×4096 bitwise 相同 0 0
64×128 bitwise 相同 0 0
7×2048 bitwise 相同 0 0
128×1024 bitwise 相同 0 0

两边 round_sf 写法不同但是同一个函数：DeepInfer 是 (int + 0x7FFFFF) & 0x7F800000（加满尾数进位 到指数域 = 向上取整到 2 的幂），FlagGems 是 exp2(ceil(log2(s)))。所以能逐位相同，不是「接近」。

算子是三层，不是一个清单

tile_kernels（DeepInfer 自己的融合算子库）和 flag_gems（通用算子库）不是替代关系。 混为一谈会导致白做很多工。实际是：

层 是什么 FlagGems 能做什么
A. 通用算子 rms_norm / softmax / silu / rotary / topk … 可替代，且这是 flag_gems 的主场（251 个 ops）
B. DeepSeek 专有融合算子 mhc_*（9 个）、engram_*、swiglu+per_token_cast 五合一、batched_transpose_weight_sf 部分可替代，但签名几乎都不 1:1
C. 引擎 kernel 后端 tilelang_ops/ 那 29 个 无关，那条线已经做完（Triton 17/29，默认路径无回退）

逐项判断（基于实际签名比对）

A 层：值得做，风险低

算子 flag_gems DeepInfer 对应 判断
per-token FP8 quant per_token_group_quant_fp8 tk_compat.quant.per_token_cast ✅ 已验证逐位相同，可直接接
rotary apply_rotary_pos_emb tk_compat.transform.apply_rotary 待 A/B。已知 DeepInfer 版有 1 bf16 ULP 差异的旧账
rms_norm rms_norm / fused_add_rms_norm 引擎内多处 待查调用点；若引擎已用 torch 原生则可直接 only_enable
grouped topk grouped_topk tk_compat.moe 内部 待 A/B。MoE 路由，影响面大
swiglu swiglu(input, quantizer=None) swiglu_forward_and_per_token_cast ⚠️ 部分：激活对得上，但 MoE 的 psum_num_tokens_per_expert / topk_weights / routed_scaling_factor 记账没有。不能整体替换

B 层：部分可做，但每个都要单独适配

算子 flag_gems 判断
mhc_pre（五路融合） mhc_pre(residual, fn, hc_scale, hc_base, eps×3, hc_post_mult_value, sinkhorn_repeat, n_splits) 对应 mhc_pre_big_fuse_fwd。返回值风格不同（DeepInfer 写 out 参数、返回 None；flag_gems 返回 tuple）
mhc_post mhc_post(x, residual, post_layer_mix, comb_res_mix) -> out 对应 mhc_post_fwd（唯一返回值的那个），最接近
hc_split_sinkhorn hc_split_sinkhorn + torch ref 对应 mhc_sinkhorn_fwd
expand / head_compute_mix / pre_apply_mix / pre_split_mixes / normw_merge flag_gems 没有 这几个是纯 broadcast/multiply，本来就不需要 kernel，保持现状
reduce_partials_and_rmsnorm flag_gems 没有 保持 torch 参考
batched_transpose_weight_sf flag_gems 没有 这是一个未结的疑点（可能解释 token 差异），别在这里动
engram_* flag_gems 没有 保持现状
dequant / cast_back flag_gems 没有 保持 torch 参考

C 层：不动

tilelang_ops/ 已经有 Triton 后端，DEEPINFER_OPS_BACKEND=triton 切换，与 flag_gems 无关。

端到端已跑通（2026-10-26）：白名单模式

DEEPINFER_ENABLE_FLAGGEMS=1 → deepinfer/flaggems_bootstrap.py 在 create_llm 里、 setup_cuda_visible_devices() 之后调 flag_gems.only_enable(include=[...])。

结果：

值
exit code 0
后端 sparse MLA = triton，FlagGems mode=allow
DP 一致 8/8
kernel 回退 无
hard failures none
回答 「中国的首都→北京」「日本的首都→东京」，带分段格式
实际派发算子 19 个、152 处调用点
注册 37 个请求名 → 96 个变体 → 43 个实际注册
警告 0

实际在跑的是 abs/amax/bincount/clamp(3变体)/cos/exp/gather/index_select/neg/nonzero/ reciprocal/relu/rsqrt/sigmoid/sin/sqrt/stack 这些——点状与 reduction 算子，正是 eager PyTorch 一次一个 launch 吃亏、Triton 能合并的地方。

为什么改成白名单，而不是「全开 + 排除坏的」

我先按「全开+排除」做了 5 轮，每轮都在完整引擎里撞出一个新问题，每轮约 5 分钟： cat(e8m0 dtype) → mul(广播形状) → topk(编译爆炸) → rms_norm(签名) → mm(编译断言)。 这条路的性质是「必须先被坑到才能修」，且中途任何状态都是不安全的。

白名单相反：任何中间状态都是安全的（没验证的算子留在 PyTorch），不完整只是性能问题， 不是正确性问题。这正好对应「能用就用、不能用就拉到」。

踩到的两个方法论错误（都已修）

  1. torch.mul 被替换后不能再当参考。第一版 ab_mul_dtype.py 在 enable() 之后同时用了 torch.mul 和 flag_gems.mul，于是「bf16 mul 0 差异」是自比自。现在所有参考值都在 enable() 之前算好并 clone。
  2. 「ok」可能根本没走 FlagGems。torch.bitwise_xor 在 FlagGems 里没有实现，所以 「bitwise_xor 通过」是 PyTorch 比 PyTorch。screen_ops.py 现在先判断该算子是否真的被注册， 否则标 skip 并说明「PyTorch answered」。square 同理（FlagGems 用别的名字）。

名字对不上是坑

FlagGems 的注册名是混合风格：add/mul 是裸名，但 eq.Scalar、lt.Tensor、where.self、 repeat_interleave.self_int、_softmax（不是 softmax）、div.true_divide、pow.pow_scalar 是 aten 风格的。只写 softmax 会静默什么都不注册。现在 _expand() 对着 flag_gems._FULL_CONFIG 自己做前缀展开，并在匹配不到时打印警告（这次就是靠它发现 softmax 应写 _softmax）。

仍未解决

  • mm 家族：隔离探测（M/N/K 扫描、NT 布局、batch、addmm）全部正确，但引擎里仍在 general_mm 撞 accumulator shape is incompatible。触发条件没定位，因此 GEMM 继续走 cuBLAS （本来也更快）。
  • token 级不一致，但是确定性的：与不用 FlagGems 的版本对比，第 1 个 token 就分叉（上一版 sparse-MLA-only 是前 15 个一致）。但同配置连跑两次逐 token 完全相同，所以差异来自 Triton kernel 与 eager PyTorch 的累加顺序，不是不确定性。回答仍然正确连贯，与 SM90 已知 2.7% logits 漂移同源，被贪心解码放大。
  • 性能完全没测。白名单里只有 19 个算子真被用到，收益未知。

建议的推进顺序

每一步都是「先 A/B 单测、通过才接线」，沿用 sparse MLA 那次的规矩。

  1. per-token quant 接入（唯一已验证的）
    • 在 tk_compat/quant.py 加一个 DEEPINFER_USE_FLAGGEMS 分支
    • 单测：直接复用 ab_quant.py 的 5 个形状 + 加 Engine 真实形状
    • 端到端：DEEPINFER_SIMULATE_NO_TILELANG=1 DEEPINFER_USE_FLAGGEMS=1
  2. rotary + rms_norm：先 A/B，再决定接哪些。rotary 有 1 ULP 旧账，必须先看清是 flag_gems 更准还是更差
  3. grouped_topk：A/B 再做。MoE 路由错了会直接影响输出分布，容错度最低
  4. mhc_post / hc_split_sinkhorn：签名最接近的两个，先做这两个而不是 mhc_pre
  5. mhc_pre：最后做，因为签名差异最大（out 参数 → 返回值）

必须先回答的两个问题

这两个不解决，后面的收益无法判断，甚至会白改：

  1. 性能到底有没有提升？ 现在全盘没有一次性能测量。tile_kernels 的 torch 参考是纯 torch 实现， 换成 Triton kernel 理论上会快，但没有数字就是猜。建议在第 1 步接入后立刻测一次 prefill/decode 的 attention 与 quant 耗时占比。
  2. batched_transpose_weight_sf 到底对不对？ 它是路径上唯一没 A/B 的算子，作用于所有权重的 scale factor，也是 token 差异的第一嫌疑。如果它本身有 bug，先修它比接 flag_gems 收益大得多。 补测要用 hugging_face/parameters.py 的真实 3-D (num_groups, -1, n) 布局。

反向清单：明确不做的

  • 不调用 enable()（无参版本）。实测它会注册全部算子，等于把引擎的数值行为整体换掉， 出了问题无法归因。只用 only_enable(include=[...])，或者干脆只 import 具体模块的函数。
  • 不替换 tilelang_ops。那条线已经完成。
  • 不为了「统一」而替换纯 broadcast 的小算子（expand / normw_merge 等）。它们本来就不是瓶颈， 换过去只增加一层出错可能。
  • 不引入 flagcx 依赖。import flag_gems 会拉 Triton TLE → dlopen libflagcx， 需要 LD_LIBRARY_PATH 里能找到 libcudart。当前部署本来就满足这一点，但不应该把这条 隐性依赖扩大化（这也是 sparse MLA 那个 kernel 内联而不依赖包的原因之一）。

环境事实（复跑用）

# flag_gems 5.0.2 已装进构建 venv
source /data/ftbuild/venv/bin/activate
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/root/.triton/nvidia/cudart/cuda_cudart-linux-x86_64-12.8.57-archive/lib
export PYTHONPATH=/data/ftbuild/repo CUDA_VISIBLE_DEVICES=3
python scripts/flaggems_port/ab_import_side_effects.py
python scripts/flaggems_port/ab_quant.py

注意：宿主机 /tmp 与容器内 /tmp 是两套，脚本要经 stdin 或 docker cp 进容器。