"""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())