diff --git a/transformer_engine/pytorch/gemm_autotune.py b/transformer_engine/pytorch/gemm_autotune.py new file mode 100644 index 000000000..86c15c1bf --- /dev/null +++ b/transformer_engine/pytorch/gemm_autotune.py @@ -0,0 +1,248 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""Opt-in on-the-fly autotuning for dense GEMM (bf16 + mxfp8, forward + backward). + +Second consumer of the pure :mod:`kernel_router` (after ``grouped_gemm_autotune``). +When ``NVTE_AUTOTUNE_KERNELS=1`` (ROCm only), each dense GEMM issued through +``general_gemm`` -- forward (TN), dgrad (NN), wgrad (NT) -- picks the fastest +backend for its shape+layout by measuring the candidates once and caching the +winner. Two regimes have a real choice: bf16 races the C++ default (hipBLASLt) +against the Triton kernel (``NVTE_USE_GEMM_TRITON=1``); mxfp8 races HipKittens (the +C++ default for mxfp8) against hipBLASLt (``NVTE_ROCM_USE_HIPBLASLT_MXFP8=1``). Both +toggles are read per call, so -- unlike the grouped mxfp8 path -- there is no +process-static to freeze the choice. fp8/fp32 have a single backend and are left +to ``general_gemm`` unchanged. + +The op-agnostic switches (``NVTE_AUTOTUNE_KERNELS`` / ``NVTE_AUTOTUNE_KERNELS_VERBOSE``) +and the low-level runtime glue (env toggling, the do_bench timer, the input-format +classifier) are shared with ``grouped_gemm_autotune``. + +Scope / safety: + +* Only plain GEMMs are routed: calls with a comm-overlap communicator (``ub``) or + an output quantizer are delegated to ``general_gemm`` unchanged -- the former + because measuring would double the collective, the latter because repeated + measurement would corrupt the output quantizer's amax. (mxfp8 quantizes the + *inputs*; the output is bf16 with no quantizer, so mxfp8 is amax-safe.) +* Timing is side-effect-free: candidates are measured into a fresh output + (``out=None``, ``accumulate=False``) so re-running them never touches the real + output or a fused ``main_grad``. Only the winner is run once with the real args. +* For mxfp8, when HipKittens does not support a shape the C++ path falls back to + hipBLASLt internally (per call, deterministic), so the ``hipkittens`` candidate + degrades to a tie with ``hipblaslt`` rather than mis-measuring. +* Unlike grouped GEMM, dense token count (M = seq*mbs) is fixed within a run, so + the key uses exact M/N/K -- no coarse token bin is needed. +* Off by default; ``NVTE_AUTOTUNE_KERNELS_VERBOSE=1`` adds per-call selection logging. +""" +from __future__ import annotations + +from dataclasses import dataclass + +from .cpp_extensions import general_gemm +from .kernel_router import AutotuneRouter + +# Shared, torch-dependent runtime glue (pending extraction into a common module). +from .grouped_gemm_autotune import ( + _FLOAT16_DTYPES, + _autotune_enabled, + _do_bench_ms, + _env, + _quant_classes, + _verbose, +) + +_GEMM_TRITON_ENV = "NVTE_USE_GEMM_TRITON" +_HIPBLASLT_MXFP8_ENV = "NVTE_ROCM_USE_HIPBLASLT_MXFP8" + + +@dataclass(frozen=True) +class GemmKey: + """GPU-free key a dense-GEMM selection is cached on. M/N/K are exact (dense + token count is stable within a run), so no coarse binning is needed.""" + + m: int + n: int + k: int + out_dtype: str + layout: str + in_format: str + + +def _in_format(operand, out_dtype): + """Classify the GEMM's input precision as ``bf16`` / ``fp8`` / ``mxfp8`` from one + operand, or ``None`` if unsupported. Mirrors the grouped classifier but takes a + single tensor (dense operands are tensors, not lists).""" + if out_dtype not in _FLOAT16_DTYPES: + return None + fp8_types, mxfp8_type, storage_type = _quant_classes() + if not isinstance(operand, storage_type): + return "bf16" if getattr(operand, "dtype", None) in _FLOAT16_DTYPES else None + q = getattr(operand, "_quantizer", None) + if isinstance(q, mxfp8_type): + return "mxfp8" + if isinstance(q, fp8_types): + return "fp8" + return None + + +def _mnk(A, B, layout): + """Logical GEMM dims from the operands, per TE's column-major BLAS convention + (see ``triton_kernels/gemm/gemm_wrapper.py``): ``m = A0 if transa else A1``, + ``k = A1 if transa else A0``, ``n = B1 if transb else B0``.""" + transa = layout[0] == "T" + transb = layout[1] == "T" + m = A.size(0) if transa else A.size(1) + k = A.size(1) if transa else A.size(0) + n = B.size(1) if transb else B.size(0) + return m, n, k + + +_last_log: str | None = None + + +def _log_selection(sel) -> None: + """One line per call: cache miss, the key, what was tried, and the winner. + Consecutive identical lines are collapsed; cache hits are silent.""" + global _last_log + if sel.from_cache: + return + key = sel.key + ks = ( + f"M={key.m} N={key.n} K={key.k} in_format={key.in_format} " + f"layout={key.layout} out_dtype={key.out_dtype}" + ) + tried = [] + for r in sel.reports: + if not r.available: + tried.append(f"{r.name}=unavailable") + elif r.error: + tried.append(f"{r.name}=ERROR({r.error})") + elif r.time_ms is None: + tried.append(f"{r.name}=rejected") + else: + tried.append(f"{r.name}={r.time_ms:.4f}ms") + line = f"[te-autotune] cache MISS [{ks}] tried: {', '.join(tried)} -> selected {sel.winner}" + if line != _last_log: + print(line, flush=True) + _last_log = line + + +@dataclass +class _GemmCall: + """Everything needed to (re)issue one ``general_gemm`` call. ``kwargs`` is the + full original keyword set (out, accumulate, layout, bias, ...).""" + + A: object + B: object + kwargs: dict + + +class _GemmBackend: + """A dense-GEMM backend reached through ``general_gemm``, selected by transiently + toggling ``NVTE_USE_GEMM_TRITON``.""" + + def __init__(self, name, env, *, formats): + self.name = name + self._env = dict(env) + self._formats = frozenset(formats) + + def available(self, key: GemmKey) -> bool: + return key.in_format in self._formats + + def prepare(self, call: _GemmCall): + # Side-effect-free timing: fresh output, accumulate off, no comm-overlap + # extra output. Only the winner is run for real via run_real(). + env = self._env + mkwargs = dict(call.kwargs) + mkwargs["out"] = None + mkwargs["accumulate"] = False + mkwargs.pop("extra_output", None) + + def run(): + with _env(**env): + general_gemm(call.A, call.B, **mkwargs) + + return run + + def run_real(self, call: _GemmCall): + with _env(**self._env): + return general_gemm(call.A, call.B, **call.kwargs) + + +_router: AutotuneRouter | None = None +_backends: dict = {} + + +def _get_router() -> AutotuneRouter: + global _router, _backends + if _router is None: + # Each candidate sets both toggles explicitly so it selects the same backend + # regardless of format: bf16 ignores the mxfp8 toggle and vice versa. + candidates = [ + _GemmBackend( + "hipblaslt", + {_GEMM_TRITON_ENV: None, _HIPBLASLT_MXFP8_ENV: "1"}, + formats=("bf16", "mxfp8"), + ), + _GemmBackend( + "triton", + {_GEMM_TRITON_ENV: "1", _HIPBLASLT_MXFP8_ENV: None}, + formats=("bf16",), + ), + _GemmBackend( + "hipkittens", + {_GEMM_TRITON_ENV: None, _HIPBLASLT_MXFP8_ENV: None}, + formats=("mxfp8",), + ), + ] + _backends = {c.name: c for c in candidates} + _router = AutotuneRouter( + candidates=candidates, + timer=_do_bench_ms, + verifier=None, + default="hipblaslt", + ) + return _router + + +def _eligible(A, B, kwargs): + """Return the input format if this dense GEMM should be autotuned, else None. + Excludes disabled autotune, comm-overlap (``ub``), and output-quantized calls. + bf16 (hipBLASLt vs Triton) and mxfp8 (hipBLASLt vs HipKittens) each have two + backends; other formats have one and are left to ``general_gemm``.""" + if not _autotune_enabled(): + return None + if kwargs.get("ub") is not None: + return None + if kwargs.get("quantization_params") is not None: + return None + fmt = _in_format(A, kwargs.get("out_dtype")) + return fmt if fmt in ("bf16", "mxfp8") else None + + +def autotuned_gemm(A, B, **kwargs): + """Drop-in for :func:`general_gemm`. + + With ``NVTE_AUTOTUNE_KERNELS=1`` and an autotunable dense GEMM (bf16 or mxfp8 + inputs, no comm-overlap, unquantized output), selects the fastest backend + (hipBLASLt / Triton for bf16, hipBLASLt / HipKittens for mxfp8) for this + shape+layout and runs it. Otherwise delegates to ``general_gemm`` unchanged, + returning its 4-tuple result. + """ + fmt = _eligible(A, B, kwargs) + if fmt is None: + return general_gemm(A, B, **kwargs) + + layout = kwargs.get("layout", "TN") + m, n, k = _mnk(A, B, layout) + key = GemmKey( + m=m, n=n, k=k, out_dtype=str(kwargs.get("out_dtype")), layout=layout, in_format=fmt + ) + call = _GemmCall(A=A, B=B, kwargs=kwargs) + sel = _get_router().select(key, call) + if _verbose(): + _log_selection(sel) + # select() measures into scratch on a miss and returns only a name on a hit; + # run the chosen backend once with the real args to produce the result. + return _backends[sel.winner].run_real(call) diff --git a/transformer_engine/pytorch/grouped_gemm_autotune.py b/transformer_engine/pytorch/grouped_gemm_autotune.py new file mode 100644 index 000000000..a91dabbab --- /dev/null +++ b/transformer_engine/pytorch/grouped_gemm_autotune.py @@ -0,0 +1,457 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""Opt-in on-the-fly autotuning for grouped GEMM (bf16/fp8/mxfp8, fwd + bwd). + +When ``NVTE_AUTOTUNE_KERNELS=1`` (ROCm only), each grouped GEMM issued through +``general_grouped_gemm`` -- forward (TN), dgrad (NN), and wgrad (NT) -- picks the +fastest backend for its shape+layout+in_format by measuring the candidates once (via +``triton.testing.do_bench``) and caching the winner. The candidates are the two +C++ backends that share the entry point and semantics (multi-stream hipBLASLt and +CK), HipKittens (mxfp8 only), plus the Triton grouped GEMM (bf16 only). The pure +selection lives in :mod:`kernel_router`; this module supplies the backend +candidates and the process-global router. + +The env vars are intentionally op-agnostic (``NVTE_AUTOTUNE_KERNELS`` / +``NVTE_AUTOTUNE_KERNELS_VERBOSE``) so the same switches govern future autotuned ops. + +Scope / safety: + +* Forward, dgrad, and wgrad grouped GEMMs are routed; each layout is a distinct + route key, measured independently. +* The C++ backends are selected by transiently toggling env vars: hipBLASLt + (multi-stream, all formats, the guaranteed floor), CK (bf16/fp8, num_gemms > 1), + and, for mxfp8 only, a single cutlass-family candidate that prefers HipKittens + (``hipkittens``). HipKittens needs 256-aligned expert dims; when they are not, + the C++ path falls back to CK internally, so that one candidate covers + HK-when-aligned and CK-otherwise. It is deliberately the *only* mxfp8 + cutlass-family candidate: the C++ HK-enable flag is a process-lifetime static + frozen on first use, so a second competing env config could not switch backends + at runtime and might pin the wrong one. Triton is a candidate for bf16 only. The + separate global ``NVTE_USE_GROUPED_GEMM_TRITON`` flag (which forces Triton for all + layouts) is independent of and mutually exclusive with this path. +* Only the *input* is quantized (fp8/mxfp8); the GEMM output is bf16 with no + output quantizer, so repeated measurement never mutates quantizer amax state. + Output-quantized calls (debug, fp8 output) are excluded and delegated. +* The Triton candidate runs TN/NN directly; NT (wgrad) runs into a fresh 3D packed + buffer (as the ``use_grouped_gemm_triton`` path allocates) and copies the result + back into the per-group output list. +* Timing is side-effect-free: candidates are measured into a scratch clone of the + output with ``accumulate=False``, so re-running them under ``do_bench`` can + never corrupt a fused ``main_grad`` (wgrad accumulation) or any real output. + Only the chosen winner is then run once into the real output with the real + arguments. +* The deferred-wgrad path (``wgrad_store``) is left untouched. +* Off by default; ``NVTE_AUTOTUNE_KERNELS_VERBOSE=1`` adds per-call selection logging + (cache hit/miss, route key, per-backend timings, winner). + +""" +from __future__ import annotations + +import contextlib +import os +from dataclasses import dataclass + +import torch +from torch.utils.cpp_extension import IS_HIP_EXTENSION + +from .cpp_extensions import general_grouped_gemm +from .kernel_router import AutotuneRouter, RouteKey, make_route_key + +# Op-agnostic autotune switches, shared by any future autotuned op. +_MASTER_ENV = "NVTE_AUTOTUNE_KERNELS" +_VERBOSE_ENV = "NVTE_AUTOTUNE_KERNELS_VERBOSE" + +# Warm-up iterations before timing, to force JIT/autotune compilation (Triton) +# out of the measured region. +_WARMUP_ITERS = 3 + + +def _autotune_enabled() -> bool: + return IS_HIP_EXTENSION and os.getenv(_MASTER_ENV, "0") == "1" + + +def _verbose() -> bool: + return os.getenv(_VERBOSE_ENV, "0") == "1" + + +_last_log: str | None = None + + +def _log_selection(sel) -> None: + """Print one line per call: cache hit/miss, the route key, what was tried + (timings or skip reason), and the winner. Consecutive identical lines are + collapsed so a steady state (repeated cache hits for one shape) logs once.""" + global _last_log + key = sel.key + ks = ( + f"G={key.num_groups} N={key.N} K={key.K} in_format={key.in_format} " + f"layout={key.layout} size_bin={key.size_bin}" + ) + if sel.from_cache: + return + line = f"[te-autotune] cache HIT [{ks}] -> {sel.winner}" + else: + tried = [] + for r in sel.reports: + if not r.available: + tried.append(f"{r.name}=unavailable") + elif r.error: + tried.append(f"{r.name}=ERROR({r.error})") + elif r.time_ms is None: + tried.append(f"{r.name}=rejected") + else: + tried.append(f"{r.name}={r.time_ms:.4f}ms") + line = ( + f"[te-autotune] cache MISS [{ks}] tried: {', '.join(tried)} " + f"-> selected {sel.winner}" + ) + if line != _last_log: + print(line, flush=True) + _last_log = line + + +@contextlib.contextmanager +def _env(**overrides): + """Temporarily set/unset env vars; None removes. The C++ grouped-GEMM + dispatch reads these per call, so each invocation sets them transiently -- + backward (which does not enter this path) is never affected.""" + prev = {k: os.environ.get(k) for k in overrides} + for k, v in overrides.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + try: + yield + finally: + for k, v in prev.items(): + if v is None: + os.environ.pop(k, None) + else: + os.environ[k] = v + + +@dataclass +class _GGCall: + """Everything needed to (re)issue one ``general_grouped_gemm`` call, for any + of the three grouped GEMMs (fprop TN, dgrad NN, wgrad NT). ``out`` is always + the list ``general_grouped_gemm`` expects; ``N``/``K`` are the weight dims, + passed by the caller so the route key is stable per layer across layouts.""" + + A: list + B: list + out: list + quantization_params: list + out_dtype: object + m_splits: list + layout: str + N: int + K: int + gemm_kwargs: dict + + @property + def num_groups(self) -> int: + return len(self.m_splits) + + +_ALL_FORMATS = ("bf16", "fp8", "mxfp8") + + +class _GroupedGemmBackend: + """A grouped-GEMM backend reached through ``general_grouped_gemm``, selected by + transiently toggling env vars. Availability is format-gated; ``needs_multi`` + requires num_groups > 1 (the C++ path only takes the grouped fast path with more + than one group).""" + + def __init__(self, name, env, *, formats, needs_multi=False): + self.name = name + self._env = dict(env) + self._formats = frozenset(formats) + self._needs_multi = needs_multi + + def available(self, key: RouteKey) -> bool: + if key.in_format not in self._formats: + return False + if self._needs_multi and key.num_groups <= 1: + return False + return True + + def prepare(self, call: _GGCall): + # Side-effect-free timing: measure into a scratch clone of the output with + # accumulate off, so repeated do_bench runs never touch the real output or + # a fused main_grad. Only the winner is run for real via run_real(). + env = self._env + scratch = [torch.empty_like(o) for o in call.out] + timing_kwargs = dict(call.gemm_kwargs) + timing_kwargs["accumulate"] = False + + def run(): + with _env(**env): + general_grouped_gemm( + call.A, call.B, scratch, call.quantization_params, call.out_dtype, + m_splits=list(call.m_splits), layout=call.layout, **timing_kwargs, + ) + + return run + + def run_real(self, call: _GGCall): + with _env(**self._env): + return general_grouped_gemm( + call.A, call.B, call.out, call.quantization_params, call.out_dtype, + m_splits=list(call.m_splits), layout=call.layout, **call.gemm_kwargs, + ) + + +def _triton_grouped_gemm(): + """Lazily import the Triton grouped GEMM; return None if unavailable.""" + try: + from transformer_engine.pytorch.triton_kernels.grouped_gemm import ( + general_grouped_gemm_triton, + ) + + return general_grouped_gemm_triton + except Exception: + return None + + +class _TritonGroupedGemmBackend: + """Triton (AITER) grouped GEMM as an autotune candidate. + + Forward (TN) and dgrad (NN) call Triton directly -- it concatenates the + per-group inputs internally (a real cost in this pre-split path, left inside + the timing). wgrad (NT) is supported partially: Triton wgrad writes a 3D + packed output (like the ``use_grouped_gemm_triton`` path allocates), so it is + run into a fresh 3D buffer and the result copied/accumulated back into the + per-group output list this path uses. An error on any shape drops it from the + ranking. + """ + + name = "triton" + + def available(self, key: RouteKey) -> bool: + return ( + key.in_format == "bf16" + and key.layout in ("TN", "NN", "NT") + and _triton_grouped_gemm() is not None + ) + + def _run_into(self, fn, call: _GGCall, out_target): + # TN/NN: Triton writes the per-group output list directly. + if call.layout != "NT": + return fn( + call.A, call.B, out_target, call.quantization_params, call.out_dtype, + m_splits=list(call.m_splits), layout=call.layout, **call.gemm_kwargs, + ) + # NT (wgrad): Triton needs a 3D packed output (G, N, K). Run into a fresh + # buffer with accumulate off, then copy (or accumulate) back into the + # per-group targets -- so it works whether they are fresh buffers or fused + # main_grads, and the copy-back cast bridges any dtype difference. + num_gemms = len(out_target) + n, k = out_target[0].shape + out3d = torch.empty((num_gemms, n, k), dtype=call.out_dtype, device=out_target[0].device) + kwargs = dict(call.gemm_kwargs) + accumulate = kwargs.pop("accumulate", False) + result = fn( + call.A, call.B, out3d, call.quantization_params, call.out_dtype, + m_splits=list(call.m_splits), layout="NT", accumulate=False, **kwargs, + ) + for i, w in enumerate(out_target): + w.add_(out3d[i]) if accumulate else w.copy_(out3d[i]) + return result + + def prepare(self, call: _GGCall): + fn = _triton_grouped_gemm() + assert fn is not None # guaranteed by available() + scratch = [torch.empty_like(o) for o in call.out] + + def run(): + self._run_into(fn, call, scratch) + + return run + + def run_real(self, call: _GGCall): + fn = _triton_grouped_gemm() + assert fn is not None # guaranteed by available() + return self._run_into(fn, call, call.out) + + +_router: AutotuneRouter | None = None +_backends: dict = {} + + +def _do_bench_ms(fn) -> float: + import triton + + # Warm up so first-call JIT/autotune compilation (Triton) is excluded from the + # timed region -- otherwise a cold compile (seconds) is charged to the first + # measured shape and permanently mis-rates the backend. do_bench's own warmup + # then runs on already-compiled kernels. Harmless for the JIT-free C++ backends. + for _ in range(_WARMUP_ITERS): + fn() + torch.cuda.synchronize() + try: + return triton.testing.do_bench(fn, warmup=100, rep=100, return_mode="median") + except TypeError: + med, _, _ = triton.testing.do_bench(fn, warmup=100, rep=100, quantiles=[0.5, 0.2, 0.8]) + return med + + +def _get_router() -> AutotuneRouter: + global _router, _backends + if _router is None: + candidates = [ + _GroupedGemmBackend( + "hipblaslt", + { + "NVTE_USE_CUTLASS_GROUPED_GEMM": None, + "NVTE_USE_CK_GROUPED_GEMM": None, + "NVTE_USE_HIPKITTENS_GROUPED_GEMM": None, + }, + formats=_ALL_FORMATS, + ), + _GroupedGemmBackend( + "ck_tile", + { + "NVTE_USE_CUTLASS_GROUPED_GEMM": "1", + "NVTE_USE_CK_GROUPED_GEMM": "1", + "NVTE_USE_HIPKITTENS_GROUPED_GEMM": None, + }, + formats=("bf16", "fp8"), + needs_multi=True, + ), + # Sole cutlass-family candidate for mxfp8. HipKittens is preferred (env + # HK=1) and the C++ path silently falls back to CK when the expert dims + # are not 256-aligned, so this one candidate covers HK-when-aligned and + # CK-otherwise. It must be the ONLY mxfp8 cutlass-family candidate: the + # C++ HK-enable flag is a process-lifetime static frozen on first use, so + # a second competing env config (e.g. a CK-forcing candidate) would not + # actually switch backends and could permanently pin the wrong choice. + _GroupedGemmBackend( + "hipkittens", + { + "NVTE_USE_CUTLASS_GROUPED_GEMM": "1", + "NVTE_USE_HIPKITTENS_GROUPED_GEMM": "1", + "NVTE_USE_CK_GROUPED_GEMM": None, + }, + formats=("mxfp8",), + needs_multi=True, + ), + _TritonGroupedGemmBackend(), + ] + _backends = {c.name: c for c in candidates} + # verifier=None: both backends are already validated by TE's grouped-GEMM + # tests, so no per-call numerics gate (which would need an fp32 reference). + _router = AutotuneRouter( + candidates=candidates, + timer=_do_bench_ms, + verifier=None, + default="hipblaslt", + ) + return _router + + +_FLOAT16_DTYPES = (torch.bfloat16, torch.float16) + +_quant_cache = None + + +def _quant_classes(): + """Lazily import the quantizer/tensor classes used to classify the input format.""" + global _quant_cache + if _quant_cache is None: + from .tensor import ( + Float8CurrentScalingQuantizer, + Float8Quantizer, + MXFP8Quantizer, + QuantizedTensorStorage, + ) + + _quant_cache = ( + (Float8Quantizer, Float8CurrentScalingQuantizer), + MXFP8Quantizer, + QuantizedTensorStorage, + ) + return _quant_cache + + +def _in_format(A, out_dtype): + """Classify the grouped GEMM as ``bf16`` / ``fp8`` / ``mxfp8`` from the input + operand, or ``None`` if unsupported (fp32, nvfp4, unknown). The GEMM output is + always bf16/fp16; the input format captures the *input* precision so a bf16 and + an fp8 GEMM of the same shape do not collide in the cache.""" + if out_dtype not in _FLOAT16_DTYPES: + return None + fp8_types, mxfp8_type, storage_type = _quant_classes() + a0 = A[0] + if not isinstance(a0, storage_type): + return "bf16" if getattr(a0, "dtype", None) in _FLOAT16_DTYPES else None + q = getattr(a0, "_quantizer", None) + if isinstance(q, mxfp8_type): + return "mxfp8" + if isinstance(q, fp8_types): + return "fp8" + return None + + +def _eligible(A, quantization_params, out_dtype): + """Return the input format (bf16/fp8/mxfp8) if this call should be autotuned, + else None. Excludes disabled autotune, output-quantized calls (debug or fp8 + output, whose quantizer amax would be corrupted by repeated measurement), and + unsupported input formats.""" + if not _autotune_enabled(): + return None + if any(q is not None for q in quantization_params): + return None + return _in_format(A, out_dtype) + + +def _key_dims(A, out, layout): + """The weight dims (N, K) for the route key, from the layout-appropriate + operand: wgrad (NT) writes them into ``out``; fprop/dgrad read them off the + weight operand ``A``. Uses ``.size()`` so it works for quantized inputs too + (the ``*TensorStorage`` classes expose ``.size()`` but not ``.shape``).""" + t = out[0] if layout == "NT" else A[0] + return t.size(0), t.size(1) + + +def autotuned_grouped_gemm(A, B, out, quantization_params, out_dtype, **kwargs): + """Drop-in for :func:`general_grouped_gemm`. + + When ``NVTE_AUTOTUNE_KERNELS=1`` and the call is an autotunable grouped GEMM + (bf16/fp8/mxfp8 inputs, unquantized output), it selects the fastest available + backend for this shape+layout+in_format and runs it. Otherwise -- disabled, + CUDA, debug, fp32, or an unsupported input format -- it delegates to + ``general_grouped_gemm`` unchanged, returning its result. + """ + in_format = _eligible(A, quantization_params, out_dtype) + if in_format is None: + return general_grouped_gemm(A, B, out, quantization_params, out_dtype, **kwargs) + + m_splits = kwargs.get("m_splits") + if m_splits is None: + return general_grouped_gemm(A, B, out, quantization_params, out_dtype, **kwargs) + + layout = kwargs.get("layout", "TN") + N, K = _key_dims(A, out, layout) + # layout/m_splits are tracked on the call; keep them out of the forwarded kwargs. + gemm_kwargs = {k: v for k, v in kwargs.items() if k not in ("layout", "m_splits")} + call = _GGCall( + A=A, + B=B, + out=out, + quantization_params=quantization_params, + out_dtype=out_dtype, + m_splits=list(m_splits), + layout=layout, + N=N, + K=K, + gemm_kwargs=gemm_kwargs, + ) + key = make_route_key(len(m_splits), tuple(m_splits), N, K, str(out_dtype), layout, in_format) + sel = _get_router().select(key, call) + if _verbose(): + _log_selection(sel) + # select() measures into scratch on a miss and returns only a name on a hit; + # run the chosen backend once into the real `out` to produce the result. + return _backends[sel.winner].run_real(call) diff --git a/transformer_engine/pytorch/kernel_router.py b/transformer_engine/pytorch/kernel_router.py new file mode 100644 index 000000000..88c0a50c8 --- /dev/null +++ b/transformer_engine/pytorch/kernel_router.py @@ -0,0 +1,168 @@ +# Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved. +# License for AMD contributions = MIT. See LICENSE for more information + +"""Pure, GPU-free kernel router for per-shape backend selection. + +This module holds no torch and no TE imports and does no I/O. It selects a +backend for a :class:`RouteKey` by measuring the available candidates once (via +an *injected* timer) and caching the winner in memory. The timer and the +(optional) numerics verifier are injected, so the selection logic is +unit-testable off-GPU and the router never depends on a specific backend or +timing method. + +It is deliberately backend-agnostic: an op wires up its own candidates (each +exposing ``name``/``available``/``prepare``), a timer, and -- optionally -- a +numerics verifier, then asks :meth:`AutotuneRouter.select` for the winner. See +``grouped_gemm_autotune.py`` for the first consumer. +""" +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Callable, Protocol, runtime_checkable + + +@dataclass(frozen=True) +class RouteKey: + """The GPU-free bundle a selection is keyed on. No live tensors. + + Structural fields (``num_groups``, ``N``, ``K``, ``out_dtype``, ``layout``, + ``in_format``) are exact and stable per layer -- ``in_format`` is the input + number format (bf16/fp8/mxfp8), which the output element type (``out_dtype``) + alone cannot convey. The token count enters only as a coarse + ``size_bin`` (small/large, see :func:`make_route_key`): the exact per-step + token count jitters with dynamic routing, so a fine key would thrash the cache + and re-measure constantly. A single 2-bin split keeps the cache stable while + still capturing the one regime where the fastest backend flips with size. + """ + + num_groups: int + N: int + K: int + out_dtype: str + layout: str + size_bin: int + in_format: str + + +# Coarse small/large split on the average per-group token count. A tunable +# heuristic, not a hard boundary: the router still *measures* both backends per +# bin, so the threshold only needs to sit near the size where the winner changes. +_LARGE_TOKENS_PER_GROUP = 2048 + + +def _size_bin(num_groups: int, m_splits) -> int: + avg_tokens = sum(m_splits) // max(num_groups, 1) + return 0 if avg_tokens < _LARGE_TOKENS_PER_GROUP else 1 + + +def make_route_key(num_groups, m_splits, N, K, out_dtype, layout, in_format) -> RouteKey: + """Derive a coarse, reuse-friendly RouteKey from the raw call parameters. Pure.""" + return RouteKey( + num_groups=num_groups, + N=N, + K=K, + out_dtype=out_dtype, + layout=layout, + size_bin=_size_bin(num_groups, m_splits), + in_format=in_format, + ) + + +@runtime_checkable +class Candidate(Protocol): + name: str + + def available(self, key: RouteKey) -> bool: ... + + def prepare(self, operands: Any) -> Callable[[], Any]: + """Return a zero-arg closure that runs the backend and returns its output. + Called once for the numerics gate (if any), then repeatedly for timing.""" + ... + + +# fn -> milliseconds (lower is better) +Timer = Callable[[Callable[[], Any]], float] +# (output, operands) -> (is_correct, sqnr_db) +Verifier = Callable[[Any, Any], "tuple[bool, float]"] + + +@dataclass +class CandidateReport: + name: str + available: bool + correct: bool | None = None + sqnr_db: float | None = None + time_ms: float | None = None # None when unavailable / incorrect / errored + error: str | None = None + + +@dataclass +class Selection: + key: RouteKey + winner: str + from_cache: bool + reports: list[CandidateReport] = field(default_factory=list) + + +class AutotuneRouter: + """Empirical per-shape selection, measured once and cached in memory. + + On a cache miss the router filters to the available candidates, optionally + runs each once through the injected verifier (an incorrect candidate is + dropped), times the survivors with the injected timer, and caches the + fastest. With no correct/available candidate it falls back to ``default`` + (the guaranteed-available floor). A ``verifier`` of ``None`` skips the + numerics gate (for use where the candidates are already validated). + """ + + def __init__( + self, + candidates: list[Candidate], + timer: Timer, + verifier: Verifier | None, + default: str, + ): + names = [c.name for c in candidates] + if default not in names: + raise ValueError(f"default {default!r} not among candidates {names}") + if len(names) != len(set(names)): + raise ValueError(f"duplicate candidate names: {names}") + self.candidates = list(candidates) + self.timer = timer + self.verifier = verifier + self.default = default + self._cache: dict[RouteKey, str] = {} + + def cached(self, key: RouteKey) -> str | None: + return self._cache.get(key) + + def select(self, key: RouteKey, operands: Any) -> Selection: + hit = self._cache.get(key) + if hit is not None: + return Selection(key, hit, from_cache=True) + reports = [self._measure(c, key, operands) for c in self.candidates] + winner = self._pick(reports) + self._cache[key] = winner + return Selection(key, winner, from_cache=False, reports=reports) + + def _measure(self, cand: Candidate, key: RouteKey, operands: Any) -> CandidateReport: + if not cand.available(key): + return CandidateReport(cand.name, available=False) + try: + fn = cand.prepare(operands) + correct, sqnr = None, None + if self.verifier is not None: + output = fn() # one real run feeds the numerics gate + correct, sqnr = self.verifier(output, operands) + if not correct: + return CandidateReport(cand.name, True, correct=False, sqnr_db=sqnr) + t = self.timer(fn) + return CandidateReport(cand.name, True, correct=correct, sqnr_db=sqnr, time_ms=t) + except Exception as exc: # a backend that errors drops out; others still rank + return CandidateReport(cand.name, True, error=repr(exc)) + + def _pick(self, reports: list[CandidateReport]) -> str: + ranked = [r for r in reports if r.time_ms is not None] + if not ranked: + return self.default + return min(ranked, key=lambda r: r.time_ms).name diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index f534da5c3..74cb9c28e 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -52,6 +52,7 @@ general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, ) +from ..grouped_gemm_autotune import autotuned_grouped_gemm from ..constants import GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload @@ -616,7 +617,9 @@ def forward( general_grouped_gemm_func = general_grouped_gemm_triton kwargs = {"m_splits_tensor": m_splits_tensor} else: - general_grouped_gemm_func = general_grouped_gemm + # Drop-in that autotunes hipBLASLt vs CK when NVTE_AUTOTUNE_KERNELS=1, + # else delegates to general_grouped_gemm unchanged. + general_grouped_gemm_func = autotuned_grouped_gemm kwargs = {} general_grouped_gemm_func( weights_fp8, @@ -1090,7 +1093,7 @@ def backward( general_grouped_gemm_func = general_grouped_gemm_triton kwargs = {"m_splits_tensor": ctx.m_splits_tensor} else: - general_grouped_gemm_func = general_grouped_gemm + general_grouped_gemm_func = autotuned_grouped_gemm kwargs = {} general_grouped_gemm_func( weights_for_dgrad, @@ -1194,7 +1197,7 @@ def backward( general_grouped_gemm_func = general_grouped_gemm_triton kwargs = {"m_splits_tensor": ctx.m_splits_tensor} else: - general_grouped_gemm_func = general_grouped_gemm + general_grouped_gemm_func = autotuned_grouped_gemm kwargs = {} grouped_gemm_wgrad = functools.partial( general_grouped_gemm_func, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index c4c9318b7..b5c4d1267 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -62,6 +62,7 @@ from ..cpp_extensions import ( general_gemm, ) +from ..gemm_autotune import autotuned_gemm from ..constants import FP8BwdTensorIdx, FP8FwdTensorIdx, GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..graph import is_graph_capturing @@ -516,7 +517,7 @@ def _linear_forward_impl( "Expected _transpose to be None or an empty tensor when transpose cache is disabled." nvtx_range_push(f"{nvtx_label}.gemm") - gemm_out, *_, reduce_scatter_out = general_gemm( + gemm_out, *_, reduce_scatter_out = autotuned_gemm( weightmat, inputmat_total, quantization_params=output_quantizer, @@ -1040,7 +1041,7 @@ def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], .. weight_for_dgrad = saved_weight if isinstance(weight_for_dgrad, QuantizedTensorStorage): weight_for_dgrad = weight_for_dgrad.dequantize(dtype=bwd_args.activation_dtype) - gemm_out, *_, reduce_scatter_out = general_gemm( + gemm_out, *_, reduce_scatter_out = autotuned_gemm( weight_for_dgrad, grad_output, layout="NN", @@ -1238,7 +1239,7 @@ def wgrad_gemm( """ nvtx_range_push(f"{nvtx_label}.wgrad_gemm") - dw, db, *_ = general_gemm(x, dy, **wgrad_gemm_kwargs) + dw, db, *_ = autotuned_gemm(x, dy, **wgrad_gemm_kwargs) nvtx_range_pop(f"{nvtx_label}.wgrad_gemm") return dw, db