diff --git a/H20-BRANCHES.md b/H20-BRANCHES.md index 9de9f99..7103241 100644 --- 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： diff --git a/H20-FLAGGEMS-PLAN.md b/H20-FLAGGEMS-PLAN.md new file mode 100644 index 0000000..b101799 --- /dev/null +++ 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` 进容器。 diff --git a/scripts/flaggems_port/ab_import_side_effects.py b/scripts/flaggems_port/ab_import_side_effects.py new file mode 100644 index 0000000..38951fa --- /dev/null +++ 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()) diff --git a/scripts/flaggems_port/ab_quant.py b/scripts/flaggems_port/ab_quant.py new file mode 100644 index 0000000..a7cb7c8 --- /dev/null +++ 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)")