mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-11-01 01:06:23 +08:00
fece5d9047
An edit of mine dropped a `return launch` from `finalize_verification_tl`. The A/B case for that kernel was not re-run, so nothing local objected; the failure surfaced as a TypeError inside the sampler of an end-to-end run, once per rank. The harness now checks that every `*_tl` symbol the selected backend exports is callable, which turns that class of mistake into one line printed before any case executes. Also records where the remaining sixteen kernels actually sit. None of them is unconditionally on the path, but the distribution is lopsided and the document had it backwards in one place: the three `candidate_*` kernels are not merely "off on SM90", they are switched off *by* `PLATFORM.is_sm90` -- and this checkpoint sets `candidate_source_layer_id=20` of 40 layers. On a card where `is_sm90` is False, that is nineteen layers on every step, which is the opposite of optional. Their harder blocker is noted alongside: that path also needs `deep_gemm.fp8_fp4_sparse_mqa_logits`, which is SM100-only, so porting those three alone would buy nothing. The four kernels that no production path reaches are listed as such, so nobody spends an afternoon on them.
286 lines
12 KiB
Python
286 lines
12 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)}")
|
|
|
|
# Every `*_tl` symbol the selected backend exports must be callable. A factory that falls off
|
|
# the end of its body returns None, and nothing notices until the engine calls it -- a stray
|
|
# edit once dropped a `return launch` and the first symptom was a TypeError eight ranks deep
|
|
# in an end-to-end run. Checking here turns that into one line before any case runs.
|
|
broken = [name for name in dir(ops) if name.endswith("_tl") and not callable(getattr(ops, name))]
|
|
if broken:
|
|
raise SystemExit(f"backend {ops._BACKEND!r} exports non-callable kernel factories: {', '.join(broken)}")
|
|
|
|
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()
|