mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-31 00:36:20 +08:00

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

4.1 KiB

Raw Blame History

分支记录（SM90 / H20 这条线）

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

一句话地图

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   6d337a8   再换掉 attention：无 tilelang 端到端
                     └── h20-sm90-flaggems       ← 当前 HEAD，全面接入 FlagGems（见 H20-FLAGGEMS-PLAN.md）

另一条独立的线：h20-sm90-lossless-requant（74ba468），做无损 requant 与 debug_dump 挪位，不在这条链上。

逐个说明

分支	HEAD	这条分支解决了什么	关键限制
`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`	`6d337a8`	attention 也换成 FlagGems 的 Triton kernel，于是「屏蔽 tilelang 也能端到端」成立	采样器的 rejection-sampling 变体仍要 TileLang（关 MTP 可绕）；不与 TileLang 版逐 token 一致
`h20-sm90-flaggems`	进行中	FlagGems 白名单启用已端到端跑通（19 算子、8/8 DP、0 回退）；per-token quant 已验证逐位相同	只覆盖点状/reduction 算子；GEMM 留 cuBLAS；token 级与不用时不同但确定性

`h20-sm90-sparse-mla-triton` 到底做了什么

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

文件	作用
`deepinfer/h20_ops/vendor/flashmla_sparse.py`	内联 PyPI `flag_gems` 5.0.2 的 sparse MLA Triton kernel（329 行，Apache-2.0，文件头记了 md5 和两处本地改动）
`deepinfer/h20_ops/sparse_mla_triton.py`	适配层：`(out, max_logits, lse)` → `(out, lse)`；lse 自然对数 → log2；两种调用约定（共享 KV / 每 query 独立 KV）
`deepinfer/h20_ops/sparse_mla.py`	分发层：`DEEPINFER_SPARSE_MLA_BACKEND=tilelang\|triton`，不设则能 import tilelang 就用它
`deepinfer/h20_ops/attn_sm90.py`	只改一行 import，两个调用点未动

验证（每步先过再往下）：

步	内容	结果
1	内联 kernel vs TileLang，7 形状	`ALL_MATCH`，≤2 bf16 ULP
2	decode 形态（fp8 cache，含 SWA+full 双缓存），5 用例	`ALL_MATCH`
3	适配层契约，5 组	`ALL_PASSED`
4	分发层 + `attn_sm90` 两个入口，4 组	`ALL_PASSED`
5	屏蔽 tilelang 端到端	`exit=0`，0 回退，8/8 DP 一致
6	有 tilelang 对照	`exit=0`，0 回退

复现：

DEEPINFER_SIMULATE_NO_TILELANG=1 DEEPINFER_KERNELS=ref DEEPINFER_OPS_STRICT=1 \
DEEPINFER_USE_NEXTN_DECODE=0 bash scripts/h20_e2e_report.sh

怎么切分支

git fetch tonyh168
git checkout h20-sm90-sparse-mla-triton      # 最新最全
git checkout h20-sm90-triton-flagtree        # 只要 engine kernel 是 Triton、attention 留 TileLang
git checkout h20-sm90-support                # 只要 SM90 能跑（attention=TileLang，工程最省）

已知未做

性能一次没测（prefill/decode 吞吐、TTFT）。所有结论都是「能跑、数值对」，不是「跑得快」。
与 TileLang 版不逐 token 一致（prompt A 前 15 个 token 相同、第 16 分叉）。控制实验已证明非采样噪声。
fused_sample 的 rejection-sampling 变体仍依赖 TileLang，故纯 Triton 需关投机解码。

4.1 KiB Raw Blame History Unescape Escape