DeepInfer/examples/neuron/device_diff.py

"""Device-vs-host differential check for the Neuron kernels.

Runs the tiny checkpoint on the NeuronCore through run_deepinfer_tiny.py and, for every call of a vendor
kernel (deep_gemm, tile_kernels, flash_mla, deep_select) or a tilelang_ops/neuron kernel, re-runs the same
pure-torch implementation on host copies of the arguments and compares the results, including arguments
mutated in place. Prints one line per kernel with the worst relative error and the argument layouts of
the worst call, so a kernel that misbehaves on the device (and only there) stands out.

    NEURON_RT_VISIBLE_CORES=63 CUDA_VISIBLE_DEVICES=63 python examples/neuron/device_diff.py \\
        --prompts short --out ~/di/neuron_diff [run_deepinfer_tiny.py args...]
"""

from __future__ import annotations

import collections
import functools
import sys
import types
from pathlib import Path
from typing import Any

import torch

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))

_STATS: dict[str, dict[str, Any]] = collections.OrderedDict()
_ACTIVE = [True]


def _to_host(x: Any) -> Any:
    if isinstance(x, torch.Tensor):
        # Keep strides: a few kernels depend on the caller's (e.g. column-major) layout.
        h = torch.empty_strided(x.shape, x.stride(), dtype=x.dtype, device="cpu").copy_(x.cpu()) if x.device.type != "cpu" else x.clone()
        # Metadata tensors carry Python state (e.g. deep_gemm's paged sparse MQA FIRST/REUSE state and KV snapshot).
        for k, v in getattr(x, "__dict__", {}).items():
            setattr(h, k, _to_host(v))
        return h
    if isinstance(x, (list, tuple)):
        return type(x)(_to_host(v) for v in x)
    if isinstance(x, dict):
        return {k: _to_host(v) for k, v in x.items()}
    return x


def _tensors(x: Any, prefix: str = "") -> list[tuple[str, torch.Tensor]]:
    if isinstance(x, torch.Tensor):
        return [(prefix, x)]
    if isinstance(x, (list, tuple)):
        return [t for i, v in enumerate(x) for t in _tensors(v, f"{prefix}[{i}]")]
    if isinstance(x, dict):
        return [t for k, v in x.items() for t in _tensors(v, f"{prefix}.{k}")]
    return []


def _rel(dev: torch.Tensor, host: torch.Tensor) -> float:
    a = dev.detach().cpu()
    b = host.detach()
    if a.shape != b.shape:
        return float("nan")
    if a.dtype in (torch.float8_e4m3fn, torch.float8_e8m0fnu) or not a.is_floating_point():
        if a.dtype in (torch.float8_e4m3fn, torch.float8_e8m0fnu):
            a, b = a.view(torch.uint8), b.view(torch.uint8)
        neq = (a != b).float().mean().item() if a.numel() else 0.0
        return neq  # fraction of differing elements / bytes
    a, b = a.float(), b.float()
    finite = torch.isfinite(b) & torch.isfinite(a)
    if not finite.all():
        if not torch.equal(torch.isfinite(a), torch.isfinite(b)):
            return float("inf")
        a, b = a[finite], b[finite]
    den = b.norm().item()
    return (a - b).norm().item() / den if den > 0 else (a - b).norm().item()


def _layout(args: Any, kwargs: Any) -> str:
    parts = []
    for name, t in _tensors(args, "a") + _tensors(kwargs, "k"):
        parts.append(f"{name}:{str(t.dtype).replace('torch.', '')}{list(t.shape)}s{list(t.stride())}")
    return " ".join(parts)[:400]


def _record(name: str, rel: float, where: str, layout: str) -> None:
    st = _STATS.setdefault(name, {"calls": 0, "worst": 0.0, "where": "", "layout": ""})
    st["calls"] += 1
    if not (rel <= st["worst"]):  # also catches nan/inf
        st["worst"], st["where"], st["layout"] = rel, where, layout


def wrap(name: str, fn: Any) -> Any:
    if getattr(fn, "_device_diff", False):
        return fn

    @functools.wraps(fn)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        if not _ACTIVE[0]:
            return fn(*args, **kwargs)
        h_args, h_kwargs = _to_host(args), _to_host(kwargs)
        layout = _layout(args, kwargs)
        out = fn(*args, **kwargs)
        _ACTIVE[0] = False  # nested kernels inside the host replay are not checked again
        try:
            h_out = fn(*h_args, **h_kwargs)
        except Exception as e:  # noqa: BLE001
            _record(name, float("nan"), f"host replay failed: {type(e).__name__}: {e}"[:200], layout)
            return out
        finally:
            _ACTIVE[0] = True
        worst, where = 0.0, ""
        for (n, d), (_, h) in zip(_tensors(out, "out"), _tensors(h_out, "out")):
            r = _rel(d, h)
            if not (r <= worst):
                worst, where = r, n
        for (n, d), (_, h) in zip(_tensors(args, "a") + _tensors(kwargs, "k"), _tensors(h_args, "a") + _tensors(h_kwargs, "k")):
            if d.device.type == "cpu":
                continue
            r = _rel(d, h)  # in-place written arguments
            if not (r <= worst):
                worst, where = r, f"inplace {n}"
        _record(name, worst, where, layout)
        return out

    wrapper._device_diff = True
    return wrapper


def wrap_factory(name: str, factory: Any) -> Any:
    """tilelang_ops: ``xxx_tl(static...)`` returns the kernel; wrap the kernel it returns."""

    @functools.wraps(factory)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        return wrap(name, factory(*args, **kwargs))

    return wrapper


def install() -> None:
    import deep_gemm
    import deep_select.interface
    import flash_mla
    import flash_mla.flash_mla_interface
    import tile_kernels
    import tile_kernels.engram
    import tile_kernels.moe
    import tile_kernels.quant
    import tile_kernels.transform

    replaced: dict[int, Any] = {}

    def patch_module(mod: types.ModuleType, prefix: str) -> None:
        for attr, val in list(vars(mod).items()):
            if attr.startswith("_") or not callable(val) or isinstance(val, type):
                continue
            if getattr(val, "__module__", "") is None or not str(getattr(val, "__module__", "")).startswith(prefix.split(".")[0]):
                continue
            w = replaced.get(id(val)) or wrap(f"{prefix}.{attr}", val)
            replaced[id(val)] = w
            setattr(mod, attr, w)

    for mod, prefix in [
        (deep_gemm, "deep_gemm"),
        (flash_mla, "flash_mla"),
        (flash_mla.flash_mla_interface, "flash_mla"),
        (deep_select.interface, "deep_select"),
        (tile_kernels.quant, "tile_kernels.quant"),
        (tile_kernels.transform, "tile_kernels.transform"),
        (tile_kernels.moe, "tile_kernels.moe"),
        (tile_kernels.engram, "tile_kernels.engram"),
    ]:
        patch_module(mod, prefix)
    for sub in ("expand_kernel", "head_compute_mix_kernel", "norm_fn_kernel", "post_kernel", "pre_apply_mix_kernel",
                "pre_split_mixes_kernel", "sinkhorn_kernel", "pre_big_fuse_kernel"):
        m = __import__(f"tile_kernels.mhc.{sub}", fromlist=["x"])
        patch_module(m, "tile_kernels.mhc")

    from deepinfer.model_executor.layers.tilelang_ops import ops

    for attr in [a for a in dir(ops) if a.endswith("_tl")]:
        setattr(ops, attr, wrap_factory(f"tl.{attr}", getattr(ops, attr)))

    # Model modules bound some kernels with `from X import f`: rebind those names too.
    import importlib

    for modname in ["deepinfer.model_executor.models.v4_modules.mhc_ops", "deepinfer.model_executor.models.v4_modules.mhc_modeling",
                    "deepinfer.model_executor.models.v4_modules.mlp", "deepinfer.model_executor.models.v4_modules.hc",
                    "deepinfer.model_executor.models.v4_modules.attn", "deepinfer.model_executor.models.v4_modules.index_attn",
                    "deepinfer.model_executor.layers.attention_backend", "deepinfer.model_executor.layers.rotary_embedding",
                    "deepinfer.model_executor.layers.sampler", "deepinfer.model_executor.layers.sampler_utils"]:
        mod = importlib.import_module(modname)
        for attr, val in list(vars(mod).items()):
            if id(val) in replaced:
                setattr(mod, attr, replaced[id(val)])


def report() -> None:
    print("\n[device_diff] kernel                                         calls   worst-rel  where / layout of worst call")
    for name, st in sorted(_STATS.items(), key=lambda kv: -(kv[1]["worst"] if kv[1]["worst"] == kv[1]["worst"] else 1e9)):
        print(f"[device_diff] {name:48s} {st['calls']:6d}  {st['worst']:9.2e}  {st['where']}  {st['layout']}")


def main() -> None:
    import run_deepinfer_tiny as R

    argv = sys.argv[1:]
    if "--device" not in argv:
        argv = ["--device", "neuron", *argv]
    R.setup_env("neuron")
    import deepinfer

    deepinfer.PLATFORM  # noqa: B018 - initialize the platform (device mapping) before patching
    install()
    sys.argv = [sys.argv[0], *argv]
    try:
        R.main()
    finally:
        report()


if __name__ == "__main__":
    main()