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