mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 01:06:23 +08:00
b4f7a42652
`h20_ops_ab.py` assumed the TileLang backend is importable, so on a vendor image it could not start at all -- which is exactly where the Triton backend needs checking. Two modes fix that: * `--dump DIR` saves the selected backend's output per case. * `--compare DIR` checks against dumps written elsewhere. Together they answer "do triton 3.5.1 and 3.7.1 agree" without either host having TileLang. Case builders that genuinely need a package the host lacks now report a skip instead of aborting the run, so partial coverage is visible rather than fatal; the import of the sampler's C++-backed inputs moved into the one builder that uses it for the same reason. This is what showed the 19 ported kernels produce element-wise identical output across triton 3.5.1, 3.6.0 and 3.7.1 -- 115 cases, no differences.
278 lines
11 KiB
Python
278 lines
11 KiB
Python
"""Compare a Triton kernel against the TileLang one it replaces, through the shared interface.
|
|
|
|
Every interface module (``gather.py``, ``norm.py``, ...) calls ``ops.<name>_tl`` and then uses the
|
|
result. That means the strongest available check is not "does Triton match a Python reference"
|
|
but "does the same public call give the same answer when the kernel behind it changes". This
|
|
script does exactly that: it swaps ``ops.<name>_tl`` between the two backends and compares the
|
|
output of the public function.
|
|
|
|
python scripts/h20_ops_ab.py # every case
|
|
python scripts/h20_ops_ab.py gather norm
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from collections.abc import Callable
|
|
from typing import Any
|
|
|
|
import torch
|
|
|
|
from deepinfer.model_executor.layers.tilelang_ops import ops
|
|
|
|
|
|
def _load_backend(name: str) -> dict[str, Any] | None:
|
|
"""Collect one backend, or None when it cannot be imported on this host.
|
|
|
|
The default backend is TileLang, which needs `tile_kernels` and a recent `tilelang`. On a host
|
|
that only carries Triton -- the whole point of the port, and the state of a vendor image --
|
|
the comparison has nothing to compare against, so the run degrades to "build and save the
|
|
selected backend's output" instead of failing to start.
|
|
"""
|
|
try:
|
|
return ops._collect(name) # noqa: SLF001 - the collector is the one place that knows the layout
|
|
except ImportError as exc:
|
|
print(f"[ops] backend {name!r} is not importable here: {str(exc).splitlines()[0]}", file=sys.stderr, flush=True)
|
|
return None
|
|
|
|
|
|
CUDA = _load_backend("cuda")
|
|
TRITON = _load_backend("triton")
|
|
BASELINE = CUDA is not None
|
|
|
|
|
|
def swap(name: str, backend: str) -> Callable[[], None]:
|
|
"""Point ``ops.<name>`` at one backend, returning a function that restores it."""
|
|
previous = getattr(ops, name)
|
|
source = CUDA if backend == "cuda" else TRITON
|
|
if source is None or name not in source:
|
|
raise AttributeError(f"backend {backend!r} does not export {name!r}")
|
|
setattr(ops, name, source[name])
|
|
|
|
def restore() -> None:
|
|
setattr(ops, name, previous)
|
|
|
|
return restore
|
|
|
|
|
|
def run_both(name: str, fn: Callable[[], Any]) -> tuple[Any, Any]:
|
|
"""Run ``fn`` once against each backend, returning (cuda result, triton result).
|
|
|
|
Without a baseline backend the first element is None and only the selected backend's output
|
|
comes back, which the caller turns into a dump or a comparison against another host's dump.
|
|
"""
|
|
if BASELINE:
|
|
restore = swap(name, "cuda")
|
|
try:
|
|
expected = fn()
|
|
finally:
|
|
restore()
|
|
else:
|
|
expected = None
|
|
|
|
restore = swap(name, "triton")
|
|
try:
|
|
actual = fn()
|
|
finally:
|
|
restore()
|
|
return expected, actual
|
|
|
|
|
|
# Tolerance follows the dtype's own precision. Anything narrower than float32 is compared at its
|
|
# representable resolution, because a different reduction order in a different compiler moves the
|
|
# last bit and nothing more: for `norm` that is 10 of 655360 elements differing by one bf16 ulp,
|
|
# while the fp8 bytes and every scale factor come out identical.
|
|
_DTYPE_TOL: dict[torch.dtype, tuple[float, float]] = {
|
|
torch.bfloat16: (1e-2, 1e-2),
|
|
torch.float16: (1e-3, 1e-3),
|
|
torch.float32: (1e-5, 1e-6),
|
|
}
|
|
|
|
|
|
def compare(
|
|
name: str,
|
|
expected: Any,
|
|
actual: Any,
|
|
*,
|
|
exact: bool,
|
|
rtol: float | None = None,
|
|
atol: float | None = None,
|
|
agreement: float | None = None,
|
|
) -> bool:
|
|
"""Report whether one case agrees; returns True on success."""
|
|
if (
|
|
agreement is not None
|
|
and isinstance(expected, torch.Tensor)
|
|
and expected.dtype
|
|
in (
|
|
torch.int32,
|
|
torch.int64,
|
|
)
|
|
):
|
|
# Random-draw outputs: require a high agreement rate rather than equality, and say what the
|
|
# rate was so a regression that halves it cannot hide behind the threshold.
|
|
same = int((expected == actual).sum())
|
|
total = expected.numel()
|
|
rate = same / max(total, 1)
|
|
if rate >= agreement:
|
|
if rate < 1.0:
|
|
print(f" note {name}: {total - same}/{total} differ (random draw on a boundary)")
|
|
return True
|
|
print(f" FAIL {name}: agreement {rate:.4f} < {agreement} ({total - same}/{total})")
|
|
return False
|
|
|
|
if isinstance(expected, torch.Tensor):
|
|
if expected.dtype != actual.dtype:
|
|
print(f" FAIL {name}: dtype {expected.dtype} vs {actual.dtype}")
|
|
return False
|
|
if expected.shape != actual.shape:
|
|
print(f" FAIL {name}: shape {tuple(expected.shape)} vs {tuple(actual.shape)}")
|
|
return False
|
|
if exact or not expected.is_floating_point():
|
|
same = torch.equal(expected, actual)
|
|
if not same:
|
|
bad = expected != actual
|
|
print(f" FAIL {name}: {int(bad.sum())}/{expected.numel()} elements differ")
|
|
idx = bad.nonzero()[0].tolist()
|
|
print(f" first at {idx}: {expected[tuple(idx)].item()} vs {actual[tuple(idx)].item()}")
|
|
return bool(same)
|
|
if expected.dtype in (torch.float8_e4m3fn, torch.float8_e5m2):
|
|
# Compare at the format's own resolution: a one-ulp difference is a rounding tie that
|
|
# moved, not a wrong value, so byte distance is the honest measure. Require the scales
|
|
# to be exact elsewhere and the bytes to agree almost everywhere.
|
|
delta = (expected.view(torch.uint8).int() - actual.view(torch.uint8).int()).abs()
|
|
frac = float((delta > 0).float().mean())
|
|
if int(delta.max()) <= 2 and frac <= 1e-3:
|
|
return True
|
|
print(f" FAIL {name}: {frac * 100:.4f}% bytes differ, max byte distance {int(delta.max())}")
|
|
return False
|
|
|
|
if rtol is None or atol is None:
|
|
default_rtol, default_atol = _DTYPE_TOL.get(expected.dtype, (1e-5, 1e-6))
|
|
rtol = default_rtol if rtol is None else rtol
|
|
atol = default_atol if atol is None else atol
|
|
close = torch.allclose(expected.float(), actual.float(), rtol=rtol, atol=atol)
|
|
if not close:
|
|
diff = (expected.float() - actual.float()).abs()
|
|
scale = expected.float().abs().clamp_min(1e-12)
|
|
print(f" FAIL {name}: max abs {diff.max().item():.3e}, max rel {(diff / scale).max().item():.3e}")
|
|
return bool(close)
|
|
|
|
if isinstance(expected, (tuple, list)):
|
|
# A case that returns several tensors is only right if every element is; recurse so each
|
|
# one gets its own report instead of collapsing into one boolean.
|
|
if len(expected) != len(actual):
|
|
print(f" FAIL {name}: {len(expected)} outputs vs {len(actual)}")
|
|
return False
|
|
ok = True
|
|
for i, (e, a) in enumerate(zip(expected, actual, strict=True)):
|
|
ok &= compare(f"{name}[{i}]", e, a, exact=exact, rtol=rtol, atol=atol, agreement=agreement)
|
|
return ok
|
|
|
|
same = expected == actual
|
|
if not same:
|
|
print(f" FAIL {name}: {expected!r} vs {actual!r}")
|
|
return bool(same)
|
|
|
|
|
|
def cases() -> dict[str, Callable[[], list[tuple[str, bool, Callable[[], Any]]]]]:
|
|
"""Registry of per-kernel case builders, filled in as kernels are ported."""
|
|
from scripts.ab_cases import CASES
|
|
|
|
return CASES
|
|
|
|
|
|
def _dump_file(root: str, kernel: str, case: str) -> str:
|
|
return os.path.join(root, f"{kernel}__{case}.pt")
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("kernels", nargs="*", help="subset of kernel names; default is all ported")
|
|
parser.add_argument("--verbose", action="store_true")
|
|
parser.add_argument(
|
|
"--dump",
|
|
metavar="DIR",
|
|
help="save the selected backend's output per case instead of comparing to the default backend. "
|
|
"Use it on a host where the default backend is not installed, then pass the directory to "
|
|
"--compare on another host to check that two Triton builds agree.",
|
|
)
|
|
parser.add_argument("--compare", metavar="DIR", help="compare against dumps written by --dump on another host")
|
|
args = parser.parse_args()
|
|
if args.dump and args.compare:
|
|
raise SystemExit("--dump and --compare are mutually exclusive")
|
|
|
|
if args.compare:
|
|
os.makedirs(args.compare, exist_ok=True)
|
|
elif args.dump:
|
|
os.makedirs(args.dump, exist_ok=True)
|
|
elif not BASELINE:
|
|
raise SystemExit(
|
|
"no baseline backend here, so there is nothing to compare against: pass --dump DIR "
|
|
"(save the outputs) or --compare DIR (check against another host's dumps)"
|
|
)
|
|
|
|
registry = cases()
|
|
selected = args.kernels or sorted(registry)
|
|
unknown = [name for name in selected if name not in registry]
|
|
if unknown:
|
|
raise SystemExit(f"no A/B cases for: {', '.join(unknown)}")
|
|
|
|
failures = 0
|
|
skipped = 0
|
|
saved = 0
|
|
for kernel in selected:
|
|
print(f"[{kernel}]")
|
|
try:
|
|
built = registry[kernel]()
|
|
except ImportError as exc:
|
|
# A case builder that needs `tile_kernels` cannot run on a Triton-only host; say so
|
|
# rather than counting it as a failure, so the report stays honest about coverage.
|
|
print(f" skip (needs a package this host lacks): {str(exc).splitlines()[0]}")
|
|
skipped += 1
|
|
continue
|
|
for case in built:
|
|
# A case is `(name, exact, fn)`, optionally with a fourth element: the fraction of
|
|
# entries allowed to differ. Only for outputs whose answer is a random draw, where the
|
|
# last bit of a cumulative mass decides which side of a boundary the token falls on;
|
|
# everything else is compared exactly or within a dtype tolerance.
|
|
case_name, exact, fn = case[:3]
|
|
agreement = case[3] if len(case) > 3 else None
|
|
try:
|
|
expected, actual = run_both(kernel, fn)
|
|
if args.dump:
|
|
torch.save(actual, _dump_file(args.dump, kernel, case_name))
|
|
saved += 1
|
|
continue
|
|
if args.compare:
|
|
path = _dump_file(args.compare, kernel, case_name)
|
|
if not os.path.exists(path):
|
|
print(f" FAIL {case_name}: no dump at {path}")
|
|
failures += 1
|
|
continue
|
|
expected = torch.load(path, weights_only=False)
|
|
except Exception as exc: # noqa: BLE001 - a raise on one backend is itself a result
|
|
lines = str(exc).splitlines()
|
|
print(f" FAIL {case_name}: {type(exc).__name__}: {lines[0][:160] if lines else ''}")
|
|
failures += 1
|
|
continue
|
|
if not compare(case_name, expected, actual, exact=exact, agreement=agreement):
|
|
failures += 1
|
|
elif args.verbose:
|
|
print(f" ok {case_name}")
|
|
|
|
print()
|
|
if args.dump:
|
|
print(f"AB_DUMPED {saved} case(s) to {args.dump}" + (f", {skipped} kernel(s) skipped" if skipped else ""))
|
|
return
|
|
if failures:
|
|
print(f"AB_FAILED {failures}" + (f" ({skipped} kernel(s) skipped)" if skipped else ""))
|
|
sys.exit(1)
|
|
print("AB_OK" + (f" ({skipped} kernel(s) skipped)" if skipped else ""))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|