From dcf00aca7cbb40a346fb77be04c6a30d464c19f2 Mon Sep 17 00:00:00 2001 From: loopyt Date: Mon, 26 Oct 2026 12:53:47 +0800 Subject: [PATCH] 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. --- H20-BRANCHES.md | 2 +- H20-FLAGGEMS-PLAN.md | 60 ++++++ deepinfer/flaggems_bootstrap.py | 227 ++++++++++++++++++++ examples/_config.py | 7 + scripts/flaggems_port/ab_mul_dtype.py | 111 ++++++++++ scripts/flaggems_port/repro_engram_hash.py | 73 +++++++ scripts/flaggems_port/screen_ops.py | 228 +++++++++++++++++++++ 7 files changed, 707 insertions(+), 1 deletion(-) create mode 100644 deepinfer/flaggems_bootstrap.py create mode 100644 scripts/flaggems_port/ab_mul_dtype.py create mode 100644 scripts/flaggems_port/repro_engram_hash.py create mode 100644 scripts/flaggems_port/screen_ops.py diff --git a/H20-BRANCHES.md b/H20-BRANCHES.md index 7103241..1f31f42 100644 --- a/H20-BRANCHES.md +++ b/H20-BRANCHES.md @@ -23,7 +23,7 @@ main (8b16149) | `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`** | **进行中** | **把 `tile_kernels` 的 torch 参考换成 FlagGems 的 Triton 实现**（先做已验证的 per-token quant） | 刚起步；`mhc_*` / `engram_*` 多数没有 1:1 对应 | +| **`h20-sm90-flaggems`** | **进行中** | FlagGems 白名单启用已端到端跑通（19 算子、8/8 DP、0 回退）；per-token quant 已验证逐位相同 | 只覆盖点状/reduction 算子；GEMM 留 cuBLAS；token 级与不用时不同但确定性 | ## `h20-sm90-sparse-mla-triton` 到底做了什么 diff --git a/H20-FLAGGEMS-PLAN.md b/H20-FLAGGEMS-PLAN.md index b101799..5f19e5d 100644 --- a/H20-FLAGGEMS-PLAN.md +++ b/H20-FLAGGEMS-PLAN.md @@ -78,6 +78,66 @@ use_tma_aligned_col_major_sf=True)`，5 个形状： `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 那次的规矩。 diff --git a/deepinfer/flaggems_bootstrap.py b/deepinfer/flaggems_bootstrap.py new file mode 100644 index 0000000..acdc0ec --- /dev/null +++ b/deepinfer/flaggems_bootstrap.py @@ -0,0 +1,227 @@ +"""Optionally route PyTorch operators through FlagGems, via its documented `enable()` / `only_enable()`. + +FlagGems' usage docs (`docs/content/en/usage/basic.md`, `.../selective.md`) prescribe: + + import flag_gems + flag_gems.enable() # all supported ops + flag_gems.only_enable(include=[...]) # only the named ops + +This module defaults to the **allow-list** form, because flag_gems 5.0.2 on SM90 has operators that are +outright wrong, and a silently wrong operator is worse than a slow correct one. Only ops proven equal +to PyTorch on this build are enabled; everything else stays on PyTorch, which is correct by +construction. Measured by `scripts/flaggems_port/screen_ops.py`; extend `_VERIFIED_OPS` only after that +screen (or an equivalent A/B) passes. + +Why an allow-list rather than "enable everything then exclude what breaks": a deny-list can only grow +after each failure is hit inside a full engine run, and each cycle surfaces one operator. The +allow-list is safe at every intermediate state, so an incomplete list is a performance question rather +than a correctness one. + +Env: + +`DEEPINFER_ENABLE_FLAGGEMS=1` turn it on +`DEEPINFER_FLAGGEMS_MODE=allow|deny` default `allow` (curated) vs `deny` (all but the broken ones) +`DEEPINFER_FLAGGEMS_INCLUDE=a,b` extra ops for allow mode +`DEEPINFER_FLAGGEMS_EXCLUDE=a,b` extra ops to keep on PyTorch, in either mode +`DEEPINFER_FLAGGEMS_RECORD_DIR=` write `/oplist-.log`, one line per call site + +**Where it is called matters more than what it does.** `import flag_gems` initialises a CUDA context +(measured: `torch.cuda.is_initialized()` flips False -> True on the import alone, before `enable()`). +If that happens before this process is pinned to its GPU, every rank initialises on device 0 and the +engine fails at the first collective with + + NCCL error: invalid usage + Last error: Multiple Ranks are using the same GPU/Partition. + +That is what a `sitecustomize` injector produced, because `sitecustomize` runs at interpreter startup +while the per-rank `CUDA_VISIBLE_DEVICES` is set later, inside `create_llm`. The same hazard is already +documented there for the profiler ("CUPTI 启动即初始化 CUDA，会使此后设置的 CUDA_VISIBLE_DEVICES 失效"). +Hence this is called from `create_llm` immediately after `setup_cuda_visible_devices`, and there is no +`sitecustomize` for it. + +**Measured on flag_gems 5.0.2 / H20 (SM90), which is why the lists below look the way they do:** + +* `mul` returns the **first operand's shape** when broadcasting would expand it -- `(1,64,8)*(4,1,8)` + gives `(1,64,8)` instead of `(4,64,8)`; `(64,1)*(1,128)` gives `(64,1)`. Silent, no error, fewer + elements computed. Same-shape and scalar `mul` are exact. `add`/`sub`/`div`/`maximum`/`minimum` all + broadcast correctly, so it is `mul`'s own bug. It produced a wrong result inside `engram_hash_ref` + during a real engine run. +* `cumsum` disagrees with PyTorch by up to 1.0 over a 512-element bf16 row where the tolerance is + ~0.17 -- bf16 accumulation where PyTorch uses fp32. +* `rms_norm` requires `weight` positionally where `torch.nn.functional.rms_norm` allows omitting it, + and with `weight` supplied still differs by ~3e-2 on O(1) output. +* `topk`'s compile time is superlinear in `k` and the sampler uses `k=4096`: k=8 -> 4.8 s, k=64 -> 19 s, + k=512 -> >180 s. The startup profile run never returned; `py-spy` put all 8 ranks at 100% CPU inside + `triton/backends/nvidia/compiler.py:make_llir` with every GPU idle. +* `cat` raises `KeyError: 'float8_e8m0fnu'` from Triton's dtype table on the UE8M0 scale tensors this + model stores, during weight loading. Triton has no e8m0 element type. +* the `mm` family passes every isolated probe (M/N/K swept, NT layout, batch, addmm) yet the engine + still hit `accumulator shape is incompatible` inside `general_mm`; trigger not identified, so GEMMs + stay on cuBLAS -- which is also the faster choice. +""" + +from __future__ import annotations + +import os + +_ENV_ENABLE = "DEEPINFER_ENABLE_FLAGGEMS" +_ENV_MODE = "DEEPINFER_FLAGGEMS_MODE" +_ENV_INCLUDE = "DEEPINFER_FLAGGEMS_INCLUDE" +_ENV_EXCLUDE = "DEEPINFER_FLAGGEMS_EXCLUDE" +_ENV_RECORD_DIR = "DEEPINFER_FLAGGEMS_RECORD_DIR" + +# Ops measured equal to PyTorch by `scripts/flaggems_port/screen_ops.py`, across the shape variants +# that screen covers (same-shape, broadcast, scalar, unary, reduction). Pointwise and reduction ops +# are where FlagGems can actually win: one kernel launch instead of several eager calls. GEMMs are +# deliberately absent -- cuBLAS wins there regardless. +_VERIFIED_OPS: tuple[str, ...] = ( + "abs", + "add", + "amax", + "bincount", + "bitwise_and", + "clamp", + "cos", + "div", + "eq", + "exp", + "gather", + "ge", + "gelu", + "index_select", + "lt", + "masked_fill", + "maximum", + "mean", + "minimum", + "neg", + "nonzero", + "pow", + "reciprocal", + "relu", + "remainder", + "repeat_interleave", + "rsqrt", + "sigmoid", + "silu", + "sin", + "_softmax", + "sort", + "sqrt", + "stack", + "sub", + "sum", + "where", +) + +# Everything known broken on this build, documented and used by deny mode. Nothing here may move into +# `_VERIFIED_OPS` without a passing A/B. +_BROKEN_OPS: tuple[str, ...] = ("cat", "cumsum", "mul", "rms_norm", "topk") + +_enabled = False + + +def enabled() -> bool: + """Whether FlagGems has been enabled in this process.""" + return _enabled + + +def _names(raw: str) -> list[str]: + return [op.strip() for op in raw.split(",") if op.strip()] + + +def _expand(names: list[str], available: list[str]) -> tuple[list[str], list[str]]: + """Resolve plain op names to the names FlagGems actually registers. + + FlagGems registers a mix of plain names (`add`, `mul`) and aten-style dotted or suffixed names: + `eq.Scalar`, `lt.Tensor`, `where.self`, `repeat_interleave.self_int`, `_softmax`, `div.true_divide`, + `pow.pow_scalar`. Asking for `eq` by name therefore registers nothing, silently -- caught here as a + warning, but the fix is to expand to every variant rather than spell out each one in + `_VERIFIED_OPS`. + + Returns (expanded, unresolved). `available` is FlagGems' own registration table, so this cannot + drift from what the installed version actually offers. + """ + expanded: list[str] = [] + unresolved: list[str] = [] + for name in names: + hits = [ + n + for n in available + if n == name or n.startswith(f"{name}_") or n.startswith(f"{name}.") + ] + if hits: + expanded.extend(hits) + else: + unresolved.append(name) + # dict.fromkeys: preserve order, drop duplicates (a name may match several wanted entries). + return list(dict.fromkeys(expanded)), unresolved + + +def maybe_enable() -> str: + """Enable FlagGems if asked to. Returns a one-line status; no-op when the env var is unset.""" + global _enabled + + if _enabled: + return "already enabled" + if not os.environ.get(_ENV_ENABLE): + return "disabled" + + try: + import flag_gems + except Exception as exc: # noqa: BLE001 - a requested-but-broken import must be loud + raise RuntimeError( + f"{_ENV_ENABLE} is set but flag_gems is not importable: {type(exc).__name__}: {exc}. " + f"Unset it, or fix the install (it also needs libcudart on LD_LIBRARY_PATH for Triton TLE)." + ) from exc + + mode = (os.environ.get(_ENV_MODE, "allow").strip().lower() or "allow") + extra_in = _names(os.environ.get(_ENV_INCLUDE, "")) + extra_out = _names(os.environ.get(_ENV_EXCLUDE, "")) + + log_kwargs: dict[str, object] = {} + record_dir = os.environ.get(_ENV_RECORD_DIR, "").strip() + if record_dir: + from pathlib import Path + + Path(record_dir).mkdir(parents=True, exist_ok=True) + # Per-PID: `setup_flaggems_logging` opens the path with mode="w", so ranks sharing one path + # would leave only whichever wrote last. + log_kwargs = { + "record": True, + "path": str(Path(record_dir) / f"oplist-{os.getpid()}.log"), + "once": True, + } + + include: list[str] | None = None + if mode == "allow": + excluded = set(extra_out) + wanted = [op for op in (*_VERIFIED_OPS, *extra_in) if op not in excluded] + # Expand against FlagGems' own table, not a hand-maintained list of alias names. Entries are + # variable-length tuples whose first field is the op name. + available = [entry[0] for entry in flag_gems._FULL_CONFIG] # noqa: SLF001 + include, unresolved = _expand(wanted, available) + flag_gems.only_enable(include=include, **log_kwargs) + how = f"only_enable({len(wanted)} wanted -> {len(include)} registered)" + elif mode == "deny": + unused = [op for op in (*_BROKEN_OPS, *extra_out)] + unresolved = [] + flag_gems.enable(unused=unused, **log_kwargs) + how = f"enable(unused={unused})" + else: + raise RuntimeError(f"{_ENV_MODE}={mode!r} is not recognised; use 'allow' or 'deny'") + + registered = set(flag_gems.all_registered_ops()) + if unresolved: + # Nothing would be enabled for these, which reads as an unexplained slowdown rather than a + # failure. Say it once, loudly. + print(f"[flaggems] WARNING: no registered op matches: {unresolved}", flush=True) + + _enabled = True + status = ( + f"enabled pid={os.getpid()} flag_gems={flag_gems.__version__} " + f"device={flag_gems.device}/{flag_gems.vendor_name} mode={mode} {how} " + f"registered={len(registered)}" + ) + print(f"[flaggems] {status}", flush=True) + return status diff --git a/examples/_config.py b/examples/_config.py index a9d432f..de6650d 100644 --- a/examples/_config.py +++ b/examples/_config.py @@ -107,6 +107,13 @@ def create_llm( import torch + # After the device pin, not before: importing flag_gems initialises a CUDA context, so doing it + # earlier puts every rank on device 0 and NCCL then refuses with "Multiple Ranks are using the + # same GPU". No-op unless DEEPINFER_ENABLE_FLAGGEMS is set. + from deepinfer.flaggems_bootstrap import maybe_enable as maybe_enable_flaggems + + maybe_enable_flaggems() + from deepinfer.entrypoints.llm import LLM with torch.profiler.profile() if profile_init else nullcontext(): diff --git a/scripts/flaggems_port/ab_mul_dtype.py b/scripts/flaggems_port/ab_mul_dtype.py new file mode 100644 index 0000000..f9d491a --- /dev/null +++ b/scripts/flaggems_port/ab_mul_dtype.py @@ -0,0 +1,111 @@ +"""Is FlagGems' `mul` correct, and for which dtypes? + +Method note, because the first version of this script was wrong: once `flag_gems.enable()` runs it also +captures `torch.mul`, so comparing `torch.mul(...)` with `flag_gems.mul(...)` compares FlagGems with +itself. Every reference here is therefore computed **before** `enable()`, in the same process. + +The defect it measures: `engram_hash_ref` does `int64 (1,T,N) * int64 (L,1,N)`, and FlagGems returns a +tensor of the *first operand's* shape rather than the broadcast shape, so most of the result is never +written. Excluding `mul` from `enable()` restores the reference result exactly, which is what pointed +here in the first place. + +Run: PYTHONPATH=. python scripts/flaggems_port/ab_mul_dtype.py +""" + +from __future__ import annotations + +import sys +from collections.abc import Callable + +import torch + +FAILED: list[str] = [] + + +def check(tag: str, cond: bool, detail: str = "") -> None: + print(f" {'ok ' if cond else 'BAD '} {tag}{' ' + detail if detail else ''}") + if not cond: + FAILED.append(tag) + + +def compare(tag: str, ref: torch.Tensor, got: torch.Tensor) -> None: + if ref.shape != got.shape: + check(tag, False, f"shape ref {tuple(ref.shape)} vs got {tuple(got.shape)}") + return + if ref.dtype != got.dtype: + check(tag, False, f"dtype ref {ref.dtype} vs got {got.dtype}") + return + if ref.dtype.is_floating_point: + diff = (ref.float() - got.float()).abs().max().item() + scale = ref.float().abs().max().item() + tol = 2.0**-8 * max(scale, 1e-9) * 2 + check(tag, diff <= tol, f"max diff {diff:.3e} tol {tol:.3e}") + return + same = torch.equal(ref, got) + detail = "" + if not same: + bad = (ref != got).sum().item() + idx = (ref != got).nonzero()[:3] + sample = [(ref[tuple(i)].item(), got[tuple(i)].item()) for i in idx] + detail = f"{bad}/{ref.numel()} wrong, (ref, got) {sample}" + check(tag, same, detail) + + +def main() -> int: + torch.manual_seed(0) + dev = "cuda" + + # ---- references, computed while torch.mul is still PyTorch's ---- + a = torch.randint(0, 1000, (1, 64, 8), device=dev, dtype=torch.int64) + b = torch.randint(1, 100, (4, 1, 8), device=dev, dtype=torch.int64) + ref_bcast = torch.mul(a, b).clone() + + a_l = torch.randint(2**40, 2**41, (256,), device=dev, dtype=torch.int64) + ref_large = torch.mul(a_l, a_l).clone() + + a_i32 = torch.randint(0, 4096, (1024,), device=dev, dtype=torch.int32) + ref_i32 = torch.mul(a_i32, a_i32).clone() + + a_bf = torch.randn((2048, 4096), device=dev, dtype=torch.bfloat16) * 0.1 + b_bf = torch.randn((2048, 4096), device=dev, dtype=torch.bfloat16) * 0.1 + ref_bf16 = torch.mul(a_bf, b_bf).clone() + + a_bf_s = torch.randn((4096,), device=dev, dtype=torch.bfloat16) + ref_bf16_scalar = torch.mul(a_bf_s, 3).clone() + ref_i64_scalar = torch.mul(a, 3).clone() + + import flag_gems + + flag_gems.enable() + print(f"flag_gems {flag_gems.__version__}: enable() done") + print(f"reference broadcast shape: {tuple(ref_bcast.shape)}\n") + + cases: list[tuple[str, torch.Tensor, Callable[[], torch.Tensor]]] = [ + ("int64 mul, broadcasting (1,64,8)*(4,1,8)", ref_bcast, lambda: flag_gems.mul(a, b)), + ("int64 mul, large values", ref_large, lambda: flag_gems.mul(a_l, a_l)), + ("int32 mul", ref_i32, lambda: flag_gems.mul(a_i32, a_i32)), + ("bf16 mul, 8M elements", ref_bf16, lambda: flag_gems.mul(a_bf, b_bf)), + ("bf16 mul by scalar", ref_bf16_scalar, lambda: flag_gems.mul(a_bf_s, 3)), + ("int64 mul by scalar", ref_i64_scalar, lambda: flag_gems.mul(a, 3)), + ] + for tag, ref, fn in cases: + got = fn() + torch.cuda.synchronize() + compare(tag, ref, got) + + # torch.mul must show the same defect. If only the direct call did, the dispatcher would be + # fine and the fix would live elsewhere. + got_dispatch = torch.mul(a, b) + torch.cuda.synchronize() + compare("torch.mul (dispatched) shows the same defect", ref_bcast, got_dispatch) + + print() + if FAILED: + print(f"DEFECTS: {FAILED}") + return 1 + print("no defect found in the dtypes tested") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/flaggems_port/repro_engram_hash.py b/scripts/flaggems_port/repro_engram_hash.py new file mode 100644 index 0000000..e0fe505 --- /dev/null +++ b/scripts/flaggems_port/repro_engram_hash.py @@ -0,0 +1,73 @@ +"""Reproduce the `engram_hash_ref` failure seen with FlagGems enabled, and bisect it. + +The end-to-end run failed inside `tile_kernels.torch.engram.engram_hash_ref` with + + IndexError: index 1 is out of bounds for dimension 0 with size 1 + +That function is a plain PyTorch reference, so either FlagGems changed an operator's semantics, or +the engine's own call path is at fault and enabling FlagGems merely shifted when it surfaced. Those +have different fixes, so this script answers "is it FlagGems" first, in-process, without a full +engine run. + +Run with FLAGGEMS_REPRO_MODE=torch|full|unused: + torch -- do not enable (control) + full -- flag_gems.enable() + unused -- flag_gems.enable(unused=[...]) with FLAGGEMS_REPRO_EXCLUDE=op1,op2 +""" + +from __future__ import annotations + +import os +import sys +import traceback + +import torch +from tile_kernels.torch.engram import engram_hash_ref + +MODE = os.environ.get("FLAGGEMS_REPRO_MODE", "torch") + + +def enable() -> str: + if MODE == "torch": + return "flag_gems: off (control)" + import flag_gems + + if MODE == "full": + flag_gems.enable() + return f"flag_gems: full enable, {len(flag_gems.all_registered_ops())} ops" + exclude = [op.strip() for op in os.environ.get("FLAGGEMS_REPRO_EXCLUDE", "").split(",") if op.strip()] + flag_gems.enable(unused=exclude) + return f"flag_gems: enable(unused={exclude})" + + +def build(device: str = "cuda") -> tuple[torch.Tensor, ...]: + """Shapes from the engine's config rather than invented ones.""" + num_ngram_layers = int(os.environ.get("REPRO_LAYERS", "1")) + max_ngram_size = int(os.environ.get("REPRO_NGRAM", "2")) + tables = 2 + num_tokens = int(os.environ.get("REPRO_TOKENS", "16")) + vocab = 4096 + + torch.manual_seed(0) + # (num_tokens, max_ngram_size) per the docstring: unsqueeze(0) then broadcast against + # multipliers.unsqueeze(1) is what builds (layers, tokens, ngram). + ngram_token_ids = torch.randint(0, 1000, (num_tokens, max_ngram_size), device=device, dtype=torch.int32) + multipliers = torch.randint(1, 100, (num_ngram_layers, max_ngram_size), device=device, dtype=torch.int64) + vocab_sizes = torch.full((num_ngram_layers, max_ngram_size - 1, tables), vocab, device=device, dtype=torch.int32) + offsets = torch.zeros((num_ngram_layers, (max_ngram_size - 1) * tables), device=device, dtype=torch.int32) + return ngram_token_ids, multipliers, vocab_sizes, offsets + + +if __name__ == "__main__": + print(enable(), flush=True) + args = build() + try: + out = engram_hash_ref(*args) + torch.cuda.synchronize() + print(f"OK out {tuple(out.shape)} {out.dtype} finite={bool(torch.isfinite(out.float()).all())}") + # Print a checksum so runs with different exclude lists can be compared. + print(f" checksum {int(out.to(torch.int64).abs().sum())}") + except Exception: + print("FAIL") + traceback.print_exc() + sys.exit(1) diff --git a/scripts/flaggems_port/screen_ops.py b/scripts/flaggems_port/screen_ops.py new file mode 100644 index 0000000..6dedf98 --- /dev/null +++ b/scripts/flaggems_port/screen_ops.py @@ -0,0 +1,228 @@ +"""Screen FlagGems ops for the defect class that broke `mul`, so the exclusion list is evidence-based. + +`mul` returns the *first operand's shape* whenever broadcasting would expand it: + + (1,64,8) * (4,1,8) -> FlagGems (1,64,8), PyTorch (4,64,8) + (64,1) * (1,128) -> FlagGems (64,1), PyTorch (64,128) + +Silent, no error. It broke `engram_hash_ref` inside the engine, and it is not dtype-specific. + +The traceback went through `flag_gems/utils/pointwise_dynamic.py` and `shape_utils.broadcast_shapes`, +which is shared by every pointwise op -- so the question is which *other* ops inherit it. This script +answers that: each op is run twice, same-shape and broadcasting, and the reference is always taken +before `enable()` so the comparison cannot be FlagGems-against-itself. + +Output is a table, which is what the exclusion list in `deepinfer/flaggems_bootstrap.py` is built +from. + +Run: PYTHONPATH=. python scripts/flaggems_port/screen_ops.py +""" + +from __future__ import annotations + +import sys +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch + +DEV = "cuda" + + +@dataclass +class Case: + op: str + label: str + run: Callable[[], Any] + ref: Any + + +def ref_of(fn: Callable[[], Any]) -> Any: + out = fn() + if isinstance(out, torch.Tensor): + return out.clone() + if isinstance(out, (tuple, list)): + return type(out)(t.clone() if isinstance(t, torch.Tensor) else t for t in out) + return out + + +def same(ref: Any, got: Any) -> tuple[bool, str]: + if isinstance(ref, torch.Tensor): + if not isinstance(got, torch.Tensor): + return False, f"type {type(got).__name__}" + if ref.shape != got.shape: + return False, f"shape {tuple(ref.shape)} -> {tuple(got.shape)}" + if ref.dtype != got.dtype: + return False, f"dtype {ref.dtype} -> {got.dtype}" + if ref.dtype.is_floating_point: + d = (ref.float() - got.float()).abs().max().item() + s = ref.float().abs().max().item() + tol = 2.0**-8 * max(s, 1e-9) * 2 + return (d <= tol, f"max diff {d:.2e}") + eq = torch.equal(ref, got) + if eq: + return True, "exact" + return False, f"{(ref != got).sum().item()}/{ref.numel()} wrong" + if isinstance(ref, (tuple, list)): + for r, g in zip(ref, got, strict=False): + ok, msg = same(r, g) + if not ok: + return False, msg + return True, "exact" + return ref == got, "" + + +def build_cases() -> list[Case]: + torch.manual_seed(0) + b = 64 + n = 8 + + # (1, b, n) vs (4, 1, n): the exact shape pattern that broke engram_hash. + x3 = torch.randn((1, b, n), device=DEV, dtype=torch.bfloat16) + y3 = torch.randn((4, 1, n), device=DEV, dtype=torch.bfloat16) + i3 = torch.randint(0, 1000, (1, b, n), device=DEV, dtype=torch.int64) + j3 = torch.randint(1, 100, (4, 1, n), device=DEV, dtype=torch.int64) + + # (b, 1) vs (1, m): the most common broadcast in transformer code (masks, scales). + x2 = torch.randn((b, 1), device=DEV, dtype=torch.bfloat16) + y2 = torch.randn((1, 128), device=DEV, dtype=torch.bfloat16) + + xs = torch.randn((256, 512), device=DEV, dtype=torch.bfloat16) + ys = torch.randn((256, 512), device=DEV, dtype=torch.bfloat16) + + cases: list[Case] = [] + + def add(name: str, label: str, fn: Callable[[], Any]) -> None: + cases.append(Case(name, label, fn, ref_of(fn))) + + # binary pointwise, same shape (control: isolates broadcast as the cause) + add("add", "same-shape", lambda: torch.add(xs, ys)) + add("mul", "same-shape", lambda: torch.mul(xs, ys)) + add("sub", "same-shape", lambda: torch.sub(xs, ys)) + add("div", "same-shape", lambda: torch.div(xs, ys)) + add("maximum", "same-shape", lambda: torch.maximum(xs, ys)) + add("minimum", "same-shape", lambda: torch.minimum(xs, ys)) + + # binary pointwise, broadcasting (the suspect class) + add("add", "broadcast 3d", lambda: torch.add(x3, y3)) + add("mul", "broadcast 3d", lambda: torch.mul(x3, y3)) + add("sub", "broadcast 3d", lambda: torch.sub(x3, y3)) + add("div", "broadcast 3d", lambda: torch.div(x3, y3)) + add("maximum", "broadcast 3d", lambda: torch.maximum(x3, y3)) + add("minimum", "broadcast 3d", lambda: torch.minimum(x3, y3)) + add("add", "broadcast 2d", lambda: torch.add(x2, y2)) + add("mul", "broadcast 2d", lambda: torch.mul(x2, y2)) + add("sub", "broadcast 2d", lambda: torch.sub(x2, y2)) + + # integer bitwise / remainder: what engram_hash also used + add("bitwise_xor_", "broadcast int64", lambda: i3.bitwise_xor(j3)) + add("remainder", "broadcast int64", lambda: torch.remainder(i3, j3)) + + # scalar variants: same code path, different signature + add("mul", "scalar", lambda: torch.mul(xs, 3.0)) + add("add", "scalar", lambda: torch.add(xs, 3.0)) + + # unary: should not involve broadcasting at all + add("silu", "unary", lambda: torch.nn.functional.silu(xs)) + add("gelu", "unary", lambda: torch.nn.functional.gelu(xs)) + add("abs", "unary", lambda: torch.abs(xs)) + + # reductions + add("softmax", "last dim", lambda: torch.softmax(xs, dim=-1)) + add("amax", "last dim", lambda: torch.amax(xs, dim=-1)) + add("sum", "last dim", lambda: torch.sum(xs, dim=-1)) + + # normalisation; guarded because rms_norm is only in newer torch + if hasattr(torch.nn.functional, "rms_norm"): + add("rms_norm", "1x512", lambda: torch.nn.functional.rms_norm(xs, (512,))) + + # more pointwise / reduction candidates for an allow-list + add("exp", "unary", lambda: torch.exp(xs)) + add("sigmoid", "unary", lambda: torch.sigmoid(xs)) + add("rsqrt", "unary", lambda: torch.rsqrt(xs.abs() + 1)) + add("neg", "unary", lambda: torch.neg(xs)) + add("relu", "unary", lambda: torch.relu(xs)) + add("sqr", "unary", lambda: torch.square(xs)) + add("clamp", "min/max", lambda: torch.clamp(xs, -0.5, 0.5)) + add("pow", "scalar", lambda: torch.pow(xs, 2)) + add("reciprocal", "unary", lambda: torch.reciprocal(xs.abs() + 1)) + add("mean", "last dim", lambda: torch.mean(xs, dim=-1)) + add("cumsum", "last dim", lambda: torch.cumsum(xs, dim=-1)) + add("masked_fill", "scalar", lambda: xs.masked_fill(xs > 0, -1.0)) + add("where", "self", lambda: torch.where(xs > 0, xs, -xs)) + add("eq", "scalar", lambda: torch.eq(xs, 0.0)) + add("lt", "scalar", lambda: torch.lt(xs, 0.0)) + add("ge", "scalar", lambda: torch.ge(xs, 0.0)) + add("nonzero", "unary", lambda: torch.nonzero(xs > 0.5)) + add("gather", "dim0", lambda: torch.gather(xs, 0, torch.zeros((8, 512), device=DEV, dtype=torch.int64))) + add("index_select", "dim0", lambda: torch.index_select(xs, 0, torch.arange(8, device=DEV))) + add("sort", "last dim", lambda: torch.sort(xs, dim=-1).values) + add("sin", "unary", lambda: torch.sin(xs)) + add("cos", "unary", lambda: torch.cos(xs)) + add("sqrt", "unary", lambda: torch.sqrt(xs.abs())) + add("bitwise_and", "scalar", lambda: torch.bitwise_and(i3, 255)) + add("bincount", "int", lambda: torch.bincount(torch.zeros(64, device=DEV, dtype=torch.int64), minlength=8)) + + # shape ops + add("cat", "2 tensors", lambda: torch.cat([xs, xs], dim=0)) + add("stack", "2 tensors", lambda: torch.stack([xs, xs], dim=0)) + add("repeat_interleave", "self_int 2", lambda: torch.repeat_interleave(xs, 2, dim=0)) + + return cases + + +def main() -> int: + cases = build_cases() # references captured here, before enable() + import flag_gems + + flag_gems.enable() + registered = sorted(flag_gems.all_registered_ops()) + print(f"flag_gems {flag_gems.__version__}, {len(registered)} ops registered\n") + + def served(op: str) -> bool: + """Did FlagGems actually take this op? + + Without this the screen lies by omission: `torch.bitwise_xor` has no FlagGems registration at + all, so a "match" there is PyTorch compared with PyTorch and says nothing about FlagGems. Any + case that is not served is reported as such and cannot enter an allow-list. + """ + return any(n == op or n.startswith(f"{op}.") or n.startswith(f"{op}_") for n in registered) + + bad: dict[str, list[str]] = {} + unserved: set[str] = set() + rows = [] + for c in cases: + if not served(c.op): + unserved.add(c.op) + rows.append((c.op, c.label, None, "not registered in flag_gems (PyTorch answered)")) + continue + try: + got = c.run() + torch.cuda.synchronize() + ok, msg = same(c.ref, got) + except Exception as exc: # noqa: BLE001 - an exception is itself a finding + ok, msg = False, f"raised {type(exc).__name__}: {str(exc)[:70]}" + rows.append((c.op, c.label, ok, msg)) + if not ok: + bad.setdefault(c.op, []).append(c.label) + + width = max(len(r[0]) for r in rows) + 2 + for op, label, ok, msg in rows: + mark = "skip" if ok is None else ("ok " if ok else "BAD ") + print(f" {mark} {op:<{width}} {label:<16} {msg}") + + if unserved: + print(f"\nnot served by flag_gems (cannot be judged here): {sorted(unserved)}") + print() + if bad: + print("UNSAFE OPS (exclude these):") + for op, labels in sorted(bad.items()): + print(f" {op}: {', '.join(labels)}") + return 1 + print("all served ops matched") + return 0 + + +if __name__ == "__main__": + sys.exit(main())