From c1ea28f5b9df1c2e389e91e54a623173394f4650 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Mon, 17 Aug 2026 17:13:51 -0500 Subject: [PATCH 1/8] autotune grouped gemm kernels --- .../pytorch/grouped_gemm_autotune.py | 265 ++++++++++++++++++ transformer_engine/pytorch/kernel_router.py | 179 ++++++++++++ .../pytorch/module/grouped_linear.py | 129 ++++++--- 3 files changed, 541 insertions(+), 32 deletions(-) create mode 100644 transformer_engine/pytorch/grouped_gemm_autotune.py create mode 100644 transformer_engine/pytorch/kernel_router.py diff --git a/transformer_engine/pytorch/grouped_gemm_autotune.py b/transformer_engine/pytorch/grouped_gemm_autotune.py new file mode 100644 index 0000000000..fc467759c3 --- /dev/null +++ b/transformer_engine/pytorch/grouped_gemm_autotune.py @@ -0,0 +1,265 @@ +# 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 bf16 grouped GEMM (forward + backward). + +When ``NVTE_AUTOTUNE=1`` (ROCm only), each bf16 grouped GEMM issued through +``general_grouped_gemm`` -- forward (TN), dgrad (NN), and wgrad (NT) -- picks +between the two C++ backends that share that entry point and semantics +(multi-stream hipBLASLt and CK) by measuring them once per shape+layout (via +``triton.testing.do_bench``) and caching the winner. 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`` / +``NVTE_AUTOTUNE_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. +* Only the two backends reachable through ``general_grouped_gemm`` are + candidates. The Triton grouped GEMM is excluded (different input layout and + backward) and stays behind its ``NVTE_USE_GROUPED_GEMM_TRITON`` flag. +* 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. +* CK requires ``num_gemms > 1`` and bf16/fp16; otherwise it is unavailable and + the router falls back to multi-stream hipBLASLt (the guaranteed floor). +* Off by default; ``NVTE_AUTOTUNE_VERBOSE=1`` adds per-call selection logging + (cache hit/miss, route key, per-backend timings, winner). + +Note: selection is per-rank. For pure GEMM a divergent per-rank choice is a perf +skew, not a hang (no collectives), but a future productionization should make the +choice rank-consistent (tune on rank 0, broadcast). +""" +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 + +_FLOAT16_KEYS = ("torch.bfloat16", "torch.float16") + +# Op-agnostic autotune switches, shared by any future autotuned op. +_MASTER_ENV = "NVTE_AUTOTUNE" +_VERBOSE_ENV = "NVTE_AUTOTUNE_VERBOSE" + + +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} dtype={key.dtype} " + f"layout={key.layout} m_bucket={key.total_m_bucket} imbal={key.imbalance_bucket}" + ) + if sel.from_cache: + line = f"[gg-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"[gg-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) + + +class _GroupedGemmBackend: + """A grouped-GEMM backend reached through ``general_grouped_gemm``, selected + by transiently toggling the CK env vars.""" + + def __init__(self, name: str, use_ck: bool): + self.name = name + self._use_ck = use_ck + + def available(self, key: RouteKey) -> bool: + if self._use_ck: + return key.num_groups > 1 and key.dtype in _FLOAT16_KEYS + return True # multi-stream hipBLASLt is the guaranteed-available floor + + def _env_overrides(self): + if self._use_ck: + return {"NVTE_USE_CUTLASS_GROUPED_GEMM": "1", "NVTE_USE_CK_GROUPED_GEMM": "1"} + return {"NVTE_USE_CUTLASS_GROUPED_GEMM": None, "NVTE_USE_CK_GROUPED_GEMM": None} + + 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(). + overrides = self._env_overrides() + scratch = [torch.empty_like(o) for o in call.out] + timing_kwargs = dict(call.gemm_kwargs) + timing_kwargs["accumulate"] = False + + def run(): + with _env(**overrides): + 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_overrides()): + 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, + ) + + +_router: AutotuneRouter | None = None +_backends: dict = {} + + +def _do_bench_ms(fn) -> float: + import triton + + 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", use_ck=False), + _GroupedGemmBackend("ck", use_ck=True), + ] + _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 + + +def maybe_autotune_grouped_gemm( + A, + B, + out, + quantization_params, + out_dtype, + *, + m_splits, + layout, + N, + K, + **gemm_kwargs, +): + """If autotune is enabled, select the fastest backend for this grouped GEMM + (identified by shape + ``layout``) and run it into ``out``. + + Returns ``(handled, result)``: ``result`` is the ``general_grouped_gemm`` + return value (wgrad uses its grad-bias output). Returns ``(False, None)`` when + disabled, so the caller runs its normal path. bf16 C++ path only. + """ + if not _autotune_enabled(): + return False, None + + 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(call.num_groups, tuple(call.m_splits), N, K, str(out_dtype), layout) + router = _get_router() + sel = router.select(key, call) + if _verbose(): + _log_selection(sel) + # select() measures into scratch on a miss and only returns a name on a hit; + # run the chosen backend once into the real `out` to produce the result. + result = _backends[sel.winner].run_real(call) + return True, result diff --git a/transformer_engine/pytorch/kernel_router.py b/transformer_engine/pytorch/kernel_router.py new file mode 100644 index 0000000000..6c80e312cb --- /dev/null +++ b/transformer_engine/pytorch/kernel_router.py @@ -0,0 +1,179 @@ +# 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 of parameters a selection is keyed on. No live tensors. + + Structural fields (``num_groups``, ``N``, ``K``, ``dtype``, ``layout``) are + exact -- they are stable per layer. The token distribution is captured + *coarsely* via ``total_m_bucket`` and ``imbalance_bucket`` (see + :func:`make_route_key`) so nearby steps -- whose exact per-expert token counts + differ every iteration -- share one cached decision instead of forcing a + re-measure on every call. This trades a little routing precision for reuse; + the bucket granularity is the knob for that trade. + """ + + num_groups: int + N: int + K: int + dtype: str + layout: str + total_m_bucket: int + imbalance_bucket: int + + +def _next_pow2(x: int) -> int: + return 1 if x <= 1 else 1 << (x - 1).bit_length() + + +def _imbalance_bucket(m_splits) -> int: + """Coefficient-of-variation of the per-group token counts, bucketed: + 0 = near-balanced, 1 = moderate skew, 2 = high skew. Multi-stream backends + (one GEMM per expert) are the most sensitive to this, so it belongs in the key.""" + n = len(m_splits) + if n <= 1: + return 0 + mean = sum(m_splits) / n + if mean == 0: + return 0 + var = sum((m - mean) ** 2 for m in m_splits) / n + cv = (var**0.5) / mean + if cv < 0.1: + return 0 + if cv < 0.5: + return 1 + return 2 + + +def make_route_key(num_groups, m_splits, N, K, dtype, layout) -> RouteKey: + """Derive a coarse, reusable RouteKey from the raw call parameters. Pure.""" + return RouteKey( + num_groups=num_groups, + N=N, + K=K, + dtype=dtype, + layout=layout, + total_m_bucket=_next_pow2(sum(m_splits)), + imbalance_bucket=_imbalance_bucket(m_splits), + ) + + +@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 f534da5c3b..0e7894fca6 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 maybe_autotune_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 @@ -618,19 +619,40 @@ def forward( else: general_grouped_gemm_func = general_grouped_gemm kwargs = {} - general_grouped_gemm_func( - weights_fp8, - inputmats, - [out], - output_quantizers, - activation_dtype, - single_output=True, - m_splits=m_splits, - bias=biases, - use_bias=use_bias, - use_split_accumulator=use_split_accumulator, - **kwargs, - ) + # Opt-in fprop autotune (NVTE_AUTOTUNE=1, ROCm bf16 path): pick + # multi-stream hipBLASLt vs CK per shape and run into `out`. Returns + # (False, None) unless enabled, leaving the normal path below unchanged. + handled = False + if general_grouped_gemm_func is general_grouped_gemm and not (fp8 or debug): + handled, _ = maybe_autotune_grouped_gemm( + weights_fp8, + inputmats, + [out], + output_quantizers, + activation_dtype, + m_splits=m_splits, + layout="TN", + N=weights_fp8[0].size(0), + K=weights_fp8[0].size(1), + single_output=True, + bias=biases, + use_bias=use_bias, + use_split_accumulator=use_split_accumulator, + ) + if not handled: + general_grouped_gemm_func( + weights_fp8, + inputmats, + [out], + output_quantizers, + activation_dtype, + single_output=True, + m_splits=m_splits, + bias=biases, + use_bias=use_bias, + use_split_accumulator=use_split_accumulator, + **kwargs, + ) output_unpadded = False @@ -1092,19 +1114,37 @@ def backward( else: general_grouped_gemm_func = general_grouped_gemm kwargs = {} - general_grouped_gemm_func( - weights_for_dgrad, - grad_output, - [dgrad], - ctx.grad_input_quantizers, - ctx.activation_dtype, - single_output=True, - layout="NN", - m_splits=ctx.m_splits, - grad=True, - use_split_accumulator=dgrad_gemm_use_split_accumulator, - **kwargs, - ) + # Opt-in dgrad autotune (NVTE_AUTOTUNE=1, ROCm bf16 path). + dgrad_handled = False + if general_grouped_gemm_func is general_grouped_gemm and not (ctx.fp8 or ctx.debug): + dgrad_handled, _ = maybe_autotune_grouped_gemm( + weights_for_dgrad, + grad_output, + [dgrad], + ctx.grad_input_quantizers, + ctx.activation_dtype, + m_splits=ctx.m_splits, + layout="NN", + N=weights_for_dgrad[0].size(0), + K=weights_for_dgrad[0].size(1), + single_output=True, + grad=True, + use_split_accumulator=dgrad_gemm_use_split_accumulator, + ) + if not dgrad_handled: + general_grouped_gemm_func( + weights_for_dgrad, + grad_output, + [dgrad], + ctx.grad_input_quantizers, + ctx.activation_dtype, + single_output=True, + layout="NN", + m_splits=ctx.m_splits, + grad=True, + use_split_accumulator=dgrad_gemm_use_split_accumulator, + **kwargs, + ) if ctx.actual_m_splits is not None and ctx.actual_m_splits != ctx.m_splits \ and not ctx.output_unpadded: @@ -1196,6 +1236,11 @@ def backward( else: general_grouped_gemm_func = general_grouped_gemm kwargs = {} + wgrad_accumulate = ( + accumulate_wgrad_into_param_main_grad + if not getattr(ctx, "origin_weights_overwrite_main_grad", False) + else False + ) grouped_gemm_wgrad = functools.partial( general_grouped_gemm_func, quantization_params=ctx.grad_weight_quantizers, @@ -1206,18 +1251,38 @@ def backward( use_bias=ctx.use_bias if grad_biases[0] is None else None, bias=biases, use_split_accumulator=wgrad_gemm_use_split_accumulator, - accumulate=( - accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "origin_weights_overwrite_main_grad", False) - else False - ), + accumulate=wgrad_accumulate, **kwargs, ) # WGRAD if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): ctx.wgrad_store.put([inputmats, grad_output, wgrad_list], grouped_gemm_wgrad) else: - _, grad_biases_, _ = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list) + # Opt-in wgrad autotune (NVTE_AUTOTUNE=1, ROCm bf16 path). The + # deferred-wgrad branch above is left untouched. + wgrad_handled = False + if general_grouped_gemm_func is general_grouped_gemm and not ( + ctx.fp8 or ctx.debug + ): + wgrad_handled, wgrad_result = maybe_autotune_grouped_gemm( + inputmats, + grad_output, + wgrad_list, + ctx.grad_weight_quantizers, + ctx.activation_dtype, + m_splits=ctx.m_splits, + layout="NT", + N=wgrad_list[0].size(0), + K=wgrad_list[0].size(1), + grad=True, + use_bias=ctx.use_bias if grad_biases[0] is None else None, + bias=biases, + use_split_accumulator=wgrad_gemm_use_split_accumulator, + accumulate=wgrad_accumulate, + ) + if not wgrad_handled: + wgrad_result = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list) + _, grad_biases_, _ = wgrad_result for i in range(ctx.num_gemms): if grad_biases[i] is None: From ca4e16502a8e7919269d805ae89d576cae65c68d Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 18 Aug 2026 10:23:52 -0500 Subject: [PATCH 2/8] rename env --- .../pytorch/grouped_gemm_autotune.py | 15 ++++++--------- .../pytorch/module/grouped_linear.py | 6 +++--- 2 files changed, 9 insertions(+), 12 deletions(-) diff --git a/transformer_engine/pytorch/grouped_gemm_autotune.py b/transformer_engine/pytorch/grouped_gemm_autotune.py index fc467759c3..95a89dbaba 100644 --- a/transformer_engine/pytorch/grouped_gemm_autotune.py +++ b/transformer_engine/pytorch/grouped_gemm_autotune.py @@ -3,7 +3,7 @@ """Opt-in on-the-fly autotuning for bf16 grouped GEMM (forward + backward). -When ``NVTE_AUTOTUNE=1`` (ROCm only), each bf16 grouped GEMM issued through +When ``NVTE_AUTOTUNE_KERNELS=1`` (ROCm only), each bf16 grouped GEMM issued through ``general_grouped_gemm`` -- forward (TN), dgrad (NN), and wgrad (NT) -- picks between the two C++ backends that share that entry point and semantics (multi-stream hipBLASLt and CK) by measuring them once per shape+layout (via @@ -11,8 +11,8 @@ :mod:`kernel_router`; this module supplies the backend candidates and the process-global router. -The env vars are intentionally op-agnostic (``NVTE_AUTOTUNE`` / -``NVTE_AUTOTUNE_VERBOSE``) so the same switches govern future autotuned ops. +The env vars are intentionally op-agnostic (``NVTE_AUTOTUNE_KERNELS`` / +``NVTE_AUTOTUNE_KERNELS_VERBOSE``) so the same switches govern future autotuned ops. Scope / safety: @@ -29,12 +29,9 @@ * The deferred-wgrad path (``wgrad_store``) is left untouched. * CK requires ``num_gemms > 1`` and bf16/fp16; otherwise it is unavailable and the router falls back to multi-stream hipBLASLt (the guaranteed floor). -* Off by default; ``NVTE_AUTOTUNE_VERBOSE=1`` adds per-call selection logging +* Off by default; ``NVTE_AUTOTUNE_KERNELS_VERBOSE=1`` adds per-call selection logging (cache hit/miss, route key, per-backend timings, winner). -Note: selection is per-rank. For pure GEMM a divergent per-rank choice is a perf -skew, not a hang (no collectives), but a future productionization should make the -choice rank-consistent (tune on rank 0, broadcast). """ from __future__ import annotations @@ -51,8 +48,8 @@ _FLOAT16_KEYS = ("torch.bfloat16", "torch.float16") # Op-agnostic autotune switches, shared by any future autotuned op. -_MASTER_ENV = "NVTE_AUTOTUNE" -_VERBOSE_ENV = "NVTE_AUTOTUNE_VERBOSE" +_MASTER_ENV = "NVTE_AUTOTUNE_KERNELS" +_VERBOSE_ENV = "NVTE_AUTOTUNE_KERNELS_VERBOSE" def _autotune_enabled() -> bool: diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 0e7894fca6..8d72951d8f 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -619,7 +619,7 @@ def forward( else: general_grouped_gemm_func = general_grouped_gemm kwargs = {} - # Opt-in fprop autotune (NVTE_AUTOTUNE=1, ROCm bf16 path): pick + # Opt-in fprop autotune (NVTE_AUTOTUNE_KERNELS=1, ROCm bf16 path): pick # multi-stream hipBLASLt vs CK per shape and run into `out`. Returns # (False, None) unless enabled, leaving the normal path below unchanged. handled = False @@ -1114,7 +1114,7 @@ def backward( else: general_grouped_gemm_func = general_grouped_gemm kwargs = {} - # Opt-in dgrad autotune (NVTE_AUTOTUNE=1, ROCm bf16 path). + # Opt-in dgrad autotune (NVTE_AUTOTUNE_KERNELS=1, ROCm bf16 path). dgrad_handled = False if general_grouped_gemm_func is general_grouped_gemm and not (ctx.fp8 or ctx.debug): dgrad_handled, _ = maybe_autotune_grouped_gemm( @@ -1258,7 +1258,7 @@ def backward( if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): ctx.wgrad_store.put([inputmats, grad_output, wgrad_list], grouped_gemm_wgrad) else: - # Opt-in wgrad autotune (NVTE_AUTOTUNE=1, ROCm bf16 path). The + # Opt-in wgrad autotune (NVTE_AUTOTUNE_KERNELS=1, ROCm bf16 path). The # deferred-wgrad branch above is left untouched. wgrad_handled = False if general_grouped_gemm_func is general_grouped_gemm and not ( From 6333b3f11c7451aee75f7b9f93c6ffc6f01030c0 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 18 Aug 2026 12:02:36 -0500 Subject: [PATCH 3/8] caorsen bins, simplify GroupedLinear --- .../pytorch/grouped_gemm_autotune.py | 71 +++++---- transformer_engine/pytorch/kernel_router.py | 47 ++---- .../pytorch/module/grouped_linear.py | 138 +++++------------- 3 files changed, 97 insertions(+), 159 deletions(-) diff --git a/transformer_engine/pytorch/grouped_gemm_autotune.py b/transformer_engine/pytorch/grouped_gemm_autotune.py index 95a89dbaba..a2ad23521e 100644 --- a/transformer_engine/pytorch/grouped_gemm_autotune.py +++ b/transformer_engine/pytorch/grouped_gemm_autotune.py @@ -71,7 +71,7 @@ def _log_selection(sel) -> None: key = sel.key ks = ( f"G={key.num_groups} N={key.N} K={key.K} dtype={key.dtype} " - f"layout={key.layout} m_bucket={key.total_m_bucket} imbal={key.imbalance_bucket}" + f"layout={key.layout} size_bin={key.size_bin}" ) if sel.from_cache: line = f"[gg-autotune] cache HIT [{ks}] -> {sel.winner}" @@ -216,29 +216,46 @@ def _get_router() -> AutotuneRouter: return _router -def maybe_autotune_grouped_gemm( - A, - B, - out, - quantization_params, - out_dtype, - *, - m_splits, - layout, - N, - K, - **gemm_kwargs, -): - """If autotune is enabled, select the fastest backend for this grouped GEMM - (identified by shape + ``layout``) and run it into ``out``. - - Returns ``(handled, result)``: ``result`` is the ``general_grouped_gemm`` - return value (wgrad uses its grad-bias output). Returns ``(False, None)`` when - disabled, so the caller runs its normal path. bf16 C++ path only. +_FLOAT16_DTYPES = (torch.bfloat16, torch.float16) + + +def _eligible(quantization_params, out_dtype) -> bool: + """True on the autotune-enabled bf16/fp16 C++ path with no quantizers (which + excludes fp8 and debug, whose quantizers are non-None).""" + return ( + _autotune_enabled() + and out_dtype in _FLOAT16_DTYPES + and all(q is None for q in quantization_params) + ) + + +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``.""" + 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 on the bf16/fp16 C++ path, it + selects the fastest backend (multi-stream hipBLASLt vs CK) for this + shape+layout and runs it. Otherwise -- disabled, CUDA, fp8/debug, or fp32 -- + it delegates to ``general_grouped_gemm`` unchanged, returning its result. """ - if not _autotune_enabled(): - return False, None + if not _eligible(quantization_params, out_dtype): + 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, @@ -251,12 +268,10 @@ def maybe_autotune_grouped_gemm( K=K, gemm_kwargs=gemm_kwargs, ) - key = make_route_key(call.num_groups, tuple(call.m_splits), N, K, str(out_dtype), layout) - router = _get_router() - sel = router.select(key, call) + key = make_route_key(len(m_splits), tuple(m_splits), N, K, str(out_dtype), layout) + sel = _get_router().select(key, call) if _verbose(): _log_selection(sel) - # select() measures into scratch on a miss and only returns a name on a hit; + # 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. - result = _backends[sel.winner].run_real(call) - return True, result + return _backends[sel.winner].run_real(call) diff --git a/transformer_engine/pytorch/kernel_router.py b/transformer_engine/pytorch/kernel_router.py index 6c80e312cb..f00bf46504 100644 --- a/transformer_engine/pytorch/kernel_router.py +++ b/transformer_engine/pytorch/kernel_router.py @@ -23,15 +23,14 @@ @dataclass(frozen=True) class RouteKey: - """The GPU-free bundle of parameters a selection is keyed on. No live tensors. + """The GPU-free bundle a selection is keyed on. No live tensors. Structural fields (``num_groups``, ``N``, ``K``, ``dtype``, ``layout``) are - exact -- they are stable per layer. The token distribution is captured - *coarsely* via ``total_m_bucket`` and ``imbalance_bucket`` (see - :func:`make_route_key`) so nearby steps -- whose exact per-expert token counts - differ every iteration -- share one cached decision instead of forcing a - re-measure on every call. This trades a little routing precision for reuse; - the bucket granularity is the knob for that trade. + exact and stable per layer. 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 @@ -39,43 +38,29 @@ class RouteKey: K: int dtype: str layout: str - total_m_bucket: int - imbalance_bucket: int + size_bin: int -def _next_pow2(x: int) -> int: - return 1 if x <= 1 else 1 << (x - 1).bit_length() +# 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 _imbalance_bucket(m_splits) -> int: - """Coefficient-of-variation of the per-group token counts, bucketed: - 0 = near-balanced, 1 = moderate skew, 2 = high skew. Multi-stream backends - (one GEMM per expert) are the most sensitive to this, so it belongs in the key.""" - n = len(m_splits) - if n <= 1: - return 0 - mean = sum(m_splits) / n - if mean == 0: - return 0 - var = sum((m - mean) ** 2 for m in m_splits) / n - cv = (var**0.5) / mean - if cv < 0.1: - return 0 - if cv < 0.5: - return 1 - return 2 +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, dtype, layout) -> RouteKey: - """Derive a coarse, reusable RouteKey from the raw call parameters. Pure.""" + """Derive a coarse, reuse-friendly RouteKey from the raw call parameters. Pure.""" return RouteKey( num_groups=num_groups, N=N, K=K, dtype=dtype, layout=layout, - total_m_bucket=_next_pow2(sum(m_splits)), - imbalance_bucket=_imbalance_bucket(m_splits), + size_bin=_size_bin(num_groups, m_splits), ) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 8d72951d8f..74cb9c28ea 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -52,7 +52,7 @@ general_grouped_gemm, general_grouped_gemm_for_grouped_tensor, ) -from ..grouped_gemm_autotune import maybe_autotune_grouped_gemm +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 @@ -617,42 +617,23 @@ 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 = {} - # Opt-in fprop autotune (NVTE_AUTOTUNE_KERNELS=1, ROCm bf16 path): pick - # multi-stream hipBLASLt vs CK per shape and run into `out`. Returns - # (False, None) unless enabled, leaving the normal path below unchanged. - handled = False - if general_grouped_gemm_func is general_grouped_gemm and not (fp8 or debug): - handled, _ = maybe_autotune_grouped_gemm( - weights_fp8, - inputmats, - [out], - output_quantizers, - activation_dtype, - m_splits=m_splits, - layout="TN", - N=weights_fp8[0].size(0), - K=weights_fp8[0].size(1), - single_output=True, - bias=biases, - use_bias=use_bias, - use_split_accumulator=use_split_accumulator, - ) - if not handled: - general_grouped_gemm_func( - weights_fp8, - inputmats, - [out], - output_quantizers, - activation_dtype, - single_output=True, - m_splits=m_splits, - bias=biases, - use_bias=use_bias, - use_split_accumulator=use_split_accumulator, - **kwargs, - ) + general_grouped_gemm_func( + weights_fp8, + inputmats, + [out], + output_quantizers, + activation_dtype, + single_output=True, + m_splits=m_splits, + bias=biases, + use_bias=use_bias, + use_split_accumulator=use_split_accumulator, + **kwargs, + ) output_unpadded = False @@ -1112,39 +1093,21 @@ 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 = {} - # Opt-in dgrad autotune (NVTE_AUTOTUNE_KERNELS=1, ROCm bf16 path). - dgrad_handled = False - if general_grouped_gemm_func is general_grouped_gemm and not (ctx.fp8 or ctx.debug): - dgrad_handled, _ = maybe_autotune_grouped_gemm( - weights_for_dgrad, - grad_output, - [dgrad], - ctx.grad_input_quantizers, - ctx.activation_dtype, - m_splits=ctx.m_splits, - layout="NN", - N=weights_for_dgrad[0].size(0), - K=weights_for_dgrad[0].size(1), - single_output=True, - grad=True, - use_split_accumulator=dgrad_gemm_use_split_accumulator, - ) - if not dgrad_handled: - general_grouped_gemm_func( - weights_for_dgrad, - grad_output, - [dgrad], - ctx.grad_input_quantizers, - ctx.activation_dtype, - single_output=True, - layout="NN", - m_splits=ctx.m_splits, - grad=True, - use_split_accumulator=dgrad_gemm_use_split_accumulator, - **kwargs, - ) + general_grouped_gemm_func( + weights_for_dgrad, + grad_output, + [dgrad], + ctx.grad_input_quantizers, + ctx.activation_dtype, + single_output=True, + layout="NN", + m_splits=ctx.m_splits, + grad=True, + use_split_accumulator=dgrad_gemm_use_split_accumulator, + **kwargs, + ) if ctx.actual_m_splits is not None and ctx.actual_m_splits != ctx.m_splits \ and not ctx.output_unpadded: @@ -1234,13 +1197,8 @@ 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 = {} - wgrad_accumulate = ( - accumulate_wgrad_into_param_main_grad - if not getattr(ctx, "origin_weights_overwrite_main_grad", False) - else False - ) grouped_gemm_wgrad = functools.partial( general_grouped_gemm_func, quantization_params=ctx.grad_weight_quantizers, @@ -1251,38 +1209,18 @@ def backward( use_bias=ctx.use_bias if grad_biases[0] is None else None, bias=biases, use_split_accumulator=wgrad_gemm_use_split_accumulator, - accumulate=wgrad_accumulate, + accumulate=( + accumulate_wgrad_into_param_main_grad + if not getattr(ctx, "origin_weights_overwrite_main_grad", False) + else False + ), **kwargs, ) # WGRAD if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): ctx.wgrad_store.put([inputmats, grad_output, wgrad_list], grouped_gemm_wgrad) else: - # Opt-in wgrad autotune (NVTE_AUTOTUNE_KERNELS=1, ROCm bf16 path). The - # deferred-wgrad branch above is left untouched. - wgrad_handled = False - if general_grouped_gemm_func is general_grouped_gemm and not ( - ctx.fp8 or ctx.debug - ): - wgrad_handled, wgrad_result = maybe_autotune_grouped_gemm( - inputmats, - grad_output, - wgrad_list, - ctx.grad_weight_quantizers, - ctx.activation_dtype, - m_splits=ctx.m_splits, - layout="NT", - N=wgrad_list[0].size(0), - K=wgrad_list[0].size(1), - grad=True, - use_bias=ctx.use_bias if grad_biases[0] is None else None, - bias=biases, - use_split_accumulator=wgrad_gemm_use_split_accumulator, - accumulate=wgrad_accumulate, - ) - if not wgrad_handled: - wgrad_result = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list) - _, grad_biases_, _ = wgrad_result + _, grad_biases_, _ = grouped_gemm_wgrad(inputmats, grad_output, wgrad_list) for i in range(ctx.num_gemms): if grad_biases[i] is None: From 6072abd815686e924737ede84197c90bfdf25b25 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 18 Aug 2026 14:17:31 -0500 Subject: [PATCH 4/8] add Triton TN/NN --- .../pytorch/grouped_gemm_autotune.py | 85 +++++++++++++++++-- 1 file changed, 77 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/grouped_gemm_autotune.py b/transformer_engine/pytorch/grouped_gemm_autotune.py index a2ad23521e..36a375e538 100644 --- a/transformer_engine/pytorch/grouped_gemm_autotune.py +++ b/transformer_engine/pytorch/grouped_gemm_autotune.py @@ -4,11 +4,12 @@ """Opt-in on-the-fly autotuning for bf16 grouped GEMM (forward + backward). When ``NVTE_AUTOTUNE_KERNELS=1`` (ROCm only), each bf16 grouped GEMM issued through -``general_grouped_gemm`` -- forward (TN), dgrad (NN), and wgrad (NT) -- picks -between the two C++ backends that share that entry point and semantics -(multi-stream hipBLASLt and CK) by measuring them once per shape+layout (via -``triton.testing.do_bench``) and caching the winner. The pure selection lives in -:mod:`kernel_router`; this module supplies the backend candidates and the +``general_grouped_gemm`` -- forward (TN), dgrad (NN), and wgrad (NT) -- picks the +fastest backend for its shape+layout 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), plus -- for TN/NN -- the Triton grouped GEMM. 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`` / @@ -18,9 +19,11 @@ * Forward, dgrad, and wgrad grouped GEMMs are routed; each layout is a distinct route key, measured independently. -* Only the two backends reachable through ``general_grouped_gemm`` are - candidates. The Triton grouped GEMM is excluded (different input layout and - backward) and stays behind its ``NVTE_USE_GROUPED_GEMM_TRITON`` flag. +* The two C++ backends (hipBLASLt, CK) are always candidates. The Triton grouped + GEMM is also a candidate for forward (TN) and dgrad (NN) when importable; it is + excluded from wgrad (NT), whose 3D packed output it cannot target here. The + separate global ``NVTE_USE_GROUPED_GEMM_TRITON`` flag (which forces Triton for + all layouts) is independent of and mutually exclusive with this path. * 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. @@ -51,6 +54,10 @@ _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" @@ -74,6 +81,7 @@ def _log_selection(sel) -> None: f"layout={key.layout} size_bin={key.size_bin}" ) if sel.from_cache: + return line = f"[gg-autotune] cache HIT [{ks}] -> {sel.winner}" else: tried = [] @@ -183,6 +191,59 @@ def run_real(self, call: _GGCall): ) +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. + + Limited to forward (TN) and dgrad (NN), whose operands Triton accepts as-is + (it concatenates the per-group inputs internally -- a real cost in this + pre-split path, so it is left inside the timing). wgrad (NT) is excluded: + Triton wgrad expects a 3D packed output this path does not allocate. An error + on any shape drops it from the ranking. + """ + + name = "triton" + + def available(self, key: RouteKey) -> bool: + return ( + key.layout in ("TN", "NN") + and key.dtype in _FLOAT16_KEYS + and _triton_grouped_gemm() is not None + ) + + 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(): + fn( + call.A, call.B, scratch, call.quantization_params, call.out_dtype, + m_splits=list(call.m_splits), layout=call.layout, **call.gemm_kwargs, + ) + + return run + + def run_real(self, call: _GGCall): + fn = _triton_grouped_gemm() + assert fn is not None # guaranteed by available() + return fn( + call.A, call.B, call.out, call.quantization_params, call.out_dtype, + m_splits=list(call.m_splits), layout=call.layout, **call.gemm_kwargs, + ) + + _router: AutotuneRouter | None = None _backends: dict = {} @@ -190,6 +251,13 @@ def run_real(self, call: _GGCall): 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: @@ -203,6 +271,7 @@ def _get_router() -> AutotuneRouter: candidates = [ _GroupedGemmBackend("hipblaslt", use_ck=False), _GroupedGemmBackend("ck", use_ck=True), + _TritonGroupedGemmBackend(), ] _backends = {c.name: c for c in candidates} # verifier=None: both backends are already validated by TE's grouped-GEMM From 98bf3f1e58ad5b91c7760f346c97cb505cf94a65 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 18 Aug 2026 14:32:58 -0500 Subject: [PATCH 5/8] add Triton NT --- .../pytorch/grouped_gemm_autotune.py | 60 +++++++++++++------ 1 file changed, 41 insertions(+), 19 deletions(-) diff --git a/transformer_engine/pytorch/grouped_gemm_autotune.py b/transformer_engine/pytorch/grouped_gemm_autotune.py index 36a375e538..0d1bb462ee 100644 --- a/transformer_engine/pytorch/grouped_gemm_autotune.py +++ b/transformer_engine/pytorch/grouped_gemm_autotune.py @@ -8,7 +8,7 @@ fastest backend for its shape+layout 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), plus -- for TN/NN -- the Triton grouped GEMM. The pure selection +CK), plus the Triton grouped GEMM. The pure selection lives in :mod:`kernel_router`; this module supplies the backend candidates and the process-global router. @@ -20,10 +20,12 @@ * Forward, dgrad, and wgrad grouped GEMMs are routed; each layout is a distinct route key, measured independently. * The two C++ backends (hipBLASLt, CK) are always candidates. The Triton grouped - GEMM is also a candidate for forward (TN) and dgrad (NN) when importable; it is - excluded from wgrad (NT), whose 3D packed output it cannot target here. The - separate global ``NVTE_USE_GROUPED_GEMM_TRITON`` flag (which forces Triton for - all layouts) is independent of and mutually exclusive with this path. + GEMM is also a candidate for all three layouts when importable: TN/NN run it + directly; NT (wgrad) runs it into a fresh 3D packed buffer (as the + ``use_grouped_gemm_triton`` path allocates) and copies/accumulates the result + back into the per-group output list this path uses. The separate global + ``NVTE_USE_GROUPED_GEMM_TRITON`` flag (which forces Triton for all layouts) is + independent of and mutually exclusive with this path. * 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. @@ -206,42 +208,62 @@ def _triton_grouped_gemm(): class _TritonGroupedGemmBackend: """Triton (AITER) grouped GEMM as an autotune candidate. - Limited to forward (TN) and dgrad (NN), whose operands Triton accepts as-is - (it concatenates the per-group inputs internally -- a real cost in this - pre-split path, so it is left inside the timing). wgrad (NT) is excluded: - Triton wgrad expects a 3D packed output this path does not allocate. An error - on any shape drops it from the ranking. + 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.layout in ("TN", "NN") + key.layout in ("TN", "NN", "NT") and key.dtype in _FLOAT16_KEYS 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(): - fn( - call.A, call.B, scratch, call.quantization_params, call.out_dtype, - m_splits=list(call.m_splits), layout=call.layout, **call.gemm_kwargs, - ) + 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 fn( - call.A, call.B, call.out, call.quantization_params, call.out_dtype, - m_splits=list(call.m_splits), layout=call.layout, **call.gemm_kwargs, - ) + return self._run_into(fn, call, call.out) _router: AutotuneRouter | None = None From 69689f3ea6c52915cc9078093b710af31cabeb0f Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Tue, 18 Aug 2026 15:57:59 -0500 Subject: [PATCH 6/8] add fp8/mxfp8 --- .../pytorch/grouped_gemm_autotune.py | 193 +++++++++++++----- transformer_engine/pytorch/kernel_router.py | 14 +- 2 files changed, 150 insertions(+), 57 deletions(-) diff --git a/transformer_engine/pytorch/grouped_gemm_autotune.py b/transformer_engine/pytorch/grouped_gemm_autotune.py index 0d1bb462ee..5147eb95ae 100644 --- a/transformer_engine/pytorch/grouped_gemm_autotune.py +++ b/transformer_engine/pytorch/grouped_gemm_autotune.py @@ -1,16 +1,16 @@ # 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 bf16 grouped GEMM (forward + backward). +"""Opt-in on-the-fly autotuning for grouped GEMM (bf16/fp8/mxfp8, fwd + bwd). -When ``NVTE_AUTOTUNE_KERNELS=1`` (ROCm only), each bf16 grouped GEMM issued through +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 by measuring the candidates once (via +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), plus the Triton grouped GEMM. The pure selection -lives in :mod:`kernel_router`; this module supplies the backend candidates and the -process-global router. +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. @@ -19,21 +19,29 @@ * Forward, dgrad, and wgrad grouped GEMMs are routed; each layout is a distinct route key, measured independently. -* The two C++ backends (hipBLASLt, CK) are always candidates. The Triton grouped - GEMM is also a candidate for all three layouts when importable: TN/NN run it - directly; NT (wgrad) runs it into a fresh 3D packed buffer (as the - ``use_grouped_gemm_triton`` path allocates) and copies/accumulates the result - back into the per-group output list this path uses. The separate global - ``NVTE_USE_GROUPED_GEMM_TRITON`` flag (which forces Triton for all layouts) is - independent of and mutually exclusive with this path. +* 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. -* CK requires ``num_gemms > 1`` and bf16/fp16; otherwise it is unavailable and - the router falls back to multi-stream hipBLASLt (the guaranteed floor). * Off by default; ``NVTE_AUTOTUNE_KERNELS_VERBOSE=1`` adds per-call selection logging (cache hit/miss, route key, per-backend timings, winner). @@ -50,8 +58,6 @@ from .cpp_extensions import general_grouped_gemm from .kernel_router import AutotuneRouter, RouteKey, make_route_key -_FLOAT16_KEYS = ("torch.bfloat16", "torch.float16") - # Op-agnostic autotune switches, shared by any future autotuned op. _MASTER_ENV = "NVTE_AUTOTUNE_KERNELS" _VERBOSE_ENV = "NVTE_AUTOTUNE_KERNELS_VERBOSE" @@ -79,7 +85,7 @@ def _log_selection(sel) -> None: global _last_log key = sel.key ks = ( - f"G={key.num_groups} N={key.N} K={key.K} dtype={key.dtype} " + 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: @@ -149,35 +155,39 @@ 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 the CK env vars.""" + """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: str, use_ck: bool): + def __init__(self, name, env, *, formats, needs_multi=False): self.name = name - self._use_ck = use_ck + self._env = dict(env) + self._formats = frozenset(formats) + self._needs_multi = needs_multi def available(self, key: RouteKey) -> bool: - if self._use_ck: - return key.num_groups > 1 and key.dtype in _FLOAT16_KEYS - return True # multi-stream hipBLASLt is the guaranteed-available floor - - def _env_overrides(self): - if self._use_ck: - return {"NVTE_USE_CUTLASS_GROUPED_GEMM": "1", "NVTE_USE_CK_GROUPED_GEMM": "1"} - return {"NVTE_USE_CUTLASS_GROUPED_GEMM": None, "NVTE_USE_CK_GROUPED_GEMM": None} + 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(). - overrides = self._env_overrides() + 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(**overrides): + 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, @@ -186,7 +196,7 @@ def run(): return run def run_real(self, call: _GGCall): - with _env(**self._env_overrides()): + 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, @@ -221,8 +231,8 @@ class _TritonGroupedGemmBackend: def available(self, key: RouteKey) -> bool: return ( - key.layout in ("TN", "NN", "NT") - and key.dtype in _FLOAT16_KEYS + key.in_format == "bf16" + and key.layout in ("TN", "NN", "NT") and _triton_grouped_gemm() is not None ) @@ -291,8 +301,42 @@ def _get_router() -> AutotuneRouter: global _router, _backends if _router is None: candidates = [ - _GroupedGemmBackend("hipblaslt", use_ck=False), - _GroupedGemmBackend("ck", use_ck=True), + _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", + { + "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} @@ -309,21 +353,64 @@ def _get_router() -> AutotuneRouter: _FLOAT16_DTYPES = (torch.bfloat16, torch.float16) +_quant_cache = None -def _eligible(quantization_params, out_dtype) -> bool: - """True on the autotune-enabled bf16/fp16 C++ path with no quantizers (which - excludes fp8 and debug, whose quantizers are non-None).""" - return ( - _autotune_enabled() - and out_dtype in _FLOAT16_DTYPES - and all(q is None for q in quantization_params) - ) + +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``.""" + 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) @@ -331,12 +418,14 @@ def _key_dims(A, out, layout): 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 on the bf16/fp16 C++ path, it - selects the fastest backend (multi-stream hipBLASLt vs CK) for this - shape+layout and runs it. Otherwise -- disabled, CUDA, fp8/debug, or fp32 -- - it delegates to ``general_grouped_gemm`` unchanged, returning its result. + 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. """ - if not _eligible(quantization_params, out_dtype): + 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") @@ -359,7 +448,7 @@ def autotuned_grouped_gemm(A, B, out, quantization_params, out_dtype, **kwargs): K=K, gemm_kwargs=gemm_kwargs, ) - key = make_route_key(len(m_splits), tuple(m_splits), N, K, str(out_dtype), layout) + 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) diff --git a/transformer_engine/pytorch/kernel_router.py b/transformer_engine/pytorch/kernel_router.py index f00bf46504..88c0a50c8f 100644 --- a/transformer_engine/pytorch/kernel_router.py +++ b/transformer_engine/pytorch/kernel_router.py @@ -25,8 +25,10 @@ class RouteKey: """The GPU-free bundle a selection is keyed on. No live tensors. - Structural fields (``num_groups``, ``N``, ``K``, ``dtype``, ``layout``) are - exact and stable per layer. The token count enters only as a coarse + 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 @@ -36,9 +38,10 @@ class RouteKey: num_groups: int N: int K: int - dtype: str + out_dtype: str layout: str size_bin: int + in_format: str # Coarse small/large split on the average per-group token count. A tunable @@ -52,15 +55,16 @@ def _size_bin(num_groups: int, m_splits) -> int: return 0 if avg_tokens < _LARGE_TOKENS_PER_GROUP else 1 -def make_route_key(num_groups, m_splits, N, K, dtype, layout) -> RouteKey: +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, - dtype=dtype, + out_dtype=out_dtype, layout=layout, size_bin=_size_bin(num_groups, m_splits), + in_format=in_format, ) From e1daca6ebea7a06f9fda55853b77db3996b656d8 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Wed, 19 Aug 2026 14:14:09 -0500 Subject: [PATCH 7/8] adjust verbose output --- transformer_engine/pytorch/grouped_gemm_autotune.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/grouped_gemm_autotune.py b/transformer_engine/pytorch/grouped_gemm_autotune.py index 5147eb95ae..a91dabbab6 100644 --- a/transformer_engine/pytorch/grouped_gemm_autotune.py +++ b/transformer_engine/pytorch/grouped_gemm_autotune.py @@ -90,7 +90,7 @@ def _log_selection(sel) -> None: ) if sel.from_cache: return - line = f"[gg-autotune] cache HIT [{ks}] -> {sel.winner}" + line = f"[te-autotune] cache HIT [{ks}] -> {sel.winner}" else: tried = [] for r in sel.reports: @@ -103,7 +103,7 @@ def _log_selection(sel) -> None: else: tried.append(f"{r.name}={r.time_ms:.4f}ms") line = ( - f"[gg-autotune] cache MISS [{ks}] tried: {', '.join(tried)} " + f"[te-autotune] cache MISS [{ks}] tried: {', '.join(tried)} " f"-> selected {sel.winner}" ) if line != _last_log: @@ -311,7 +311,7 @@ def _get_router() -> AutotuneRouter: formats=_ALL_FORMATS, ), _GroupedGemmBackend( - "ck", + "ck_tile", { "NVTE_USE_CUTLASS_GROUPED_GEMM": "1", "NVTE_USE_CK_GROUPED_GEMM": "1", From c749262398d8eb0ce3043c07c2a2972fe066f124 Mon Sep 17 00:00:00 2001 From: Matthias Diener Date: Fri, 21 Aug 2026 16:46:30 -0500 Subject: [PATCH 8/8] autotune GEMM --- transformer_engine/pytorch/gemm_autotune.py | 248 ++++++++++++++++++++ transformer_engine/pytorch/module/linear.py | 7 +- 2 files changed, 252 insertions(+), 3 deletions(-) create mode 100644 transformer_engine/pytorch/gemm_autotune.py diff --git a/transformer_engine/pytorch/gemm_autotune.py b/transformer_engine/pytorch/gemm_autotune.py new file mode 100644 index 0000000000..86c15c1bff --- /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/module/linear.py b/transformer_engine/pytorch/module/linear.py index 556348fde9..d0363f51bd 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -60,6 +60,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 @@ -514,7 +515,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, @@ -1038,7 +1039,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", @@ -1236,7 +1237,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