docs(h20): plan the FlagGems integration, with the first operator already proven

Opens `h20-sm90-flaggems` to replace `tile_kernels`'s torch reference with FlagGems' Triton kernels. The point is not "runs without TileLang" -- that landed already. It is speed: the compat layer is plain torch today, so every activation in the hot path is going through eager PyTorch. Wrote the plan down because the tempting framing is wrong. `tile_kernels` and `flag_gems` are not substitutes; they overlap in three layers and only the first is a clean fit: * generic ops (quant, rotary, rms_norm, topk) -- FlagGems' home turf, 251 ops * DeepSeek-specific fused ops (mhc_*, engram_*, the five-way swiglu+quant) -- partial overlap, and almost none of the signatures match 1:1 * the engine kernel backend (tilelang_ops) -- already done, unrelated De-risked the highest-value one first, since per-token quantisation runs on every activation. `ab_quant.py`, 5 shapes: `per_token_group_quant_fp8(scale_ue8m0=True, column_major_scales=True)` against `per_token_cast(fmt="e4m3", 128, round_sf=True, use_packed_ue8m0=False, use_tma_aligned_col_major_sf=True)` -- **quantised values bitwise identical, scales and dequantised values exactly zero difference**. The two `round_sf` spellings are the same function: DeepInfer does `(int + 0x7FFFFF) & 0x7F800000`, which carries into the exponent field, i.e. ceiling to a power of two; FlagGems does `exp2(ceil(log2(s)))`. Same function, so the test demands exactness rather than tolerance and reports any difference. `ab_import_side_effects.py` pins the mechanism that makes this safe: import is inert, `only_enable(include=[...])` touches exactly what it names. Recorded also that the package root pulls Triton TLE -> libflagcx, so selective imports or vendoring stay the default posture rather than reaching for `enable()`. Two questions the plan puts ahead of the work, because without answers the payoff is unknowable: nothing on this line has ever been measured for speed, and `batched_transpose_weight_sf` is still the only path operator never A/B'd -- it rewrites every weight's scale factor and is the leading suspect for the token divergence. Fixing that could be worth more than any op swap here. Also names what is deliberately off the table: `enable()` with no arguments, swapping tilelang_ops again, and churning broadcast-only micro-ops that were never the bottleneck.
2026-10-30 15:57:20 +08:00 · 2026-10-25 14:38:08 +08:00 · 2026-10-25 14:38:08 +08:00 · 32d3d2c6bf
commit 32d3d2c6bf
parent 3486b93ac9
4 changed files with 335 additions and 5 deletions
--- a/H20-BRANCHES.md
+++ b/H20-BRANCHES.md
@ -1,6 +1,6 @@
 # 分支记录（SM90 / H20 这条线）

-截止 2026-10-25 14:30。所有分支都已推到 `tonyh168/DeepInfer`。
+截止 2026-10-25 14:40。所有分支都已推到 `tonyh168/DeepInfer`。

 ## 一句话地图

@ -9,8 +9,8 @@ main (8b16149)
 └── h20-sm90-support            5004c77   TileLang PR 版：SM90 attention 用 TileLang 实现
      └── h20-sm90-triton        73b8090   把引擎 kernel 换成 Triton 后端
           └── h20-sm90-triton-flagtree   c6a65fe   同上，在 flagtree 镜像里重编 vendor 库
-                └── h20-sm90-sparse-mla-triton   f25e5d8   ← 当前 HEAD
-                                                再换掉 attention：无 tilelang 端到端
+                └── h20-sm90-sparse-mla-triton   6d337a8   再换掉 attention：无 tilelang 端到端
+                     └── h20-sm90-flaggems       ← 当前 HEAD，全面接入 FlagGems（见 H20-FLAGGEMS-PLAN.md）
 ```

 另一条独立的线：`h20-sm90-lossless-requant`（74ba468），做无损 requant 与 debug_dump 挪位，不在这条链上。
@ -22,9 +22,10 @@ main (8b16149)
 | `h20-sm90-support` | `5004c77` | 上游 PR #180 的基础：FlashMLA 是 SM100-only，SM90 上按 `patch_imports` 钩子换成自己实现的 attention | attention 仍是 TileLang |
 | `h20-sm90-triton` | `73b8090` | 引擎侧 29 个 kernel 加 Triton 实现（17 个可用），`DEEPINFER_OPS_BACKEND=triton` 切换 | 融合算子库 `tile_kernels` 仍要 TileLang |
 | `h20-sm90-triton-flagtree` | `c6a65fe` | 在 flagtree 镜像（torch 2.11 + cu13）里重编 `deepinfer._C` / `deep_gemm` / `deep_ep` / `deep_select`；`tile_kernels` 加 compat 层走 torch 参考，从而不 import TileLang；修掉参考实现在 sync 守卫下抛错的问题 | attention 仍要 TileLang |
-| **`h20-sm90-sparse-mla-triton`** | **`f25e5d8`** | **attention 也换成 FlagGems 的 Triton kernel，于是「屏蔽 tilelang 也能端到端」成立** | 采样器的 rejection-sampling 变体仍要 TileLang（关 MTP 可绕）；不与 TileLang 版逐 token 一致 |
+| `h20-sm90-sparse-mla-triton` | `6d337a8` | **attention 也换成 FlagGems 的 Triton kernel，于是「屏蔽 tilelang 也能端到端」成立** | 采样器的 rejection-sampling 变体仍要 TileLang（关 MTP 可绕）；不与 TileLang 版逐 token 一致 |
+| **`h20-sm90-flaggems`** | **进行中** | **把 `tile_kernels` 的 torch 参考换成 FlagGems 的 Triton 实现**（先做已验证的 per-token quant） | 刚起步；`mhc_*` / `engram_*` 多数没有 1:1 对应 |

-## `f25e5d8`（当前版本）到底做了什么
+## `h20-sm90-sparse-mla-triton` 到底做了什么

 新增三个文件、改一处 import：

--- a/H20-FLAGGEMS-PLAN.md
+++ b/H20-FLAGGEMS-PLAN.md
@ -0,0 +1,127 @@
+# 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 无关。
+
+## 建议的推进顺序
+
+每一步都是「先 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 内联而不依赖包的原因之一）。
+
+## 环境事实（复跑用）
+
+```bash
+# 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` 进容器。
--- a/scripts/flaggems_port/ab_import_side_effects.py
+++ b/scripts/flaggems_port/ab_import_side_effects.py
@ -0,0 +1,86 @@
+"""What does importing `flag_gems` actually do to the process?
+
+This exists because the answer was assumed wrongly once. The earlier belief was that importing the
+package registers ATen implementations process-wide, which would make any dependency on it a
+numerical risk for the whole engine. That is not what happens: import creates the `aten_lib` handle
+and pulls the op modules in, but registration only occurs when `enable()` or `only_enable()` is
+called.
+
+The distinction decides the architecture, so it is measured rather than read:
+
+* `import flag_gems`               -> ATen dispatch tables must be unchanged
+* `only_enable(include=[...])`     -> exactly the named ops change, nothing else
+* `enable()`                       -> everything changes (not run here, and not wanted)
+
+Also records the import's real cost, which is the reason to stay selective: the package root
+star-imports every op module and loads Triton TLE, which dlopen's `libflagcx`. That needs
+`LD_LIBRARY_PATH` to find `libcudart`, and it is a native dependency the engine does not otherwise
+have.
+
+Run: PYTHONPATH=. python scripts/flaggems_port/ab_import_side_effects.py
+"""
+
+from __future__ import annotations
+
+import hashlib
+import sys
+
+import torch
+
+# A spread of ops: FlagGems provides all of these, so any registration would be visible here.
+OPS = ["aten::rms_norm", "aten::silu", "aten::softmax", "aten::topk", "aten::lerp"]
+
+
+def snapshot() -> dict[str, str]:
+    """Hash each op's dispatch table. Content, not length -- registration rewrites entries."""
+    out = {}
+    for op in OPS:
+        try:
+            table = torch._C._dispatch_dump_table(op)
+        except Exception as exc:  # noqa: BLE001 - op may not exist on this build
+            table = f"unavailable: {exc}"
+        out[op] = hashlib.md5(table.encode()).hexdigest()[:12]
+    return out
+
+
+def report(tag: str, before: dict[str, str], after: dict[str, str]) -> bool:
+    changed = [op for op in OPS if before[op] != after[op]]
+    print(f"\n[{tag}]")
+    for op in OPS:
+        mark = "changed" if before[op] != after[op] else "same"
+        print(f"  {op:16} {before[op]} -> {after[op]}  {mark}")
+    return bool(changed)
+
+
+def main() -> int:
+    failures = []
+
+    before = snapshot()
+    import flag_gems
+
+    print(f"flag_gems {flag_gems.__version__} on {flag_gems.vendor_name}/{flag_gems.device}")
+
+    # 1. Import alone must not touch anything.
+    if report("after import flag_gems (expect: nothing changed)", before, snapshot()):
+        failures.append("import registered ATen ops")
+
+    # 2. only_enable must touch exactly what it names.
+    before_narrow = snapshot()
+    flag_gems.only_enable(include=["rms_norm"])
+    after_narrow = snapshot()
+    changed = [op for op in OPS if before_narrow[op] != after_narrow[op]]
+    print("\n[only_enable(include=['rms_norm'])]")
+    print(f"  changed: {changed or 'none'}")
+    if changed != ["aten::rms_norm"]:
+        failures.append(f"only_enable changed {changed}, expected only aten::rms_norm")
+
+    print()
+    if failures:
+        print(f"UNEXPECTED: {failures}")
+        return 1
+    print("AS_EXPECTED: import is inert; only_enable is scoped to the ops it names")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
--- a/scripts/flaggems_port/ab_quant.py
+++ b/scripts/flaggems_port/ab_quant.py
@ -0,0 +1,116 @@
+"""Is FlagGems' per-token FP8 quant a drop-in for `tile_kernels`'s, on the SM90 recipe?
+
+This is the highest-value candidate in the FlagGems inventory: per-token quantisation runs on every
+activation, in attention and in the MLP, so it is the one where replacing a torch reference with a
+Triton kernel actually buys something.
+
+Recipe on SM90 (`get_fp8_cast_params`):
+
+    fmt=e4m3, num_per_channels=128, round_sf=True,
+    use_packed_ue8m0=False, use_tma_aligned_col_major_sf=True
+
+The two `round_sf` implementations look different but should be the same function. DeepInfer does
+`(int_bits + 0x7FFFFF) & 0x7F800000` -- add the maximum mantissa and mask to the exponent field,
+which carries into the next exponent unless the mantissa was already zero, i.e. ceiling to a power of
+two. FlagGems does `exp2(ceil(log2(s)))`, the same thing spelled arithmetically. Both then divide by
+that scale. So this should match bit for bit, not merely closely -- which is why the comparison
+insists on exactness and reports any difference rather than tolerating it.
+
+Run: PYTHONPATH=. python scripts/flaggems_port/ab_quant.py
+"""
+
+from __future__ import annotations
+
+import sys
+
+import torch
+
+from deepinfer.tk_compat.quant import per_token_cast
+
+DEV = "cuda"
+GROUP = 128
+FAILED: list[str] = []
+
+
+def flag_gems_quant(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Call FlagGems' op without importing the package root.
+
+    `import flag_gems` star-imports every op module and loads Triton TLE (which needs libflagcx).
+    Pulling the single module in through the package still runs the package `__init__`, so this goes
+    through the already-installed package but touches only what it must -- see the note in
+    `ab_import_side_effects.py` for why importing the root is safe anyway.
+    """
+    from flag_gems.ops.per_token_group_quant_fp8 import per_token_group_quant_fp8
+
+    return per_token_group_quant_fp8(x, GROUP, scale_ue8m0=True, column_major_scales=True)
+
+
+def reference_quant(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    return per_token_cast(
+        x,
+        "e4m3",
+        GROUP,
+        round_sf=True,
+        use_packed_ue8m0=False,
+        use_tma_aligned_col_major_sf=True,
+    )
+
+
+def compare(tag: str, a: torch.Tensor, b: torch.Tensor, exact: bool) -> None:
+    if a.shape != b.shape:
+        print(f"  DIFF {tag}: shape {tuple(a.shape)} vs {tuple(b.shape)}")
+        FAILED.append(f"{tag}:shape")
+        return
+    if a.dtype != b.dtype:
+        print(f"  DIFF {tag}: dtype {a.dtype} vs {b.dtype}")
+        FAILED.append(f"{tag}:dtype")
+        return
+    if exact:
+        same = torch.equal(a.view(torch.uint8) if a.dtype.itemsize == 1 else a, b.view(torch.uint8) if b.dtype.itemsize == 1 else b)
+        print(f"  {'ok  ' if same else 'DIFF'} {tag}: bitwise identical = {same}")
+        if not same:
+            FAILED.append(tag)
+        return
+    diff = (a.float() - b.float()).abs()
+    mx = diff.max().item()
+    scale = a.float().abs().max().item()
+    print(f"  {'ok  ' if mx == 0 else 'DIFF'} {tag}: max abs diff {mx:.3e} (max|a| {scale:.4f})")
+    if mx != 0:
+        FAILED.append(tag)
+
+
+def case(tag: str, tokens: int, dim: int, seed: int) -> None:
+    torch.manual_seed(seed)
+    x = (torch.randn((tokens, dim), device=DEV, dtype=torch.bfloat16) * 0.5).contiguous()
+    print(f"\n[{tag}] x {tuple(x.shape)} {x.dtype}")
+
+    q_ref, s_ref = reference_quant(x)
+    q_fg, s_fg = flag_gems_quant(x)
+    torch.cuda.synchronize()
+
+    compare("quantised values", q_ref, q_fg, exact=True)
+    compare("scales", s_ref, s_fg, exact=False)
+
+    # Dequantising both ways is the property that matters downstream: a matching scale with
+    # mismatched values would still be wrong after the GEMM.
+    deq_ref = q_ref.to(torch.float32) * s_ref.float().repeat_interleave(GROUP, dim=-1)
+    deq_fg = q_fg.to(torch.float32) * s_fg.float().repeat_interleave(GROUP, dim=-1)
+    compare("dequantised", deq_ref, deq_fg, exact=False)
+
+
+CASES = [
+    ("typical activation", 512, 4096, 0),
+    ("single token", 1, 4096, 1),
+    ("narrow", 64, 128, 2),
+    ("odd token count", 7, 2048, 3),
+    ("with small values", 128, 1024, 4),
+]
+
+if __name__ == "__main__":
+    for c in CASES:
+        case(*c)
+    print()
+    if FAILED:
+        print(f"MISMATCH {len(FAILED)}: {FAILED}")
+        sys.exit(1)
+    print(f"ALL_MATCH ({len(CASES)} cases)")