DeepInfer/deepinfer/flaggems_bootstrap.py

"""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=<dir>`  write `<dir>/oplist-<pid>.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