Kernels:

kernels-community
/

finegrained-kernels

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on Sep 16

Commit

4a5c820

verified ·

1 Parent(s): 64a16f9

Uploaded using `kernel-builder`.

Browse files

Files changed (24) hide show

build/torch-rocm/__init__.py +48 -0
build/torch-rocm/_ops.py +38 -0
build/torch-rocm/backward.py +1052 -0
build/torch-rocm/batched.py +1947 -0
build/torch-rocm/bayesian_autotuner.py +940 -0
build/torch-rocm/compat.py +498 -0
build/torch-rocm/descriptors.py +216 -0
build/torch-rocm/epilogue.py +639 -0
build/torch-rocm/finegrained_kernels/__init__.py +26 -0
build/torch-rocm/formats.py +357 -0
build/torch-rocm/grouped.py +0 -0
build/torch-rocm/loading/__init__.py +17 -0
build/torch-rocm/loading/scales.py +325 -0
build/torch-rocm/loading/tiles.py +585 -0
build/torch-rocm/matmul.py +1970 -0
build/torch-rocm/metadata.json +48 -0
build/torch-rocm/metadata.json.sigstore +1 -0
build/torch-rocm/mma.py +367 -0
build/torch-rocm/moe.py +864 -0
build/torch-rocm/norm.py +156 -0
build/torch-rocm/pruners.py +1213 -0
build/torch-rocm/quant.py +955 -0
build/torch-rocm/scheduling.py +403 -0
build/torch-rocm/swizzle.py +170 -0

build/torch-rocm/__init__.py ADDED Viewed

	@@ -0,0 +1,48 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from .matmul import matmul_2d
+from .batched import matmul_batched
+from .grouped import matmul_grouped
+from .moe import moe_fused_batched, moe_fused_grouped
+from . import backward
+from .formats import get_supported_act_fns
+from .norm import get_supported_norms, rms_norm_rows
+from .swizzle import swizzle_mx_scales, unswizzle_mx_scales
+from .quant import mxfp4_act_quant, mxfp8_act_quant, nvfp4_act_quant
+__all__ = [
+    # the three GEMM dispatchers (the weight format is read off the tensors; `activation_format`
+    # names the activations', the gate|up fusion rides as kwargs)
+    "matmul_2d",
+    "matmul_batched",
+    "matmul_grouped",
+    # the fused MoE forwards over them
+    "moe_fused_batched",
+    "moe_fused_grouped",
+    # importing it registers the dgrad formulas on the ops, so a forward call differentiates
+    "backward",
+    "get_supported_act_fns",
+    # the per-expert output norm a model can hand the forwards by name, and its standalone op
+    "get_supported_norms",
+    "rms_norm_rows",
+    # load-time helpers: the swizzled scale layout the SM100 scaled-MMA reads, and the row-wise
+    # quantizers a loader uses to quantize weights into the group formats
+    "swizzle_mx_scales",
+    "unswizzle_mx_scales",
+    "mxfp8_act_quant",
+    "mxfp4_act_quant",
+    "nvfp4_act_quant",
+]

build/torch-rocm/_ops.py ADDED Viewed

	@@ -0,0 +1,38 @@

+import torch
+def get_backend() -> str:
+    """Detect the backend by inspecting torch."""
+    import torch
+    if hasattr(torch, "neuron"):
+        # Needs to be sorted before specific Torch builds, since Neuron
+        # extension can be loaded into e.g. CUDA Torch builds.
+        return "neuron"
+    elif torch.version.cuda is not None:
+        return "cuda"
+    elif torch.version.hip is not None:
+        return "rocm"
+    elif torch.backends.mps.is_available():
+        return "metal"
+    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+        return "xpu"
+    else:
+        return "cpu"
+def _find_ops_name() -> str:
+    kernel_name = "finegrained_kernels"
+    unique_id = "1348465"
+    backend = get_backend()
+    return f"_{kernel_name}_{backend}_{unique_id}"
+_OPS_NAME = _find_ops_name()
+ops = getattr(torch.ops, _OPS_NAME)
+def add_op_namespace_prefix(op_name: str) -> str:
+    """
+    Prefix op by namespace.
+    """
+    return f"{_OPS_NAME}::{op_name}"

build/torch-rocm/backward.py ADDED Viewed

	@@ -0,0 +1,1052 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Backward for the quantized ops: dgrad only, registered on the forward ops themselves.
+``torch.library.register_autograd`` on the ops themselves rather than an ``autograd.Function``
+wrapper: the ops take the RAW activation (they quantize it in one pass internally), so the op
+boundary is exactly where a usable gradient lives, and hanging autograd there means an ordinary
+``matmul_2d`` call differentiates — through compile, export and AOTAutograd, with no second
+entry point to keep in sync. ``compile_time_only_triton_op``'s dispatch routes to the op
+whenever an input wants a gradient, so this fires in eager as well.
+dgrad ONLY: ``dA = dY @ B``, contracting N. The quantized weight never receives a gradient, and
+not merely by omission — autograd accumulates into a leaf, and an E4M3 leaf has no ``add``
+("ufunc_add_CUDA not implemented for Float8_e4m3fn"), so an fp8/fp4 tensor cannot be an autograd
+parameter at all. A frozen quantized base with trainable high-precision adapters in front of it
+(QLoRA) is therefore the shape this serves, and the only gradient it needs is the one that flows
+THROUGH the weight to whatever precedes it.
+Block-FP8 and per-tensor formats reach that product on the forward kernel via transposed views
+(their scale grids subdivide both axes). MX/NVFP4 cannot be reoriented by a view — group scales
+bind to K and packed E2M1 packs along K — so they contract N against the weight's natural (N, K)
+tile, dequantized per-tile in-register; see ``backward.py``. Either way the FORWARD-oriented
+weight is the only copy that exists.
+── the gradient products ──
+Forward contracts K (``C = A @ B.T``, both operands K-contiguous). dgrad contracts N::
+    dX[m, k] = sum_n dY[m, n] * W[n, k]
+The weight is STORED ``(N, K)`` K-contiguous, so the ``[BN, BK]`` tile this reduction wants is
+already a natural contiguous tile — n leads, k is contiguous. Nothing is transposed and no second
+quantized copy is kept: the tile is dequantized in-register and fed to a plain ``tl.dot``, exactly
+as the weight-only forward's ``dot`` arm does. That is what makes a 4-bit frozen weight
+differentiable at no memory cost, which is the whole QLoRA case.
+The group axis stays K, which here is the OUTPUT axis rather than the contraction axis. That is
+why ``dot_scaled`` cannot serve this pass — the tensor core's scaled MMA requires the scaled axis
+to be the one it reduces over — so dgrad is plain-dot only. The tuner already prefers ``dot`` over
+``dot_scaled`` on the gate|up tile, so little is given up.
+Format genericity lives in ``_dgrad_weight_tile`` alone. A format is fully described to this
+kernel by three numbers plus the value dtype:
+    SCALE_GROUP_K    scale extent along K   32 (MX)   128 (block-FP8)
+    SCALE_ROW_DIV    rows sharing a scale    1 (MX)   128 (block-FP8)
+    WEIGHT_VALUES_PER_BYTE  2 for packed E2M1, else 1
+so MX (per-row K-groups) and block-FP8 (128x128 blocks) are the same code path at different
+divisors, and no new arm is needed to add a format whose scales tile the same way.
+"""
+import torch
+import triton
+import triton.language as tl
+from .bayesian_autotuner import bayesian_autotune
+from .compat import (
+    add_op_namespace_prefix,
+    compile_time_only_triton_op,
+    compile_time_only_triton_wrap,
+    device_context,
+    get_accelerator_autotuning_configs,
+    sm_count,
+)
+from .descriptors import maybe_descriptor, rebind_dgrad_descriptors, rebind_dgrad_grouped_descriptors
+from .matmul import matmul_2d
+from .pruners import (
+    PATH_ANCHOR_AXES,
+    block_within_dim_pruner,
+    config_filter,
+    descriptor_box_pruner,
+    compose_pruners,
+    smem_pruner,
+    warp_spec_compile_guard_pruner,
+)
+from .formats import global_scale_stride, normalize_global_scale
+from .quant import e2m1_cols_to_bf16, e2m1_cols_to_f32
+from .loading.scales import decode_group_scale
+from .loading.tiles import load_grouped_act_tile, operand_tile_ptrs
+from .scheduling import build_tile_layout, resolve_grouped_tile
+# Internal: these are the gradient PRODUCTS, not API. Backward hangs off the forward ops via
+# ``register_autograd`` (the table at the bottom of this file), so an ordinary forward call
+# differentiates and there is
+# no second entry point to keep in sync. Nothing here is re-exported from the package root.
+@triton.jit
+def _dgrad_load_weight(
+    b_ptrs, b_descriptor, row0, n_off, kb_off,
+    BLOCK_ROWS: tl.constexpr, BLOCK_COLS: tl.constexpr,
+    B_MEMORY_MODE: tl.constexpr, GROUPED: tl.constexpr = False,
+):
+    """One ``[BN, BK_bytes]`` weight tile: the explicit pointer tile, or the host-TMA box at the
+    same absolute offset. The box is a natural contiguous window — unlike the forward's gated
+    tile there is no interleaved span to express — so the descriptor arm needs no doubling.
+    Single return; only the taken arm compiles."""
+    if B_MEMORY_MODE == "pointer":
+        w = tl.load(b_ptrs)
+    elif GROUPED:  # (E, N, K_bytes) box — the expert leads
+        w = tl.reshape(b_descriptor.load([row0, n_off, kb_off]), BLOCK_ROWS, BLOCK_COLS)
+    else:
+        w = b_descriptor.load([n_off, kb_off])
+    return w
+@triton.jit
+def _dgrad_weight_tile(
+    w,
+    w_scale,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    """One natural ``[BLOCK_SIZE_N, BLOCK_SIZE_K]`` weight tile, dequantized. The K-oriented
+    counterpart of ``mx_weight_upcast`` — the group axis is already last here, so unlike the
+    forward no transpose is needed.
+    The scale broadcast happens INSIDE the multiply: the real (contiguous) weight reshapes to
+    ``[BN, ng, g]`` and the ``[BN, ng, 1]`` scale broadcasts across each group's columns.
+    Materializing the broadcast first and reshaping it does not lower — the same trap
+    ``mx_weight_upcast`` documents.
+    ``OUT_DTYPE`` picks the arm, and both are this codebase's measured verdicts:
+    bf16 for the MMA consumers (exact — E2M1/E4M3 codes and power-of-two UE8M0 scales all fit in
+    bf16, and the multiply is a cheap exponent shift), fp32 for the scalar reduce, which widens
+    anyway and for which the bf16/E4M3 hop is pure cost. Single return; only the taken arm
+    compiles."""
+    NG: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K
+    if OUT_DTYPE == tl.float32:  # scalar-reduce arm: unpack straight to fp32, no narrow hop
+        vals = e2m1_cols_to_f32(w) if w.dtype == tl.uint8 else w.to(tl.float32)
+        sc = decode_group_scale(w_scale).to(tl.float32)[:, :, None]
+    else:  # MMA arm: dequantize in bf16, exactly as the forward's weight upcast does
+        vals = e2m1_cols_to_bf16(w) if w.dtype == tl.uint8 else w.to(tl.bfloat16)
+        sc = decode_group_scale(w_scale).to(tl.bfloat16)[:, :, None]
+    tile = vals.reshape(BLOCK_SIZE_N, NG, SCALE_GROUP_K) * sc
+    return tile.reshape(BLOCK_SIZE_N, BLOCK_SIZE_K).to(OUT_DTYPE)
+@bayesian_autotune(
+    # Plain-dot loop over the N reduction — same physics as the full-precision 2D kernel (tile +
+    # WS), no memory-mode axis: the dequantize is in-tile so both operands read through pointers,
+    # and no COMPUTE_MODE axis because the scaled MMA cannot reduce over an unscaled axis.
+    get_accelerator_autotuning_configs(
+        tune_block_nk=True,
+        tune_block_m=True,
+        warp_spec=True,
+        a_memory_modes=("descriptor", "pointer"),
+        b_memory_modes=("descriptor", "pointer"),
+        pre_hook=rebind_dgrad_descriptors,
+    ),
+    ["N", "K", "m_bit_length"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    finite_check_args=("C",),
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            block_within_dim_pruner("N"),  # exact N-loop: no tail mask, constant-stride advance
+            config_filter(  # the scale row advances by BN // SCALE_ROW_DIV — must be whole
+                lambda c, a: c.kwargs["BLOCK_SIZE_N"] % max(a.get("SCALE_ROW_DIV", 1), 1) == 0
+            ),
+            config_filter(  # BK below one scale group gives NG=0 -> tl.arange(0, 0)
+                lambda c, a: c.kwargs["BLOCK_SIZE_K"] % max(a.get("SCALE_GROUP_K", 1), 1) == 0
+            ),
+            warp_spec_compile_guard_pruner(),
+            descriptor_box_pruner("BLOCK_SIZE_K"),
+            smem_pruner("BLOCK_SIZE_K"),
+        )
+    },
+)
+@triton.jit
+def dgrad_matmul_2d_kernel(
+    A,  # (M, N) upstream gradient, raw bf16/fp16
+    ADescriptor,  # host TMA descriptor over A, box (BM, BN); read iff A_MEMORY_MODE != "pointer"
+    B,  # (N, K) quantized weight, K-contiguous — read in its FORWARD orientation
+    BDescriptor,  # host TMA descriptor over B, box (BN, BK_bytes); read iff B_MEMORY_MODE != "pointer"
+    Bs,  # (N // SCALE_ROW_DIV, K // SCALE_GROUP_K) weight scales
+    BsGlobal,  # fp32 NVFP4 second-level weight global; folded on the accumulator, None folds out
+    C,  # (M, K) output gradient
+    # Shape
+    M,
+    N,
+    K,
+    m_bit_length,  # autotune key only (log2 M bucket); unused in body
+    # Strides
+    stride_a_m,
+    stride_a_n,
+    stride_b_n,
+    stride_b_k,
+    stride_bs_n,
+    stride_bs_k,
+    stride_c_m,
+    stride_c_k,
+    # Format
+    SCALE_GROUP_K: tl.constexpr,
+    SCALE_ROW_DIV: tl.constexpr,
+    WEIGHT_VALUES_PER_BYTE: tl.constexpr,
+    # Meta-parameters
+    A_MEMORY_MODE: tl.constexpr,
+    B_MEMORY_MODE: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    WARP_SPEC: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_k = tl.program_id(1)
+    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_k = pid_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
+    offs_kb = tl.minimum(
+        pid_k * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE)
+        + tl.arange(0, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE),
+        K // WEIGHT_VALUES_PER_BYTE - 1,
+    )
+    offs_kg = tl.minimum(
+        (pid_k * BLOCK_SIZE_K) // SCALE_GROUP_K + tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K),
+        K // SCALE_GROUP_K - 1,
+    )
+    m_mask = offs_m < M
+    k_mask = offs_k < K
+    # BLOCK_SIZE_N divides N (block_within_dim_pruner), so the reduction loop is EXACT: no tail
+    # mask on the activation, no clamp on the weight rows, and the addresses advance by a constant
+    # stride instead of being rebuilt from offsets each step — the forward's loop shape.
+    offs_n0 = tl.arange(0, BLOCK_SIZE_N)
+    # SWAP_AB is the WEIGHT's viewpoint, inverted for activations: True = rows-major [BM, BN]
+    dy_ptrs = operand_tile_ptrs(A, offs_m, offs_n0, stride_a_m, stride_a_n, A_MEMORY_MODE, True)
+    b_ptrs = operand_tile_ptrs(B, offs_n0, offs_kb, stride_b_n, stride_b_k, B_MEMORY_MODE, True)
+    bs_ptrs = (
+        Bs + (offs_n0[:, None] // SCALE_ROW_DIV) * stride_bs_n + offs_kg[None, :] * stride_bs_k
+    )
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K), dtype=tl.float32)
+    for n in tl.range(0, tl.cdiv(N, BLOCK_SIZE_N), warp_specialize=WARP_SPEC):
+        n_off = n * BLOCK_SIZE_N
+        dy = load_grouped_act_tile(
+            dy_ptrs, ADescriptor, pid_m * BLOCK_SIZE_M, n_off, m_mask, 0, A_MEMORY_MODE
+        )
+        # the tile the reduction wants is the weight's NATURAL layout: n leads, k contiguous
+        w = _dgrad_load_weight(
+            b_ptrs, BDescriptor, 0, n_off,
+            pid_k * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE),
+            BLOCK_SIZE_N, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE, B_MEMORY_MODE,
+        )
+        w_scale = tl.load(bs_ptrs)
+        accumulator += tl.dot(
+            dy.to(tl.bfloat16),
+            _dgrad_weight_tile(
+                w, w_scale, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, tl.bfloat16
+            ),
+        )
+        dy_ptrs += BLOCK_SIZE_N * stride_a_n
+        b_ptrs += BLOCK_SIZE_N * stride_b_n
+        bs_ptrs += (BLOCK_SIZE_N // SCALE_ROW_DIV) * stride_bs_n
+    if BsGlobal is not None:  # NVFP4's second level — one matrix, so one value
+        accumulator = accumulator * tl.load(BsGlobal)
+    tl.store(
+        C + offs_m[:, None] * stride_c_m + offs_k[None, :] * stride_c_k,
+        accumulator.to(C.dtype.element_ty),
+        mask=m_mask[:, None] & k_mask[None, :],
+    )
+@compile_time_only_triton_op("finegrained::dgrad_matmul_2d", mutates_args=())
+def dgrad_matmul_2d(
+    dY: torch.Tensor,
+    W: torch.Tensor,
+    Ws: torch.Tensor,
+    scale_group_k: int,
+    scale_row_div: int,
+    output_dtype: torch.dtype | None = None,
+    w_global_scale: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """``dX = dY @ W`` against the FORWARD-oriented quantized weight — the gradient a frozen
+    quantized base weight passes to whatever trains in front of it (LoRA adapters, earlier
+    layers).
+    ``W`` is ``(N, K)`` exactly as the forward holds it; nothing is transposed or re-quantized.
+    ``scale_group_k``/``scale_row_div`` describe the scale grid (MX: ``32, 1``; block-FP8:
+    ``128, 128``), which is the only thing that differs between formats."""
+    assert dY.ndim == 2 and W.ndim == 2, f"expected 2D operands, got {dY.ndim}D and {W.ndim}D"
+    M, N = dY.shape
+    values_per_byte = 2 if W.dtype == torch.uint8 else 1
+    K = W.shape[1] * values_per_byte
+    assert W.shape[0] == N, f"weight rows {W.shape[0]} != dY's N {N}"
+    assert K % scale_group_k == 0, (
+        f"K={K} must be a whole number of scale groups ({scale_group_k})"
+    )
+    dX = torch.empty(M, K, device=dY.device, dtype=output_dtype or dY.dtype)
+    grid = lambda META: (  # noqa: E731
+        triton.cdiv(M, META["BLOCK_SIZE_M"]),
+        triton.cdiv(K, META["BLOCK_SIZE_K"]),
+    )
+    with device_context(dY.device):
+        # built INSIDE the device context: the tensormap is created against the CURRENT context,
+        # so constructing it before the switch fails as "invalid device context" whenever the
+        # caller's device is not the active one. Boxes are placeholders; the pre_hook rebinds
+        # them to the tuned tile per config.
+        a_descriptor = maybe_descriptor(dY, [1, 128])
+        b_descriptor = maybe_descriptor(W, [1, 128])
+        compile_time_only_triton_wrap(dgrad_matmul_2d_kernel)[grid](
+            dY,
+            a_descriptor,
+            W,
+            b_descriptor,
+            Ws,
+            w_global_scale,
+            dX,
+            M,
+            N,
+            K,
+            max(int(M).bit_length(), 1),
+            dY.stride(0),
+            dY.stride(1),
+            W.stride(0),
+            W.stride(1),
+            Ws.stride(0),
+            Ws.stride(1),
+            dX.stride(0),
+            dX.stride(1),
+            SCALE_GROUP_K=scale_group_k,
+            SCALE_ROW_DIV=scale_row_div,
+            WEIGHT_VALUES_PER_BYTE=values_per_byte,
+        )
+    return dX
+@bayesian_autotune(
+    # Persistent grouped loop, plain dot over the N reduction. No memory-mode axis (the
+    # dequantize is in-tile, so both operands read through pointers) and no COMPUTE_MODE axis
+    # (the scaled MMA cannot reduce over an unscaled axis — see the module docstring).
+    get_accelerator_autotuning_configs(
+        tune_block_nk=True,
+        tune_block_m=True,
+        warp_spec=True,
+        a_memory_modes=("descriptor", "pointer"),
+        b_memory_modes=("descriptor", "pointer"),
+        pre_hook=rebind_dgrad_grouped_descriptors,
+    ),
+    ["N", "K", "tokens_per_expert_bit_length"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    finite_check_args=("C",),
+    # C is ACCUMULATED into (tl.atomic_add: several routed rows share a source token), so the
+    # kernel is not idempotent. Benchmarking calls each config many times against one buffer, and
+    # without this every repetition adds again — the tune reads a correct kernel as 1.7e5x wrong.
+    reset_to_zero=("C",),
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            block_within_dim_pruner("N"),  # exact N-loop: no tail mask, constant-stride advance
+            config_filter(  # the scale row advances by BN // SCALE_ROW_DIV — must be whole
+                lambda c, a: c.kwargs["BLOCK_SIZE_N"] % max(a.get("SCALE_ROW_DIV", 1), 1) == 0
+            ),
+            config_filter(  # BK below one scale group gives NG=0 -> tl.arange(0, 0)
+                lambda c, a: c.kwargs["BLOCK_SIZE_K"] % max(a.get("SCALE_GROUP_K", 1), 1) == 0
+            ),
+            warp_spec_compile_guard_pruner(),
+            descriptor_box_pruner("BLOCK_SIZE_K"),
+            smem_pruner("BLOCK_SIZE_K"),
+        )
+    },
+)
+@triton.jit
+def dgrad_matmul_grouped_kernel(
+    A,  # (S, N) upstream gradient in the forward's OUTPUT row order
+    ADescriptor,  # host TMA descriptor over A; box (1|BM, BN) — 1 row when gathering
+    B,  # (E, N, K) quantized weights, K-contiguous — the forward orientation
+    BDescriptor,  # host TMA descriptor over B, box (1, BN, BK_bytes)
+    Bs,  # (E, N // SCALE_ROW_DIV, K // SCALE_GROUP_K) weight scales
+    BsGlobal,  # fp32 NVFP4 weight global, per tensor or per expert; folded on the accumulator
+    stride_bs_global,  # expert stride of BsGlobal (0 broadcasts a scalar)
+    C,  # (S, K) output gradient, written in the forward's INPUT row order
+    GatherIdx,  # (S,) sorted position -> source row of the forward's A
+    ScatterIdx,  # (S,) sorted position -> destination row of the forward's C
+    ExpertStart,  # (NUM_EXPERTS_POW2 + 1,) cumulative row starts, S sentinel
+    # Shape
+    S,
+    N,
+    K,
+    tokens_per_expert_bit_length,  # autotune key only; unused in body
+    # Strides
+    stride_a_m,
+    stride_a_n,
+    stride_b_e,
+    stride_b_n,
+    stride_b_k,
+    stride_bs_e,
+    stride_bs_n,
+    stride_bs_k,
+    stride_c_m,
+    stride_c_k,
+    # Format
+    SCALE_GROUP_K: tl.constexpr,
+    SCALE_ROW_DIV: tl.constexpr,
+    WEIGHT_VALUES_PER_BYTE: tl.constexpr,
+    # Meta-parameters
+    A_MEMORY_MODE: tl.constexpr,
+    NUM_EXPERTS_POW2: tl.constexpr,
+    NUM_SMS: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    WARP_SPEC: tl.constexpr,
+    B_MEMORY_MODE: tl.constexpr,
+    ROUTED_OUT: tl.constexpr = False,
+):
+    """Grouped dgrad. The tile resolver is reused verbatim with K-tiles in its N-tile slot —
+    it is generic in that second axis, so ``offs_bn`` comes back as this kernel's K offsets and
+    ``pid_n`` as its K-tile index (GATE=False: the gradient of a gate|up GEMM is an UNGATED
+    product at the doubled N extent, so no gated arm exists here).
+    The two row maps SWAP roles against the forward: the forward gathered A at ``in_row`` and
+    scattered C to ``out_row``, so its gradient reads ``A`` at ``out_row`` and writes ``C`` at
+    ``in_row``. Both are already real parameters, which is why routing needs no new machinery."""
+    start_pid = tl.program_id(axis=0)
+    exp_start, freqs, tile_start_excl, total_m_tiles, e_offs = build_tile_layout(
+        ExpertStart, NUM_EXPERTS_POW2, BLOCK_SIZE_M
+    )
+    num_k_tiles = tl.cdiv(K, BLOCK_SIZE_K)
+    for tile_id in tl.range(start_pid, total_m_tiles * num_k_tiles, NUM_SMS):
+        pid_k, _, expert_id64, in_row, out_row, row_mask, offs_k, row0, _, m_start = (
+            resolve_grouped_tile(
+                tile_id,
+                num_k_tiles,
+                exp_start,
+                freqs,
+                tile_start_excl,
+                e_offs,
+                GatherIdx,
+                ScatterIdx,
+                BLOCK_SIZE_K,
+                BLOCK_SIZE_M,
+                False,
+            )
+        )
+        offs_kb = tl.minimum(
+            pid_k * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE)
+            + tl.arange(0, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE),
+            K // WEIGHT_VALUES_PER_BYTE - 1,
+        )
+        offs_kg = tl.minimum(
+            (pid_k * BLOCK_SIZE_K) // SCALE_GROUP_K + tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K),
+            K // SCALE_GROUP_K - 1,
+        )
+        b_base = B + expert_id64 * stride_b_e
+        bs_base = Bs + expert_id64 * stride_bs_e
+        # exact N-loop (BLOCK_SIZE_N divides N): maskless weight, constant-stride advance
+        offs_n0 = tl.arange(0, BLOCK_SIZE_N)
+        # SWAP_AB is the WEIGHT's viewpoint and activation callers pass it INVERTED: False builds
+        # [K, rows], so both of these want True to get the rows-major tiles the dot expects
+        # ([BM, BN] and [BN, BK]). Only the pointer arm is affected — the descriptor arm returns
+        # the scalar base and reads its box from the descriptor.
+        dy_ptrs = operand_tile_ptrs(
+            A, out_row, offs_n0, stride_a_m, stride_a_n, A_MEMORY_MODE, True
+        )
+        b_ptrs = operand_tile_ptrs(
+            b_base, offs_n0, offs_kb, stride_b_n, stride_b_k, B_MEMORY_MODE, True
+        )
+        bs_ptrs = (
+            bs_base
+            + (offs_n0[:, None] // SCALE_ROW_DIV) * stride_bs_n
+            + offs_kg[None, :] * stride_bs_k
+        )
+        accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K), dtype=tl.float32)
+        for n in tl.range(0, tl.cdiv(N, BLOCK_SIZE_N), warp_specialize=WARP_SPEC):
+            n_off = n * BLOCK_SIZE_N
+            # the forward gathered A at in_row and scattered C to out_row, so its gradient READS
+            # at out_row — ScatterIdx is this pass's gather map. The box anchor must be out_row's
+            # tile start too: the resolver's `m_start` is min(in_row) — the FORWARD's read side —
+            # and with a gather-only launch (ScatterIdx None, in_row = the accumulation targets)
+            # it anchors the box at a TOKEN id, reading the wrong dY rows.
+            dy = load_grouped_act_tile(
+                dy_ptrs, ADescriptor, tl.min(out_row).to(tl.int32), n_off, row_mask, out_row,
+                A_MEMORY_MODE, ScatterIdx is not None,
+            )
+            w = _dgrad_load_weight(
+                b_ptrs, BDescriptor, row0, n_off,
+                pid_k * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE),
+                BLOCK_SIZE_N, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE, B_MEMORY_MODE, True,
+            )
+            w_scale = tl.load(bs_ptrs)
+            accumulator += tl.dot(
+                dy.to(tl.bfloat16),
+                _dgrad_weight_tile(
+                    w, w_scale, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, tl.bfloat16
+                ),
+            )
+            dy_ptrs += BLOCK_SIZE_N * stride_a_n
+            b_ptrs += BLOCK_SIZE_N * stride_b_n
+            bs_ptrs += (BLOCK_SIZE_N // SCALE_ROW_DIV) * stride_bs_n
+        # Fallback when the rows are NOT token-major (no scatter map): accumulate straight into
+        # the source token. `in_row` is that token, and with top_k > 1 several routed rows share
+        # one, so a store would keep whichever tile landed last and drop the rest — invisible to
+        # any test whose gather map is the identity (S == M), which is how it survived one.
+        #
+        # The FORWARD reduces without atomics (`weighted_reduce`: (S, H) -> (T, H)) — but its
+        # routed-space tensor already exists, because the down GEMM's natural output IS (S, H), and
+        # its reduce also applies routing weights and skips EP sentinels, which an atomic could not
+        # express. dgrad is the transpose: its natural output is SOURCE space, so mirroring the
+        # forward would mean materializing a (S, K) buffer purely to avoid atomics.
+        #
+        # Measured and rejected: store-into-(S, K) + index_add_ ran 4038 vs 4181us at DSV3 top_k=8
+        # (-3.4%), 1278 vs 1294 at GPT-OSS top_k=4, and SLOWER at decode (411 vs 404) — to buy that
+        # you pay a top_k x fp32 intermediate, ~940MB at DSV3 shapes. top_k-way contention is not
+        # the bottleneck here; the GEMM is.
+        if BsGlobal is not None:  # NVFP4's second level, this tile's expert
+            accumulator = accumulator * tl.load(BsGlobal + expert_id64 * stride_bs_global)
+        if ROUTED_OUT:
+            # Plain store at the TOKEN-MAJOR routed row (`out_row`), then the caller sums the
+            # top_k rows per token. This is the FORWARD's shape — it stores routed rows and
+            # reduces them with `weighted_reduce` — and it is 1.40x faster than atomics at DSV3
+            # top_k=8 (3003 vs 4215us), 1.13x at GPT-OSS top_k=4, break-even at decode. top_k-way
+            # atomic contention is the difference; the reduce itself is nearly free.
+            tl.store(
+                C + out_row[:, None] * stride_c_m + offs_k[None, :] * stride_c_k,
+                accumulator,
+                mask=row_mask[:, None] & (offs_k[None, :] < K),
+            )
+        else:
+            tl.atomic_add(
+                C + in_row[:, None] * stride_c_m + offs_k[None, :] * stride_c_k,
+                accumulator,
+                mask=row_mask[:, None] & (offs_k[None, :] < K),
+            )
+@compile_time_only_triton_op("finegrained::dgrad_matmul_grouped", mutates_args=())
+def dgrad_matmul_grouped(
+    dY: torch.Tensor,
+    W: torch.Tensor,
+    Ws: torch.Tensor,
+    expert_start: torch.Tensor,
+    scale_group_k: int,
+    scale_row_div: int,
+    gather_idx: torch.Tensor | None = None,
+    scatter_idx: torch.Tensor | None = None,
+    output_dtype: torch.dtype | None = None,
+    num_input_rows: int | None = None,
+    w_global_scale: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Routed ``dX = dY @ W[expert]`` over expert-sorted positions — the MoE counterpart of
+    ``dgrad_matmul_2d``.
+    ``gather_idx`` / ``scatter_idx`` are the FORWARD's maps, passed unchanged: the kernel swaps
+    their roles itself (reads ``dY`` at the forward's scatter destination, writes ``dX`` at its
+    gather source), so callers hand over the same two tensors the forward got."""
+    assert dY.ndim == 2 and W.ndim == 3, f"expected (S,N) and (E,N,K), got {dY.shape}, {W.shape}"
+    S, N = dY.shape
+    values_per_byte = 2 if W.dtype == torch.uint8 else 1
+    K = W.shape[2] * values_per_byte
+    assert W.shape[1] == N, f"weight N {W.shape[1]} != dY's N {N}"
+    assert K % scale_group_k == 0, (
+        f"K={K} must be a whole number of scale groups ({scale_group_k})"
+    )
+    num_experts = W.shape[0]
+    num_sms = sm_count(dY.device.index)
+    # (num_input_rows, K), ZERO-initialised: the kernel accumulates each routed row's gradient
+    # into its source token, so the buffer is indexed in SOURCE space and must start empty.
+    # fp32 because that is what tl.atomic_add takes.
+    rows = S if num_input_rows is None else num_input_rows
+    # Token-major store + a top_k sum beats atomics whenever the scatter map makes the rows
+    # token-major (see the kernel). It costs a top_k x fp32 intermediate, which is why the atomic
+    # arm stays as the fallback for gather-only launches.
+    routed_out = scatter_idx is not None and rows > 0 and S % rows == 0
+    dX = torch.zeros(S if routed_out else rows, K, device=dY.device, dtype=torch.float32)
+    with device_context(dY.device):
+        # inside the context — see dgrad_matmul_2d; placeholder boxes, rebound by the pre_hook
+        a_descriptor = maybe_descriptor(dY, [1, 128])
+        b_descriptor = maybe_descriptor(W, [1, 1, 128])
+        compile_time_only_triton_wrap(dgrad_matmul_grouped_kernel)[(num_sms,)](
+            dY,
+            a_descriptor,
+            W,
+            b_descriptor,
+            Ws,
+            w_global_scale,
+            global_scale_stride(w_global_scale),
+            dX,
+            gather_idx,
+            scatter_idx,
+            expert_start,
+            S,
+            N,
+            K,
+            max(int(max(S // max(num_experts, 1), 1)).bit_length(), 1),
+            dY.stride(0),
+            dY.stride(1),
+            W.stride(0),
+            W.stride(1),
+            W.stride(2),
+            Ws.stride(0),
+            Ws.stride(1),
+            Ws.stride(2),
+            dX.stride(0),
+            dX.stride(1),
+            SCALE_GROUP_K=scale_group_k,
+            SCALE_ROW_DIV=scale_row_div,
+            WEIGHT_VALUES_PER_BYTE=values_per_byte,
+            NUM_EXPERTS_POW2=triton.next_power_of_2(num_experts),
+            NUM_SMS=num_sms,
+            ROUTED_OUT=routed_out,
+        )
+    if routed_out:
+        dX = dX.view(rows, S // rows, K).sum(1)
+    return dX.to(output_dtype or dY.dtype)
+@triton.jit
+def _glu_grad(g, u, dh, ACT_FN: tl.constexpr, SWIGLU_ALPHA: tl.constexpr,
+              SWIGLU_LIMIT: tl.constexpr):
+    """(dgate, dup) for one tile, mirroring ``epilogue.glu`` arm for arm.
+    The alpha arm is ``(u + 1) * g * sigmoid(ALPHA * g)`` — the ``+1`` is load-bearing and easy to
+    drop: without it the gate gradient is wrong on exactly the GPT-OSS / MiniMax path and nowhere
+    else. ``SWIGLU_LIMIT`` clamps gate ABOVE and up to ``[-LIMIT, LIMIT]``, so the gradient is zero
+    wherever the forward saturated. Single return; only the taken arms compile."""
+    gc = tl.minimum(g, SWIGLU_LIMIT) if SWIGLU_LIMIT is not None else g
+    uc = (
+        tl.minimum(tl.maximum(u, -SWIGLU_LIMIT), SWIGLU_LIMIT)
+        if SWIGLU_LIMIT is not None
+        else u
+    )
+    if SWIGLU_ALPHA is not None:
+        sig = tl.sigmoid(gc * SWIGLU_ALPHA)
+        act = gc * sig
+        dact = sig + gc * SWIGLU_ALPHA * sig * (1.0 - sig)
+        dg = dh * (uc + 1.0) * dact
+    elif ACT_FN == "silu":
+        sig = tl.sigmoid(gc)
+        act = gc * sig
+        dg = dh * uc * (sig + gc * sig * (1.0 - sig))
+    elif ACT_FN == "relu":
+        act = tl.maximum(gc, 0.0)
+        dg = dh * uc * (gc > 0.0).to(tl.float32)
+    else:  # gelu, exact via erf — the forward's form
+        cdf = 0.5 * (1.0 + tl.erf(gc * 0.7071067811865476))
+        act = gc * cdf
+        dg = dh * uc * (cdf + gc * tl.exp(-0.5 * gc * gc) * 0.3989422804014327)
+    du = dh * act
+    if SWIGLU_LIMIT is not None:  # saturated inputs pass no gradient
+        dg = tl.where(g > SWIGLU_LIMIT, 0.0, dg)
+        du = tl.where(tl.abs(u) > SWIGLU_LIMIT, 0.0, du)
+    return dg, du
+@triton.jit
+def _glu_backward_kernel(
+    Z,  # (M, 2*INTER) pre-activation, gate at even columns / up at odd
+    DH,  # (M, INTER) gradient w.r.t. the GLU output
+    DZ,  # (M, 2*INTER) output, interleaved to match Z
+    M, INTER,
+    stride_z_m, stride_z_n, stride_dh_m, stride_dh_n, stride_dz_m, stride_dz_n,
+    ACT_FN: tl.constexpr, SWIGLU_ALPHA: tl.constexpr, SWIGLU_LIMIT: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_I: tl.constexpr,
+):
+    """Elementwise GLU backward over the interleaved layout. Gate ``j`` sits at column ``2j`` and
+    up at ``2j+1``, so one tile reads both halves with a stride-2 pair of loads and writes them
+    back the same way — no split, no concatenate, no materialized halves."""
+    offs_m = tl.program_id(0) * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_i = tl.program_id(1) * BLOCK_SIZE_I + tl.arange(0, BLOCK_SIZE_I)
+    mask = (offs_m[:, None] < M) & (offs_i[None, :] < INTER)
+    z_gate = Z + offs_m[:, None] * stride_z_m + (2 * offs_i[None, :]) * stride_z_n
+    g = tl.load(z_gate, mask=mask, other=0.0).to(tl.float32)
+    u = tl.load(z_gate + stride_z_n, mask=mask, other=0.0).to(tl.float32)
+    dh = tl.load(
+        DH + offs_m[:, None] * stride_dh_m + offs_i[None, :] * stride_dh_n,
+        mask=mask, other=0.0,
+    ).to(tl.float32)
+    dg, du = _glu_grad(g, u, dh, ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT)
+    dz_gate = DZ + offs_m[:, None] * stride_dz_m + (2 * offs_i[None, :]) * stride_dz_n
+    tl.store(dz_gate, dg.to(DZ.dtype.element_ty), mask=mask)
+    tl.store(dz_gate + stride_dz_n, du.to(DZ.dtype.element_ty), mask=mask)
+def glu_backward(
+    dH: torch.Tensor,
+    Z: torch.Tensor,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+) -> torch.Tensor:
+    """Gradient of the fused GLU epilogue: ``dH`` (grad w.r.t. ``H = f(Z_gate) * Z_up``) and the
+    2I-wide pre-activation ``Z`` -> ``dZ``, interleaved.
+    One pass over Z, no materialized halves. The GEMM that consumes ``dZ`` is then an UNGATED
+    product at the doubled N extent — the forward kernel with no gate flag, because
+    ``pid_n * 2BN == 2 * pid_n * BN`` — so no gated backward kernel exists or is needed."""
+    assert dH.ndim == 2 and Z.ndim == 2, f"expected 2D, got {dH.shape}, {Z.shape}"
+    M, INTER = dH.shape
+    assert Z.shape == (M, 2 * INTER), f"Z must be (M, 2I) = {(M, 2 * INTER)}, got {tuple(Z.shape)}"
+    dZ = torch.empty_like(Z)
+    grid = lambda META: (  # noqa: E731
+        triton.cdiv(M, META["BLOCK_SIZE_M"]),
+        triton.cdiv(INTER, META["BLOCK_SIZE_I"]),
+    )
+    with device_context(Z.device):
+        compile_time_only_triton_wrap(_glu_backward_kernel)[grid](
+            Z, dH, dZ, M, INTER,
+            Z.stride(0), Z.stride(1), dH.stride(0), dH.stride(1), dZ.stride(0), dZ.stride(1),
+            ACT_FN=act_fn, SWIGLU_ALPHA=swiglu_alpha, SWIGLU_LIMIT=swiglu_limit,
+            BLOCK_SIZE_M=32, BLOCK_SIZE_I=64, num_warps=4,
+        )
+    return dZ
+@bayesian_autotune(
+    # Neither a BLOCK_SIZE_M nor a WARP_SPEC axis. One row per program (the rows are not
+    # expert-sorted, so a tile cannot span them) and the reduce is FMA rather than MMA, so there is
+    # no M tile to choose — and WS, which needs ``(BN if swapped else BM) >= 64`` of M work, has no
+    # M operand to satisfy it. Emitting either axis would search configs that cannot compile.
+    get_accelerator_autotuning_configs(tune_block_nk=True),
+    ["N", "K", "S"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    finite_check_args=("C",),
+    # No pruners. The smem model is fit to MMA loops — ``num_stages`` buffers of an A tile and a B
+    # tile, keyed on BLOCK_SIZE_M — and this loop has neither: the activation is a ``[BN]`` vector
+    # and the reduce is FMA, so there is no M dimension for it to reason about. A config that
+    # genuinely will not fit raises OutOfResources, which the tuner already inf-forgives.
+)
+@triton.jit
+def dgrad_matmul_batched_kernel(
+    A,  # (S, N) upstream gradient
+    B,  # (E, N, K) quantized weights, K-contiguous — the forward orientation
+    Bs,  # (E, N // SCALE_ROW_DIV, K // SCALE_GROUP_K) weight scales
+    BsGlobal,  # fp32 NVFP4 weight global, per tensor or per expert; folded on the accumulator
+    stride_bs_global,  # expert stride of BsGlobal (0 broadcasts a scalar)
+    C,  # (S, K) output gradient
+    ExpertIds,  # (S,) which expert each row routed to
+    # Shape
+    S,
+    N,
+    K,
+    # Strides
+    stride_a_m,
+    stride_a_n,
+    stride_b_e,
+    stride_b_n,
+    stride_b_k,
+    stride_bs_e,
+    stride_bs_n,
+    stride_bs_k,
+    stride_c_m,
+    stride_c_k,
+    # Format
+    SCALE_GROUP_K: tl.constexpr,
+    SCALE_ROW_DIV: tl.constexpr,
+    WEIGHT_VALUES_PER_BYTE: tl.constexpr,
+    # Meta-parameters
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    """Batched dgrad: ``dX[s] = dY[s] @ B[ExpertIds[s]]``.
+    Unlike the grouped form the rows are NOT expert-sorted — each row names its own expert — so a
+    tile must cover one expert. ``BLOCK_SIZE_M`` is pinned to 1 by the wrapper for that reason,
+    which is also the shape the batched forward runs at."""
+    pid_m = tl.program_id(0)
+    pid_k = tl.program_id(1)
+    row_ok = pid_m < S
+    expert_id = tl.load(ExpertIds + pid_m, mask=row_ok, other=0).to(tl.int64)
+    offs_k = pid_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
+    offs_kb = tl.minimum(
+        pid_k * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE)
+        + tl.arange(0, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE),
+        K // WEIGHT_VALUES_PER_BYTE - 1,
+    )
+    offs_kg = tl.minimum(
+        (pid_k * BLOCK_SIZE_K) // SCALE_GROUP_K + tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K),
+        K // SCALE_GROUP_K - 1,
+    )
+    b_base = B + expert_id * stride_b_e
+    bs_base = Bs + expert_id * stride_bs_e
+    # CUDA-core FMA reduce, NOT tl.dot: one row per program is an M=1 GEVM, and feeding the MMA a
+    # broadcast [1, BN] operand runs it emulated at a fraction of peak. The forward reaches the same
+    # conclusion for its decode path (``mx_weight_only_scalar_swapped``) — the scalar arm wins at
+    # M=1 and the tuner picks it there. Same reduction, no MMA, no padded M dimension.
+    accumulator = tl.zeros((BLOCK_SIZE_K,), dtype=tl.float32)
+    for n in tl.range(0, tl.cdiv(N, BLOCK_SIZE_N)):
+        offs_n = n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+        dy = tl.load(
+            A + pid_m * stride_a_m + offs_n * stride_a_n,
+            mask=row_ok & (offs_n < N),
+            other=0.0,
+        ).to(tl.float32)
+        offs_nc = tl.minimum(offs_n, N - 1)  # clamp: dy is already zeroed past N
+        w = tl.load(b_base + offs_nc[:, None] * stride_b_n + offs_kb[None, :] * stride_b_k)
+        w_scale = tl.load(
+            bs_base
+            + (offs_nc[:, None] // SCALE_ROW_DIV) * stride_bs_n
+            + offs_kg[None, :] * stride_bs_k
+        )
+        wdeq = _dgrad_weight_tile(
+            w, w_scale, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, tl.float32
+        )
+        accumulator += tl.sum(dy[:, None] * wdeq, axis=0)
+    if BsGlobal is not None:  # NVFP4's second level, this row's expert
+        accumulator = accumulator * tl.load(BsGlobal + expert_id * stride_bs_global)
+    tl.store(
+        C + pid_m * stride_c_m + offs_k * stride_c_k,
+        accumulator.to(C.dtype.element_ty),
+        mask=row_ok & (offs_k < K),
+    )
+@compile_time_only_triton_op("finegrained::dgrad_matmul_batched", mutates_args=())
+def dgrad_matmul_batched(
+    dY: torch.Tensor,
+    W: torch.Tensor,
+    Ws: torch.Tensor,
+    expert_ids: torch.Tensor,
+    scale_group_k: int,
+    scale_row_div: int,
+    output_dtype: torch.dtype | None = None,
+    w_global_scale: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """``dX[s] = dY[s] @ W[expert_ids[s]]`` — the batched (per-row expert) counterpart of
+    ``dgrad_matmul_grouped``. Registered because the batched experts forward is a SELECTABLE
+    dispatch in transformers, not a decode-only path: without this, finetuning against it
+    silently produces no gradient."""
+    assert dY.ndim == 2 and W.ndim == 3, f"expected (S,N) and (E,N,K), got {dY.shape}, {W.shape}"
+    S, N = dY.shape
+    values_per_byte = 2 if W.dtype == torch.uint8 else 1
+    K = W.shape[2] * values_per_byte
+    assert W.shape[1] == N, f"weight N {W.shape[1]} != dY's N {N}"
+    dX = torch.empty(S, K, device=dY.device, dtype=output_dtype or dY.dtype)
+    grid = lambda META: (S, triton.cdiv(K, META["BLOCK_SIZE_K"]))  # noqa: E731
+    with device_context(dY.device):
+        compile_time_only_triton_wrap(dgrad_matmul_batched_kernel)[grid](
+            dY, W, Ws, w_global_scale, global_scale_stride(w_global_scale), dX, expert_ids, S, N, K,
+            dY.stride(0), dY.stride(1),
+            W.stride(0), W.stride(1), W.stride(2),
+            Ws.stride(0), Ws.stride(1), Ws.stride(2),
+            dX.stride(0), dX.stride(1),
+            SCALE_GROUP_K=scale_group_k,
+            SCALE_ROW_DIV=scale_row_div,
+            WEIGHT_VALUES_PER_BYTE=values_per_byte,
+        )
+    return dX
+# ── gradient products, per format family ─────────────────────────────────────────
+# What each format's dgrad actually computes. ``backward.py`` holds only the table saying which
+# op gets which of these — compute lives here, attachment lives there.
+def _dgrad_mx(dY, B, Bs, b_global, out_dtype):
+    """dA for an MX/NVFP4 weight: contract N on the forward-oriented weight (see backward.py).
+    The scale grid is DERIVED FROM SHAPES, not from a format constant — ``Bs`` is
+    ``(N // row_div, K // group)``, so the divisors read straight off it. That covers group-32
+    UE8M0 (MX) and group-16 E4M3 (NVFP4) without either being named here, and cannot drift when a
+    format's block changes.
+    Containers are normalized, not reinterpreted: packed E2M1 reaches the kernel as ``int8`` or
+    ``uint8`` depending on the producer, and UE8M0 as ``float8_e8m0fnu`` or ``uint8``. The kernel
+    keys its unpack/decode off ``uint8``, and decoding UE8M0 bytes as floats is off by orders of
+    magnitude — the failure mode is a non-finite dX, which the tuner's numerics veto turns into an
+    all-inf tune rather than a wrong answer.
+    The NVFP4 second-level global rides into the kernel, which folds it on the accumulator —
+    one value here, a 2D weight being one matrix."""
+    if B.dtype == torch.int8:
+        B = B.view(torch.uint8)
+    if Bs.dtype == torch.float8_e8m0fnu:
+        Bs = Bs.view(torch.uint8)
+    values_per_byte = 2 if B.dtype == torch.uint8 else 1
+    K = B.shape[-1] * values_per_byte
+    group = K // Bs.shape[-1]
+    row_div = max(B.shape[-2] // Bs.shape[-2], 1)
+    dX = dgrad_matmul_2d(
+        dY.contiguous(), B, Bs, group, row_div, output_dtype=torch.float32,
+        w_global_scale=normalize_global_scale(b_global, 1),
+    )
+    return dX.to(out_dtype)
+def _dgrad_transposed_weight(dY, B, Bs, out_dtype):
+    """dA on the QUANTIZED weight, reoriented by transposition. Valid when the scale does not
+    bind a group to K: a per-tensor scalar is orientation-free, and a BLOCK grid subdivides both
+    axes, so transposing it is the same reindexing as transposing the tile.
+    The scale grid is made CONTIGUOUS, not just transposed: reading it column-major cost 79us on
+    a 4096x8192x6144 dgrad (405.7 -> 326.5us) while the copy is ~3KB, because every tile re-reads
+    it strided. The weight itself stays a view — materializing it saves a further 47us but costs
+    211us to produce, so that only pays when the caller already keeps both orientations."""
+    Bs_t = Bs if Bs.numel() == 1 else Bs.t().contiguous()
+    return matmul_2d(dY.contiguous(), B.t(), None, Bs_t, output_dtype=out_dtype)
+def _dgrad_grouped(ctx, dY, B, Bs, expert_start, gather_idx, scatter_idx, b_global):
+    """Routed dgrad. GATE is not a case to handle: the gradient of a gate|up GEMM is an UNGATED
+    product at the doubled N extent, which is what a 2N-row weight already is. The forward's two
+    row maps are passed through unchanged — the kernel swaps their roles itself. The NVFP4
+    global folds in the kernel, which already resolves each tile's expert."""
+    if B.dtype == torch.int8:
+        B = B.view(torch.uint8)
+    if Bs.dtype == torch.float8_e8m0fnu:
+        Bs = Bs.view(torch.uint8)
+    values_per_byte = 2 if B.dtype == torch.uint8 else 1
+    K = B.shape[-1] * values_per_byte
+    dX = dgrad_matmul_grouped(
+        dY.contiguous(), B, Bs, expert_start,
+        K // Bs.shape[-1], max(B.shape[-2] // Bs.shape[-2], 1),
+        gather_idx=gather_idx, scatter_idx=scatter_idx, output_dtype=torch.float32,
+        num_input_rows=getattr(ctx, "a_rows", None),
+        w_global_scale=normalize_global_scale(b_global, B.shape[0]),
+    )
+    return dX.to(ctx.a_dtype)
+__all__: list[str] = []
+def _register(name: str, saved: tuple[int, ...], dgrad):
+    """Attach a backward to one op. ``saved`` names the positional inputs backward needs — they
+    differ per op, e.g. the MX signature is ``(A, B, As, Bs, ...)`` so its weight scale is index
+    3, not 2.
+    Everything but the activation gets None — weights and scales take no gradient, the rest is
+    non-differentiable config."""
+    def setup_context(ctx, inputs, output):
+        ctx.save_for_backward(*(inputs[i] for i in saved))
+        ctx.a_dtype = inputs[0].dtype
+        # routed dgrad accumulates into SOURCE tokens, which is a different row count from the
+        # routed rows the op consumed whenever a gather map is in play
+        ctx.a_rows = inputs[0].shape[0]
+    def backward(ctx, *grads):
+        # torch matches backward's return against `tree_flatten(args, is_leaf=not_list_of_tensor)`
+        # of the args AS INVOKED — which setup_context cannot see (it gets the bound args, with
+        # defaults filled). `needs_input_grad` is sized to that same flattened list plus the
+        # trailing metadata arg, so it is the one place the right count is observable. torch
+        # appends the metadata's own None itself, so this returns exactly that many.
+        n_grads = len(ctx.needs_input_grad)
+        if not ctx.needs_input_grad[0]:
+            return (None,) * n_grads
+        # the ops return list[Tensor] (a second entry under output_format), so autograd hands the
+        # gradient back in that shape — unwrap to the gradient of the primary output
+        dY = grads[0]
+        while isinstance(dY, (list, tuple)):
+            dY = dY[0]
+        return (dgrad(ctx, dY, *ctx.saved_tensors),) + (None,) * (n_grads - 1)
+    torch.library.register_autograd(
+        add_op_namespace_prefix(name), backward, setup_context=setup_context
+    )
+# block-FP8 (128x128 / 1x128, fp32 or UE8M0 scales): the grid subdivides both axes, so dgrad
+# reuses the quantized weight through a transposed view.
+_register(
+    "w8a8_block_dynamic_fp8_matmul", saved=(1, 2),
+    dgrad=lambda ctx, dY, B, Bs: _dgrad_transposed_weight(dY, B, Bs, ctx.a_dtype),
+)
+# per-tensor FP8: a scalar scale is orientation-free, so the same route needs no reorientation.
+_register(
+    "w8a8_tensor_dynamic_fp8_matmul", saved=(1, 2),
+    dgrad=lambda ctx, dY, B, Bs: _dgrad_transposed_weight(dY, B, Bs, ctx.a_dtype),
+)
+# MX / NVFP4: per-row group scales bind to K and packed E2M1 packs along K, so the weight cannot
+# be reoriented by a VIEW — but it does not need to be. The reduction runs over N against the
+# weight's natural (N, K) tile, dequantized per-tile in-register, so the FORWARD-oriented weight
+# serves dgrad as-is: no transposed copy, no second quantized orientation, no fp32 materialization.
+_register(
+    "mx_dynamic_matmul", saved=(1, 3, 13),
+    dgrad=lambda ctx, dY, B, Bs, b_global: _dgrad_mx(dY, B, Bs, b_global, ctx.a_dtype),
+)
+# Weight-only (W4A16 / W8A16) — the QLoRA format: a frozen 4-bit base weight with raw bf16
+# activations. Same route; the weight takes no gradient, so nothing here needs a master copy.
+_register(
+    "mx_weight_only_matmul_2d", saved=(1, 2, 9),
+    dgrad=lambda ctx, dY, B, Bs, b_global: _dgrad_mx(dY, B, Bs, b_global, ctx.a_dtype),
+)
+def _dgrad_batched(ctx, dY, B, Bs, expert_ids, b_global):
+    """Per-row-expert dgrad. Same shape-derived scale grid as the others; the only difference from
+    ``_dgrad_grouped`` is that the expert comes from a per-row index rather than sorted ranges."""
+    if B.dtype == torch.int8:
+        B = B.view(torch.uint8)
+    if Bs.dtype == torch.float8_e8m0fnu:
+        Bs = Bs.view(torch.uint8)
+    values_per_byte = 2 if B.dtype == torch.uint8 else 1
+    K = B.shape[-1] * values_per_byte
+    dX = dgrad_matmul_batched(
+        dY.contiguous(), B, Bs, expert_ids,
+        K // Bs.shape[-1], max(B.shape[-2] // Bs.shape[-2], 1),
+        output_dtype=torch.float32,
+        w_global_scale=normalize_global_scale(b_global, B.shape[0]),
+    )
+    return dX.to(ctx.a_dtype)
+# Routed MoE, weight-only — QLoRA on an MoE model. Positional slots: B=1, Bs=2, expert_start=3,
+# gather_idx=10, scatter_idx=11, b_global_scale=12.
+_register(
+    "mx_weight_only_matmul_grouped", saved=(1, 2, 3, 10, 11, 12),
+    dgrad=_dgrad_grouped,
+)
+# Routed MoE, block-FP8. Same product as the weight-only grouped op — a 128x128 block grid is
+# the per-row K-group layout at wider divisors, which `_dgrad_grouped` derives from `Bs.shape` —
+# but it is a DIFFERENT op, and only registering the weight-only one left the grouped block-FP8
+# path with no grad_fn at all. That is the path a PEFT LoRA on fp8 experts actually runs.
+# Positional slots: B=1, Bs=3, expert_start=4, gather_idx=14, scatter_idx=15; no global scale.
+_register(
+    "w8a8_block_dynamic_fp8_matmul_grouped", saved=(1, 3, 4, 14, 15),
+    dgrad=lambda ctx, dY, B, Bs, es, gi, si: _dgrad_grouped(ctx, dY, B, Bs, es, gi, si, None),
+)
+# The batched experts forward is a SELECTABLE dispatch in transformers, not a decode-only path, so
+# it needs a formula too — without one, finetuning against `batched_mm` silently trains nothing.
+# Rows are not expert-sorted here: slot 3 is per-row `expert_ids`, not `expert_start`.
+_register(
+    "mx_weight_only_matmul_batched", saved=(1, 2, 3, 12),
+    dgrad=_dgrad_batched,
+)

build/torch-rocm/batched.py ADDED Viewed

	@@ -0,0 +1,1947 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait
+from .bayesian_autotuner import bayesian_autotune
+from .compat import add_op_namespace_prefix, FP8_DTYPE, MX_SCALE_GROUP_K, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, tl_dtype, decode_pdl
+from .descriptors import rebind_batched_mx_bs_descriptor
+from .formats import check_activation_format, global_scale_stride, normalize_global_scale, e2m1_as_uint8, expert_weight_shape, is_mx, mx_scale_family, normalize_per_expert_scale, resolve_activation_format, resolve_output_dtype, ue8m0_as_uint8, validate_dense_operands, weight_block_size, weight_format
+from .epilogue import fused_glu
+from .quant import MX_ACT_QUANT, fp8_act_quant_block_dynamic, fp8_act_quant_tensor_wide
+from .swizzle import swizzled_scale_descriptor
+from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot
+from .loading.tiles import (
+    advance_ptrs,
+    load_act_block_dynamic,
+    load_act_mx,
+    load_act_plain,
+    load_act_static,
+    load_weight_block_dynamic,
+    load_weight_mx,
+    load_weight_plain,
+    load_weight_static,
+    operand_tile_ptrs,
+    oriented_tile_ptrs,
+    weight_tile_ptrs,
+)
+from .epilogue import acc_finalize, acc_init, add_bias, bias_strides, gemm_epilogue
+from .pruners import PATH_ANCHOR_AXES, fp8_dot_warp_pruner, dot_scaled_staging_pruner, block_fits_dim_pruner, block_within_dim_pruner, compose_pruners, mx_config_pruner, require_moe_dims_aligned, scale_subblock_pruner, smem_pruner, swizzled_scale_config_pruner, weight_only_swap_scope_pruner
+@triton.jit
+def expert_setup(
+    A,
+    B,
+    C,
+    Bs,
+    ExpertIds,
+    GatherIdx,
+    ScatterIdx,
+    stride_a_m,
+    stride_b_e,
+    stride_c_m,
+    stride_bs_e,
+    stride_eid,
+    ADVANCE_BS: tl.constexpr = True,
+):
+    """Per-(row, expert) prologue shared by the batched kernels: read the program
+    ids, look up the routed expert, and advance the A/B/C/Bs base pointers to this
+    row's slice. Returns ``(batch_id, pid_n, expert_id, A, B, C, Bs, in_row, out_row)`` — the
+    resolved source/destination rows fold the gather/scatter out of the kernel bodies.
+    ``ADVANCE_BS=False`` leaves ``Bs`` at the buffer base (the mx scale leaf applies the expert
+    offset itself — its swizzled path indexes by 128-row block, not the row-major expert stride).
+    ``A``'s source row is ``GatherIdx[batch_id]`` when ``GatherIdx`` is not None (the gate_up reading
+    unexpanded activations, many-to-one for top_k > 1) else ``batch_id``; ``C``'s destination
+    row is ``ScatterIdx[batch_id]`` when ``ScatterIdx`` is not None else ``batch_id`` — the same virtual
+    gather/scatter ``matmul_grouped`` does, so the routed rows need no materialized copy.
+    The caller must early-return on the EP sentinel (``expert_id >= num_experts``)
+    before any load — the pointer arithmetic itself is harmless, only the loads on a
+    non-local expert would be out of bounds."""
+    batch_id = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    # Cast to int64 to prevent overflow on expert_id * stride_b_e.
+    expert_id = tl.load(ExpertIds + batch_id * stride_eid).to(tl.int64)
+    in_row = tl.load(GatherIdx + batch_id).to(tl.int64) if GatherIdx is not None else batch_id
+    out_row = tl.load(ScatterIdx + batch_id).to(tl.int64) if ScatterIdx is not None else batch_id
+    A = A + in_row * stride_a_m
+    B = B + expert_id * stride_b_e
+    C = C + out_row * stride_c_m
+    if ADVANCE_BS:
+        Bs = Bs + expert_id * stride_bs_e
+    return batch_id, pid_n, expert_id, A, B, C, Bs, in_row, out_row
+@triton.jit
+def store_row(
+    C,
+    accumulator,
+    pid_n,
+    stride_c_n,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+):
+    """Output epilogue shared by the batched kernels (``C`` already advanced to the
+    row). The fake-batch trick aliases all ``BLOCK_SIZE_M`` lanes to the same C row,
+    so a plain store would issue ``BLOCK_SIZE_M`` duplicate-address writes — benign on
+    NVIDIA WGMMA (last-write-wins of identical bytes) but hardware-undefined on Intel
+    XPU, where it corrupts the output. Mask so only lane 0 stores; the accumulator
+    rows are mathematically identical (same A row × same B), so lane 0 is correct."""
+    c = accumulator.to(C.dtype.element_ty)
+    offs_cm = tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    # offs_cm[:, None] * 0: broadcast to a [BM, BN] pointer tile (all rows alias the one C row)
+    # so the lane-0 mask below has a row axis to select; the M stride is deliberately 0.
+    c_ptrs = C + offs_cm[:, None] * 0 + stride_c_n * offs_cn[None, :]
+    tl.store(c_ptrs, c, mask=(offs_cm == 0)[:, None])
+@bayesian_autotune(
+    get_accelerator_autotuning_configs(swap_ab=True, tune_block_n=True),
+    # one winner per (shape, requant): requant narrows the legal tiles to the quant block, so a
+    # shared entry could replay a sub-block winner on a launch whose output scale needs the block.
+    # GATE stacks a 2*BN-wide dot — a distinct config space that must not share a winner.
+    # BLOCK_N/BLOCK_K (the launch-pinned quant block) bound the legal tiles the same way —
+    # two checkpoints at one shape with different blocks must not share a winner.
+    ["N", "K", "S", "OUTPUT_FORMAT", "GATE", "BLOCK_N", "BLOCK_K"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            scale_subblock_pruner(),
+            fp8_dot_warp_pruner(),
+        )
+    },
+)
+@triton.jit
+def w8a8_block_dynamic_fp8_matmul_batched_kernel(
+    A,  # (S, K) E4M3 activations (pre-quantized once by the wrapper)
+    As,  # (S, K // BLOCK_SIZE_K) fp32 per-row, per-K-block activation scales
+    B,  # (num_experts, N, K) FP8 weights; under GATE the (num_experts, 2N, K) gate|up stack
+    Bs,  # (num_experts, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales (2N under GATE)
+    C,  # (S, N) output; under an OUTPUT_FORMAT the FP8-requantized intermediate
+    Cs,  # (S, N // BLOCK_SIZE_N) per-(row, block) output scale; written iff OUTPUT_FORMAT
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    ExpertIds,  # (S,) — which expert each batch element routes to
+    GatherIdx,  # (S,) int — batch_id -> source row of A; read only when not None
+    ScatterIdx,  # (S,) int — batch_id -> destination row of C; read only when not None
+    # Shape
+    S,
+    N,
+    K,
+    # Strides
+    stride_a_m,
+    stride_a_k,
+    stride_as_m,
+    stride_b_e,
+    stride_b_k,
+    stride_b_n,
+    stride_bs_e,
+    stride_bs_k,
+    stride_bs_n,
+    stride_c_m,
+    stride_c_n,
+    stride_cs_m,
+    stride_cs_n,
+    stride_bias_e,
+    stride_bias_n,
+    stride_eid,
+    num_experts,
+    # Meta-parameters. BLOCK_N/BLOCK_K are the QUANT block (one scale per (BLOCK_N, BLOCK_K) tile of
+    # the weight); BLOCK_SIZE_* are the tuned COMPUTE tile. The N tile may subdivide the quant block
+    # (BLOCK_SIZE_N <= BLOCK_N); the K tile is the quant block (one scale step per K iteration).
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    SWAP_AB: tl.constexpr = False,
+    # Gate|up fusion epilogue (GATE=False -> plain batched GEMM, every arm below folds out)
+    GATE: tl.constexpr = False,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    # the output format name (None | "fp8")
+    OUTPUT_FORMAT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
+    PDL: tl.constexpr = False,
+):
+    """Block-scale batched FP8 expert matmul kernel.
+    Each program handles one routed token row and one N-tile, looking up the
+    owning expert from ``ExpertIds``. Activations arrive pre-quantized (one wrapper
+    pass — an inline quant would repeat per N-tile and pay a per-tile amax reduction).
+    ``SWAP_AB`` (tuner axis, M=1 decode): load the weight output-rows-major ``[BN, BK]`` and put
+    those rows in the MMA M dim, padding the single token to the N=16 atom; column 0 of the
+    ``[BN, 16]`` accumulator is the result. No-swap keeps the token in M (padded to 16).
+    ``GATE`` fuses the gate|up projection: ``B`` is the ``(E, 2N, K)`` stack (gate rows [0,N),
+    interleaved rows), run as two dots (the decode-validated form), SwiGLU-combined, and — under
+    an ``OUTPUT_FORMAT`` — FP8-requantized into ``C`` + a per-(row, block) scalar ``Cs``. Every gate arm
+    folds out at compile time; ``GATE=False`` is the plain GEMM, bit-identical."""
+    if PDL:
+        gdc_wait()
+    batch_id, pid_n, expert_id, A, B, C, Bs, in_row, out_row = expert_setup(
+        A,
+        B,
+        C,
+        Bs,
+        ExpertIds,
+        GatherIdx,
+        ScatterIdx,
+        stride_a_m,
+        stride_b_e,
+        stride_c_m,
+        stride_bs_e,
+        stride_eid,
+    )
+    # EP sentinel: row routed to a non-local expert; output is left uninit.
+    if expert_id >= num_experts:
+        return
+    # One scale per quant block, broadcast over the tile, so an N tile narrower than BLOCK_N just
+    # reads its own block: tile pid_n sits in block (pid_n * BLOCK_SIZE_N) // BLOCK_N, which is pid_n
+    # when the tile IS the block. Narrower tiles multiply the N grid (see scale_subblock_pruner).
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    offs_bn = pid_n * n_width + tl.arange(0, n_width)
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = operand_tile_ptrs(A, tl.arange(0, BLOCK_SIZE_M) * 0, offs_k, stride_a_m, stride_a_k, "pointer", True)
+    as_ptrs = As + in_row * stride_as_m + tl.zeros((BLOCK_SIZE_M,), tl.int32)
+    # One gate|up weight tile + one block-scale pointer, like every other kernel: the n_width
+    # span shares a tile + a single dot, each row's scale block following from its global row.
+    b_ptrs = weight_tile_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, SWAP_AB)
+    bs_ptr = Bs + (pid_n * n_width // BLOCK_N) * stride_bs_n
+    bs_off = ((pid_n * n_width % BLOCK_N + tl.arange(0, n_width)) // BLOCK_N) * stride_bs_n
+    acc = acc_init("dot", BLOCK_SIZE_M, n_width, SWAP_AB)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a, a_s = load_act_block_dynamic(
+            a_ptrs, as_ptrs, None, None, 0, 0, 0, 0, 0, "pointer", False, False, False
+        )
+        w, b_s = load_weight_block_dynamic(
+            b_ptrs, b_ptrs, bs_ptr + bs_off, None, 0, 0, 0, 0, 0, 0,
+            GATE, False, "pointer", SWAP_AB, BLOCK_SIZE_N, BLOCK_SIZE_K,
+        )
+        acc = block_dynamic_dot(acc, a, a_s, w, b_s, BLOCK_SIZE_K, SWAP_AB, False, True)
+        a_ptrs, as_ptrs, b_ptrs, bs_ptr, _, _ = advance_ptrs(
+            a_ptrs, as_ptrs, b_ptrs, bs_ptr, b_ptrs, bs_ptr,
+            BLOCK_SIZE_K * stride_a_k, 1, BLOCK_SIZE_K * stride_b_k, stride_bs_k,
+            "pointer", "pointer", True, True, False,
+        )
+    if PDL:
+        gdc_launch_dependents()
+    gemm_epilogue(
+        C, Cs, acc, out_row, pid_n, 0, out_row, 1, stride_c_n, stride_cs_m, stride_cs_n,
+        BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, OUTPUT_FORMAT, BLOCK_K,
+        ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+        COMPUTE_MODE="dot", SWAP_AB=SWAP_AB, FAKE_BATCH=True,
+        Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n,
+        global_row=expert_id,
+    )
+@bayesian_autotune(
+    get_accelerator_autotuning_configs(swap_ab=True, tune_block_n=True),
+    # keyed like the block-dynamic sibling: requant narrows the legal tiles to the quant block
+    ["N", "K", "S", "OUTPUT_FORMAT", "GATE", "BLOCK_N", "BLOCK_K"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            scale_subblock_pruner(),
+            fp8_dot_warp_pruner(),
+        )
+    },
+)
+@triton.jit
+def w8a8_block_static_fp8_matmul_batched_kernel(
+    A,  # (S, K) E4M3 activations (pre-quantized against the static scale by the wrapper)
+    As,  # scalar — static per-tensor activation scale (calibration-time)
+    B,  # (num_experts, N, K) FP8 weights; under GATE the (num_experts, 2N, K) gate|up stack
+    Bs,  # (num_experts, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales (2N under GATE)
+    C,  # (S, N) output; under an OUTPUT_FORMAT the FP8-requantized intermediate
+    Cs,  # (S, N // BLOCK_SIZE_N) per-(row, block) output scale; written iff OUTPUT_FORMAT
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    ExpertIds,  # (S,) — which expert each batch element routes to
+    GatherIdx,  # (S,) int — batch_id -> source row of A; read only when not None
+    ScatterIdx,  # (S,) int — batch_id -> destination row of C; read only when not None
+    # Shape
+    S,
+    N,
+    K,
+    # Strides
+    stride_a_m,
+    stride_a_k,
+    stride_b_e,
+    stride_b_k,
+    stride_b_n,
+    stride_bs_e,
+    stride_bs_k,
+    stride_bs_n,
+    stride_c_m,
+    stride_c_n,
+    stride_cs_m,
+    stride_cs_n,
+    stride_bias_e,
+    stride_bias_n,
+    stride_eid,
+    num_experts,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    SWAP_AB: tl.constexpr = False,
+    # Gate|up fusion epilogue (GATE=False -> plain batched GEMM, every arm below folds out)
+    GATE: tl.constexpr = False,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    OUTPUT_FORMAT: tl.constexpr = None,  # None | "fp8" (per-(row, block) requant of the intermediate)
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
+    PDL: tl.constexpr = False,
+):
+    """Block-scale batched FP8 expert matmul with a static (per-tensor) activation scale — the
+    block-dynamic batched sibling (one program per routed token + N-tile, fake-batch decode,
+    ``SWAP_AB``, ``GATE`` gate|up fusion) with the 2D ``block_static`` format: ``A`` arrives
+    pre-quantized against the calibrated scalar, per-block weight scales apply per-K-tile
+    (``accumulate`` ``"static"``, ``FAKE_BATCH``), and the scalar activation scale multiplies the
+    accumulator once after the loop. bf16 GLU output only (no fused requant). GATE=False is the plain GEMM."""
+    a_s_static = tl.load(As)  # per-tensor static activation scale, applied post-loop
+    if PDL:
+        gdc_wait()
+    batch_id, pid_n, expert_id, A, B, C, Bs, in_row, out_row = expert_setup(
+        A,
+        B,
+        C,
+        Bs,
+        ExpertIds,
+        GatherIdx,
+        ScatterIdx,
+        stride_a_m,
+        stride_b_e,
+        stride_c_m,
+        stride_bs_e,
+        stride_eid,
+    )
+    # EP sentinel: row routed to a non-local expert; output is left uninit.
+    if expert_id >= num_experts:
+        return
+    # the N tile may subdivide the quant block — see the dynamic sibling / scale_subblock_pruner
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    offs_bn = pid_n * n_width + tl.arange(0, n_width)
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = operand_tile_ptrs(A, tl.arange(0, BLOCK_SIZE_M) * 0, offs_k, stride_a_m, stride_a_k, "pointer", True)
+    # One gate|up weight tile over the interleaved rows + one block-scale pointer;
+    # each weight row's scale block follows from its global row.
+    b_ptrs = weight_tile_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, SWAP_AB)
+    bs_ptr = Bs + (pid_n * n_width // BLOCK_N) * stride_bs_n
+    bs_off = ((pid_n * n_width % BLOCK_N + tl.arange(0, n_width)) // BLOCK_N) * stride_bs_n
+    acc = acc_init("dot", BLOCK_SIZE_M, n_width, SWAP_AB)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a, _ = load_act_static(a_ptrs, 0, 0, 0, None, 0, 0.0, "pointer", False)  # pre-quantized E4M3 token (fake-batch replicated)
+        w, b_s = load_weight_static(
+            b_ptrs, b_ptrs, bs_ptr + bs_off, None, 0, 0, 0, 0, 0, 0,
+            GATE, False, "pointer", SWAP_AB, BLOCK_SIZE_N, BLOCK_SIZE_K,
+        )
+        acc = static_dot(acc, a, w, b_s, SWAP_AB, BLOCK_SIZE_K, True)
+        a_ptrs, _, b_ptrs, bs_ptr, _, _ = advance_ptrs(
+            a_ptrs, a_ptrs, b_ptrs, bs_ptr, b_ptrs, bs_ptr,
+            BLOCK_SIZE_K * stride_a_k, 0, BLOCK_SIZE_K * stride_b_k, stride_bs_k,
+            "pointer", "pointer", False, True, False,
+        )
+    acc = acc * a_s_static
+    if PDL:
+        gdc_launch_dependents()
+    gemm_epilogue(
+        C, Cs, acc, out_row, pid_n, 0, out_row, 1, stride_c_n, stride_cs_m, stride_cs_n,
+        BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, OUTPUT_FORMAT, BLOCK_K,
+        ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+        COMPUTE_MODE="dot", SWAP_AB=SWAP_AB, FAKE_BATCH=True,
+        Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n,
+        global_row=expert_id,
+    )
+@bayesian_autotune(
+    # S (routed rows) keyed like the block-dynamic/mx batched siblings — decode re-tunes per batch.
+    get_accelerator_autotuning_configs(tune_block_nk=True, swap_ab=True),
+    ["N", "K", "S"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    # BLOCK_SIZE_K/N are tuned axes; the K-loop is maskless and the N-tile store is
+    # row-masked only — veto non-dividing tiles on both.
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            fp8_dot_warp_pruner(),
+            block_within_dim_pruner("K"),
+            block_within_dim_pruner("N", "BLOCK_SIZE_N"),
+        )
+    },
+)
+@triton.jit
+def w8a8_tensor_dynamic_fp8_matmul_batched_kernel(
+    A,  # (S, K) pre-quantized FP8 activations
+    As,  # (S,) per-token activation scales
+    B,  # (num_experts, N, K) FP8 weight matrices
+    Bs,  # (num_experts, 1, 1) per-tensor weight scales
+    C,  # (S, N) output
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    ExpertIds,  # (S,) — which expert each batch element routes to
+    GatherIdx,  # (S,) int — batch_id -> source row of A; read only when not None
+    ScatterIdx,  # (S,) int — batch_id -> destination row of C; read only when not None
+    # Shape
+    S,
+    N,
+    K,
+    # Strides
+    stride_a_m,
+    stride_a_k,
+    stride_as_m,
+    stride_b_e,
+    stride_b_k,
+    stride_b_n,
+    stride_bs_e,
+    stride_c_m,
+    stride_c_n,
+    stride_bias_e,
+    stride_bias_n,
+    stride_eid,
+    num_experts,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SWAP_AB: tl.constexpr = False,
+    PDL: tl.constexpr = False,
+):
+    """Tensor-scale batched FP8 expert matmul kernel.
+    Activations are already quantized; the kernel applies per-token activation
+    scales and per-expert tensor weight scales.
+    ``SWAP_AB`` (tuner axis, M=1 decode): weight output rows in the MMA M dim (``B`` as ``[BN, BK]``,
+    single token padded to N=16); column 0 of the ``[BN, 16]`` accumulator is the result. Both
+    scales are per-token/per-tensor scalars, applied once after the loop, orientation-agnostic."""
+    if PDL:
+        gdc_wait()
+    batch_id, pid_n, expert_id, A, B, C, Bs, in_row, out_row = expert_setup(
+        A,
+        B,
+        C,
+        Bs,
+        ExpertIds,
+        GatherIdx,
+        ScatterIdx,
+        stride_a_m,
+        stride_b_e,
+        stride_c_m,
+        stride_bs_e,
+        stride_eid,
+    )
+    # EP sentinel: row routed to a non-local expert; output is left uninit.
+    if expert_id >= num_experts:
+        return
+    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = operand_tile_ptrs(A, tl.arange(0, BLOCK_SIZE_M) * 0, offs_k, stride_a_m, stride_a_k, "pointer", True)
+    b_ptrs = oriented_tile_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, SWAP_AB)
+    b_s = tl.load(Bs)
+    a_s = tl.load(As + in_row * stride_as_m)
+    accumulator = acc_init("dot", BLOCK_SIZE_M, BLOCK_SIZE_N, SWAP_AB)
+    for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a, _ = load_act_plain(a_ptrs, 0, 0, 0, None, 0, "pointer", False)
+        b, _ = load_weight_plain(
+            b_ptrs, b_ptrs, 0, 0, 0, False, False, "pointer", SWAP_AB, BLOCK_SIZE_N, BLOCK_SIZE_K
+        )
+        accumulator = accumulator + fp8_dot(a, b, SWAP_AB, BLOCK_SIZE_K)
+        a_ptrs, _, b_ptrs, _, _, _ = advance_ptrs(
+            a_ptrs, a_ptrs, b_ptrs, b_ptrs, b_ptrs, b_ptrs,
+            BLOCK_SIZE_K * stride_a_k, 0, BLOCK_SIZE_K * stride_b_k, 0,
+            "pointer", "pointer", False, False, False,
+        )
+    accumulator = acc_finalize(accumulator, "dot", BLOCK_SIZE_N, SWAP_AB) * a_s * b_s
+    # this split keeps its own store path (ungated, per-tensor dequant) but takes the same bias
+    accumulator = add_bias(
+        accumulator, Bias, stride_bias_e, stride_bias_n, expert_id, pid_n, BLOCK_SIZE_N
+    )
+    if PDL:
+        gdc_launch_dependents()
+    store_row(C, accumulator, pid_n, stride_c_n, BLOCK_SIZE_M, BLOCK_SIZE_N)
+# The MXFP4/MXFP8 (and packed-activation) splits key themselves — the tuner appends every tensor
+# arg's dtype to its cache key. BLOCK_SIZE_M is always 1 here (per-token decode), so plain `dot`
+# is excluded (only scalar / dot_scaled-swap are emitted); the swapped dot helper stays
+# implemented for future shapes but is not fielded. Swap verdicts are B200 (sm_100) — re-measure
+# on H100 or the target device before inheriting.
+@bayesian_autotune(
+    get_accelerator_autotuning_configs(
+        mx=True,
+        tune_block_nk=True,
+        compute_modes=("dot_scaled", "scalar"),
+        swap_ab=True,
+        pre_hook=rebind_batched_mx_bs_descriptor,
+    ),
+    # ACTIVATION_FORMAT keys the inline act-quant grid: A stays raw bf16 under every
+    # format, so the tuner's dtype-appended key can't split W4A8 from W4A4 itself.
+    # SWIZZLED_SCALES keys the weight-scale load: it constrains the config space (BK % 128 == 0)
+    # and picks a different optimum (full-block descriptor vs the pointer scale), so the swizzled
+    # and un-swizzled launches of one shape must not share a tune.
+    # GATE keys the tune AND the GATE-conditional prunes (the packed-fp4 silent-zeros
+    # fence in mx_config_pruner runs once per key, so a GATE=False-first tune must not
+    # hand its winner to a GATE=True launch at the same shape).
+    # OUTPUT_FORMAT keys the requant epilogue explicitly (the C/Cs dtype append splits it
+    # incidentally today; this kernel's requant Cs is always row-major — no SWIZZLED_OUT axis).
+    ["N", "K", "S", "ACTIVATION_FORMAT", "SWIZZLED_SCALES", "GATE", "OUTPUT_FORMAT"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    # BK-within-K + the sm_10x MMA-shape guards (swapped dot_scaled needs BN >= 128 for the
+    # native scaled-MMA; smaller-BN swap configs never win and mislead the TPE).
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            mx_config_pruner("K", "N"), swizzled_scale_config_pruner(allow_gate_subblock=True), smem_pruner(),
+            dot_scaled_staging_pruner(),
+        )
+    },
+)
+@triton.jit
+def mx_dynamic_matmul_batched_kernel(
+    A,  # (S, K) activations: raw BF16/FP16 (inline-quant) or E4M3 (pre-quantized, As set)
+    As,  # (S, K // SCALE_GROUP_K) UE8M0 act scales; None ⇒ inline-quant, read iff not None
+    B,  # (num_experts, N, K[/2]); under GATE the (num_experts, 2N, K[/2]) gate|up stack
+    Bs,  # (num_experts, N, K // SCALE_GROUP_K) UE8M0 weight scales (2N under GATE)
+    BSDescriptor,  # host TMA descriptor over Bs when SWIZZLED (BN=128 bulk load); dummy otherwise
+    C,  # (S, N[/2]) output; under an OUTPUT_FORMAT the MX-requantized intermediate
+    Cs,  # (S, N // SCALE_GROUP_K) UE8M0 output scale; written iff OUTPUT_FORMAT
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    AsGlobal,  # (1,) fp32 NVFP4 activation global g_a — SOLELY normalizes the inline raw-A quant (A/g_a); read iff not None
+    AsBsGlobal,  # (num_experts,) fp32 NVFP4 WEIGHT global g_b — recovers on the accumulator together with AsGlobal (in-register product); read iff not None
+    CsGlobal,  # (1,) fp32 NVFP4 output global (next proj's provided input_scale); normalizes the requant; read iff not None
+    ExpertIds,  # (S,) — which expert each routed row uses
+    GatherIdx,  # (S,) int — batch_id -> source row of A; read only when not None
+    ScatterIdx,  # (S,) int — batch_id -> destination row of C; read only when not None
+    # Shape
+    S,
+    N,
+    K,
+    # Strides
+    stride_a_m,
+    stride_a_k,
+    stride_as_m,
+    stride_b_e,
+    stride_b_k,
+    stride_b_n,
+    stride_bs_e,
+    stride_bs_k,
+    stride_bs_n,
+    stride_c_m,
+    stride_c_n,
+    stride_cs_m,
+    stride_cs_n,
+    stride_bias_e,
+    stride_bias_n,
+    stride_as_global,  # expert stride of AsGlobal (0 broadcasts a shared scalar)
+    stride_bs_global,  # expert stride of AsBsGlobal (``global_scale_stride``)
+    stride_cs_global,  # expert stride of CsGlobal (0 broadcasts a shared scalar)
+    stride_eid,
+    num_experts,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    COMPUTE_MODE: tl.constexpr,
+    SWAP_AB: tl.constexpr = False,
+    # Gate|up fusion epilogue (GATE=False -> plain batched GEMM, every arm below folds out)
+    GATE: tl.constexpr = False,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    # the output format name (None | "mxfp8" | "mxfp4" | "nvfp4")
+    OUTPUT_FORMAT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
+    ACTIVATION_FORMAT: tl.constexpr = "mxfp8",
+    # SWIZZLED_SCALES: Bs arrives pre-swizzled (SWIZZLE_32_4_4) — the checkpoint stores one layout,
+    # shared with the grouped (prefill) kernel. Read via load_weight's per-expert scale leaf off the
+    # single Bs pointer (+ BSDescriptor for the BN=128 bulk load); un-swizzled Bs takes the affine
+    # arm in the same leaf. The op never swizzles — a 3D caller runs un-swizzled at no penalty.
+    SWIZZLED_SCALES: tl.constexpr = False,
+    PDL: tl.constexpr = False,
+):
+    """Unified batched microscaled expert matmul (MXFP8/MXFP4/NVFP4, W4A8/W4A4) with
+    fused act quant.
+    One routed row + one N-tile per program; expert looked up from ``ExpertIds``. ``A`` is
+    quantized to E4M3 per K-group inline (UE8M0 scale). The weight dtype picks the
+    weight format (2 = packed E2M1 / MXFP4, 1 = unpacked E4M3 / MXFP8); ``COMPUTE_MODE``
+    picks ``tl.dot_scaled`` (native M=128) vs the scalar CUDA-core reduce (wins at decode).
+    ``SWAP_AB`` (tuner axis, M=1 decode): weight output rows in the MMA M dim (``B`` as ``[BN, BK]``,
+    single token padded to N=16); column 0 of the ``[BN, 16]`` accumulator is the result. dot_scaled
+    uses the swapped scaled-MMA; scalar reduces over K with the weight output-rows-major.
+    """
+    if PDL:
+        gdc_wait()
+    batch_id, pid_n, expert_id, A, B, C, Bs, in_row, out_row = expert_setup(
+        A,
+        B,
+        C,
+        Bs,
+        ExpertIds,
+        GatherIdx,
+        ScatterIdx,
+        stride_a_m,
+        stride_b_e,
+        stride_c_m,
+        stride_bs_e,
+        stride_eid,
+        ADVANCE_BS=False,  # scale leaf applies the per-expert offset (swizzled indexes by block)
+    )
+    # EP sentinel: row routed to a non-local expert; output is left uninit.
+    if expert_id >= num_experts:
+        return
+    # each operand's format is its dtype: uint8 = packed E2M1 (two values per byte, W4A4
+    # for A / MXFP4 for B — it also keys the autotune cache), else E4M3
+    ACT_VALUES_PER_BYTE: tl.constexpr = 2 if A.dtype.element_ty == tl.uint8 else 1
+    WEIGHT_VALUES_PER_BYTE: tl.constexpr = 2 if B.dtype.element_ty == tl.uint8 else 1
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    # Non-128 N: the partial last N-tile's pointer rows wrap into B (offs_bn % N) so the load
+    # never reads past the expert's N rows; the wrapped columns' output is masked off (N_COLS)
+    # in the epilogue. Inert when N % BLOCK_SIZE_N == 0 (the affine arm's BN|N veto), so it is
+    # load-bearing only for the swizzled arm, whose scale rides pid_n's block index, not offs_bn.
+    offs_bn = (pid_n * n_width + tl.arange(0, n_width)) % (2 * N if GATE else N)
+    offs_kb = tl.arange(0, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE)
+    offs_sf = tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K)
+    offs_ka = tl.arange(0, BLOCK_SIZE_K // ACT_VALUES_PER_BYTE)
+    a_ptrs = operand_tile_ptrs(A, tl.arange(0, BLOCK_SIZE_M) * 0, offs_ka, stride_a_m, stride_a_k, "pointer", True)
+    # As is not None ⇒ pre-quantized: A is E4M3 + As UE8M0 (the down reading a requantized
+    # intermediate). Else A is raw, quantized inline onto ACTIVATION_FORMAT's grid (gate_up / plain —
+    # packed E2M1 under fp4, one act row per program so the quant is decode-free); As stays None.
+    if As is not None:  # build the scale pointers only when the scale is read
+        as_ptrs = (
+            As
+            + in_row * stride_as_m
+            + tl.arange(0, BLOCK_SIZE_M)[:, None] * 0
+            + offs_sf[None, :]
+        )
+    else:
+        as_ptrs = a_ptrs  # dead placeholder so advance_ptrs can take it unconditionally
+    # GATE reads one weight tile spanning the interleaved gate|up rows, oriented by
+    # SWAP_AB; load_weight reads value + scale (swizzled/un-swizzled hidden) off these pointers.
+    b_ptrs = weight_tile_ptrs(
+        B, offs_bn, offs_kb, stride_b_n, stride_b_k, SWAP_AB
+    )
+    accumulator = acc_init(COMPUTE_MODE, BLOCK_SIZE_M, n_width, SWAP_AB)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a, a_scale = load_act_mx(
+            a_ptrs, as_ptrs, AsGlobal, expert_id * stride_as_global, None, None, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
+            "pointer", False, False, False,
+            BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K, ACTIVATION_FORMAT,
+        )
+        b, b_s = load_weight_mx(
+            b_ptrs, b_ptrs, Bs, None, BSDescriptor, Bs, 0, 0, 0, 0, expert_id, pid_n, k, N, K,
+            stride_bs_e, stride_bs_n, stride_bs_k,
+            GATE, False, True, "pointer", SWAP_AB, SWIZZLED_SCALES,
+            BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, WEIGHT_VALUES_PER_BYTE,
+        )
+        accumulator = mx_compute(
+            accumulator, a, a_scale, b, b_s, COMPUTE_MODE,
+            n_width, BLOCK_SIZE_K, SCALE_GROUP_K, SWAP_AB,
+        )
+        a_ptrs, as_ptrs, b_ptrs, _, _, _ = advance_ptrs(
+            a_ptrs, as_ptrs, b_ptrs, b_ptrs, b_ptrs, b_ptrs,
+            (BLOCK_SIZE_K // ACT_VALUES_PER_BYTE) * stride_a_k,
+            BLOCK_SIZE_K // SCALE_GROUP_K,
+            (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE) * stride_b_k,
+            0,
+            "pointer", "pointer", As is not None, False, False,
+        )
+    # NVFP4 two-level: block e4m3 scales rode through the reduce; recover the combined per-tensor
+    if PDL:
+        gdc_launch_dependents()
+    gemm_epilogue(
+        C, Cs, accumulator, out_row, pid_n, 0, out_row, 1, stride_c_n, stride_cs_m, stride_cs_n,
+        BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, OUTPUT_FORMAT, SCALE_GROUP_K,
+        ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+        COMPUTE_MODE=COMPUTE_MODE, SWAP_AB=SWAP_AB, FAKE_BATCH=True, N_COLS=N,
+        CsGlobal=CsGlobal,
+        stride_cs_global=stride_cs_global,
+        GlobalScale=AsBsGlobal,
+        stride_global_e=stride_bs_global,
+        GlobalScaleA=AsGlobal,
+        stride_global_a=stride_as_global,
+        global_row=expert_id,
+        Bias=Bias,
+        stride_bias_e=stride_bias_e,
+        stride_bias_n=stride_bias_n,
+    )
+@bayesian_autotune(
+    # weight-only plain bf16 dot (fp4/fp8 weight upcast per-group in-loop), plus the swapped
+    # decode reduce (weight output rows lead the tile, no MMA — see the scope pruner).
+    get_accelerator_autotuning_configs(
+        tune_block_nk=True, compute_modes=("dot", "dot_scaled", "scalar"), swap_ab=True
+    ),
+    ["N", "K", "S", "GATE"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    prune_configs_by={
+        # No BK-within-K veto: this kernel's tail tile slides back in bounds (see the K-loop),
+        # so any BK the shape can hold is legal — that is the point, it lets a K=2880 decode
+        # take a 256-byte burst instead of 32. BK must still FIT (the clamp needs K >= BK).
+        "early_config_prune": compose_pruners(
+            block_fits_dim_pruner("K"),
+            block_within_dim_pruner("N", "BLOCK_SIZE_N"),
+            mx_config_pruner("K", "N", block_within_k=False),  # dot_scaled shape gates
+            weight_only_swap_scope_pruner(),
+        )
+    },
+)
+@triton.jit
+def mx_weight_only_matmul_batched_kernel(
+    A,  # (rows, K) raw BF16/FP16 activations — NOT quantized
+    B,  # (num_experts, N, K[/2]) MXFP4/NVFP4/MXFP8 weights; 2N under GATE
+    Bs,  # (num_experts, N, K // SCALE_GROUP_K) group scales — UE8M0 (MX) or E4M3 (NVFP4)
+    BsGlobal,  # (num_experts,) fp32 NVFP4 per-expert global — recovers on the accumulator; read iff not None
+    C,
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    ExpertIds,
+    GatherIdx,
+    ScatterIdx,
+    S,
+    N,
+    K,
+    stride_a_m,
+    stride_a_k,
+    stride_b_e,
+    stride_b_k,
+    stride_b_n,
+    stride_bs_e,
+    stride_bs_n,
+    stride_bs_k,
+    stride_c_m,
+    stride_c_n,
+    stride_bias_e,
+    stride_bias_n,
+    stride_bs_global,  # expert stride of BsGlobal (``global_scale_stride``)
+    stride_eid,
+    num_experts,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    WEIGHT_VALUES_PER_BYTE: tl.constexpr,
+    COMPUTE_MODE: tl.constexpr = "dot",
+    SWAP_AB: tl.constexpr = False,
+    GATE: tl.constexpr = False,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
+    PDL: tl.constexpr = False,
+):
+    """weight-only batched (decode) expert matmul: raw bf16 activations against MXFP4/MXFP8 weights upcast
+    to bf16 in-loop (unpack + per-group group-scale), plain ``tl.dot``. One routed row + one N-tile
+    per program (expert from ``ExpertIds``). Pointer/affine. ``SWAP_AB`` (tuner axis, M=1 decode):
+    weight output rows lead the tile and the CUDA-core reduce replaces the MMA. ``GATE`` fuses gate|up."""
+    if PDL:
+        gdc_wait()
+    batch_id, pid_n, expert_id, A, B, C, Bs, in_row, out_row = expert_setup(
+        A, B, C, Bs, ExpertIds, GatherIdx, ScatterIdx,
+        stride_a_m, stride_b_e, stride_c_m, stride_bs_e, stride_eid, ADVANCE_BS=False,
+    )
+    if expert_id >= num_experts:  # EP sentinel: non-local expert, output left uninit
+        return
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    # non-128 N: the last N-tile's rows run past B (N=320, n_width=256 -> tile 2 wants rows
+    # 512..767 of 640), so wrap them back into the tensor. The wrapped columns are masked off by
+    # the epilogue's N_COLS. An unwrapped read is out of bounds and faults once the allocation
+    # layout puts an unmapped page there — invisible in a short run, an illegal access in a long one.
+    offs_bn = (pid_n * n_width + tl.arange(0, n_width)) % (2 * N if GATE else N)
+    offs_kb = tl.arange(0, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE)
+    a_ptrs = operand_tile_ptrs(
+        A, tl.arange(0, BLOCK_SIZE_M) * 0, tl.arange(0, BLOCK_SIZE_K),
+        stride_a_m, stride_a_k, "pointer", True,
+    )
+    b_ptrs = weight_tile_ptrs(
+        B, offs_bn, offs_kb, stride_b_n, stride_b_k, SWAP_AB
+    )
+    accumulator = acc_init(COMPUTE_MODE, BLOCK_SIZE_M, n_width, SWAP_AB)
+    # Each trip indexes off the BASE tile by the scalar ``k_start`` rather than advancing a
+    # pointer tile in place: an advanced tile stays live across the loop, and at [2*BN, BK/2]
+    # int64 that is ~128 registers per thread of addresses alone — enough to halve resident
+    # CTAs and starve the weight stream (measured n_regs 128 with spills vs 55 for the same
+    # math indexed from a base). The clamp below already computes ``k_start``, so the walk is
+    # the same addresses either way.
+    # The K-loop is maskless, which normally pins BK to a divisor of K — leaving a
+    # non-power-of-two K (gpt-oss 2880) with BK=64, a 32-byte burst per weight row over 45
+    # trips. Instead the LAST tile slides back to end at K (``k_start`` clamp): every load
+    # stays in bounds, the tile re-reads positions the previous trip already accumulated, and
+    # zeroing those activation lanes cancels them for every arm (a zero operand contributes
+    # nothing to the dot, the scaled MMA, or the reduce). BK | K makes the clamp a no-op, so
+    # aligned shapes keep today's pointer walk bit-for-bit.
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        k_start = tl.minimum(k * BLOCK_SIZE_K, K - BLOCK_SIZE_K)
+        overlap = k * BLOCK_SIZE_K - k_start
+        a, _ = load_act_plain(
+            a_ptrs + k_start * stride_a_k, 0, 0, 0, None, 0, "pointer", False
+        )
+        if overlap > 0:  # only the slid tail trip pays the select
+            a = tl.where(tl.arange(0, BLOCK_SIZE_K)[None, :] >= overlap, a, 0.0)
+        b_k = b_ptrs + (k_start // WEIGHT_VALUES_PER_BYTE) * stride_b_k
+        w, w_s = load_weight_mx(
+            b_k, b_k, Bs, None, None, Bs, 0, 0, 0, 0, expert_id, pid_n, k, N, K,
+            stride_bs_e, stride_bs_n, stride_bs_k,
+            GATE, False, True, "pointer", SWAP_AB, False,
+            BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, WEIGHT_VALUES_PER_BYTE,
+            k_col_off=-(overlap // SCALE_GROUP_K),
+        )
+        accumulator = mx_weight_only_compute(
+            accumulator, a, w, w_s, COMPUTE_MODE, BLOCK_SIZE_K, n_width, SCALE_GROUP_K,
+            SWAP_AB,
+        )
+    if PDL:
+        gdc_launch_dependents()
+    gemm_epilogue(
+        C, None, accumulator, out_row, pid_n, 0, out_row, 1, stride_c_n, 1, 1,
+        BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, BLOCK_SIZE_K,
+        ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+        COMPUTE_MODE=COMPUTE_MODE, SWAP_AB=SWAP_AB, FAKE_BATCH=True, N_COLS=N,
+        GlobalScale=BsGlobal,
+        stride_global_e=stride_bs_global,
+        global_row=expert_id,
+        Bias=Bias,
+        stride_bias_e=stride_bias_e,
+        stride_bias_n=stride_bias_n,
+    )
+@bayesian_autotune(
+    get_accelerator_autotuning_configs(tune_block_nk=True, swap_ab=True),
+    # S (routed rows) keyed like the fp8/mx batched siblings — decode re-tunes per batch;
+    # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide).
+    ["N", "K", "S", "GATE"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    # BLOCK_SIZE_K/N are tuned axes; the K-loop is maskless and the N-tile store is
+    # row-masked only — veto non-dividing tiles on both.
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            block_within_dim_pruner("K"),
+            block_within_dim_pruner("N", "BLOCK_SIZE_N"),
+        )
+    },
+)
+@triton.jit
+def full_precision_matmul_batched_kernel(
+    A,  # (rows, K) BF16/FP16 activations
+    B,  # (num_experts, N, K) weights in A's dtype; under GATE the (num_experts, 2N, K) gate|up stack
+    C,  # (S, N) output; under GATE the GLU intermediate
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    ExpertIds,  # (S,) — which expert each batch element routes to
+    GatherIdx,  # (S,) int — batch_id -> source row of A; read only when not None
+    ScatterIdx,  # (S,) int — batch_id -> destination row of C; read only when not None
+    # Shape
+    S,
+    N,
+    K,
+    # Strides
+    stride_a_m,
+    stride_a_k,
+    stride_b_e,
+    stride_b_k,
+    stride_b_n,
+    stride_c_m,
+    stride_c_n,
+    stride_bias_e,
+    stride_bias_n,
+    stride_eid,
+    num_experts,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SWAP_AB: tl.constexpr = False,
+    # Gate|up fusion epilogue (GATE=False -> plain batched GEMM). No requant arm: the
+    # full-precision chain has no quantized intermediate — down consumes the GLU output as is.
+    GATE: tl.constexpr = False,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
+    PDL: tl.constexpr = False,
+):
+    """Full-precision batched expert matmul kernel: plain ``tl.dot`` over unquantized
+    BF16/FP16 activations and weights, fp32 accumulation, no scales anywhere. ``GATE``
+    computes gate|up as ONE stacked tile + dot (straight-line, both orientations) and
+    applies the ``ACT_FN``/SwiGLU ``glu``. ``SWAP_AB`` (tuner axis, M=1 decode): weight
+    output rows in the MMA M dim, the single token padded to the N=16 atom."""
+    if PDL:
+        gdc_wait()
+    batch_id, pid_n, expert_id, A, B, C, _, in_row, out_row = expert_setup(
+        A,
+        B,
+        C,
+        None,  # no scales on the full-precision path
+        ExpertIds,
+        GatherIdx,
+        ScatterIdx,
+        stride_a_m,
+        stride_b_e,
+        stride_c_m,
+        0,
+        stride_eid,
+        ADVANCE_BS=False,
+    )
+    # EP sentinel: row routed to a non-local expert; output is left uninit.
+    if expert_id >= num_experts:
+        return
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    offs_bn = pid_n * n_width + tl.arange(0, n_width)
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    a_ptrs = operand_tile_ptrs(A, tl.arange(0, BLOCK_SIZE_M) * 0, offs_k, stride_a_m, stride_a_k, "pointer", True)
+    # GATE reads one weight tile spanning the interleaved gate|up rows;
+    # GATE=False -> the plain oriented tile.
+    b_ptrs = weight_tile_ptrs(
+        B, offs_bn, offs_k, stride_b_n, stride_b_k, SWAP_AB
+    )
+    accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, SWAP_AB)
+    for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a, _ = load_act_plain(a_ptrs, 0, 0, 0, None, 0, "pointer", False)
+        w, _ = load_weight_plain(
+            b_ptrs, b_ptrs, 0, 0, 0, GATE, False, "pointer", SWAP_AB, BLOCK_SIZE_N, BLOCK_SIZE_K
+        )
+        accumulator = accumulator + fp8_dot(a, w, SWAP_AB, BLOCK_SIZE_K)
+        a_ptrs, _, b_ptrs, _, _, _ = advance_ptrs(
+            a_ptrs, a_ptrs, b_ptrs, b_ptrs, b_ptrs, b_ptrs,
+            BLOCK_SIZE_K * stride_a_k, 0, BLOCK_SIZE_K * stride_b_k, 0,
+            "pointer", "pointer", False, False, False,
+        )
+    if PDL:
+        gdc_launch_dependents()
+    gemm_epilogue(
+        C, None, accumulator, out_row, pid_n, 0, out_row, 1, stride_c_n, 1, 1,
+        BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, BLOCK_SIZE_K,
+        ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+        COMPUTE_MODE="dot", SWAP_AB=SWAP_AB, FAKE_BATCH=True,
+        Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n,
+        global_row=expert_id,
+    )
+GATE_UNSTACK_MAX_S = 16  # the unstacked-gate decode band (see the dispatch in the wrapper)
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("w8a8_block_dynamic_fp8_matmul_batched"),
+    mutates_args=(),
+    opaque=True,
+)
+def w8a8_block_dynamic_fp8_matmul_batched(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor | None,
+    Bs: torch.Tensor,
+    expert_ids: torch.Tensor,
+    block_size: list[int],
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    activation_format: str | None = None,
+    quantize_output: bool = False,
+    output_dtype: torch.dtype | None = None,
+    gather_idx: torch.Tensor | None = None,
+    scatter_idx: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    """Block-scale batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T; activations
+    quantized offline in one pass. The ``gate``/``act_fn``/``swiglu_*``/``requant``/``output_dtype``
+    flags are the dispatcher's gate|up fusion kwargs (torch custom ops take only primitives —
+    ``matmul_batched`` unpacks the bundle). ``gather_idx``/``scatter_idx`` map the source row of A
+    / destination row of C per program (None = row s). Returns ``[C]``, or ``[C, Cs]`` under
+    ``requant``.
+    A:  (rows, K) raw bf16/fp16 activations — rows addressed via ``gather_idx``
+    B:  (num_experts, N, K) FP8 weights; under ``gate`` the (num_experts, 2N, K) gate|up stack
+    Bs: (num_experts, N // block_n, K // block_k) per-block weight scales (2N under gate)
+    """
+    validate_dense_operands(A, B)
+    output_dtype = resolve_output_dtype(output_dtype, A, As)
+    # S is the routed-row count (one program per expert_id); A may hold fewer rows when
+    # gather_idx maps many programs to one source row (gate_up reading unexpanded hidden).
+    K = A.shape[1]
+    S = expert_ids.shape[0]
+    num_experts, n_rows, N = expert_weight_shape(B, gate)
+    assert len(block_size) == 2, (
+        f"block_size must be [block_n, block_k], got {block_size}"
+    )
+    block_n, block_k = block_size[0], block_size[1]
+    require_moe_dims_aligned(N, K, block_n, block_k)
+    assert Bs.shape == (num_experts, n_rows // block_n, K // block_k), (
+        f"Bs shape {tuple(Bs.shape)} != expected ({num_experts}, {n_rows // block_n}, {K // block_k})"
+    )
+    bs_u8 = ue8m0_as_uint8(Bs)
+    # Offline quant wins here even at decode. An inline quant would rerun once per N-tile
+    # of the (S x N-tiles) grid, and block-FP8 quant is an fp32 amax+div per element, so
+    # the redundant work outweighs the extra launch down to T=1 (inline only edges ahead
+    # near T=64). UE8M0 quant is ~free per pass, which is why the MX kernels do it inline.
+    assert activation_format in (None, "fp8"), (
+        f"block-dynamic activations are E4M3 ('fp8'), got {activation_format!r}"
+    )
+    output_format = "fp8" if quantize_output else None
+    requant = output_format is not None
+    # the requantized intermediate's scale groups follow gate_up's block_n, and the
+    # down consumes per-block_k — a non-square block format would misalign them
+    assert not requant or block_size[0] == block_size[1], (
+        f"the fused 'fp8' requant needs square quant blocks, got {block_size}"
+    )
+    # The decode band runs the ``gate`` contract UNSTACKED: this same op without the gate
+    # (one plain GEMM over the stacked weight), then the one-kernel ``fused_glu`` — bit-
+    # identical rounding, and exactly the unfused-reference order, so ``simulate_unfused``
+    # needs no carve-out. Holding the gate AND up tiles per CTA doubles the smem footprint
+    # and halves occupancy precisely where the launch is weight-bandwidth-bound, and the
+    # fp8 ``tl.dot`` — unlike the MX/NVFP4 scaled MMA, whose wide M operand earns the
+    # native instruction — gains nothing back: DSV3 shape, same 235MB weight read, stacked
+    # 69.8µs vs unstacked 49.8µs at S=8, a wash by S=32. A requant is the same offline
+    # quant the raw activation gets below, applied to the GLU output — quantizing the bf16
+    # intermediate (the unfused-reference order) where the stacked epilogue quantizes its
+    # fp32 accumulator, a sub-quantum difference consistent with the band's semantics.
+    if gate and S <= GATE_UNSTACK_MAX_S:
+        [gate_up] = w8a8_block_dynamic_fp8_matmul_batched(
+            A, B, As, Bs, expert_ids, block_size,
+            activation_format=activation_format, output_dtype=output_dtype,
+            gather_idx=gather_idx, scatter_idx=scatter_idx,
+        )
+        out = fused_glu(gate_up, act_fn, swiglu_alpha, swiglu_limit,
+                        quant_group=block_n if requant else None,
+                        use_ue8m0=bs_u8.dtype == torch.uint8)
+        return list(out) if requant else [out]
+    # A raw (As is None) -> quantize here (offline); else pre-quantized (As given, e.g. the
+    # requantized intermediate handed to the down projection).
+    if As is None:
+        A_q, A_s = fp8_act_quant_block_dynamic(
+            A, block_k, use_ue8m0=bs_u8.dtype == torch.uint8
+        )
+    else:
+        A_q, A_s = A, As
+    if requant:
+        C = A.new_empty(S, N, dtype=FP8_DTYPE)
+        # UE8M0 model (ue8m0 weights) -> UE8M0 intermediate scales (whole-model contract);
+        # the kernel infers the requant format from this dtype. fp32 weights keep fp32.
+        cs_dtype = bs_u8.dtype  # uint8 (UE8M0) or float32 — the whole-model scale format
+        Cs = torch.empty(S, N // block_n, device=A.device, dtype=cs_dtype)
+    else:
+        C = A.new_empty(S, N, dtype=output_dtype)
+        Cs = None  # unread without an OUTPUT_FORMAT; strides literal below
+    # the N tile is tuned (it may subdivide the block scale — see scale_subblock_pruner)
+    grid = lambda meta: (S, triton.cdiv(N, meta["BLOCK_SIZE_N"]))  # noqa: E731
+    with device_context(A.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(w8a8_block_dynamic_fp8_matmul_batched_kernel)[
+            grid
+        ](
+            A_q,
+            A_s,
+            B,
+            bs_u8,
+            C,
+            Cs,
+            bias,
+            expert_ids,
+            gather_idx,  # None = A is expert-sorted; read only when not None (folds at trace time)
+            scatter_idx,  # None = C is expert-sorted; read only when not None (folds at trace time)
+            S,
+            N,
+            K,
+            A_q.stride(0),
+            A_q.stride(1),
+            A_s.stride(0),
+            B.stride(0),
+            B.stride(2),
+            B.stride(1),
+            bs_u8.stride(0),
+            bs_u8.stride(2),
+            bs_u8.stride(1),
+            C.stride(0),
+            C.stride(1),
+            Cs.stride(0) if requant else 1,
+            Cs.stride(1) if requant else 1,
+            bias_stride_e,
+            bias_stride_n,
+            expert_ids.stride(0),
+            BLOCK_SIZE_K=block_k,
+            BLOCK_N=block_n,
+            BLOCK_K=block_k,
+            num_experts=num_experts,
+            GATE=gate,
+            ACT_FN=act_fn,
+            SWIGLU_ALPHA=swiglu_alpha,
+            SWIGLU_LIMIT=swiglu_limit,
+            OUTPUT_FORMAT=output_format,
+            SIMULATE_UNFUSED=simulate_unfused,
+            INTERMEDIATE_DTYPE=tl_dtype(output_dtype),
+            PDL=decode_pdl(),
+            launch_pdl=decode_pdl(),
+        )
+    return [C, Cs] if requant else [C]
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("w8a8_block_static_fp8_matmul_batched"),
+    mutates_args=(),
+    opaque=True,
+)
+def w8a8_block_static_fp8_matmul_batched(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor,
+    Bs: torch.Tensor,
+    expert_ids: torch.Tensor,
+    block_size: list[int],
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    activation_format: str | None = None,
+    quantize_output: bool = False,
+    output_dtype: torch.dtype | None = None,
+    gather_idx: torch.Tensor | None = None,
+    scatter_idx: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    """Block-scale batched FP8 matmul with a static (per-tensor calibrated) activation scale — the
+    block-dynamic batched sibling with the 2D ``block_static`` format. ``A`` is raw here: the op
+    quantizes it against the scalar ``As`` (offline), the kernel applies the per-block weight scales
+    in the K-loop and the scalar once post-loop. Returns ``[C]``, or ``[C, Cs]`` under
+    ``output_format="fp8"`` (the per-row output scale is independent of the per-tensor input scale).
+    A:  (rows, K) raw bf16/fp16 activations — rows addressed via ``gather_idx``
+    B:  (num_experts, N, K) FP8 weights; under ``gate`` the (num_experts, 2N, K) gate|up stack
+    As: scalar / (1,) — the calibrated per-tensor (static) activation scale
+    Bs: (num_experts, N // block_n, K // block_k) per-block weight scales (2N under gate)
+    """
+    validate_dense_operands(A, B)
+    output_dtype = resolve_output_dtype(output_dtype, A, None)
+    K = A.shape[1]
+    S = expert_ids.shape[0]
+    num_experts, n_rows, N = expert_weight_shape(B, gate)
+    assert len(block_size) == 2, (
+        f"block_size must be [block_n, block_k], got {block_size}"
+    )
+    block_n, block_k = block_size[0], block_size[1]
+    require_moe_dims_aligned(N, K, block_n, block_k)
+    assert Bs.shape == (num_experts, n_rows // block_n, K // block_k), (
+        f"Bs shape {tuple(Bs.shape)} != expected ({num_experts}, {n_rows // block_n}, {K // block_k})"
+    )
+    assert activation_format in (None, "fp8"), (
+        f"block-static activations are E4M3 ('fp8'), got {activation_format!r}"
+    )
+    output_format = "fp8" if quantize_output else None
+    requant = output_format is not None
+    assert not requant or block_n == block_k, (
+        f"the fused 'fp8' requant needs square quant blocks, got {block_size}"
+    )
+    As = As.reshape(1).to(torch.float32)
+    bs_u8 = ue8m0_as_uint8(Bs)
+    # Pre-quantize the raw activations against the calibrated scalar (offline; the kernel folds
+    # the scalar back post-loop).
+    A_q = (A.to(torch.float32) / As).to(FP8_DTYPE)
+    if requant:
+        C = A.new_empty(S, N, dtype=FP8_DTYPE)
+        Cs = torch.empty(S, N // block_n, device=A.device, dtype=bs_u8.dtype)
+    else:
+        C = A.new_empty(S, N, dtype=output_dtype)
+        Cs = None  # unread without an OUTPUT_FORMAT; strides literal below
+    # the N tile is tuned (it may subdivide the block scale — see scale_subblock_pruner)
+    grid = lambda meta: (S, triton.cdiv(N, meta["BLOCK_SIZE_N"]))  # noqa: E731
+    with device_context(A.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(w8a8_block_static_fp8_matmul_batched_kernel)[
+            grid
+        ](
+            A_q,
+            As,
+            B,
+            bs_u8,
+            C,
+            Cs,
+            bias,
+            expert_ids,
+            gather_idx,  # None = A is expert-sorted; read only when not None (folds at trace time)
+            scatter_idx,  # None = C is expert-sorted; read only when not None (folds at trace time)
+            S,
+            N,
+            K,
+            A_q.stride(0),
+            A_q.stride(1),
+            B.stride(0),
+            B.stride(2),
+            B.stride(1),
+            bs_u8.stride(0),
+            bs_u8.stride(2),
+            bs_u8.stride(1),
+            C.stride(0),
+            C.stride(1),
+            Cs.stride(0) if requant else 1,
+            Cs.stride(1) if requant else 1,
+            bias_stride_e,
+            bias_stride_n,
+            expert_ids.stride(0),
+            num_experts,
+            BLOCK_SIZE_K=block_k,
+            BLOCK_N=block_n,
+            BLOCK_K=block_k,
+            GATE=gate,
+            ACT_FN=act_fn,
+            SWIGLU_ALPHA=swiglu_alpha,
+            SWIGLU_LIMIT=swiglu_limit,
+            OUTPUT_FORMAT=output_format,
+            SIMULATE_UNFUSED=simulate_unfused,
+            INTERMEDIATE_DTYPE=tl_dtype(output_dtype),
+            PDL=decode_pdl(),
+            launch_pdl=decode_pdl(),
+        )
+    return [C, Cs] if requant else [C]
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("w8a8_tensor_dynamic_fp8_matmul_batched"),
+    mutates_args=(),
+    opaque=True,
+)
+def w8a8_tensor_dynamic_fp8_matmul_batched(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor | None,
+    Bs: torch.Tensor,
+    expert_ids: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+    gather_idx: torch.Tensor | None = None,
+    scatter_idx: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Tensor-scale batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T. ``A`` raw
+    (``As`` None) -> quantized here (offline, per-token); else pre-quantized (``As`` given).
+    ``gather_idx``/``scatter_idx`` map the source row of A / destination row of C per program
+    (None = row s).
+    A:  (rows, K) raw or pre-quantized FP8 activations — rows addressed via ``gather_idx``
+    B:  (num_experts, N, K) FP8 expert weights
+    As: (rows,) per-token scales, or None when A is raw
+    Bs: (num_experts,) or (num_experts, 1, 1) per-expert weight scales
+    """
+    validate_dense_operands(A, B)
+    output_dtype = resolve_output_dtype(output_dtype, A, As)
+    K = A.shape[1]
+    S = expert_ids.shape[0]
+    num_experts, N, _ = B.shape
+    # Normalize Bs to (num_experts, 1, 1)
+    Bs = normalize_per_expert_scale(Bs, num_experts)
+    bs_u8 = ue8m0_as_uint8(Bs)
+    if As is None:
+        qA, As = fp8_act_quant_tensor_wide(A, K)
+    else:
+        qA = A
+    C = A.new_empty(S, N, dtype=output_dtype)
+    def grid(META):
+        return (S, triton.cdiv(N, META["BLOCK_SIZE_N"]))
+    with device_context(A.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(w8a8_tensor_dynamic_fp8_matmul_batched_kernel)[
+            grid
+        ](
+            qA,
+            As,
+            B,
+            bs_u8,
+            C,
+            bias,
+            expert_ids,
+            gather_idx,  # None = A is expert-sorted; read only when not None (folds at trace time)
+            scatter_idx,  # None = C is expert-sorted; read only when not None (folds at trace time)
+            S,
+            N,
+            K,
+            qA.stride(0),
+            qA.stride(1),
+            As.stride(0),
+            B.stride(0),
+            B.stride(2),
+            B.stride(1),
+            bs_u8.stride(0),
+            C.stride(0),
+            C.stride(1),
+            bias_stride_e,
+            bias_stride_n,
+            expert_ids.stride(0),
+            num_experts=num_experts,
+            PDL=decode_pdl(),
+            launch_pdl=decode_pdl(),
+        )
+    return C
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("mx_dynamic_matmul_batched"), mutates_args=(), opaque=True
+)
+def mx_dynamic_matmul_batched(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor | None,
+    Bs: torch.Tensor,
+    expert_ids: torch.Tensor,
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    activation_format: str | None = None,
+    quantize_output: bool = False,
+    output_dtype: torch.dtype | None = None,
+    gather_idx: torch.Tensor | None = None,
+    scatter_idx: torch.Tensor | None = None,
+    a_global_scale: torch.Tensor | None = None,
+    b_global_scale: torch.Tensor | None = None,
+    output_global_scale: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    """Batched MX matmul ``C[s] = A[s] @ B[expert_ids[s]].T``; activations quantized
+    inline in the kernel (decode: one act row per program, inline is free). The
+    ``gate``/``act_fn``/``swiglu_*``/``requant``/``output_dtype`` flags are ``matmul_batched``'s
+    gate|up fusion kwargs. ``gather_idx``/``scatter_idx`` map the
+    source row of A / destination row of C per program (None = row s). Returns ``[C]``, or
+    ``[C, Cs]`` under ``requant``. Weight format is detected from ``B.dtype``: ``int8`` → packed
+    E2M1 (MXFP4, ``B`` is ``(num_experts, N, K//2)``); ``float8_e4m3fn`` → unpacked E4M3 (MXFP8);
+    both use UE8M0 group-32 scales; under ``gate`` the (num_experts, 2N, K[/2]) gate|up stack.
+    A:  (rows, K) activations — raw bf16/fp16/fp32 (inline-quant) or pre-quantized E4M3
+    expert_ids: (S,) which expert each routed row uses
+    """
+    assert A.ndim == 2 and B.ndim == 3 and Bs.ndim in (3, 5)  # 5D = pre-swizzled
+    assert expert_ids.ndim == 1
+    # A raw (As None) -> quantized inline in the kernel (decode-free UE8M0); pre-quantized
+    # (As given, e.g. the down reading a requantized intermediate) -> loaded with its scales
+    # (the kernel folds on As is None).
+    # the kernel quantizes raw A inline on this grid (fp4 formats pack in-register);
+    # NVFP4 batched runs on the software arms — decode grid BM <= 16 < the native
+    # mxf4nvf4 M=128 staging (scalar / swap-scalar column-unpack + E4M3 scale decode)
+    activation_format = resolve_activation_format(activation_format, B, Bs)
+    output_format = activation_format if quantize_output else None
+    requant = quantize_output
+    if As is not None:
+        assert (As.dtype == torch.float8_e4m3fn) == (Bs.dtype == torch.float8_e4m3fn), (
+            f"activation scales ({As.dtype}) must match the weight scale family ({Bs.dtype})"
+        )
+    pre_quantized = As is not None
+    assert B.dtype in (torch.int8, torch.float8_e4m3fn), (
+        f"B must be int8 (packed E2M1) or float8_e4m3fn (E4M3), got {B.dtype}"
+    )
+    WEIGHT_VALUES_PER_BYTE = NIBBLES_PER_BYTE if B.dtype == torch.int8 else 1
+    # NVFP4 raw activations pre-quantize HERE, above the layout derivation below — the quant
+    # packs A to E2M1, so ACT_VALUES_PER_BYTE and K must be read off the packed tensor. Offline
+    # beats the inline arm at every M for NVFP4 (unlike the UE8M0 formats, whose exponent-only
+    # inline quant is free): its E4M3 block scale costs an amax divide plus a global normalize
+    # per group, and the decode grid re-runs that in every N-tile program against the SAME
+    # routed row. Measured bit-identical, GLM decode gate_up 34.2 -> 31.9us, down 18.9 -> 16.9us.
+    # The act scale stays affine under a swizzled weight: load_act_mx reads it off as_ptrs
+    # row-major, SWIZZLED_SCALES governs only the weight side.
+    if As is None and activation_format == "nvfp4" and A.dtype not in (torch.int8, torch.float8_e4m3fn):
+        # a per-expert g_a quantizes each routed row against ITS expert's global; the rows are
+        # routed slots, so the row -> expert map is this op's own `expert_ids`
+        A, As = MX_ACT_QUANT["nvfp4"](
+            A,
+            global_scale=normalize_global_scale(a_global_scale, B.shape[0]),
+            expert_index=expert_ids,
+        )
+        pre_quantized = True
+    # int8 A = caller-provided packed-E2M1 activations (W4A4, native mxf4 MMA): K is two
+    # values per stored byte and the scales are mandatory (nothing left to quantize).
+    ACT_VALUES_PER_BYTE = NIBBLES_PER_BYTE if A.dtype == torch.int8 else 1
+    if ACT_VALUES_PER_BYTE == NIBBLES_PER_BYTE:
+        assert As is not None, "packed-E2M1 activations need their UE8M0 scales (As)"
+    output_dtype = resolve_output_dtype(output_dtype, A, As)
+    K = A.shape[1] * ACT_VALUES_PER_BYTE
+    S = expert_ids.shape[0]
+    num_experts, n_rows, N = expert_weight_shape(B, gate)
+    K_b = B.shape[2]
+    assert K == WEIGHT_VALUES_PER_BYTE * K_b, (
+        f"K (={K}) must equal {WEIGHT_VALUES_PER_BYTE} * B.shape[2] (={K_b})"
+    )
+    # Bs arrives either row-major (num_experts, n_rows, K // scale_group) — read affine — or
+    # already SWIZZLE_32_4_4 (5D: 1, num_experts * n_rows // 128, cols // 4, 2, 256), the shared
+    # checkpoint layout swizzled once at model load (the deployment contract, no per-call
+    # rearrange). The format is the scale dtype (E4M3 = NVFP4 group-16, UE8M0 = MX group-32); the
+    # swizzled cols encode (K // scale_group) // 4.
+    swizzled_scales = Bs.ndim == 5
+    scale_group = mx_scale_family(Bs, K)
+    if not swizzled_scales:
+        assert Bs.shape == (num_experts, n_rows, K // scale_group), (
+            f"Bs shape {tuple(Bs.shape)} != ({num_experts}, {n_rows}, {K // scale_group})"
+        )
+    else:
+        # the artifact's expert count and per-expert row blocks must match THIS weight stack — a
+        # wrong-layer artifact with a matching K would otherwise dequantize with garbage scales silently
+        assert Bs.shape[:2] == (num_experts, triton.cdiv(n_rows, 128)), (
+            f"swizzled Bs carries {tuple(Bs.shape[:2])} (experts, 128-row blocks), expected "
+            f"{(num_experts, triton.cdiv(n_rows, 128))} for ({n_rows}, K) weights — wrong artifact"
+        )
+    a_u8 = e2m1_as_uint8(A)
+    as_u8 = ue8m0_as_uint8(As)  # None when raw (A quantized inline)
+    b_u8 = e2m1_as_uint8(B)
+    bs_u8 = ue8m0_as_uint8(Bs)
+    # The op never swizzles: Bs is read in whatever layout it arrives (format-agnostic — MX or
+    # NVFP4). A pre-swizzled SWIZZLE_32_4_4 Bs (5D, the shared checkpoint layout the grouped kernel
+    # also consumes) takes the fast descriptor/gather path; a row-major (3D) Bs takes the affine
+    # path at no penalty. Callers swizzle once at load (public swizzle_mx_scales) to opt into perf.
+    # The descriptor is built only on the swizzled path — the un-swizzled arm never reads it (None).
+    bs_descriptor = swizzled_scale_descriptor(bs_u8) if swizzled_scales else None
+    # Requant scales are written ROW-MAJOR (never SWIZZLE_32_4_4), unlike the grouped/2D
+    # requant which fuse the swizzle in-epilogue. This is deliberate, not a gap: batched is the
+    # decode kernel — one distinct routed row per program (FAKE_BATCH replicates it across the BM
+    # lanes), so it never forms the 128-distinct-row MMA tile the tcgen05 swizzled scaled-MMA
+    # fast path needs. Swizzled scales give decode no speedup, and the 128-row swizzle block
+    # can't be written from a one-row-per-program grid without cross-program collisions. The
+    # down projection reads this row-major intermediate directly (its As is row-major too).
+    if output_format in ("mxfp4", "nvfp4"):
+        # packed E2M1 intermediate (nibble pairs along N) + group scales (UE8M0 for MX,
+        # E4M3 for NVFP4) — feeds a W4A4 down as-is
+        assert N % (2 * scale_group) == 0, (
+            f"N (={N}) must be a multiple of {2 * scale_group} to pack E2M1 pairs"
+        )
+        C = a_u8.new_empty((S, N // 2), dtype=torch.int8)
+        Cs = torch.empty(
+            S,
+            N // scale_group,
+            device=a_u8.device,
+            dtype=bs_u8.dtype,  # UE8M0 -> uint8, NVFP4 -> e4m3 (the binder-safe weight-scale dtype)
+        )
+    elif requant:
+        C = a_u8.new_empty((S, N), dtype=FP8_DTYPE)
+        Cs = torch.empty(S, N // MX_SCALE_GROUP_K, device=a_u8.device, dtype=torch.uint8)
+    else:
+        C = a_u8.new_empty((S, N), dtype=output_dtype)
+        Cs = None  # unread without an OUTPUT_FORMAT; strides literal below
+    def grid(META):
+        return (S, triton.cdiv(N, META["BLOCK_SIZE_N"]))
+    # NVFP4 accumulator correction: the per-expert g_a·g_b product folded onto the fp32 accumulator.
+    # g_b per expert and g_a scalar go down SEPARATELY; the kernel multiplies them in-register
+    # (the old host-side g_a*g_b product was a torch launch per GEMM per call)
+    input_global_scale = normalize_global_scale(b_global_scale, B.shape[0])
+    act_global_scale = normalize_global_scale(a_global_scale, B.shape[0])
+    out_global_scale = normalize_global_scale(output_global_scale, B.shape[0])
+    with device_context(a_u8.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(mx_dynamic_matmul_batched_kernel)[grid](
+            a_u8,
+            as_u8,  # None when raw (A quantized inline)
+            b_u8,
+            bs_u8,
+            bs_descriptor,
+            C,
+            Cs,
+            bias,
+            act_global_scale,  # AsGlobal: g_a per tensor or per expert (acc, and the inline-quant arm's A/g_a)
+            input_global_scale,  # AsBsGlobal = g_b per expert (acc; g_a folds in-kernel via AsGlobal)
+            out_global_scale,  # CsGlobal: requant output normalization (next proj's provided input_scale); None folds out
+            expert_ids,
+            gather_idx,  # None = A is expert-sorted; read only when not None (folds at trace time)
+            scatter_idx,  # None = C is expert-sorted; read only when not None (folds at trace time)
+            S,
+            N,
+            K,
+            a_u8.stride(0),
+            a_u8.stride(1),
+            as_u8.stride(0) if pre_quantized else 1,
+            b_u8.stride(0),
+            b_u8.stride(2),
+            b_u8.stride(1),
+            bs_u8.stride(0),
+            bs_u8.stride(2),
+            bs_u8.stride(1),
+            C.stride(0),
+            C.stride(1),
+            Cs.stride(0) if requant else 1,
+            Cs.stride(1) if requant else 1,
+            bias_stride_e,
+            bias_stride_n,
+            global_scale_stride(act_global_scale),
+            global_scale_stride(input_global_scale),
+            global_scale_stride(out_global_scale),
+            expert_ids.stride(0),
+            SCALE_GROUP_K=scale_group,
+            num_experts=num_experts,
+            SWIZZLED_SCALES=swizzled_scales,
+            ACTIVATION_FORMAT=activation_format,
+            GATE=gate,
+            ACT_FN=act_fn,
+            SWIGLU_ALPHA=swiglu_alpha,
+            SWIGLU_LIMIT=swiglu_limit,
+            OUTPUT_FORMAT=output_format,
+            SIMULATE_UNFUSED=simulate_unfused,
+            INTERMEDIATE_DTYPE=tl_dtype(output_dtype),
+            PDL=decode_pdl(),
+            launch_pdl=decode_pdl(),
+        )
+    return [C, Cs] if requant else [C]
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("full_precision_matmul_batched"),
+    mutates_args=(),
+    opaque=True,
+)
+def full_precision_matmul_batched(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    expert_ids: torch.Tensor,
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    activation_format: str | None = None,
+    quantize_output: bool = False,
+    output_dtype: torch.dtype | None = None,
+    gather_idx: torch.Tensor | None = None,
+    scatter_idx: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    """Full-precision (BF16/FP16) batched matmul: C[s] = A[s] @ B[expert_ids[s]].T — no
+    quantization anywhere, fp32 accumulation. ``gate``/``act_fn``/``swiglu_*``/
+    ``simulate_unfused`` are the dispatcher's gate|up fusion kwargs (GLU only; ``requant`` is
+    meaningless without a quantized format). ``gather_idx``/``scatter_idx`` map the source
+    row of A / destination row of C per program (None = row s).
+    A:  (rows, K) BF16/FP16 activations — rows addressed via ``gather_idx``
+    B:  (num_experts, N, K) expert weights in A's dtype; under ``gate`` the (num_experts, 2N, K) stack
+    """
+    validate_dense_operands(A, B)
+    assert A.dtype == B.dtype and A.dtype in (torch.bfloat16, torch.float16), (
+        f"full-precision path needs matching BF16/FP16 A and B, got {A.dtype} / {B.dtype}"
+    )
+    assert activation_format in (None, "bf16") and not quantize_output, (
+        "the full-precision path quantizes nothing — no quantized activation format or requant applies"
+    )
+    output_dtype = resolve_output_dtype(output_dtype, A, None)
+    K = A.shape[1]
+    S = expert_ids.shape[0]
+    num_experts, _, N = expert_weight_shape(B, gate)
+    C = A.new_empty(S, N, dtype=output_dtype)
+    def grid(META):
+        return (S, triton.cdiv(N, META["BLOCK_SIZE_N"]))
+    with device_context(A.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(full_precision_matmul_batched_kernel)[grid](
+            A,
+            B,
+            C,
+            bias,
+            expert_ids,
+            gather_idx,  # None = A is expert-sorted; read only when not None (folds at trace time)
+            scatter_idx,  # None = C is expert-sorted; read only when not None (folds at trace time)
+            S,
+            N,
+            K,
+            A.stride(0),
+            A.stride(1),
+            B.stride(0),
+            B.stride(2),
+            B.stride(1),
+            C.stride(0),
+            C.stride(1),
+            bias_stride_e,
+            bias_stride_n,
+            expert_ids.stride(0),
+            num_experts=num_experts,
+            GATE=gate,
+            ACT_FN=act_fn,
+            SWIGLU_ALPHA=swiglu_alpha,
+            SWIGLU_LIMIT=swiglu_limit,
+            SIMULATE_UNFUSED=simulate_unfused,
+            INTERMEDIATE_DTYPE=tl_dtype(output_dtype),
+            PDL=decode_pdl(),
+            launch_pdl=decode_pdl(),
+        )
+    return [C]
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("mx_weight_only_matmul_batched"),
+    mutates_args=(),
+    opaque=True,
+)
+def mx_weight_only_matmul_batched(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    expert_ids: torch.Tensor,
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    output_dtype: torch.dtype | None = None,
+    gather_idx: torch.Tensor | None = None,
+    scatter_idx: torch.Tensor | None = None,
+    b_global_scale: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    """weight-only batched (decode) matmul: ``C[s] = A[s] @ B[expert_ids[s]].T`` with raw bf16/fp16
+    activations against MXFP4/NVFP4/MXFP8 weights upcast to bf16 in-loop — the ``matmul_ogs`` format.
+    Affine (3D) weight scales — UE8M0 group-32 (MX) or E4M3 group-16 (NVFP4, whose per-expert
+    fp32 ``b_global_scale`` ``(E,)`` recovers on the accumulator; ``None`` = single-level).
+    Returns ``[C]``."""
+    assert A.dtype in (torch.bfloat16, torch.float16, torch.float32), (
+        f"weight-only takes a raw bf16/fp16/fp32 activation, got {A.dtype}"
+    )
+    assert Bs.ndim == 3, f"weight-only batched takes affine (3D) weight scales, got ndim={Bs.ndim}"
+    # The ``gate`` contract runs UNSTACKED at every S: this same op without the gate (one
+    # plain GEMM over the stacked weight), then the one-kernel ``fused_glu`` — bit-identical
+    # rounding, the unfused-reference order, so ``simulate_unfused`` rides through. The
+    # upcast weight feeds a plain bf16 dot, so the stacked tile buys no native MMA and only
+    # halves occupancy on a weight-bandwidth-bound loop; unlike the block-FP8 band this
+    # never inverts (GPT-OSS shape: 48.0->39.7µs at S=4, still ahead at S=512; 2026-08-06).
+    if gate:
+        # fp32 intermediate = the exact GEMM accumulators, so the GLU keeps the gated
+        # epilogue's FUSED-order rounding (bf16 operands would drift steep-sigmoid elements
+        # past the weight-only tests' exact-ish tolerances). Under ``simulate_unfused`` the
+        # caller is asking for the UNFUSED order instead, where the reference lands its gate_up
+        # in the activation dtype before the GLU — carrying fp32 there leaves the two a bf16 ULP
+        # apart (256.0 absolute at magnitude 2^15), which the parity tolerance rejects.
+        inter_dtype = (
+            resolve_output_dtype(output_dtype, A, None) if simulate_unfused else torch.float32
+        )
+        [gate_up] = mx_weight_only_matmul_batched(
+            A, B, Bs, expert_ids, output_dtype=inter_dtype,
+            gather_idx=gather_idx, scatter_idx=scatter_idx, b_global_scale=b_global_scale,
+        )
+        return [fused_glu(gate_up, act_fn, swiglu_alpha, swiglu_limit,
+                          out_dtype=resolve_output_dtype(output_dtype, A, None))]
+    output_dtype = resolve_output_dtype(output_dtype, A, None)
+    K = A.shape[1]
+    S = expert_ids.shape[0]
+    WEIGHT_VALUES_PER_BYTE = 2 if B.dtype == torch.int8 else 1
+    num_experts, rows, K_b = B.shape
+    assert K == WEIGHT_VALUES_PER_BYTE * K_b, (
+        f"K ({K}) must equal {WEIGHT_VALUES_PER_BYTE} * B.shape[2] ({K_b})"
+    )
+    N = rows // 2 if gate else rows
+    scale_group = mx_scale_family(Bs, K)
+    b_global_scale = normalize_global_scale(b_global_scale, num_experts)
+    C = A.new_empty(S, N, dtype=output_dtype)
+    b_u8 = e2m1_as_uint8(B)
+    bs_u8 = ue8m0_as_uint8(Bs)
+    def grid(META):
+        return (S, triton.cdiv(N, META["BLOCK_SIZE_N"]))
+    with device_context(A.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(mx_weight_only_matmul_batched_kernel)[grid](
+            A,
+            b_u8,
+            bs_u8,
+            b_global_scale,
+            C,
+            bias,
+            expert_ids,
+            gather_idx,
+            scatter_idx,
+            S,
+            N,
+            K,
+            A.stride(0),
+            A.stride(1),
+            b_u8.stride(0),
+            b_u8.stride(2),
+            b_u8.stride(1),
+            bs_u8.stride(0),
+            bs_u8.stride(1),
+            bs_u8.stride(2),
+            C.stride(0),
+            C.stride(1),
+            bias_stride_e,
+            bias_stride_n,
+            global_scale_stride(b_global_scale),
+            expert_ids.stride(0),
+            num_experts=num_experts,
+            SCALE_GROUP_K=scale_group,
+            WEIGHT_VALUES_PER_BYTE=WEIGHT_VALUES_PER_BYTE,
+            GATE=gate,
+            ACT_FN=act_fn,
+            SWIGLU_ALPHA=swiglu_alpha,
+            SWIGLU_LIMIT=swiglu_limit,
+            SIMULATE_UNFUSED=simulate_unfused,
+            INTERMEDIATE_DTYPE=tl_dtype(output_dtype),
+            PDL=decode_pdl(),
+            launch_pdl=decode_pdl(),
+        )
+    return [C]
+def matmul_batched(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor | None = None,
+    Bs: torch.Tensor | None = None,
+    *,
+    expert_ids: torch.Tensor,
+    bias: torch.Tensor | None = None,  # (E, N_out) per-expert output bias; (N_out,) for 2D
+    activation_format: str | None = None,
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    quantize_output: bool = False,
+    simulate_unfused: bool = False,
+    output_dtype: torch.dtype | None = None,
+    gather_idx: torch.Tensor | None = None,
+    scatter_idx: torch.Tensor | None = None,
+    a_global_scale: torch.Tensor | None = None,
+    b_global_scale: torch.Tensor | None = None,
+    output_global_scale: torch.Tensor | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """Batched matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4, or full-precision). Routes one
+    program per routed row (``expert_ids`` gives its expert).
+    ``As`` marks ``A`` as already quantized (framework-precomputed scales, or a requantized
+    intermediate handed to the down projection); a per-tensor scalar ``As`` is instead the static
+    (calibrated) activation scale for block-scale FP8 weights — the op quantizes raw ``A`` against
+    it; ``None`` = raw ``A``, quantized dynamically by the op in ``activation_format`` (``None`` =
+    the weights' own format — fp8/E4M3 offline for bd/tensor, inline for MX; ``"bf16"`` = weight-only;
+    ``"mxfp4"`` = packed E2M1 activations). ``Bs`` ``None`` = unquantized BF16/FP16 weights.
+    ``quantize_output`` requantizes the output into ``activation_format`` — the return is then
+    ``(C, Cs)``. ``gate``/``act_fn``/``swiglu_alpha``/``swiglu_limit`` fuse the gate|up GLU into the
+    epilogue. ``As``/``Bs`` are each a bare block-scale tensor; the two-level NVFP4
+    second-level scales ride the separate ``a_global_scale``/``b_global_scale`` (fp32 per-tensor,
+    weights per-expert ``(E,)``; from ``nvfp4_quantize_two_level``), and the op folds ``g_a · g_b``
+    onto the accumulator. The activation global ``g_a`` is CALIBRATED (the checkpoint's
+    ``input_scale``): ``a_global_scale=g_a`` with a raw ``A`` has the op quantize ``A / g_a`` per
+    block, and rides a pre-quantized ``As`` the same way. Under an NVFP4 requant the
+    fused requant normalizes the GLU intermediate by the PROVIDED ``output_global_scale`` (the next
+    proj's calibrated ``input_scale``) before the block quant and returns ``[C, Cs]``; the down
+    consumes it as ``As=Cs, a_global_scale=output_global_scale``. ``gather_idx``/``scatter_idx`` (each None or a ``(S,)`` map)
+    address the source row of ``A`` / destination row of ``C`` per program — None means row
+    ``s``; the gather lets the gate_up read unexpanded activations (source row
+    ``s // num_top_k``) with no copy.
+    Routes by what the weight tensors themselves say (there is no ``block_size``
+    parameter — the quantization block is derived from the scale shape,
+    ``weight_block_size``):
+    - ``Bs`` None → ``full_precision_matmul_batched`` (plain dot, no scales anywhere).
+    - MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0
+      group-32 ``Bs`` → ``mx_dynamic_matmul_batched``.
+    - one scale per expert (``Bs`` ``(E,)``/``(E, 1, 1)``) →
+      ``w8a8_tensor_dynamic_fp8_matmul_batched``.
+    - block scales (``Bs`` ``(E, N/bn, K/bk)``) → ``w8a8_block_dynamic_fp8_matmul_batched``.
+    """
+    check_activation_format(activation_format, quantize_output)
+    assert (a_global_scale is None and b_global_scale is None) or (
+        Bs is not None and weight_format(B, Bs) == "nvfp4"
+    ), "two-level globals (a_global_scale / b_global_scale) are NVFP4-only"
+    assert output_global_scale is None or (
+        Bs is not None and quantize_output and resolve_activation_format(activation_format, B, Bs) == "nvfp4"
+    ), (
+        "output_global_scale is the NVFP4 requant second level — it requires quantize_output=True on NVFP4 "
+        "(the epilogue would otherwise normalize by it with nothing downstream to compensate)"
+    )
+    if As is not None and As.numel() == 1:
+        # static (per-tensor calibrated) activation quant: a per-tensor scalar As for block-scale FP8
+        # weights — the caller hands raw A, the op quantizes it against the scalar (As IS the scale).
+        assert Bs is not None and not is_mx(B, Bs) and weight_block_size(B, Bs) is not None, (
+            "a per-tensor scalar As (static activation scale) needs block-scale FP8 weights"
+        )
+        out = w8a8_block_static_fp8_matmul_batched(
+            A,
+            B,
+            As,
+            Bs,
+            expert_ids,
+            weight_block_size(B, Bs),
+            gate,
+            act_fn,
+            swiglu_alpha,
+            swiglu_limit,
+            simulate_unfused,
+            activation_format,
+            quantize_output,
+            output_dtype,
+            gather_idx,
+            scatter_idx,
+            bias=bias,
+        )
+        return out[0] if len(out) == 1 else tuple(out)
+    if Bs is None:
+        assert As is None, (
+            "the full-precision path (Bs=None) takes no activation scales"
+        )
+        out = full_precision_matmul_batched(
+            A,
+            B,
+            expert_ids,
+            gate,
+            act_fn,
+            swiglu_alpha,
+            swiglu_limit,
+            simulate_unfused,
+            activation_format,
+            quantize_output,
+            output_dtype,
+            gather_idx,
+            scatter_idx,
+            bias=bias,
+        )
+    elif is_mx(B, Bs) and activation_format == "bf16":  # weight-only: raw bf16 acts, MX weight upcast in-MMA
+        assert As is None and a_global_scale is None and not quantize_output, (
+            "weight-only (activation_format='bf16') takes a raw activation, no As/global/requant"
+        )
+        assert output_global_scale is None, (
+            "weight-only has no requant epilogue for output_global_scale to normalize"
+        )
+        out = mx_weight_only_matmul_batched(
+            A,
+            B,
+            Bs,
+            expert_ids,
+            gate,
+            act_fn,
+            swiglu_alpha,
+            swiglu_limit,
+            simulate_unfused,
+            output_dtype,
+            gather_idx,
+            scatter_idx,
+            b_global_scale=b_global_scale,
+            bias=bias,
+        )
+    elif is_mx(B, Bs):
+        out = mx_dynamic_matmul_batched(
+            A,
+            B,
+            As,
+            Bs,
+            expert_ids,
+            gate,
+            act_fn,
+            swiglu_alpha,
+            swiglu_limit,
+            simulate_unfused,
+            activation_format,
+            quantize_output,
+            output_dtype,
+            gather_idx,
+            scatter_idx,
+            a_global_scale,
+            b_global_scale,
+            output_global_scale,
+            bias=bias,
+        )
+    elif (block_size := weight_block_size(B, Bs)) is None:
+        assert not gate, (
+            "the batched op has no tensor-wide gate|up fusion (grouped and 2D support it)"
+        )
+        assert activation_format in (None, "fp8") and not quantize_output, (
+            "tensor-wide supports neither packed activations nor a fused requant"
+        )
+        out = w8a8_tensor_dynamic_fp8_matmul_batched(
+            A, B, As, Bs, expert_ids, output_dtype, gather_idx, scatter_idx, bias=bias
+        )
+    else:
+        out = w8a8_block_dynamic_fp8_matmul_batched(
+            A,
+            B,
+            As,
+            Bs,
+            expert_ids,
+            block_size,
+            gate,
+            act_fn,
+            swiglu_alpha,
+            swiglu_limit,
+            simulate_unfused,
+            activation_format,
+            quantize_output,
+            output_dtype,
+            gather_idx,
+            scatter_idx,
+            bias=bias,
+        )
+    # bd/mx/full-precision ops return a list ([C] or [C, Cs]); the tensor op returns a bare tensor.
+    if isinstance(out, (list, tuple)):
+        return out[0] if len(out) == 1 else tuple(out)
+    return out

build/torch-rocm/bayesian_autotuner.py ADDED Viewed

	@@ -0,0 +1,940 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Bayesian-optimization autotuner: benches a budgeted sample of the config grid instead of
+the full grid, via a Tree-structured Parzen Estimator (TPE) over the config *dimensions*
+(``num_warps`` / ``num_stages`` / tile sizes / flags). After a short random seed phase it
+models the good (top-``gamma`` by measured time) vs bad per-dimension value densities and
+benches the unmeasured config that maximizes ``l(x)/g(x)`` — the Expected-Improvement proxy
+— then a coordinate-descent pass polishes the best. Subclasses Triton's ``Autotuner`` and
+uses only the Python stdlib (no external optimizer dependency).
+Each new key's search is warm-started from the most recently tuned key's best config (nearby
+workloads share tile-shape preferences). Grids smaller than ``n_trials`` defer to the stock
+exhaustive bench-all.
+"""
+from __future__ import annotations
+import hashlib
+import logging
+import json
+import math
+import os
+import traceback
+import random
+import time
+from collections import defaultdict
+from contextlib import contextmanager
+import torch
+from typing import Dict, List
+from triton.runtime.autotuner import (
+    Autotuner,
+    Config,
+    JITFunction,
+    driver,
+    get_cache_invalidating_env_vars,
+    get_cache_manager,
+    knobs,
+    triton_key,
+)
+logger = logging.getLogger(__name__)
+# Env-gated host-execution counter (FINEGRAINED_AUTOTUNE_CALLCOUNT=<path>): every
+# BayesianAutotuner.run() call increments its kernel's tally, dumped to <path> at exit.
+# Discriminates cudagraph REPLAY (host code runs only at capture) from per-token
+# re-execution — the decode-gap forensics tool.
+_CALL_COUNTS = None
+if os.environ.get("FINEGRAINED_AUTOTUNE_CALLCOUNT"):
+    import atexit
+    _CALL_COUNTS = {}
+    def _dump_call_counts(path=os.environ["FINEGRAINED_AUTOTUNE_CALLCOUNT"]):
+        with open(path, "w") as f:
+            json.dump(_CALL_COUNTS, f, indent=1, sort_keys=True)
+    atexit.register(_dump_call_counts)
+class BayesianAutotuner(Autotuner):
+    """Drop-in replacement for ``triton.runtime.autotuner.Autotuner`` that
+    benches ~``n_trials`` configs per key via TPE Bayesian optimization +
+    coordinate-descent refinement, instead of the full grid."""
+    def __init__(
+        self,
+        *args,
+        n_trials: int = 80,
+        max_failures: int | None = None,
+        n_startup_trials: int = 12,
+        gamma: float = 0.25,
+        refine: bool = True,
+        max_refine_iters: int = 5,
+        log_path: str | None = None,
+        path_anchor_axes: tuple[str, ...] = (),
+        finite_check_args: tuple[str, ...] = (),
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.fn_name = getattr(self.fn, "__name__", str(self.fn))
+        # Bayesian trial budget — the per-decorator default, overridable via the
+        # FINEGRAINED_AUTOTUNE_TRIALS env var (quick sweeps / exhaustive runs without touching
+        # the decorators; set it >= grid size to fall back to stock exhaustive bench-all).
+        self.n_trials = int(os.environ.get("FINEGRAINED_AUTOTUNE_TRIALS") or n_trials)
+        # A tune aborts once this many configs have failed to compile/run without the measured
+        # budget being met. Defaults to n_trials: a grid that cannot land its measurements without
+        # that many rejects has a pruner gap, and the fix is to fence the dead region rather than
+        # to tolerate the compiles. Its own axis so a kernel can raise it deliberately.
+        self.max_failures = int(
+            os.environ.get("FINEGRAINED_AUTOTUNE_MAX_FAILURES")
+            or max_failures
+            or self.n_trials
+        )
+        self.n_startup_trials = n_startup_trials
+        # top fraction of measured configs the TPE treats as "good"
+        self.gamma = gamma
+        self.refine = refine
+        self.max_refine_iters = max_refine_iters
+        # Branch axes (by kwarg name) that relocate the TILE optimum — the caller's
+        # declaration, since it is configuration knowledge (the tuner never hardcodes
+        # axis names). Each distinct value combination of these axes gets one guaranteed
+        # max-tile anchor; empty = median anchors only.
+        self.path_anchor_axes = tuple(path_anchor_axes)
+        # Output arg NAMES (a declaration — the tuner never hardcodes kernel details) whose
+        # values are checked for NaN/inf after each config's bench: a config that writes
+        # non-finite output scores inf and can never be crowned. The tuner times, it never
+        # validates — this veto closes the NaN half of the "wrong-answer configs win
+        # tunes" class at negligible cost (one isfinite pass per benched config).
+        self.finite_check_args = tuple(finite_check_args)
+        # JSONL log of every benched (key, config, ms) — set here or via the
+        # FINEGRAINED_AUTOTUNE_LOG env var. Analyse offline to prune bad configs.
+        self.log_path = log_path or os.environ.get("FINEGRAINED_AUTOTUNE_LOG")
+        # steady-state launch support: key extraction without dict builds (eager decode
+        # is host-bound; the cache-hit path of run() is on the per-token critical path)
+        self._arg_position = {name: i for i, name in enumerate(self.arg_names)}
+        self._arg_name_set = frozenset(self.arg_names)
+        self._dtype_tag_cache: Dict[tuple, str] = {}
+        self._dtype_str_cache: Dict[torch.dtype, str] = {}
+    def _dtype_str(self, dtype) -> str:
+        s = self._dtype_str_cache.get(dtype)
+        if s is None:
+            s = self._dtype_str_cache[dtype] = str(dtype)
+        return s
+    # substrings marking a COMPILE-stage failure — the only class safe to memoize on
+    # disk (a benching/CUDA error can be transient or sticky-context contamination;
+    # persisting one would permanently fence a healthy config for this source version)
+    _COMPILE_FAILURE_MARKS = (
+        "PassManager",
+        "CompilationError",
+        "MLIR",
+        "ConvertTritonGPUToLLVM",
+        "TritonGPUAccelerateMatmul",
+    )
+    def _bench(self, *args, config, **meta):
+        """Score any failing config as inf instead of raising — a compile failure is data
+        for the search, not a fatal error. Stock Triton forgives only OutOfResources, but
+        e.g. Triton 3.7.1's ``warp_specialize`` raises RuntimeError at unsupported
+        (shape, config) combos; unguarded, one such config kills the whole tune when a
+        small grid falls through to the stock exhaustive path below. Every failure is
+        recorded and ``_report_bench_failures`` reports every distinct failure —
+        inf-scoring must not silently hide a broken path behind a healthy one.
+        (Stock already inf's OutOfResources / CompileTimeAssertionFailure / PTXASError
+        internally without reaching this handler — those are deliberate guard classes,
+        e.g. our own ``tl.static_assert`` fences; the reporter covers the UNEXPECTED
+        failure classes stock would otherwise let kill the tune.)
+        Compile-stage failures are memoized on disk keyed by everything that determines
+        compilation (kernel source hash + config + constexpr arg values + tensor dtypes +
+        invalidating env), so the next fresh tune of ANY shape skips the doomed compile
+        outright (~40-66 wasted compiles per nvfp4 tune before this). The memo dies with
+        the kernel source (``fn.cache_key``); bench-stage errors are never persisted."""
+        sig = self._compile_signature(config, meta)
+        memo = self._failed_compile_memo()
+        if sig is not None and sig in memo:
+            self._failures.append((config, memo[sig] + "  [memoized]"))
+            return [float("inf")] * 3
+        try:
+            timing = super()._bench(*args, config=config, **meta)
+            bad = self._nonfinite_output(meta)
+            if bad is not None:
+                self._failures.append((config, f"non-finite output in {bad} (numerics veto)"))
+                return [float("inf")] * 3
+            return timing
+        except Exception as e:
+            # triton's CompilationError carries only the location + source snippet in str(e);
+            # the actual reason lives in the chained cause, so without it every compile failure
+            # reads as an anonymous caret. Keep the head of the location and the cause verbatim.
+            msg = str(e)[:250]
+            # triton nests a CompilationError per @jit frame, each carrying only its own location
+            # and source — the actual reason is at the BOTTOM of the chain. Walk to it, else every
+            # compile failure reads as an anonymous caret and cannot be fenced.
+            root, seen = e, 0
+            while (root.__cause__ or root.__context__) is not None and seen < 12:
+                root = root.__cause__ or root.__context__
+                seen += 1
+            if root is not e:
+                msg += f" || root: {type(root).__name__}: {str(root)[-300:]}"
+            err = f"{type(e).__name__}: {msg}"
+            if os.environ.get("FINEGRAINED_AUTOTUNE_TRACEBACKS"):
+                logger.warning(
+                    "[autotune] %s failed on %s:\n%s",
+                    self.fn_name,
+                    config,
+                    "".join(traceback.format_exception(type(e), e, e.__traceback__)),
+                )
+            self._failures.append((config, err))
+            if sig is not None and any(m in err for m in self._COMPILE_FAILURE_MARKS):
+                memo[sig] = err
+                self._persist_failed_compile_memo()
+            return [float("inf")] * 3
+    def _nonfinite_output(self, meta=None) -> str | None:
+        """Name of the first declared output arg holding NaN/inf after a bench, else None.
+        fp8 tensors byte-check the E4M3 NaN pattern (no ``isfinite`` kernel for fp8);
+        integer/packed outputs are skipped (every bit pattern is a value)."""
+        if not self.finite_check_args:
+            return None
+        nargs = {**(getattr(self, "nargs", None) or {}), **(meta or {})}
+        for name in self.finite_check_args:
+            v = nargs.get(name)
+            if not isinstance(v, torch.Tensor) or v.numel() == 0:
+                continue
+            if v.dtype == torch.float8_e4m3fn:
+                flat = v.reshape(-1).view(torch.uint8)
+                bad = bool((((flat & 0x7F) == 0x7F)).any())
+            elif v.is_floating_point():
+                bad = bool((~v.reshape(-1).float().isfinite()).any())
+            else:
+                continue
+            if bad:
+                return name
+        return None
+    def _compile_signature(self, config, meta=None) -> str | None:
+        """Hash of the compile determinants for one config: the config itself plus the
+        constexpr argument VALUES and tensor argument dtypes from this launch (both feed
+        specialization — e.g. GATE flips arms, a uint8 A packs the loads). Source/env
+        live in the memo FILE's key, not here.
+        ``meta`` carries the launch kwargs: every constexpr is kwarg-passed at the call
+        sites, so ``self.nargs`` (positional-only) alone would hash GATE/format flags as
+        None and collide memo entries across arms."""
+        nargs = {**(getattr(self, "nargs", None) or {}), **(meta or {})}
+        if not nargs:
+            return None
+        fn = self.fn
+        while not isinstance(fn, JITFunction):
+            fn = fn.fn
+        parts = [str(sorted(config.all_kwargs().items()))]
+        for param in fn.params:
+            v = nargs.get(param.name)
+            if hasattr(v, "dtype"):
+                parts.append(f"{param.name}:{v.dtype}")
+            elif param.is_constexpr and isinstance(
+                v, (bool, int, float, str, type(None))
+            ):
+                parts.append(f"{param.name}={v}")
+            elif isinstance(v, int) and not isinstance(v, bool):
+                # non-constexpr ints feed Triton's specialization (==1 / %16 divisibility),
+                # which is (shape, config)-dependent for WS lowering failures — hashing the
+                # bucket keeps one shape's PassManager failure from fencing every shape
+                parts.append(
+                    f"{param.name}:i{1 if v == 1 else (16 if v % 16 == 0 else 0)}"
+                )
+        return hashlib.sha256("-".join(parts).encode("utf-8")).hexdigest()
+    def _failed_compile_memo(self) -> dict:
+        """The per-(kernel source, backend, env) failed-compile dict, loaded from
+        Triton's on-disk cache once per autotuner instance."""
+        if getattr(self, "_failed_memo", None) is not None:
+            return self._failed_memo
+        try:
+            from triton.compiler.compiler import make_backend
+            fn = self.fn
+            while not isinstance(fn, JITFunction):
+                fn = fn.fn
+            group = hashlib.sha256(
+                "-".join(
+                    [
+                        triton_key(),
+                        make_backend(driver.active.get_current_target()).hash(),
+                        fn.cache_key,
+                        str(sorted(get_cache_invalidating_env_vars().items())),
+                    ]
+                ).encode("utf-8")
+            ).hexdigest()
+            self._failed_memo_cache = get_cache_manager(group)
+            self._failed_memo_file = f"{fn.__name__[:150]}.failed_compiles.json"
+            path = self._failed_memo_cache.get_file(self._failed_memo_file)
+            self._failed_memo = json.load(open(path)) if path else {}
+        except Exception:
+            self._failed_memo_cache = None
+            self._failed_memo = {}
+        return self._failed_memo
+    def _persist_failed_compile_memo(self):
+        if getattr(self, "_failed_memo_cache", None) is None:
+            return
+        try:
+            self._failed_memo_cache.put(
+                json.dumps(self._failed_memo), self._failed_memo_file, binary=False
+            )
+        except Exception:
+            pass  # persistence is best-effort; the in-memory memo still holds
+    def _report_bench_failures(self):
+        """After every tune, report every UNIQUE failure — a failure is never silent:
+        inf-scoring keeps the search alive, but a human must see what broke (e.g. a code
+        change that kills one compute path would otherwise silently degrade into "the other
+        path wins"). Distinct errors are deduped with a count and an example config; the
+        JSONL autotune log has per-config detail."""
+        if not self._failures:
+            return
+        by_err = defaultdict(list)
+        for c, err in self._failures:
+            by_err[err].append(c)
+        for err, cfgs in by_err.items():
+            c = cfgs[0]
+            example = ", ".join(f"{k}={v}" for k, v in c.kwargs.items())
+            logger.warning(
+                "[autotune] %s: %d config(s) failed to compile/run — %s  (e.g. %s, w%d s%d)",
+                self.fn_name,
+                len(cfgs),
+                err,
+                example,
+                c.num_warps,
+                c.num_stages,
+            )
+    def run(self, *args, **kwargs):
+        if _CALL_COUNTS is not None:
+            _CALL_COUNTS[self.fn_name] = _CALL_COUNTS.get(self.fn_name, 0) + 1
+        self._failures = []
+        # Small grid → defer to parent (stock exhaustive bench-all). Steady state
+        # still short-circuits the parent's per-call dict builds: a crowned config
+        # with no pre_hook launches directly. The lookup key must be STOCK format
+        # (key values + untagged dtype strings) — it is the parent that populates
+        # this branch's cache entries.
+        if len(self.configs) <= 1 or self.n_trials >= len(self.configs):
+            config = self.configs[0] if len(self.configs) == 1 else None
+            if config is None and self.cache:
+                n_positional = len(args)
+                key = []
+                for k in self.keys:
+                    i = self._arg_position.get(k)
+                    if i is not None and i < n_positional:
+                        key.append(args[i])
+                    elif k in kwargs and k in self._arg_name_set:
+                        key.append(kwargs[k])
+                for i in range(n_positional):
+                    dtype = getattr(args[i], "dtype", None)
+                    if dtype is not None:
+                        key.append(self._dtype_str(dtype))
+                for k, v in kwargs.items():
+                    if k in self._arg_name_set and self._arg_position[k] >= n_positional:
+                        dtype = getattr(v, "dtype", None)
+                        if dtype is not None:
+                            key.append(self._dtype_str(dtype))
+                config = self.cache.get(tuple(key))
+            if config is not None and config.pre_hook is None:
+                self.best_config = config
+                config_kwargs = getattr(config, "_all_kwargs_memo", None)
+                if config_kwargs is None:
+                    config_kwargs = config._all_kwargs_memo = config.all_kwargs()
+                return self.fn.run(*args, **kwargs, **config_kwargs)
+            ret = super().run(*args, **kwargs)
+            self._report_bench_failures()
+            return ret
+        # Key extraction runs on every launch and eager decode is host-bound, so it
+        # avoids dict builds: positional args by precomputed index, then kwargs —
+        # the same (arg_names-prefix, kwargs-order) iteration the dict merge gave,
+        # keeping keys byte-compatible with existing disk crown caches.
+        arg_names = self.arg_names
+        n_positional = len(args)
+        tags = self._dtype_tag_cache
+        key = []
+        for k in self.keys:
+            i = self._arg_position.get(k)
+            if i is not None and i < n_positional:
+                key.append(args[i])
+            elif k in kwargs and k in self._arg_name_set:
+                key.append(kwargs[k])
+        # name-tagged, with an explicit token for absent optional tensors: a positional
+        # dtype list lets a gather-only and a scatter-only launch alias to one entry,
+        # while descriptor-arm legality (descriptor_box_pruner) hangs off GatherIdx
+        # presence — the aliased replay would read wrong rows without ever re-pruning
+        for i in range(n_positional):
+            v = args[i]
+            dtype = getattr(v, "dtype", None) if v is not None else None
+            if v is not None and dtype is None:
+                continue
+            tag_key = (arg_names[i], dtype)
+            tag = tags.get(tag_key)
+            if tag is None:
+                tag = tags[tag_key] = f"{tag_key[0]}:{dtype}"
+            key.append(tag)
+        for k, v in kwargs.items():
+            if k not in self._arg_name_set or self._arg_position[k] < n_positional:
+                continue
+            dtype = getattr(v, "dtype", None) if v is not None else None
+            if v is not None and dtype is None:
+                continue
+            tag_key = (k, dtype)
+            tag = tags.get(tag_key)
+            if tag is None:
+                tag = tags[tag_key] = f"{tag_key[0]}:{dtype}"
+            key.append(tag)
+        key = tuple(key)
+        first_launch = key not in self.cache
+        if first_launch:
+            self.nargs = dict(zip(self.arg_names, args))
+            pruned = self.prune_configs(kwargs)
+            def benchmark():
+                t0 = time.time()
+                self.cache[key] = self._bayesian_search(pruned, args, kwargs, key)
+                self._report_bench_failures()
+                self.bench_time = time.time() - t0
+                if knobs.autotuning.print:
+                    fn_name = getattr(self.fn, "__name__", str(self.fn))
+                    print(
+                        f"[bayesian-autotune] {fn_name} tuned "
+                        f"{len(self.configs_timings)} configs in {self.bench_time:.1f}s — "
+                        f"key={key}, best={self.cache[key].all_kwargs()}"
+                    )
+                full_nargs = {**self.nargs, **kwargs, **self.cache[key].all_kwargs()}
+                self.pre_hook(full_nargs, reset_only=True)
+            if self.cache_results:
+                self.check_disk_cache(key, pruned, benchmark)
+            else:
+                benchmark()
+        config = self.cache[key]
+        self.best_config = config
+        config_kwargs = getattr(config, "_all_kwargs_memo", None)
+        if config_kwargs is None:
+            config_kwargs = config._all_kwargs_memo = config.all_kwargs()
+        if config.pre_hook is not None:
+            if not first_launch:  # tuned path already built nargs for prune_configs
+                self.nargs = dict(zip(self.arg_names, args))
+            config.pre_hook({**self.nargs, **kwargs, **config_kwargs})
+        ret = self.fn.run(*args, **kwargs, **config_kwargs)
+        if first_launch and self.finite_check_args:
+            # First launch after a key miss: the one regime where a crowned launch has been
+            # observed returning non-finite output while every subsequent identical launch
+            # is correct (large multi-device models; root cause open — see the DSV3
+            # investigation record). Detect via the declared output args and relaunch once,
+            # loudly. Sentinel/padding rows can false-trigger (uninitialized by contract);
+            # the cost is one duplicate launch per key, never a dropped config.
+            bad = self._nonfinite_output(kwargs)
+            if bad is not None:
+                logger.warning(
+                    "[autotune] %s: non-finite %s on the FIRST launch of key=%s — relaunching once",
+                    self.fn_name, bad, key,
+                )
+                ret = self.fn.run(*args, **kwargs, **config.all_kwargs())
+        self.nargs = None
+        return ret
+    def _bayesian_search(self, configs: List[Config], args, kwargs, key) -> Config:
+        timings: Dict[int, float] = {}
+        sigs = [tuple(sorted(c.all_kwargs().items())) for c in configs]
+        def bench_idx(idx: int) -> float:
+            if idx in timings:
+                return timings[idx]
+            try:
+                ms = self._bench(*args, config=configs[idx], **kwargs)
+                if isinstance(ms, (tuple, list)):
+                    ms = ms[0]
+                timings[idx] = float(ms)
+            except Exception:
+                timings[idx] = float("inf")
+            self._log_result(key, configs[idx], timings[idx])
+            return timings[idx]
+        # Distinct values per config dimension, for the TPE's Laplace smoothing.
+        dim_vals: Dict = defaultdict(set)
+        for sig in sigs:
+            for d, v in sig:
+                dim_vals[d].add(v)
+        # Seed phase: one BASIN ANCHOR per (COMPUTE_MODE, SWAP_AB) group, the most recent
+        # cached key's best (warm start), then seeded-random fill. The anchors guarantee every
+        # categorical basin gets at least one real measurement — without them the TPE's
+        # per-dimension model can write off a whole axis it never saw succeed (two dsv4 tunes
+        # shipped 25-60% slow winners because their random seeds only sampled a basin's dead
+        # configs), and coordinate descent can't recover a winner two coupled flips away.
+        n_startup = max(2, min(self.n_startup_trials, self.n_trials))
+        anchors = self._basin_anchor_indices(configs)
+        order = list(range(len(configs)))
+        random.Random(0).shuffle(order)
+        warm_idx = self._warm_start_index(configs)
+        head = anchors + ([warm_idx] if warm_idx is not None else [])
+        order = list(dict.fromkeys(head + order))
+        for idx in order[: max(n_startup, len(head))]:
+            bench_idx(idx)
+        # TPE: split measured configs into good (top-gamma) / bad, build per-dimension value
+        # densities for each, and bench the unmeasured config maximizing log l(x) - log g(x)
+        # (Expected-Improvement proxy), updating the model after each measurement.
+        # inf (failed-to-compile) configs are EXCLUDED from the densities: a compile failure is
+        # evidence about that one joint shape (usually shared memory), not about its dimension
+        # values — counting them as "bad" buried SWAP_AB under a wall of BN=256 smem failures
+        # and made the tuner ship a 53µs winner while the 41µs swap config sat unbenched.
+        # They don't consume the trial budget either (n_trials = MEASURED configs; a failure
+        # stays in ``timings`` only as a skip-list entry) — the compile it burned is the one
+        # cost that can't be refunded here, which is what the smem/compile-guard pruners avoid.
+        warned_failures = False
+        while sum(1 for t in timings.values() if t != float("inf")) < self.n_trials:
+            # Failures don't consume the trial budget (above), which is right for a handful of
+            # smem/regs rejects but unbounded if a launch is broadly broken: the loop keeps
+            # picking, each pick costs an O(configs) TPE rescan, and the tune degenerates into a
+            # silent multi-hour spin instead of a diagnosable error. Cap the failures at the same
+            # budget as the measurements and fail loudly.
+            if (
+                not warned_failures
+                and sum(1 for t in timings.values() if t == float("inf")) >= self.max_failures
+            ):
+                # Loud, once, and non-fatal. Aborting here would break tunes that DO find a winner
+                # after many rejects (mxfp4 W4A4 at a non-128 N burns >100), trading a silent
+                # inefficiency for a hard failure. Warning keeps the waste visible — this is how
+                # ~100 dead compiles per mx tune (an IndexError in load_act_mx's inline-quant arm)
+                # were found at all — without gating the kernel on a diagnostic.
+                warned_failures = True
+                self._report_bench_failures()
+                logger.warning(
+                    "[autotune] %s: %d config(s) scored inf before the trial budget was met "
+                    "(%d of them UNEXPECTED failures — the rest are deliberate guard rejections: "
+                    "OutOfResources and tl.static_assert fences, which stock inf's internally). "
+                    "The unexpected ones are dead compile time every tune, and a pruner gap "
+                    "worth closing.",
+                    self.fn_name,
+                    sum(1 for t in timings.values() if t == float("inf")),
+                    len(self._failures),
+                )
+            ranked = sorted(
+                (i for i, t in timings.items() if t != float("inf")), key=timings.get
+            )
+            if not ranked:  # nothing compiled yet — keep seeding in shuffled order
+                nxt = next((i for i in order if i not in timings), None)
+                if nxt is None:
+                    break
+                bench_idx(nxt)
+                continue
+            n_good = max(1, round(self.gamma * len(ranked)))
+            good_c: Dict = defaultdict(lambda: defaultdict(int))
+            bad_c: Dict = defaultdict(lambda: defaultdict(int))
+            for j, i in enumerate(ranked):
+                tgt = good_c if j < n_good else bad_c
+                for d, v in sigs[i]:
+                    tgt[d][v] += 1
+            n_bad = len(ranked) - n_good
+            best_i, best_score = None, -math.inf
+            for i in range(len(configs)):
+                if i in timings:
+                    continue
+                score = 0.0
+                for d, v in sigs[i]:
+                    V = len(dim_vals[d])
+                    lp = (good_c[d][v] + 1.0) / (
+                        n_good + V
+                    )  # P(v | good), Laplace-smoothed
+                    gp = (bad_c[d][v] + 1.0) / (n_bad + V)  # P(v | bad)
+                    score += math.log(lp) - math.log(gp)
+                if score > best_score:
+                    best_score, best_i = score, i
+            if best_i is None:
+                break
+            bench_idx(best_i)
+        # Coordinate-descent refinement: try single-dim perturbations around
+        # the current best until no neighbor improves.
+        if self.refine:
+            for _ in range(self.max_refine_iters):
+                best_idx = min(timings, key=timings.get)
+                best_sig = sigs[best_idx]
+                best_ms = timings[best_idx]
+                improved = False
+                for i, s in enumerate(sigs):
+                    if i in timings:
+                        continue
+                    diff = sum(1 for (_, a), (_, b) in zip(best_sig, s) if a != b)
+                    if diff != 1:
+                        continue
+                    if bench_idx(i) < best_ms:
+                        improved = True
+                if not improved:
+                    break
+        # Crown RUNOFF: the search's single-shot timings carry enough noise to flip crowns
+        # between same-code tune rolls (gpt-oss prefill grouped rolled 1727-2150us on identical
+        # grids; the tuner's one RNG site is seeded, so noise is the only roll source). Re-time
+        # the top candidates back-to-back — same warmed process, adjacent in time, so they see
+        # the same clocks — and crown the winner of THAT comparison. Costs ~3 extra benches.
+        finite = sorted((t, i) for i, t in timings.items() if t != float("inf"))
+        if len(finite) > 1:
+            for _, i in finite[:3]:
+                timings.pop(i, None)
+                bench_idx(i)  # memoizing closure: re-benches and re-records the popped entry
+        self.configs_timings = {configs[i]: t for i, t in timings.items()}
+        best = min(timings, key=timings.get)
+        if timings[best] == float("inf"):
+            # crowning an arbitrary broken config would persist it as `best` on disk and
+            # every later healthy process would replay it without re-tuning — the exact
+            # sticky-contaminated-context scenario inf-forgiveness exists to survive
+            self._report_bench_failures()
+            raise RuntimeError(
+                f"[bayesian-autotune] {self.fn_name}: every benched config failed "
+                "(all-inf tune) — nothing correct to crown or persist; the context may "
+                "be trap-contaminated (retry in a fresh process) or the grid is broken "
+                "for this launch (see the failure report above)."
+            )
+        return configs[best]
+    def _log_result(self, key, config: Config, ms: float):
+        """Append one ``(key, config, ms)`` record as JSONL for offline analysis —
+        e.g. pruning configs that are consistently far off the per-key best. ``inf`` ms
+        marks a config that failed to compile/run (out of resources, etc.)."""
+        if not self.log_path:
+            return
+        try:
+            rec = {
+                "timestamp": time.time(),
+                "fn_name": getattr(self.fn, "__name__", str(self.fn)),
+                "key": list(key),
+                "kwargs": config.kwargs,
+                "num_warps": config.num_warps,
+                "num_stages": config.num_stages,
+                "ms": ms,
+            }
+            with open(self.log_path, "a") as f:
+                f.write(json.dumps(rec, default=str) + "\n")
+        except Exception:
+            pass
+    def _basin_anchor_indices(self, configs: List[Config]) -> List[int]:
+        """One representative config index per (COMPUTE_MODE, SWAP_AB) basin — the MEDIAN in
+        tile-sort order, a mid-sized tile with mid warps/stages. Not the smallest: a basin's
+        minimal corner (min BK x 2 warps) can be latency-bound pathological (a 131µs anchor in
+        a basin whose peak is 41µs re-poisons the axis it was meant to protect). Coordinate
+        descent climbs BN/BK/warps/stages from wherever the TPE lands within the basin.
+        Returns [] when the grid has no such axes (single basin)."""
+        # Basin axes are DERIVED, not declared: a config kwarg with string or boolean
+        # values is a branch axis (different code path — compute mode, operand swap,
+        # warp specialization), which partitions the grid into disjoint performance
+        # basins. Numeric kwargs (tiles/warps/stages) are ordinal — the TPE's densities
+        # and coordinate descent handle those.
+        basin_axes = sorted(
+            {
+                k
+                for c in configs
+                for k, v in c.kwargs.items()
+                if isinstance(v, (bool, str))
+            }
+        )
+        if not basin_axes:
+            return []
+        # A basin = the compute-path axes (constexprs selecting different compiled code).
+        groups: Dict = {}
+        for i, c in enumerate(configs):
+            basin = tuple(c.kwargs.get(k) for k in basin_axes)
+            groups.setdefault(basin, []).append(i)
+        if len(groups) <= 1:
+            return []
+        def tile_order(i):
+            return (
+                configs[i].kwargs.get("BLOCK_SIZE_N", 0),
+                configs[i].kwargs.get("BLOCK_SIZE_K", 0),
+                configs[i].num_warps,
+                configs[i].num_stages,
+            )
+        # Median anchor per basin, PLUS one max-tile anchor per declared PATH-axis combo
+        # (``path_anchor_axes``). The median alone left the big-tile end unsampled: TPE
+        # densities lock onto the first-timed mid-tile region and coordinate descent cannot
+        # cross DIAGONAL ridges — dsv4 mxfp8 prefill gate_up's BN=128/BK=256 winner (840us)
+        # needs BN and warps to move together, so descent from the BN=64 crown rejects the
+        # lone BN step and n_trials 100/200/400 all missed it; one guaranteed big-tile
+        # sample fixed it (crown 958 -> 880 at the default budget). The max anchor rides
+        # the declared COARSE grouping, NOT the full branch-axis cross: per-basin max
+        # anchors on the weight-only grouped grid (32 basins: modes x WS x PS x 2 memory
+        # modes each way) put 64 forced trials into a 100-trial budget and starved the
+        # search — gpt-oss prefill regressed 3463 -> 4099us on a quiet box before this was
+        # scoped. Which axes qualify is the KERNEL's declaration (configuration knowledge);
+        # axes absent from this grid are ignored.
+        anchors = []
+        coarse_seen = set()
+        coarse_axes = [k for k in self.path_anchor_axes if k in basin_axes]
+        for basin, idxs in groups.items():
+            ordered = sorted(idxs, key=tile_order)
+            anchors.append(ordered[len(ordered) // 2])
+            if not self.path_anchor_axes:
+                continue  # no declaration -> median anchors only
+            # declared axes absent from THIS grid still coarse-group to () — one global
+            # max-tile anchor, the big-tile guarantee without per-basin anchor spam
+            coarse = tuple(basin[basin_axes.index(k)] for k in coarse_axes)
+            if coarse not in coarse_seen and ordered[-1] != anchors[-1]:
+                coarse_seen.add(coarse)
+                anchors.append(ordered[-1])
+        return anchors
+    def _warm_start_index(self, configs: List[Config]):
+        """Return the index in ``configs`` matching the most recently cached
+        key's best config (or ``None`` if no prior tune or no match in the
+        current pruned list)."""
+        if not self.cache:
+            return None
+        # Python 3.7+ dicts preserve insertion order; last entry = most recent tune.
+        prev_best = next(reversed(self.cache.values()))
+        prev_kwargs = prev_best.all_kwargs()
+        for i, c in enumerate(configs):
+            if c.all_kwargs() == prev_kwargs:
+                return i
+        return None
+    def check_disk_cache(self, tuning_key, configs, bench_fn):
+        """Persist the tuned best config to Triton's on-disk cache so a later run (or process)
+        skips the whole minutes-long search+compile — most of which is Triton codegen, not
+        benching. Keyed exactly like Triton's own autotune cache (triton version + backend +
+        kernel source hash + invalidating env + key + the full config grid), so any of those
+        changing re-tunes.
+        Unlike Triton's stock version we do NOT bail when configs carry a ``pre_hook`` (ours
+        always do): we persist each config's ``all_kwargs()`` — serializable and pre_hook-free —
+        and on a hit re-match it to the live ``Config`` object, so the pre_hook is preserved from
+        code rather than lost in (de)serialization. A hit launches only the winning config."""
+        if not tuning_key:
+            bench_fn()
+            return False
+        from triton.compiler.compiler import make_backend
+        fn = self.fn
+        while not isinstance(fn, JITFunction):
+            fn = fn.fn
+        env_vars = get_cache_invalidating_env_vars()
+        cache_key = [
+            triton_key(),
+            make_backend(driver.active.get_current_target()).hash(),
+            fn.cache_key,
+            str(sorted(env_vars.items())),
+            # Key NAMES and the search budget are determinants too: tuning_key holds
+            # values only (renaming a key or swapping two same-typed keys would collide),
+            # and FINEGRAINED_AUTOTUNE_TRIALS is not in Triton's invalidating-env list —
+            # without n_trials a 3-trial exploratory winner gets replayed at 100 trials.
+            str(self.keys),
+            f"n_trials={self.n_trials},n_startup_trials={self.n_startup_trials},gamma={self.gamma}",
+            str(tuning_key),
+        ] + [str(c) for c in configs]  # str(Config) is pre_hook-free and process-stable
+        cache_key = hashlib.sha256("-".join(cache_key).encode("utf-8")).hexdigest()
+        cache = get_cache_manager(cache_key)
+        file_name = f"{fn.__name__[:150]}.bayes_autotune.json"
+        keylog = os.environ.get("FINEGRAINED_AUTOTUNE_KEYLOG")
+        if keylog:
+            with open(keylog, "a") as f:
+                f.write(json.dumps({
+                    "fn": fn.__name__,
+                    "hit": cache.get_file(file_name) is not None,
+                    "tuning_key": str(tuning_key),
+                    "keys": str(self.keys),
+                    "n_configs": len(configs),
+                    "env": str(sorted(env_vars.items())),
+                    "fn_cache_key": fn.cache_key[:16],
+                    "hash": cache_key[:16],
+                }) + "\n")
+        # signature -> live Config (carries the pre_hook); used to re-match the cached winner
+        by_sig = {tuple(sorted(c.all_kwargs().items())): c for c in configs}
+        def load_crown() -> bool:
+            path = cache.get_file(file_name)
+            if not path:
+                return False
+            try:
+                with open(path) as f:
+                    data = json.load(f)
+                best = by_sig.get(tuple(sorted(data["best"].items())))
+                if best is None:
+                    return False
+                self.cache[tuning_key] = best
+                self.configs_timings = {
+                    by_sig[s]: t
+                    for kw, t in data["timings"]
+                    for s in (tuple(sorted(kw.items())),)
+                    if s in by_sig
+                }
+                return True
+            except Exception:
+                return False  # corrupt/stale cache file → re-tune
+        if load_crown():
+            return True
+        # Cross-process tuning lock. Without it, TP/EP ranks sharing this cache dir all tune the
+        # same key CONCURRENTLY and each crowns its own winner: the search is seeded-random and
+        # near-ties break on measurement noise, so ranks end up on different tiles, round
+        # differently, and greedy decoding turns that into different tokens (measured: an 8-rank
+        # MXFP8 run where 3 ranks emitted a different token). Serializing means one rank tunes and
+        # the rest read the same crown — identical configs, and ~world_size less tuning wall-clock.
+        # Best-effort: any lock failure falls through to tuning locally (correct, just redundant).
+        with _tuning_lock(cache, file_name):
+            if load_crown():  # another rank tuned this key while we waited
+                return True
+            bench_fn()
+            # publish INSIDE the lock: a rank released to an empty cache would tune again
+            try:
+                cache.put(
+                    json.dumps(
+                        {
+                            "best": self.cache[tuning_key].all_kwargs(),
+                            "timings": [
+                                (c.all_kwargs(), t) for c, t in self.configs_timings.items()
+                            ],
+                        }
+                    ),
+                    file_name,
+                    binary=False,
+                )
+            except Exception:
+                pass
+        return False
+@contextmanager
+def _tuning_lock(cache, file_name: str):
+    """Serialize tuning of ONE key across every process sharing the Triton cache directory.
+    Held while a key is tuned and its crown published, so concurrent TP/EP ranks converge on the
+    same config instead of each crowning its own (see ``check_disk_cache``). Advisory ``flock`` on
+    a sidecar file next to the crown; entirely best-effort — no cache dir, no ``fcntl``, or a lock
+    that cannot be taken within ``FINEGRAINED_TUNE_LOCK_TIMEOUT`` (default 30 min, since the holder
+    is doing a real tune) all fall through to tuning locally. Losing the lock costs redundant work,
+    never correctness."""
+    lock_file = None
+    try:
+        import fcntl
+        cache_dir = getattr(cache, "cache_dir", None)
+        if cache_dir:
+            os.makedirs(cache_dir, exist_ok=True)
+            deadline = time.time() + float(
+                os.environ.get("FINEGRAINED_TUNE_LOCK_TIMEOUT") or 1800
+            )
+            lock_file = open(os.path.join(cache_dir, f"{file_name}.lock"), "w")
+            while True:
+                try:
+                    fcntl.flock(lock_file, fcntl.LOCK_EX | fcntl.LOCK_NB)
+                    break
+                except BlockingIOError:
+                    if time.time() > deadline:
+                        logger.warning(
+                            "[autotune] %s: tuning lock held elsewhere for >%.0fs — tuning locally",
+                            file_name,
+                            time.time() - (deadline - 1800),
+                        )
+                        lock_file.close()
+                        lock_file = None
+                        break
+                    time.sleep(1.0)
+    except Exception:
+        lock_file = None
+    try:
+        yield
+    finally:
+        if lock_file is not None:
+            try:
+                import fcntl
+                fcntl.flock(lock_file, fcntl.LOCK_UN)
+            finally:
+                lock_file.close()
+def bayesian_autotune(
+    configs,
+    key,
+    *,
+    n_trials: int = 80,
+    max_failures: int | None = None,
+    n_startup_trials: int = 12,
+    gamma: float = 0.25,
+    refine: bool = True,
+    max_refine_iters: int = 5,
+    log_path: str | None = None,
+    path_anchor_axes: tuple[str, ...] = (),
+    finite_check_args: tuple[str, ...] = (),
+    cache_results: bool = True,
+    reset_to_zero=None,
+    restore_value=None,
+    **kwargs,
+):
+    """Decorator mirroring ``@triton.autotune``. Extra kwargs:
+    n_trials:                 successfully measured configs per key (TPE budget; configs
+                              that fail to compile/run are skipped without consuming it)
+    n_startup_trials:         random seed configs before the TPE model kicks in
+    gamma:                    top fraction of measured configs treated as "good"
+    refine, max_refine_iters: coordinate-descent refinement after the TPE
+    path_anchor_axes:         branch axes (kwarg names) whose values relocate the tile
+                              optimum — one guaranteed max-tile anchor per value combo
+                              (see _basin_anchor_indices); the kernel declares them, the
+                              tuner stays independent of configuration details
+    max_failures:             abort the tune once this many configs fail to compile/run before
+                              the measured budget is met (default n_trials; env override
+                              FINEGRAINED_AUTOTUNE_MAX_FAILURES). Without it a broadly broken
+                              grid spins: failures don't consume the trial budget, so the search
+                              keeps picking and each pick costs an O(configs) rescan.
+    log_path:                 JSONL of benched configs (or FINEGRAINED_AUTOTUNE_LOG)
+    cache_results:            persist the tuned best config to disk (on by default) so later
+                              runs skip the search+compile — see BayesianAutotuner.check_disk_cache"""
+    def decorator(fn):
+        return BayesianAutotuner(
+            fn,
+            fn.arg_names,
+            configs,
+            key,
+            reset_to_zero,
+            restore_value,
+            n_trials=n_trials,
+            max_failures=max_failures,
+            n_startup_trials=n_startup_trials,
+            gamma=gamma,
+            refine=refine,
+            max_refine_iters=max_refine_iters,
+            log_path=log_path,
+            path_anchor_axes=path_anchor_axes,
+            finite_check_args=finite_check_args,
+            cache_results=cache_results,
+            **kwargs,
+        )
+    return decorator

build/torch-rocm/compat.py ADDED Viewed

	@@ -0,0 +1,498 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import contextvars
+import functools
+from contextlib import contextmanager
+import os
+import torch
+import triton
+import triton.language as tl
+from torch.library import triton_op, wrap_triton
+# Programmatic dependent launch for the decode chain (batched GEMMs, GLU, act quant, reduce): each
+# kernel waits on its grid dependency at the top and releases its dependents just before the
+# epilogue. The LATE trigger is load-bearing: released early, the next kernel's CTAs co-reside and
+# steal issue slots (-15us on a single-wave chain); released late they only hide launch latency
+# (GLM/DSV4/GPT-OSS decode -4.5..-4.9%, bit-identical). Prefill grids are multi-wave and never
+# take it. FINEGRAINED_PDL=0 disables.
+DECODE_PDL = os.environ.get("FINEGRAINED_PDL", "1") == "1"
+def decode_pdl() -> bool:
+    """PDL for this launch: the flag, and not under torch.compile — dynamo's TTIR access analysis
+    cannot read the ``griddepcontrol`` inline asm and would mark every input mutated (extra copies,
+    no fusion). Deployment decode runs eager launches under cudagraphs, where PDL applies."""
+    return DECODE_PDL and not torch.compiler.is_compiling()
+# kernel-builder generates ``_ops.py`` into the built variant (the op namespace carries the build's
+# unique id) and refuses to build over a tracked one, so the source tree has none: an unbuilt
+# checkout (tests, bench, FINEGRAINED_KERNELS_PATH) registers its ops under the plain package name.
+try:
+    from ._ops import add_op_namespace_prefix
+except ImportError:
+    def add_op_namespace_prefix(name: str) -> str:
+        return f"finegrained_kernels::{name}"
+# ── Format constants ──────────────────────────────────────────────────────────
+# FP8 (E4M3) is the main format for weights and activations;
+FP8_DTYPE = torch.float8_e4m3fn
+# FP4 (E2M1) packs two 4-bit nibbles per byte. MX formats (MXFP4 weights, MXFP8
+# E4M3 weights/activations) share one UE8M0 scale per 32-element K-group — the OCP
+# MX block size, consumed by ``tl.dot_scaled``. Format constants, not tunables.
+NIBBLES_PER_BYTE = 2
+MX_SCALE_GROUP_K = 32
+NVFP4_SCALE_GROUP_K = 16
+# ── Host-side helpers ─────────────────────────────────────────────────────────
+# set ONLY while an opaque op's registered fake impl runs (shape inference); read by
+# compile_time_only_triton_wrap to no-op the kernel launches within
+_SKIP_LAUNCHES = contextvars.ContextVar("finegrained_kernels_skip_launches", default=False)
+# Plain mirror of the contextvar, read on the launch path. Dynamo cannot trace `ContextVar.get`
+# ("Unsupported method call ... method `get` of class `ContextVar`"), which broke tracing of every
+# non-opaque op; a module-level bool it can constant-fold. The contextvar stays authoritative for
+# set/reset (nesting, async isolation) — this only mirrors it for the read.
+_SKIP_LAUNCHES_MIRROR = False
+def compile_time_only_triton_wrap(kernel):
+    """``wrap_triton`` while torch.compile is tracing (required to capture a raw Triton
+    launch into the graph), the bare kernel in eager. Every kernel launch goes through
+    this: eager skips ``wrap_triton``'s per-call dispatch overhead, and it dodges a
+    correctness trap — eager ``wrap_triton`` of a stock ``@triton.autotune`` kernel that
+    has ``prune_configs_by`` re-runs the FULL tune on every call (torch's wrapper rebuilds
+    a fresh autotuner around the pruned configs per call, ``triton_kernel_wrap.py``:
+    ``autotune(configs=pruned, key=[])``), so nothing ever lands in the original
+    ``Autotuner.cache`` (~2.3s/call for ``_mx_act_quant_kernel``) and the mid-capture
+    tune invalidates cudagraph capture. Pruner-less kernels and the ``bayesian_autotune``
+    kernels reuse their cache through the wrapper (measured) — but eager never needs the
+    wrapper at all.
+    Inside an OPAQUE op's fake impl (see ``compile_time_only_triton_op``) the launch is
+    skipped outright: the op function has already allocated the correctly-shaped
+    outputs, which is all fake mode needs. The skip is an EXPLICIT contextvar, never
+    ambient fake-mode detection — ``triton_op``'s capture also runs under FakeTensor
+    mode, and skipping there would compile graphs with the kernel MISSING (uninitialized
+    outputs; found as identical garbage across formats, 2026-07-16). The skip is also
+    load-bearing in the other direction — probed: without it the opaque ops' fake impls
+    reach the launcher with FakeTensors ("RuntimeError when making fake tensor call";
+    ``is_compiling`` is still true during fake prop, so they'd take the wrap_triton
+    branch and hit the pre_hook assert that made them opaque in the first place)."""
+    if _SKIP_LAUNCHES_MIRROR:
+        class _NoLaunch:
+            def __getitem__(self, grid):
+                return lambda *args, **kwargs: None
+        return _NoLaunch()
+    return wrap_triton(kernel) if torch.compiler.is_compiling() else kernel
+def compile_time_only_triton_op(name, mutates_args=(), opaque=False):
+    """``@triton_op`` under torch.compile (the registered custom op is what dynamo
+    captures), the plain function in eager: the torch.library op dispatch stack costs
+    ~160µs per eager call — the dominant decode CPU cost in the eager-breakdown probe —
+    and eager needs none of it. The op is still registered at import time, so compiled
+    callers and ``torch.ops`` introspection see it unchanged. Sibling of
+    ``compile_time_only_triton_wrap`` (same rule one level down, at the kernel launch)."""
+    def decorator(fn):
+        if opaque:
+            # kernels whose configs carry a pre_hook (TMA descriptor rebinders) cannot
+            # be traced by triton_op's capture; register OPAQUE instead — the fake impl
+            # is the op fn itself, whose launches no-op under FakeTensor mode (see
+            # compile_time_only_triton_wrap), leaving just the shape-correct allocations
+            op = torch.library.custom_op(name, mutates_args=mutates_args)(fn)
+            @functools.wraps(fn)
+            def fake_impl(*args, **kwargs):
+                global _SKIP_LAUNCHES_MIRROR
+                token = _SKIP_LAUNCHES.set(True)
+                prev, _SKIP_LAUNCHES_MIRROR = _SKIP_LAUNCHES_MIRROR, True
+                try:
+                    return fn(*args, **kwargs)
+                finally:
+                    _SKIP_LAUNCHES.reset(token)
+                    _SKIP_LAUNCHES_MIRROR = prev
+            op.register_fake(fake_impl)
+        else:
+            op = triton_op(name, mutates_args=mutates_args)(fn)
+        @functools.wraps(fn)
+        def dispatch(*args, **kwargs):
+            # The op path is also the AUTOGRAD path (`register_autograd` hangs off the op, and
+            # these ops take the RAW activation, so that is where a usable gradient exists).
+            # Calling `fn` skips the dispatcher entirely, so a differentiable caller would
+            # silently get no gradient — route through the op whenever a tensor wants one.
+            # Training runs at prefill shapes where the ~160us dispatch is noise; inference
+            # runs under no_grad and keeps the bare-function fast path.
+            if torch.compiler.is_compiling() or (
+                torch.is_grad_enabled()
+                and any(getattr(a, "requires_grad", False) for a in args)
+            ):
+                return op(*args, **kwargs)
+            return fn(*args, **kwargs)
+        dispatch._triton_op = op
+        return dispatch
+    return decorator
+@contextmanager
+def device_context(device: torch.device):
+    """Context manager that sets the active device for any backend (cuda, xpu, etc.)."""
+    backend = getattr(torch, device.type, None)
+    if backend is not None and hasattr(backend, "device"):
+        with backend.device(device):
+            yield
+    else:
+        yield
+@functools.lru_cache(maxsize=8)
+def sm_count(device_index: int) -> int:
+    """Processor (CUDA SM / XPU Xe-core) count for a device (a constant) — cached so callers
+    don't re-query it each launch. Read from Triton's active driver so it works on any backend,
+    with a CUDA fallback."""
+    try:
+        return triton.runtime.driver.active.utils.get_device_properties(device_index)[
+            "multiprocessor_count"
+        ]
+    except Exception:
+        active_device = get_active_device_type()
+        if active_device == "cuda":
+            return torch.cuda.get_device_properties(device_index).multi_processor_count
+        elif active_device == "xpu":
+            return torch.xpu.get_device_properties(device_index).multi_processor_count
+        else:
+            raise RuntimeError(
+                f"Unsupported device type {active_device} for sm_count; only cuda/xpu are supported."
+            )
+def tl_dtype(dtype: torch.dtype) -> tl.dtype:
+    """The ``tl`` dtype matching a torch dtype (``torch.bfloat16`` → ``tl.bfloat16``) — the
+    attribute names line up, so no table. For passing a tensor's dtype as a kernel constexpr
+    when the kernel can't read it off a pointer argument."""
+    return getattr(tl, str(dtype).removeprefix("torch."))
+def get_active_device_type() -> str:
+    """Active torch device type for the current Triton backend (``"cuda"``, ``"xpu"``, ...).
+    Falls back to ``"cuda"`` when no driver is loaded — Triton raises
+    ``RuntimeError: 0 active drivers ([])`` on driverless build boxes, and the
+    autotune-config builder is evaluated at module-import time under the
+    ``@triton.autotune`` decorator (no kernel launches there, so the default is
+    only used to shape the config list).
+    """
+    try:
+        return triton.runtime.driver.active.get_active_torch_device().type
+    except RuntimeError:
+        return "cuda"
+def resolve_memory_modes(memory_modes):
+    """Resolve the generic ``"descriptor"`` B_MEMORY_MODE to the device's flavor: host-built
+    (NVIDIA TMA) on CUDA, device-built in-kernel tensormap on XPU. Installs the XPU
+    scratch allocator device-built tensormaps need."""
+    descriptor = (
+        "device_descriptor" if get_active_device_type() == "xpu" else "host_descriptor"
+    )
+    memory_modes = tuple(descriptor if m == "descriptor" else m for m in memory_modes)
+    if "device_descriptor" in memory_modes:
+        # device-side descriptors build in-kernel tensormaps that need a scratch
+        # allocator — on whichever accelerator is active (XPU always; CUDA when a kernel
+        # requests "device_descriptor" explicitly to skip the host-descriptor plumbing).
+        device_type = get_active_device_type()
+        triton.set_allocator(
+            lambda size, alignment, stream: torch.empty(
+                size, dtype=torch.int8, device=device_type
+            )
+        )
+    return memory_modes
+def get_accelerator_autotuning_configs(
+    *,
+    mx: bool = False,
+    tune_block_m: bool = False,
+    tune_block_n: bool = False,
+    tune_block_nk: bool = False,
+    swap_ab: bool = False,
+    warp_spec: bool = False,
+    packed_schedule: bool = False,
+    compute_modes=None,
+    a_memory_modes=None,
+    b_memory_modes=None,
+    pre_hook=None,
+):
+    """Autotune search grid for the current accelerator — the single generator for every
+    GEMM kernel. One union span serves both families: num_warps {2,4,8,16} x num_stages
+    {2..6} (XPU: warps {8,16}); the tile comes from the caller's ``block_size``, or
+    ``tune_block_nk`` spans BN {32,64,128,256} x BK {64,128,256,512}. ``mx=True`` (group-32
+    scale formats, MXFP4/MXFP8) requires ``compute_modes`` (BK is coupled per mode) and
+    ``tune_block_nk`` (MX kernels have no caller block_size — the tile is always tuned).
+    ``compute_modes`` emits the ``COMPUTE_MODE`` axis (``None`` — the default — emits no
+    axis, like ``b_memory_modes``), tuner-selected per workload (token count is in the
+    key): ``"dot_scaled"`` native group-32 scaled MMA (wide BK {128,256}; wins once the
+    grid saturates, ~S>=32), ``"dot"`` fp8 ``tl.dot`` + per-group software rescale
+    (under ``mx`` its BK is pinned to the 32-group), ``"scalar"`` CUDA-core FMA reduce
+    (no MMA M->16 pad; wins memory-bound decode, only emit where implemented — it is
+    pinned to BM=1 here because a bigger-BM scalar config at BN == BM COMPILES and
+    silently computes garbage the tuner would time and pick).
+    Axes, crossed per flag:
+    - ``tune_block_m``: BLOCK_SIZE_M in {16,32,64,128} instead of a fixed
+      launch-time choice.
+    - ``tune_block_n``: independent BLOCK_SIZE_N in {64,128,256} for the block-scale
+      kernels whose BK is pinned by the caller (activation scale groups are per
+      block_k) but whose N tile may grow past the scale granularity.
+    - ``swap_ab``: marks the single-token decode GEVM kernels — crosses in the coupled
+      decode ``(BLOCK_SIZE_M, SWAP_AB)`` pairs (see below). BK=512 earns its union-span
+      seat here: at decode's tiny memory-parallelism-bound grids, longer per-iteration
+      bursts are the one in-block lever that helps (+12% on the dsv4 swap GEMV; deeper
+      num_stages and split-K both measured WORSE). Prefill tunes carry the same rows as
+      trial-budget dilution (descriptor_box_pruner drops the illegal e4m3 boxes).
+    - ``warp_spec`` (CUDA only): warp-specialize the K-loop, crossed over every mode —
+      pair it with ``warp_spec_compile_guard_pruner``, which owns the can't-compile
+      regions (including dot_scaled + WS, a PassManager failure). Compile support is
+      (shape, config)-dependent on Triton 3.7.1, so it is a tuner axis — failures score
+      inf and self-prune; where it compiles it is both faster and (bd grouped gate_up)
+      load-bearing for correctness.
+    - ``a_memory_modes``: the A_MEMORY_MODE activation-load axis (descriptor legal only
+      without a gather — the tile's rows are the contiguous sorted positions; the pruner
+      fences descriptor rows when GatherIdx is passed).
+    - ``b_memory_modes``: the B_MEMORY_MODE weight-load axis — pass
+      ``("descriptor", "pointer")`` to let the tuner pick the device's tensor-descriptor
+      flavor vs explicit pointers; ``None`` (default) emits no axis. In the plain-dot
+      kernels the descriptor arm is the SWAPPED loop (see the bd 2D kernel).
+    A given ``pre_hook`` is attached to every config and must self-guard on
+    ``B_MEMORY_MODE``.
+    Winner census over 743 tunes (B200, all kernels/shapes), useful when reading tuner
+    logs: every value wins somewhere — MX lives in stages 5/6 (69 wins) and never picked
+    warps 16; plain-dot picked warps 16 (19 wins, e.g. the bd gate_up decode) and warps 2
+    (24 wins, the tensor batched decode 9/11) and concentrated in stages 3/4.
+    """
+    assert not (mx and compute_modes is None), (
+        "the MX grid couples BLOCK_SIZE_K to COMPUTE_MODE — pass compute_modes"
+    )
+    assert not (mx and not tune_block_nk), (
+        "MX kernels have no caller block_size — their (BN, BK) tile is always tuned"
+    )
+    is_xpu = get_active_device_type() == "xpu"
+    # ONE union span for both families — the per-family sets were stock-tuner-era cost
+    # caps; the Bayesian tuner samples a fixed trial budget, values measured dead for
+    # one family (MX warps=16, MX BK=64; plain-dot stages 5/6 were simply never tested)
+    # just never win, and can't-compile configs (smem) score inf and stay out of the
+    # TPE densities. The 743-tune winner census (see docstring) maps where winners
+    # actually live per family. BK=512 is in the union span for every tune_block_nk
+    # kernel — its win is decode-only (see the swap_ab docstring bullet), and prefill
+    # rows are accepted trial-budget dilution.
+    #
+    # Axes measured DEAD and deliberately absent (do not re-add without fresh evidence):
+    #   num_ctas          slower everywhere + miscompiles with TMA at num_stages<=4
+    #   flatten (tl.range)  proxy +17% never survived e2e (<=1%; gate_up regressed)
+    #   persistent 2D     grid-stride 2D prototype's win was an inline-vs-helper artifact
+    #   batched WARP_SPEC  decode loops are launch/latency-bound; WS never won a batched
+    #                     cell, so the axis is not emitted at the six batched sites
+    num_warps = [8, 16] if is_xpu else [2, 4, 8, 16]
+    num_stages = [2, 3, 4, 5, 6]
+    bn_span = (128,) if is_xpu else (32, 64, 128, 256)
+    bk_span = (128,) if is_xpu else (64, 128, 256, 512)
+    # no tuned tile -> one empty meta-dict (the tile comes from the launch kwargs)
+    blocks = (
+        [{"BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk} for bn in bn_span for bk in bk_span]
+        if tune_block_nk
+        else [{}]
+    )
+    if compute_modes is not None:
+        # generic COMPUTE_MODE axis; MX "dot"'s BK is STRUCTURALLY the UE8M0 scale group
+        # (one software rescale per K-step — a format constraint, not a tile choice), so
+        # it overrides the wide-BK span and the dedup below collapses the copies.
+        blocks = [
+            {
+                **b,
+                "COMPUTE_MODE": mode,
+                **({"BLOCK_SIZE_K": MX_SCALE_GROUP_K} if mx and mode == "dot" else {}),
+            }
+            for b in blocks
+            for mode in compute_modes
+        ]
+    if tune_block_m:
+        # BM=256 is OPEN, not disproven (2026-07-29). A plain fp8 GEMM at the dsv4 attn prefill
+        # shape clearly prefers BM=BN=256 (350-367us vs 420-478us at 128, 7 interleaved repeats,
+        # no overlap), so the basin is real for an unscaled loop. Our kernel measured 523us vs
+        # 495us against it — but the run-to-run spread at BM=128 on this box is 13.9%, so a 5.8%
+        # difference decides nothing. Left out because nothing has shown it WINNING here and the
+        # axis widens every tune 25%; re-add only with repeated tunes that clear the noise floor.
+        blocks = [{**b, "BLOCK_SIZE_M": bm} for b in blocks for bm in (16, 32, 64, 128)]
+    if tune_block_n:
+        blocks = [{**b, "BLOCK_SIZE_N": bn} for b in blocks for bn in (64, 128, 256)]
+    if warp_spec and get_active_device_type() == "cuda":
+        blocks = [{**b, "WARP_SPEC": ws} for b in blocks for ws in (False, True)]
+    if packed_schedule:
+        # grouped M-tile resolution axis: packed-table loads vs the E-wide register-resident
+        # layout — pure register relief, bit-identical tiles; wins only where the inline
+        # resolve is pressure-walled (dsv4 BM128/BK256 regs 244->187, +8%), so the tuner picks
+        blocks = [{**b, "PACKED_SCHEDULE": ps} for b in blocks for ps in (False, True)]
+    # (descriptor-mode can't-win/can't-serve regions are fenced per kernel: descriptor_box_pruner,
+    # matched_memory_modes_pruner, descriptor_needs_prequant_pruner)
+    if a_memory_modes is not None:
+        blocks = [
+            {**b, "A_MEMORY_MODE": mm}
+            for b in blocks
+            for mm in resolve_memory_modes(a_memory_modes)
+        ]
+    if b_memory_modes is not None:
+        blocks = [
+            {**b, "B_MEMORY_MODE": mm}
+            for b in blocks
+            for mm in resolve_memory_modes(b_memory_modes)
+        ]
+    # The coupled decode (BLOCK_SIZE_M, SWAP_AB) pairs — the single source of the swap/BM
+    # coupling. At M=1 the MMA 16-atom is filled either by replicating the token (non-swap
+    # BM=16, ~40% over the degenerate BM=1 on plain tl.dot) or by putting the weight's
+    # output rows in M (swap, BM=1). Swap REQUIRES BM=1 (its rhs is the single token
+    # padded to the N=16 atom), so (16, swap) never exists; the tuner picks. Prefill
+    # kernels leave SWAP_AB unset — the prefill swap arms measured losers everywhere
+    # (see the bd 2D kernel's tuner note).
+    if swap_ab and tune_block_m:
+        # tuned-BM kernels (the 2D mx GEMM): the prefill BM sweep keeps no-swap; the swap
+        # decode rows are APPENDED at their structural BM=1 (dedup collapses the copies)
+        blocks = [{**b, "SWAP_AB": False} for b in blocks] + [
+            {**b, "BLOCK_SIZE_M": 1, "SWAP_AB": True} for b in blocks
+        ]
+    elif swap_ab:
+        blocks = [
+            {**b, "BLOCK_SIZE_M": bm, "SWAP_AB": sw}
+            for b in blocks
+            for bm, sw in ((1, False), (16, False), (1, True))
+        ]
+    # scalar's BM is STRUCTURALLY 1 (a single-row GEVM reduce — a bigger-BM scalar
+    # config at BN == BM even COMPILES and silently computes garbage the tuner would
+    # time and pick); the dedup below collapses the pinned copies of the BM sweep.
+    blocks = [
+        {**b, "BLOCK_SIZE_M": 1}
+        if b.get("COMPUTE_MODE") == "scalar" and "BLOCK_SIZE_M" in b
+        else b
+        for b in blocks
+    ]
+    configs = [
+        triton.Config(b, num_warps=w, num_stages=s, pre_hook=pre_hook)
+        for b in blocks
+        for w in num_warps
+        for s in num_stages
+    ]
+    # dedup: the scalar BM=1 pin and the MX "dot" BK override collapse swept axes onto
+    # one value, leaving identical copies that would each be benched.
+    return list({str(c): c for c in configs}.values())
+@functools.cache
+def is_sm10x() -> bool:
+    """Whether the CURRENT device is Blackwell-datacenter (sm_10x, TMEM scaled-MMA) —
+    the target the ``dot_scaled`` compiler-bug guards are scoped to. False off-CUDA and on
+    driverless build boxes (``torch.cuda.is_available()`` short-circuits before the
+    capability query, which would otherwise raise); cached (pruners call this per tune, and
+    a process is pinned to one device)."""
+    return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10
+@functools.cache
+def is_sm90() -> bool:
+    """Whether the CURRENT device is Hopper (sm_90) — the target the measured NVIDIA-backend
+    MMA-lowering guards are scoped to. A POSITIVE arch test on purpose: ``not is_sm10x()`` is
+    also true on XPU, ROCm and driverless build boxes, none of which run the CUDA codegen those
+    guards were probed against, so scoping by negation would restrict backends no one measured.
+    Same driverless-safe short-circuit and caching as ``is_sm10x``."""
+    return torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0)
+@functools.lru_cache(maxsize=None)
+def sm_shared_memory_limit() -> int:
+    """Max dynamic shared memory per block (bytes) for the CURRENT device — the cap
+    Triton reports as the 'Hardware limit' on an ``out of resource: shared memory``
+    failure (~232 KB on B200, ~227 KB on H100, much less on older/consumer parts).
+    Queried from the driver so the prune adapts to the hardware instead of hardcoding
+    one GPU; cached (a process is pinned to one device)."""
+    dev = get_active_device_type()
+    device_index = (
+        torch.cuda.current_device()
+        if dev == "cuda"
+        else torch.xpu.current_device()
+        if dev == "xpu"
+        else 0
+    )
+    try:
+        return triton.runtime.driver.active.utils.get_device_properties(device_index)[
+            "max_shared_mem"
+        ]
+    except Exception:
+        if dev == "cuda":
+            return torch.cuda.get_device_properties(
+                device_index
+            ).shared_memory_per_block_optin
+        elif dev == "xpu":
+            return torch.xpu.get_device_properties(
+                device_index
+            ).shared_memory_per_block_optin
+        else:
+            raise RuntimeError(
+                f"Unsupported device type {dev} for sm_shared_memory_limit; only cuda/xpu are supported."
+            )

build/torch-rocm/descriptors.py ADDED Viewed

	@@ -0,0 +1,216 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Host-TMA descriptors for the kernels' descriptor memory arms. ``maybe_descriptor`` builds one
+where the operand layout can back it; the ``rebind_*`` functions are the autotuner per-config
+``pre_hook``s that size each descriptor's box to the tuned tile (they MUST mutate ``block_shape`` in
+place — a rebind never reaches the launch) and are no-ops for the pointer arms, whose descriptor
+argument is a dead placeholder."""
+import torch
+from triton.tools.tensor_descriptor import TensorDescriptor
+def rebind_operand_box(nargs, mode_key, desc_key, rows, cols):
+    """Set one operand's host-TMA box to ``[rows, cols]`` in place — MUST mutate (a rebind never
+    reaches the launch). No-op for a pointer config, whose descriptor is a dead int placeholder."""
+    desc = nargs[desc_key]
+    # None = the operand's layout cannot back a descriptor (see `maybe_descriptor`); those
+    # configs are pruned, so reaching here with one would mean a pruner gap, not a rebind target
+    if nargs.get(mode_key, "pointer") != "pointer" and not isinstance(desc, int) and desc is not None:
+        desc.block_shape = [rows, cols]
+def maybe_descriptor(t: torch.Tensor, box: list[int]) -> TensorDescriptor | None:
+    """A host-TMA descriptor over ``t``, or None when the layout cannot back one — a TMA box
+    needs a unit-stride innermost dim, which a transposed view (as a backward pass hands in to
+    contract over the other axis) does not have. ``None`` is the same signal the pointer arms
+    already take for an unused descriptor, and ``descriptor_box_pruner`` drops the configs that
+    would read it, so the tuner is left with the pointer arms the kernel serves stride-generally."""
+    if t is None or (t.ndim >= 2 and t.stride(-1) != 1):
+        return None
+    return TensorDescriptor.from_tensor(t, box)
+def rebind_bd_descriptors(nargs):
+    """Per-config pre_hook: set the A and B host-TMA boxes to the tuned tile over the
+    ``(rows, K)`` matrices — ``[BLOCK_SIZE_M, block_k]`` and ``[BLOCK_SIZE_N, block_k]``."""
+    rebind_operand_box(nargs, "A_MEMORY_MODE", "ADescriptor", nargs["BLOCK_SIZE_M"], nargs["block_k"])
+    rebind_operand_box(nargs, "B_MEMORY_MODE", "BDescriptor", (2 if nargs.get("GATE") else 1) * nargs["BLOCK_SIZE_N"], nargs["block_k"])
+def rebind_weight_only_descriptors(nargs):
+    """Per-config pre_hook for the weight-only kernels — set the A/B host-TMA boxes to the tuned tile:
+    ``[BM, BK]`` over the (M, K) bf16 activation and ``[BN, BK // WEIGHT_VALUES_PER_BYTE]`` over the
+    packed (N, K_bytes) weight (bytes: uint8 = packed E2M1, two values/byte). No scale boxes — weight-only
+    scales are always affine (never swizzled), read through the pointer arm."""
+    wvpb = 2 if nargs["B"].dtype == torch.uint8 else 1
+    bk = nargs["BLOCK_SIZE_K"]
+    rebind_operand_box(nargs, "A_MEMORY_MODE", "ADescriptor", nargs["BLOCK_SIZE_M"], bk)
+    rebind_operand_box(nargs, "B_MEMORY_MODE", "BDescriptor", (2 if nargs.get("GATE") else 1) * nargs["BLOCK_SIZE_N"], bk // wvpb)
+def rebind_mx_descriptors(nargs):
+    """Per-config pre_hook for the MX kernel — set the A/B host-TMA boxes to the tuned tile
+    in BYTES over the (rows, K_bytes) packed matrices: ``[BM, BK // ACT_VALUES_PER_BYTE]`` and
+    ``[BN, BK // WEIGHT_VALUES_PER_BYTE]`` (values-per-byte read off the operand dtype; uint8 =
+    packed E2M1), plus the SWIZZLE_32_4_4 scale boxes on the offline (pre-quantized A) path."""
+    avpb = 2 if nargs["A"].dtype == torch.uint8 else 1
+    wvpb = 2 if nargs["B"].dtype == torch.uint8 else 1
+    bk = nargs["BLOCK_SIZE_K"]
+    rebind_operand_box(nargs, "A_MEMORY_MODE", "ADescriptor", nargs["BLOCK_SIZE_M"], bk // avpb)
+    rebind_operand_box(nargs, "B_MEMORY_MODE", "BDescriptor", (2 if nargs.get("GATE") else 1) * nargs["BLOCK_SIZE_N"], bk // wvpb)
+    # SWIZZLE_32_4_4 scale boxes: [1, BLOCK//128, (BK // SCALE_GROUP_K) // 4, 2, 256]. Only where the
+    # scale is actually swizzled — else the SA/SB descriptor is a dummy aliased to the operand
+    # descriptor, and stamping a scale box would clobber its [BM, BK] operand box. The weight is
+    # swizzled iff SWIZZLED_SCALES; the act only when it was also offline-quantized (E4M3 / packed
+    # E2M1) — inline (raw bf16 A) stays affine even under a swizzled weight.
+    if nargs["SWIZZLED_SCALES"]:
+        # Sub-128 tiles read scales via the per-row pointer gather, leaving the descriptor
+        # unread — clamp its box to one block so it keeps a valid, non-degenerate shape
+        # (a 0-block box traps the descriptor-encoding pass; same clamp as the grouped hook).
+        rep_k = max((nargs["BLOCK_SIZE_K"] // nargs["SCALE_GROUP_K"]) // 4, 1)
+        bn_blocks = max(1, ((2 if nargs.get("GATE") else 1) * nargs["BLOCK_SIZE_N"]) // 128)
+        nargs["BSDescriptor"].block_shape = [1, bn_blocks, rep_k, 2, 256]
+        if nargs["A"].dtype in (torch.float8_e4m3fn, torch.uint8):
+            bm_blocks = max(nargs["BLOCK_SIZE_M"] // 128, 1)
+            nargs["ASDescriptor"].block_shape = [1, bm_blocks, rep_k, 2, 256]
+    # swizzled requant output (Cs is a descriptor): store tile [1, 1, rep_n, 2, 256], rep_n per config
+    if nargs["CSDescriptor"] is not None:
+        rep_n = (nargs["BLOCK_SIZE_N"] // nargs["SCALE_GROUP_K"]) // 4
+        nargs["CSDescriptor"].block_shape = [1, 1, rep_n, 2, 256]
+def rebind_batched_mx_bs_descriptor(nargs):
+    """Per-config pre_hook: size the swizzled weight-scale descriptor box to the tile's 128-row
+    blocks (doubled under GATE, whose tile spans 2*BN interleaved rows). BN<128 (fp8 scalar)
+    pointer-gathers instead and never reads the descriptor. Only under SWIZZLED_SCALES; the
+    un-swizzled path keeps its dummy box."""
+    if not nargs.get("SWIZZLED_SCALES"):
+        return
+    rep = max(1, ((2 if nargs.get("GATE") else 1) * nargs["BLOCK_SIZE_N"]) // 128)
+    rep_k = (nargs["BLOCK_SIZE_K"] // nargs["SCALE_GROUP_K"]) // 4
+    nargs["BSDescriptor"].block_shape = [1, rep, rep_k, 2, 256]
+def rebind_grouped_weight_descriptor(nargs):
+    """Per-config pre_hook: set the MX weight descriptor box to the tuned
+    ``[1, (2 if GATE else 1) * BLOCK_SIZE_N, BLOCK_SIZE_K // values_per_byte]`` over the
+    ``(E, 2N|N, K_bytes)`` weight view. MUST mutate ``block_shape`` in place — a rebind never reaches
+    the launch. No-op for pointer configs (they never read the descriptor)."""
+    if nargs.get("B_MEMORY_MODE", "pointer") == "pointer" or isinstance(
+        nargs["BDescriptor"], int
+    ):
+        return
+    values_per_byte = 2 if nargs["B"].dtype == torch.uint8 else 1
+    nargs["BDescriptor"].block_shape = [
+        1,
+        (2 if nargs.get("GATE") else 1) * nargs["BLOCK_SIZE_N"],
+        nargs["BLOCK_SIZE_K"] // values_per_byte,
+    ]
+def rebind_grouped_act_descriptor(nargs):
+    """Per-config pre_hook: set the activation descriptor box to the tuned
+    ``[BLOCK_SIZE_M, BLOCK_SIZE_K // act_values_per_byte]`` over the ``(rows, K_bytes)``
+    activation matrix. In-place mutate; no-op for pointer-A configs."""
+    if nargs.get("A_MEMORY_MODE", "pointer") == "pointer" or isinstance(
+        nargs["ADescriptor"], int
+    ):
+        return
+    act_values_per_byte = 2 if nargs["A"].dtype == torch.uint8 else 1
+    # tma gather4 loads N independent rows: descriptor_gather requires a 1-row box;
+    # the contiguous (no-gather) arm loads the whole [BM, BK_bytes] tile in one box
+    nargs["ADescriptor"].block_shape = [
+        1 if nargs.get("GatherIdx") is not None else nargs["BLOCK_SIZE_M"],
+        nargs["BLOCK_SIZE_K"] // act_values_per_byte,
+    ]
+def rebind_grouped_descriptors(nargs):
+    """Composite pre_hook: both weight and activation descriptor boxes."""
+    rebind_grouped_weight_descriptor(nargs)
+    rebind_grouped_act_descriptor(nargs)
+def build_grouped_operand_descriptors(a_operand, b_operand):
+    """Operand host-TMA descriptors for a grouped launch: A box ``[16, 64]``, B box
+    ``[1, 128, 64]`` over the ``(E, 2N|N, K_bytes)`` weight view. Placeholder boxes, re-bound to
+    the tuned tile per config by ``rebind_grouped_descriptors``."""
+    return (
+        TensorDescriptor.from_tensor(a_operand, block_shape=[16, 64]),
+        TensorDescriptor.from_tensor(b_operand, block_shape=[1, 128, 64]),
+    )
+def rebind_grouped_mx_descriptors(nargs):
+    """MX composite pre_hook: the operand boxes plus the two SWIZZLE_32_4_4 scale boxes
+    ``[1, BLOCK // 128, (BK // SCALE_GROUP_K) // 4, 2, 256]`` over the swizzled
+    ``(1, rows // 128, cols // 4, 2, 256)`` views. Activation box is BM // 128 (BM pinned 128);
+    the weight box is ``(2 if GATE else 1) * BN // 128`` — the stacked gate|up tile is one 2*BN
+    block (BN pinned 128 under GATE by ``swizzled_scales_bm_pruner``). Mutate in place. Both scale
+    descriptors are None on the un-swizzled arm (affine read — one SWIZZLED_SCALES flag governs both
+    operands), so their boxes are skipped."""
+    rebind_grouped_descriptors(nargs)
+    # BK below 4 scale groups (BK=64 rows the pruner rejects) still passes through here when a
+    # descriptor is unread — clamp like bn_blocks below so it keeps a valid, non-degenerate
+    # shape (a 0-block box traps the descriptor-encoding pass).
+    rep_k = max((nargs["BLOCK_SIZE_K"] // nargs["SCALE_GROUP_K"]) // 4, 1)
+    if nargs["ASDescriptor"] is not None:
+        nargs["ASDescriptor"].block_shape = [1, nargs["BLOCK_SIZE_M"] // 128, rep_k, 2, 256]
+    if nargs["BSDescriptor"] is not None:
+        # One bulk-load spans (2 if GATE) * BN//128 blocks: GATE reads the block-interleaved gate|up
+        # pair ([g,u] adjacent) as one 2*BN tile. BN<128 (non-gate, non-128 N) reads via the per-row
+        # pointer gather instead — clamp the box to one block so the (unread) descriptor keeps a valid,
+        # non-degenerate shape (a 0-block box traps the descriptor-encoding pass).
+        bn_blocks = max(1, ((2 if nargs.get("GATE") else 1) * nargs["BLOCK_SIZE_N"]) // 128)
+        nargs["BSDescriptor"].block_shape = [1, bn_blocks, rep_k, 2, 256]
+    # Swizzled requant output (Cs is a descriptor): the store tile is [1, 1, rep_n, 2, 256], and
+    # rep_n = (BN // SCALE_GROUP_K) // 4 depends on the tuned BLOCK_SIZE_N and the group size — so
+    # rebind per config, else nvfp4 (group-16 -> rep_n=2) mismatches the [.,.,1,.,.] build default.
+    if nargs["CSDescriptor"] is not None:
+        rep_n = (nargs["BLOCK_SIZE_N"] // nargs["SCALE_GROUP_K"]) // 4
+        nargs["CSDescriptor"].block_shape = [1, 1, rep_n, 2, 256]
+def rebind_dgrad_descriptors(nargs):
+    """Per-config pre_hook: set the dgrad boxes to the tuned tile — ``[BM, BN]`` over the (M, N)
+    gradient and ``[BN, BK // WEIGHT_VALUES_PER_BYTE]`` over the (N, K_bytes) weight. Both are
+    natural contiguous windows in the FORWARD-oriented operands, which is why dgrad can take a
+    descriptor at all: a transposed view has no unit-stride innermost dim (see
+    ``maybe_descriptor``). Scales stay affine on the pointer arm, as weight-only does."""
+    wvpb = 2 if nargs["B"].dtype == torch.uint8 else 1
+    rebind_operand_box(
+        nargs, "A_MEMORY_MODE", "ADescriptor", nargs["BLOCK_SIZE_M"], nargs["BLOCK_SIZE_N"]
+    )
+    rebind_operand_box(
+        nargs, "B_MEMORY_MODE", "BDescriptor",
+        nargs["BLOCK_SIZE_N"], nargs["BLOCK_SIZE_K"] // wvpb,
+    )
+def rebind_dgrad_grouped_descriptors(nargs):
+    """Grouped dgrad boxes. The activation box is 1 row when the pass gathers (tma gather4 needs a
+    1-row box) and ``[BM, BN]`` otherwise — ScatterIdx is the gather map here, since dgrad reads
+    dY at the forward's SCATTER destination. The weight box is ``[1, BN, BK_bytes]`` over
+    (E, N, K_bytes): the expert leads, so this descriptor takes THREE offsets."""
+    wvpb = 2 if nargs["B"].dtype == torch.uint8 else 1
+    gathering = nargs.get("ScatterIdx") is not None
+    rebind_operand_box(
+        nargs, "A_MEMORY_MODE", "ADescriptor",
+        1 if gathering else nargs["BLOCK_SIZE_M"], nargs["BLOCK_SIZE_N"],
+    )
+    desc = nargs["BDescriptor"]
+    if nargs.get("B_MEMORY_MODE", "pointer") != "pointer" and not isinstance(desc, int) and desc is not None:
+        desc.block_shape = [1, nargs["BLOCK_SIZE_N"], nargs["BLOCK_SIZE_K"] // wvpb]

build/torch-rocm/epilogue.py ADDED Viewed

	@@ -0,0 +1,639 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait
+from .compat import compile_time_only_triton_wrap, decode_pdl, device_context
+from .mma import MMA_N_ATOM
+from .quant import fp8_act_quant_inline, mx_act_quant_inline
+from .loading.scales import apply_global_scale
+@triton.jit
+def store_masked(
+    C,
+    accumulator,
+    pid_m,
+    pid_n,
+    M,
+    N,
+    stride_c_m,
+    stride_c_n,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+):
+    """Shared output epilogue of the kernels below: cast the fp32 accumulator to
+    ``C``'s dtype and store the ``(BLOCK_SIZE_M, BLOCK_SIZE_N)`` tile at the swizzled
+    ``(pid_m, pid_n)``, masked to the ``(M, N)`` bounds. (A descriptor-store arm was
+    measured an EXACT tie at the store-heavy gap shape and dropped, 2026-07-16 —
+    stores are fire-and-forget, TMA has nothing to hide there. B200-only verdict;
+    re-measure on H100 or the target device.)"""
+    c = accumulator.to(C.dtype.element_ty)
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = C + stride_c_m * offs_cm[:, None] + stride_c_n * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, c, mask=c_mask)
+@triton.jit
+def store_masked_oriented(
+    C,
+    accumulator,
+    pid_m,
+    pid_n,
+    M,
+    N,
+    stride_c_m,
+    stride_c_n,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    SWAP_AB: tl.constexpr,
+):
+    """``store_masked`` with the (m, n) roles swapped under ``SWAP_AB``, where the
+    accumulator is ``(BLOCK_SIZE_N, BLOCK_SIZE_M)``. Only the taken branch compiles."""
+    if SWAP_AB:
+        store_masked(
+            C,
+            accumulator,
+            pid_n,
+            pid_m,
+            N,
+            M,
+            stride_c_n,
+            stride_c_m,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_M,
+        )
+    else:
+        store_masked(
+            C,
+            accumulator,
+            pid_m,
+            pid_n,
+            M,
+            N,
+            stride_c_m,
+            stride_c_n,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+        )
+@triton.jit
+def acc_init(
+    COMPUTE_MODE: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    SWAP_AB: tl.constexpr,
+):
+    """Zero accumulator shaped for the layout: swapped scalar reduces into ``[1, N]``; any other
+    swapped mode keeps the ``[N, M]`` MMA acc (weight rows in M, act tile in N — padded up to
+    the ``MMA_N_ATOM`` when the caller is the single-token decode GEVM, col 0 taken after the
+    K-loop); no-swap uses ``[M, N]``. ``COMPUTE_MODE`` matters only under swap — kernels with no
+    mode axis (fp8 ``tl.dot``) pass ``"dot"``. ``N`` is the weight-output tile (``BLOCK_SIZE_H``
+    for the fp8 down projection). Single return: only the taken branch compiles."""
+    if SWAP_AB and COMPUTE_MODE == "scalar":
+        acc = tl.zeros((1, BLOCK_SIZE_N), dtype=tl.float32)
+    elif SWAP_AB and BLOCK_SIZE_M < MMA_N_ATOM:
+        acc = tl.zeros((BLOCK_SIZE_N, MMA_N_ATOM), dtype=tl.float32)
+    elif SWAP_AB:
+        acc = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32)
+    else:
+        acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    return acc
+@triton.jit
+def acc_finalize(
+    acc, COMPUTE_MODE: tl.constexpr, ROWS: tl.constexpr, SWAP_AB: tl.constexpr
+):
+    """Bookend to ``acc_init``: when the acc was built as the persistent ``[ROWS, MMA_N_ATOM]`` MMA
+    tile (any swapped non-scalar mode), collapse the padded token dim to column 0 → ``[1, ROWS]``.
+    Swapped scalar (already ``[1, ROWS]``) and no-swap pass through unchanged. ``COMPUTE_MODE``
+    matches ``acc_init``'s (fp8 ``tl.dot`` kernels, which have no mode axis, pass ``"dot"``)."""
+    if SWAP_AB and COMPUTE_MODE != "scalar":
+        # take column 0: the padded token dim collapses back to the single real token
+        acc = tl.reshape(
+            tl.sum(acc * (tl.arange(0, MMA_N_ATOM)[None, :] == 0), axis=1), (1, ROWS)
+        )
+    return acc
+def bias_strides(bias: torch.Tensor | None) -> tuple[int, int]:
+    """(expert, column) strides for a per-expert output bias, for ``add_bias``. A 1-D bias is the
+    single-expert 2D case, so its expert stride is 0. ``None`` gives (0, 0) — the kernel arm folds
+    out at trace time and never reads them."""
+    if bias is None:
+        return 0, 0
+    if bias.dim() == 1:
+        return 0, bias.stride(0)
+    return bias.stride(0), bias.stride(1)
+@triton.jit
+def add_bias(acc, Bias, stride_bias_e, stride_bias_n, bias_row, pid_n,
+             BLOCK_SIZE_N: tl.constexpr, N_COLS: tl.constexpr = 0,
+             SWAP_AB: tl.constexpr = False):
+    """Add this expert's slice of a per-expert output bias to an output tile. Gated callers pass
+    the doubled extent — an interleaved gate|up tile is an ungated tile of width ``2*BN``.
+    ``None`` folds out at trace time.
+    Two shapes reach here. Post-``acc_finalize`` tiles are ``[rows, BN]`` for every orientation
+    (``SWAP_AB=False``, what ``gemm_epilogue`` passes) — finalize is what collapses the swapped
+    MMA accumulator, so adding before it would broadcast against the wrong axis. The 2D kernels'
+    ungated arm stores straight from the accumulator with no finalize, and there ``SWAP_AB``
+    leaves it transposed as ``(BN, BM)``, so the bias varies down axis 0 instead."""
+    if Bias is not None:
+        offs_bias = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+        ptrs = Bias + bias_row * stride_bias_e + offs_bias * stride_bias_n
+        if N_COLS > 0:  # dense N isn't BN-aligned; the tail columns are off the end of the bias
+            bias = tl.load(ptrs, mask=offs_bias < N_COLS, other=0.0)
+        else:
+            bias = tl.load(ptrs)
+        if SWAP_AB:
+            out = acc + bias.to(acc.dtype)[:, None]
+        else:
+            out = acc + bias.to(acc.dtype)[None, :]
+    else:
+        out = acc
+    return out
+@triton.jit
+def split_gate_up(flat, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, SWAP_AB: tl.constexpr):
+    """De-interleave a finalized ``[rows, 2*BN]`` gate|up tile into the ``(gate, up)`` pair, each
+    ``[rows, BN]`` (rows = 1 under swap, else BM). The N axis alternates gate/up per column, so
+    this is a trailing-axis split. The caller finalizes (and biases) before calling."""
+    rows: tl.constexpr = 1 if SWAP_AB else BLOCK_SIZE_M
+    g, u = tl.split(tl.reshape(flat, (rows, BLOCK_SIZE_N, 2)))
+    return g, u
+@triton.jit
+def glu(
+    gate,
+    up,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.float32,
+):
+    """Gated linear unit on the gate/up matmul accumulators. ``SWIGLU_LIMIT`` clamps gate above and up
+    to ``[-LIMIT, LIMIT]``; ``SWIGLU_ALPHA`` gives the clamped/scaled SwiGLU ``(up + 1) * gate * sigmoid(ALPHA *
+    gate)`` (GPT-OSS / MiniMax), else ``ACT_FN(gate) * up`` (``ACT_FN`` in {silu, gelu, relu}, gelu exact
+    via erf). ``SIMULATE_UNFUSED`` rounds each materialized value through ``INTERMEDIATE_DTYPE`` (the dtype the unfused path lands intermediates in) to match the
+    unfused (separate-kernel) path, where every intermediate lands in that dtype."""
+    g = gate
+    u = up
+    if SIMULATE_UNFUSED:
+        g = g.to(INTERMEDIATE_DTYPE).to(tl.float32)
+        u = u.to(INTERMEDIATE_DTYPE).to(tl.float32)
+    if SWIGLU_LIMIT is not None:
+        g = tl.minimum(g, SWIGLU_LIMIT)
+        u = tl.minimum(tl.maximum(u, -SWIGLU_LIMIT), SWIGLU_LIMIT)
+    if SWIGLU_ALPHA is not None:
+        gate_scaled = g * SWIGLU_ALPHA
+        if SIMULATE_UNFUSED:
+            gate_scaled = gate_scaled.to(INTERMEDIATE_DTYPE).to(tl.float32)
+        sig = tl.sigmoid(gate_scaled)
+        if SIMULATE_UNFUSED:
+            sig = sig.to(INTERMEDIATE_DTYPE).to(tl.float32)
+        act = g * sig
+        u = u + 1.0
+    elif ACT_FN == "silu":
+        sig = tl.sigmoid(g)
+        # SIMULATE_UNFUSED must be bit-exact vs the unfused ``apply_glu`` (``g * torch.sigmoid(g)``),
+        # where torch.sigmoid returns bf16 — i.e. the sigmoid is rounded before the multiply. Round
+        # it here to match (the fp32 sigmoid otherwise flips e4m3 requant bits, ~35% on MXFP4 down).
+        if SIMULATE_UNFUSED:
+            sig = sig.to(INTERMEDIATE_DTYPE).to(tl.float32)
+        act = g * sig
+    elif ACT_FN == "gelu":
+        if SIMULATE_UNFUSED:
+            # Bit-match the unfused ``apply_glu`` gelu ``0.5 * g * (1 + erf(g * c))``, which rounds
+            # to bf16 at every torch op (input to erf, erf, 1+erf, 0.5*g, final mul). Rounding only
+            # a subset diverges (~0.7 rel) on the MX requant — round each op, like torch does.
+            gc = (g * 0.7071067811865476).to(INTERMEDIATE_DTYPE).to(tl.float32)
+            e = tl.erf(gc).to(INTERMEDIATE_DTYPE).to(tl.float32)
+            one_plus = (1.0 + e).to(INTERMEDIATE_DTYPE).to(tl.float32)
+            half_g = (0.5 * g).to(INTERMEDIATE_DTYPE).to(tl.float32)
+            act = half_g * one_plus
+        else:
+            act = 0.5 * g * (1.0 + tl.erf(g * 0.7071067811865476))
+    elif ACT_FN == "relu":
+        act = tl.maximum(g, 0.0)
+    else:
+        tl.static_assert(
+            False, "unsupported ACT_FN; expected 'silu', 'gelu', or 'relu'"
+        )
+    if SIMULATE_UNFUSED:
+        act = act.to(INTERMEDIATE_DTYPE).to(tl.float32)
+        u = u.to(INTERMEDIATE_DTYPE).to(tl.float32)
+    gated = act * u
+    if SIMULATE_UNFUSED:
+        gated = gated.to(INTERMEDIATE_DTYPE).to(tl.float32)
+    return gated
+@triton.jit
+def _glu_kernel(
+    GateUp,  # (S, 2*INTER) interleaved [g0,u0,g1,u1,...] GEMM output; fp32 = exact accumulators
+    Out,  # (S, INTER) activation-dtype GLU result; the E4M3 tensor under the requant arm
+    QuantScales,  # block-FP8 requant arm (``fused_glu(quant_group=...)``): per-``QUANT_GROUP``
+    #             scales (fp32, or UE8M0 exponent bytes); ``None`` folds the arm out — the GLU
+    #             result then rounds through bf16 (the reference two-kernel order) and quantizes
+    n_out,
+    INTER: tl.constexpr,
+    ACT_FN: tl.constexpr,
+    SWIGLU_ALPHA: tl.constexpr,
+    SWIGLU_LIMIT: tl.constexpr,
+    BLOCK: tl.constexpr,
+    QUANT_GROUP: tl.constexpr = 128,
+    UE8M0: tl.constexpr = False,
+    PDL: tl.constexpr = False,
+):
+    """Fused GLU over a contiguous (S, 2*INTER) interleaved gate|up tensor -> (S, INTER), covering every
+    ``apply_glu`` variant (silu/gelu/relu, ``SWIGLU_ALPHA``/``SWIGLU_LIMIT`` clamped-SwiGLU
+    — ``None`` folds the arm out, the in-kernel ``glu``'s convention). Rounds through the
+    output dtype after each op exactly where the torch chain does (each torch tensor op
+    materializes), so the two are bit-identical."""
+    if PDL:
+        gdc_wait()
+    offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n_out
+    # gate|up columns alternate (gate 2j, up 2j+1) — the interleaved row order the GEMM emits
+    base = (offs // INTER) * 2 * INTER + 2 * (offs % INTER)
+    # fp32 input = exact GEMM accumulators (the unstacked fused-order path): compute the
+    # GLU in fp32 and round ONCE at the store — the gated epilogue's math. Narrow inputs
+    # keep the torch-chain per-op rounding (bit-identical to apply_glu).
+    FUSED_ORDER: tl.constexpr = GateUp.dtype.element_ty == tl.float32
+    dt = tl.float32 if FUSED_ORDER else (
+        tl.bfloat16 if QuantScales is not None else Out.dtype.element_ty)
+    g = tl.load(GateUp + base, mask=mask).to(tl.float32)
+    u = tl.load(GateUp + base + 1, mask=mask).to(tl.float32)
+    if SWIGLU_LIMIT is not None:
+        g = tl.minimum(g, SWIGLU_LIMIT)
+        u = tl.clamp(u, -SWIGLU_LIMIT, SWIGLU_LIMIT)
+    if SWIGLU_ALPHA is not None:  # (up + 1) * gate * sigmoid(alpha * gate)
+        sig = tl.sigmoid((g * SWIGLU_ALPHA).to(dt).to(tl.float32)).to(dt).to(tl.float32)
+        act = (g * sig).to(dt).to(tl.float32)
+        out = ((u + 1.0).to(dt).to(tl.float32) * act).to(dt)
+    elif ACT_FN == "silu":
+        sig = tl.sigmoid(g).to(dt).to(tl.float32)
+        act = (g * sig).to(dt).to(tl.float32)
+        out = (act * u).to(dt)
+    elif ACT_FN == "gelu":  # exact via erf, torch's evaluation order
+        e = tl.erf((g * 0.7071067811865476).to(dt).to(tl.float32)).to(dt).to(tl.float32)
+        d = (1.0 + e).to(dt).to(tl.float32)
+        a = (0.5 * g).to(dt).to(tl.float32)
+        out = ((a * d).to(dt).to(tl.float32) * u).to(dt)
+    else:  # relu
+        out = (tl.maximum(g, 0.0) * u).to(dt)
+    if QuantScales is None:
+        tl.store(Out + offs, out.to(Out.dtype.element_ty), mask=mask)
+    else:
+        grp = tl.reshape(out.to(tl.float32), (BLOCK // QUANT_GROUP, QUANT_GROUP))
+        amax = tl.max(tl.abs(grp), axis=1)
+        if UE8M0:
+            sc = tl.exp2(tl.ceil(tl.log2(tl.maximum(amax, 1e-10) / 448.0)))
+        else:
+            sc = tl.maximum(amax, 1e-10) / 448.0
+        q = tl.clamp(grp / sc[:, None], -448.0, 448.0)
+        tl.store(Out + offs, tl.reshape(q, (BLOCK,)).to(tl.float8e4nv), mask=mask)
+        soffs = tl.program_id(0) * (BLOCK // QUANT_GROUP) + tl.arange(0, BLOCK // QUANT_GROUP)
+        smask = soffs * QUANT_GROUP < n_out
+        if UE8M0:
+            tl.store(QuantScales + soffs, (tl.log2(sc) + 127.0).to(tl.uint8), mask=smask)
+        else:
+            if PDL:
+                gdc_launch_dependents()
+            tl.store(QuantScales + soffs, sc, mask=smask)
+def fused_glu(
+    gate_up: torch.Tensor,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    quant_group: int | None = None,
+    use_ue8m0: bool = False,
+    out_dtype: torch.dtype | None = None,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    """GLU over the interleaved (S, 2*INTER) gate|up GEMM output — the one-kernel ``_glu_kernel``
+    on a contiguous tensor, the torch ``apply_glu`` on any fallback layout; bit-identical
+    either way. ``quant_group`` additionally block-FP8-requantizes the result per group in
+    the same kernel (the unstacked decode band's handoff to the down projection) and
+    returns ``(q, scales)`` — scales fp32, or UE8M0 exponent bytes under ``use_ue8m0``;
+    non-SiLU/clamped variants take the two-step order (same bits — the fused arm rounds
+    through bf16 first for exactly this equivalence)."""
+    inter_dim = gate_up.shape[-1] // 2
+    if quant_group is not None:
+        if act_fn == "silu" and swiglu_alpha is None and swiglu_limit is None and gate_up.is_contiguous():
+            S = gate_up.shape[0]
+            q = torch.empty(S, inter_dim, dtype=torch.float8_e4m3fn, device=gate_up.device)
+            sc = torch.empty(S, inter_dim // quant_group,
+                             dtype=torch.uint8 if use_ue8m0 else torch.float32,
+                             device=gate_up.device)
+            n = q.numel()
+            with device_context(gate_up.device):
+                compile_time_only_triton_wrap(_glu_kernel)[(triton.cdiv(n, 1024),)](
+                    gate_up, q, sc, n, inter_dim, "silu", None, None, BLOCK=1024,
+                    QUANT_GROUP=quant_group, UE8M0=use_ue8m0,
+                    PDL=decode_pdl(),
+                    launch_pdl=decode_pdl(),
+                )
+            return q, sc
+        from .quant import fp8_act_quant_block_dynamic  # deferred: module import order
+        inter = fused_glu(gate_up, act_fn, swiglu_alpha, swiglu_limit)
+        return fp8_act_quant_block_dynamic(inter, quant_group, use_ue8m0=use_ue8m0)
+    if gate_up.is_contiguous() and act_fn in ("silu", "gelu", "relu"):
+        out = torch.empty(*gate_up.shape[:-1], inter_dim,
+                          dtype=out_dtype or gate_up.dtype, device=gate_up.device)
+        n = out.numel()
+        with device_context(gate_up.device):
+            compile_time_only_triton_wrap(_glu_kernel)[(triton.cdiv(n, 1024),)](
+                gate_up, out, None, n, inter_dim, act_fn, swiglu_alpha, swiglu_limit, BLOCK=1024,
+                PDL=decode_pdl(),
+                launch_pdl=decode_pdl(),
+            )
+        return out
+    gate, up = gate_up[..., 0::2], gate_up[..., 1::2]
+    return apply_glu(gate, up, act_fn, swiglu_alpha, swiglu_limit)
+def apply_glu(
+    gate: torch.Tensor,
+    up: torch.Tensor,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+) -> torch.Tensor:
+    """Host-side (torch) gated linear unit — the unfused path's activation, mirroring the triton
+    ``glu``. ``swiglu_limit`` clamps gate above / up to ``[-limit, limit]``; ``swiglu_alpha`` gives
+    the clamped/scaled SwiGLU ``(up + 1) * gate * sigmoid(alpha * gate)`` (GPT-OSS / MiniMax), else
+    ``act_fn(gate) * up`` (``act_fn`` in {silu, gelu, relu}, gelu exact via erf)."""
+    if swiglu_limit is not None:
+        gate = gate.clamp(max=swiglu_limit)
+        up = up.clamp(min=-swiglu_limit, max=swiglu_limit)
+    if swiglu_alpha is not None:
+        return (up + 1.0) * (gate * torch.sigmoid(gate * swiglu_alpha))
+    if act_fn == "silu":
+        act = gate * torch.sigmoid(gate)
+    elif act_fn == "gelu":
+        act = 0.5 * gate * (1.0 + torch.erf(gate * 0.7071067811865476))
+    elif act_fn == "relu":
+        act = gate.clamp(min=0.0)
+    else:
+        raise ValueError(
+            f"unsupported act_fn {act_fn!r}; expected 'silu', 'gelu', or 'relu'"
+        )
+    return act * up
+@triton.jit
+def split_gate_up_glu(
+    flat,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    SWAP_AB: tl.constexpr,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.float32,
+):
+    """Gate|up epilogue in one step: split the finalized stacked tile into its (gate, up) pair
+    (``split_gate_up``) and apply the ``ACT_FN``/SwiGLU gated linear unit (``glu``), returning
+    the combined intermediate. See those two for the orientation and activation details."""
+    gate, up = split_gate_up(flat, BLOCK_SIZE_M, BLOCK_SIZE_N, SWAP_AB)
+    return glu(
+        gate,
+        up,
+        ACT_FN,
+        SWIGLU_ALPHA,
+        SWIGLU_LIMIT,
+        SIMULATE_UNFUSED,
+        INTERMEDIATE_DTYPE,
+    )
+@triton.jit
+def _store_out(
+    C, acc, out_row, pid_n, row_mask, stride_c_m, stride_c_n,
+    BLOCK_SIZE_M: tl.constexpr, WIDTH: tl.constexpr, FAKE_BATCH: tl.constexpr,
+    N_COLS: tl.constexpr = 0,
+):
+    """Cast + store one output tile of N-width ``WIDTH`` (halved when the format packs nibble
+    pairs). ``FAKE_BATCH`` (batched decode): the BM lanes alias one C row (``C`` pre-advanced), so
+    a plain store would duplicate-write the same bytes (hardware-undefined on Intel XPU) — mask to
+    lane 0 (the replicated rows are identical). Else a real scatter to global rows ``out_row`` under
+    ``row_mask``. ``N_COLS`` > 0 also masks the column tail (the 2D dense output isn't ``BN``-aligned
+    like the ``N % BN == 0`` grouped/batched MoE outputs); 0 skips it. Single return."""
+    c = acc.to(C.dtype.element_ty)
+    offs_cm = tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * WIDTH + tl.arange(0, WIDTH)
+    col_ok = (offs_cn < N_COLS) if N_COLS > 0 else (offs_cn >= 0)
+    if FAKE_BATCH:
+        c_ptrs = C + offs_cm[:, None] * 0 + stride_c_n * offs_cn[None, :]
+        tl.store(c_ptrs, c, mask=(offs_cm == 0)[:, None] & col_ok[None, :])
+    else:
+        c_ptrs = C + stride_c_m * out_row[:, None] + stride_c_n * offs_cn[None, :]
+        tl.store(c_ptrs, c, mask=row_mask[:, None] & col_ok[None, :])
+@triton.jit
+def _epilogue_requant_fp8(
+    C, Cs, out, out_row, pid_n, row_mask, stride_c_m, stride_c_n, stride_cs_m, stride_cs_n,
+    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, FAKE_BATCH: tl.constexpr,
+    N_COLS: tl.constexpr,
+):
+    """Requantize the SwiGLU intermediate to fp8 with a per-(row, N-tile) scalar scale. UE8M0
+    intermediate scales under a UE8M0 model (inferred from Cs's dtype) keep the down proj's activation
+    scales power-of-two so its dot_scaled arm fires. FAKE_BATCH collapses the replicated rows to one
+    scalar per (row, N-tile)."""
+    q, q_s = fp8_act_quant_inline(out, UE8M0=Cs.dtype.element_ty == tl.uint8)
+    _store_out(C, q, out_row, pid_n, row_mask, stride_c_m, stride_c_n, BLOCK_SIZE_M, BLOCK_SIZE_N, FAKE_BATCH, N_COLS)
+    cs_ptr = Cs + out_row * stride_cs_m + pid_n * stride_cs_n
+    if FAKE_BATCH:  # replicated rows -> one scalar per (row, N-tile)
+        tl.store(cs_ptr, tl.max(q_s))
+    else:
+        tl.store(cs_ptr, q_s, mask=row_mask)
+@triton.jit
+def _epilogue_requant_mx(
+    C, Cs, out, out_row, pid_n, pid_m, row_mask, stride_c_m, stride_c_n, stride_cs_m, stride_cs_n,
+    CSDescriptor, CsGlobal, cs_global_row,
+    BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, SCALE_GROUP_K: tl.constexpr,
+    OUTPUT_FORMAT: tl.constexpr, SWIZZLED_OUT: tl.constexpr, FAKE_BATCH: tl.constexpr,
+    N_COLS: tl.constexpr,
+):
+    """Requantize the SwiGLU intermediate to MX group-``SCALE_GROUP_K`` (mxfp8 UE8M0 / mxfp4 / nvfp4
+    E4M3 — the fp4 formats pack nibble pairs so ``C`` halves). ``SWIZZLED_OUT`` writes ``Cs`` straight
+    into the down proj's SWIZZLE_32_4_4 descriptor at block ``(pid_m, pid_n)`` (BM/BN pinned 128), else
+    a row-major affine store. NVFP4 two-level: normalize by the next proj's provided input_scale
+    (``CsGlobal[cs_global_row]`` — per tensor or this expert's) before the block quant. FAKE_BATCH
+    collapses replicated rows via the row-max."""
+    if CsGlobal is not None:
+        # NVFP4 two-level requant: normalize the fp32 GLU intermediate by the NEXT proj's provided
+        # (calibrated) input_scale before the block quant — the canonical two-step. The down folds it
+        # back via its As pair ([Cs, g_out]); nothing is computed at runtime.
+        out = out / tl.load(CsGlobal + cs_global_row).to(tl.float32)
+    q, q_s = mx_act_quant_inline(out, BLOCK_SIZE_M, BLOCK_SIZE_N, SCALE_GROUP_K, OUTPUT_FORMAT)
+    width: tl.constexpr = BLOCK_SIZE_N if OUTPUT_FORMAT == "mxfp8" else BLOCK_SIZE_N // 2
+    _store_out(C, q, out_row, pid_n, row_mask, stride_c_m, stride_c_n, BLOCK_SIZE_M, width, FAKE_BATCH,
+               N_COLS if OUTPUT_FORMAT == "mxfp8" else N_COLS // 2)
+    if SWIZZLED_OUT:
+        # group scales straight into the down proj's SWIZZLE_32_4_4 layout (inverse of
+        # load_swizzled_scale) at block (pid_m, pid_n) — BM/BN pinned 128.
+        REP_K_CS: tl.constexpr = (BLOCK_SIZE_N // SCALE_GROUP_K) // 4
+        sw = (
+            q_s.reshape(1, 4, 32, REP_K_CS, 4)
+            .trans(0, 3, 2, 1, 4)
+            .reshape(1, 1, REP_K_CS, 2, 256)
+        )
+        CSDescriptor.store([0, pid_m, pid_n * REP_K_CS, 0, 0], sw)
+    else:
+        offs_sc = pid_n * (BLOCK_SIZE_N // SCALE_GROUP_K) + tl.arange(0, BLOCK_SIZE_N // SCALE_GROUP_K)
+        if FAKE_BATCH:
+            # replicated rows -> the row-max IS the row's scale (f32-exact for UE8M0 exponent bytes
+            # and E4M3 values alike)
+            tl.store(
+                Cs + out_row * stride_cs_m + offs_sc[None, :] * stride_cs_n,
+                tl.reshape(tl.max(q_s.to(tl.float32), axis=0), (1, BLOCK_SIZE_N // SCALE_GROUP_K)),
+            )
+        else:
+            tl.store(
+                Cs + out_row[:, None] * stride_cs_m + offs_sc[None, :] * stride_cs_n,
+                q_s,
+                mask=row_mask[:, None],
+            )
+@triton.jit
+def gemm_epilogue(
+    C,
+    Cs,  # row-major requant-scale pointer (None under SWIZZLED_OUT — CSDescriptor writes instead)
+    acc,
+    out_row,
+    pid_n,
+    pid_m,
+    row_mask,
+    stride_c_m,
+    stride_c_n,
+    stride_cs_m,
+    stride_cs_n,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    GATE: tl.constexpr,
+    OUTPUT_FORMAT: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    ACT_FN: tl.constexpr,
+    SWIGLU_ALPHA: tl.constexpr,
+    SWIGLU_LIMIT: tl.constexpr,
+    SIMULATE_UNFUSED: tl.constexpr,
+    INTERMEDIATE_DTYPE: tl.constexpr,
+    COMPUTE_MODE: tl.constexpr = "dot",
+    SWAP_AB: tl.constexpr = False,
+    SWIZZLED_OUT: tl.constexpr = False,
+    FAKE_BATCH: tl.constexpr = False,
+    N_COLS: tl.constexpr = 0,  # >0 masks the column tail (2D dense N isn't BN-aligned); 0 = no mask
+    CSDescriptor=0,  # SWIZZLE_32_4_4 requant-scale descriptor; read only under SWIZZLED_OUT (else dummy)
+    CsGlobal=None,  # fp32 NVFP4 output global (the NEXT proj's provided input_scale), per tensor or per expert; normalizes the requant, None folds out
+    stride_cs_global=0,  # expert stride of CsGlobal (0 broadcasts a scalar)
+    GlobalScale=None,  # fp32 NVFP4 weight global g_b, per tensor or per expert; applied PRE-GLU, None folds out
+    stride_global_e=0,  # expert stride of GlobalScale (global_scale_stride; 0 broadcasts a scalar)
+    global_row=0,  # expert id indexing the globals and the bias (0 for the 2D case)
+    GlobalScaleA=None,  # fp32 NVFP4 activation global g_a, per tensor or per expert; multiplied in-register with GlobalScale
+    stride_global_a=0,  # expert stride of GlobalScaleA (0 broadcasts a scalar)
+    PreAct=None,  # (M, 2N) pre-activation buffer; written iff not None — the GLU backward's Z
+    stride_pa_m=0,
+    stride_pa_n=0,
+    Bias=None,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; row is `global_row`
+    stride_bias_e=0,
+    stride_bias_n=0,
+):
+    """Unified output epilogue for grouped (a real scatter tile) and batched (fake-batch decode:
+    one token replicated across the BM lanes) GEMMs. Plain: cast + store the accumulator. ``GATE``:
+    split the stacked gate|up accumulator + SwiGLU (``split_gate_up_glu``); ``OUTPUT_FORMAT`` — the
+    ``activation_format`` vocabulary — then requantizes into ``C`` + ``Cs``: ``"fp8"`` (per-(row, N-tile)
+    scalar), or MX group-``SCALE_GROUP_K`` (UE8M0/E4M3 — the fp4 formats pack nibble pairs so ``C``
+    halves), with ``SWIZZLED_OUT`` writing the MX ``Cs`` straight into the down proj's SWIZZLE_32_4_4
+    descriptor at block ``(pid_m, pid_n)`` (the tcgen05 fast path, BM/BN pinned 128). ``FAKE_BATCH``
+    shims the store: value masks to lane 0 (``C`` pre-advanced), the scale collapses the replicated
+    rows with ``tl.max``; else a real BM-row scatter (``out_row`` + ``row_mask``).
+    ``COMPUTE_MODE``/``SWAP_AB`` orient the decode GLU/finalize (grouped passes ``"dot"``/no-swap,
+    both no-ops there). Every arm is constexpr-pruned."""
+    # The globals fold on the FINALIZED tile: a swapped accumulator still carries the token dim
+    # padded to the MMA atom, so folding before the collapse costs 16x the multiplies.
+    if GATE:
+        # Finalize + bias once, up here rather than inside the split: the gate|up bias belongs to
+        # the pre-activation, so PreAct below has to see it too.
+        flat = add_bias(
+            apply_global_scale(
+                acc_finalize(acc, COMPUTE_MODE, 2 * BLOCK_SIZE_N, SWAP_AB),
+                GlobalScale, stride_global_e, GlobalScaleA, stride_global_a, global_row,
+            ),
+            Bias, stride_bias_e, stride_bias_n, global_row, pid_n,
+            2 * BLOCK_SIZE_N, 2 * N_COLS if N_COLS > 0 else 0,
+        )
+        # SAVE_PREACT: the GLU backward needs Z, the 2*BN-wide pre-activation. Interleaved, that
+        # tile IS an ungated output tile at doubled width (pid_n * 2BN == 2 * pid_n * BN), so the
+        # ordinary store serves it — no separate path, no gated store arm.
+        if PreAct is not None:
+            _store_out(
+                PreAct, flat,
+                out_row, pid_n, row_mask, stride_pa_m, stride_pa_n,
+                BLOCK_SIZE_M, 2 * BLOCK_SIZE_N, FAKE_BATCH,
+                2 * N_COLS if N_COLS > 0 else 0,
+            )
+        out = split_gate_up_glu(
+            flat, BLOCK_SIZE_M, BLOCK_SIZE_N, SWAP_AB,
+            ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+        )
+        if OUTPUT_FORMAT == "fp8":
+            _epilogue_requant_fp8(
+                C, Cs, out, out_row, pid_n, row_mask, stride_c_m, stride_c_n, stride_cs_m, stride_cs_n,
+                BLOCK_SIZE_M, BLOCK_SIZE_N, FAKE_BATCH, N_COLS,
+            )
+        elif OUTPUT_FORMAT is not None:  # "mxfp8" | "mxfp4" | "nvfp4"
+            _epilogue_requant_mx(
+                C, Cs, out, out_row, pid_n, pid_m, row_mask, stride_c_m, stride_c_n, stride_cs_m,
+                stride_cs_n, CSDescriptor, CsGlobal, global_row * stride_cs_global,
+                BLOCK_SIZE_M, BLOCK_SIZE_N, SCALE_GROUP_K, OUTPUT_FORMAT, SWIZZLED_OUT, FAKE_BATCH, N_COLS,
+            )
+        else:  # bf16 (unquantized) SwiGLU output
+            _store_out(C, out, out_row, pid_n, row_mask, stride_c_m, stride_c_n, BLOCK_SIZE_M, BLOCK_SIZE_N, FAKE_BATCH, N_COLS)
+    else:  # plain GEMM: cast + store the accumulator (no fused requant on the non-gate path)
+        acc = add_bias(
+            apply_global_scale(
+                acc_finalize(acc, COMPUTE_MODE, BLOCK_SIZE_N, SWAP_AB),
+                GlobalScale, stride_global_e, GlobalScaleA, stride_global_a, global_row,
+            ),
+            Bias, stride_bias_e, stride_bias_n, global_row, pid_n, BLOCK_SIZE_N, N_COLS,
+        )
+        _store_out(C, acc, out_row, pid_n, row_mask, stride_c_m, stride_c_n, BLOCK_SIZE_M, BLOCK_SIZE_N, FAKE_BATCH, N_COLS)

build/torch-rocm/finegrained_kernels/__init__.py ADDED Viewed

	@@ -0,0 +1,26 @@

+import ctypes
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+def _import_from_path(file_path: Path) -> ModuleType:
+    # We cannot use the module name as-is, after adding it to `sys.modules`,
+    # it would also be used for other imports. So, we make a module name that
+    # depends on the path for it to be unique using the hex-encoded hash of
+    # the path.
+    path_hash = "{:x}".format(ctypes.c_size_t(hash(file_path.absolute())).value)
+    module_name = path_hash
+    spec = importlib.util.spec_from_file_location(module_name, file_path)
+    if spec is None:
+        raise ImportError(f"Cannot load spec for {module_name} from {file_path}")
+    module = importlib.util.module_from_spec(spec)
+    if module is None:
+        raise ImportError(f"Cannot load module {module_name} from spec")
+    sys.modules[module_name] = module
+    spec.loader.exec_module(module)  # type: ignore
+    return module
+globals().update(vars(_import_from_path(Path(__file__).parent.parent / "__init__.py")))

build/torch-rocm/formats.py ADDED Viewed

	@@ -0,0 +1,357 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+import triton
+from .compat import MX_SCALE_GROUP_K, NIBBLES_PER_BYTE, NVFP4_SCALE_GROUP_K
+def ue8m0_as_uint8(scale: torch.Tensor | None) -> torch.Tensor | None:
+    """View UE8M0 (``float8_e8m0fnu``) weight scales as ``uint8`` for the Triton
+    binder, which doesn't recognize the dtype; kernels decode ``2^(exp-127)``
+    inline. fp32 (non-UE8M0) scales pass through unchanged; ``None`` (an absent
+    optional scale) passes through as ``None`` — kernels take the arg as None."""
+    if scale is None:
+        return None
+    return scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale
+def e2m1_as_uint8(weight: torch.Tensor) -> torch.Tensor:
+    """View an ``int8``-stored MXFP4 (packed E2M1) weight as ``uint8`` — a zero-cost
+    reinterpret. ``tl.dot_scaled`` requires the packed rhs as ``uint8``, so do it once here
+    instead of casting in-kernel at every load. E4M3 (MXFP8) weights pass through unchanged."""
+    return weight.view(torch.uint8) if weight.dtype == torch.int8 else weight
+# UE8M0 group-32 scales arrive either as ``float8_e8m0fnu`` or as raw ``uint8`` — the same 8
+# exponent bits, and a common on-disk encoding (e.g. group-32 "mxfp8" checkpoints store the
+# scale tensor as uint8). Both are valid MX scales: ``ue8m0_as_uint8`` reinterprets to uint8
+# and the kernels decode ``2^(exp-127)`` inline, so the detectors accept either dtype.
+UE8M0_SCALE_DTYPES = (torch.float8_e8m0fnu, torch.uint8)
+def _shapes_match(weight: torch.Tensor, scale: torch.Tensor, group: int) -> bool:
+    """Shape leg of the family predicates: matching leading dims and a last dim of
+    one scale per ``group`` unpacked values. Early-return ``if``s, not an ``and``
+    chain: callers compare predicate results (``is_x(gate) != is_x(down)``), and an
+    ``and`` chain hands them a lazy SymBool under dynamo — the resulting nested
+    symbolic Eq crashes ``evaluate_expr``. Control flow forces each comparison to a
+    real bool (weight shapes are static parameters, so the guards are correct)."""
+    packed = NIBBLES_PER_BYTE if weight.dtype == torch.int8 else 1
+    if scale.shape[:-1] != weight.shape[:-1]:
+        return False
+    if scale.shape[-1] != (weight.shape[-1] * packed) // group:
+        return False
+    return True
+def is_mxfp8(weight: torch.Tensor, scale: torch.Tensor) -> bool:
+    """MXFP8 weight/scale pair: E4M3 weights with UE8M0 group-32 scales — last dim
+    ``scale.shape[-1] == weight.shape[-1] // MX_SCALE_GROUP_K``, matching leading dims.
+    Works for 2D ``(N, K)`` and 3D ``(E, N, K)`` weights. The group-32 layout is what
+    separates MXFP8 from 128-block FP8 (which may also carry UE8M0 scales)."""
+    return (
+        weight.dtype == torch.float8_e4m3fn
+        and scale.dtype in UE8M0_SCALE_DTYPES
+        and _shapes_match(weight, scale, MX_SCALE_GROUP_K)
+    )
+def is_mxfp4(weight: torch.Tensor, scale: torch.Tensor) -> bool:
+    """MXFP4 weight/scale pair: packed E2M1 weights (``int8``, two codes/byte) with
+    UE8M0 group-32 scales — ``scale.shape[-1] == weight.shape[-1] * NIBBLES_PER_BYTE //
+    MX_SCALE_GROUP_K`` (unpacked K = ``2 * K_half``), matching leading dims. 2D or 3D."""
+    return (
+        weight.dtype == torch.int8
+        and scale.dtype in UE8M0_SCALE_DTYPES
+        and _shapes_match(weight, scale, MX_SCALE_GROUP_K)
+    )
+def is_nvfp4(weight: torch.Tensor, scale: torch.Tensor) -> bool:
+    """NVFP4 weight/scale pair: packed E2M1 weights (``int8``, two codes/byte) with E4M3
+    group-16 block scales — the scale DTYPE is the format carrier (UE8M0 = MX, E4M3 = NV), and
+    the group falls out of the shape. This predicate reads the block scale; the per-tensor
+    second-level global is a separate ``b_global_scale`` argument (``nvfp4_quantize_two_level``).
+    Layout-agnostic questions ("is this NVFP4, swizzled or not?") go through ``weight_format``,
+    which carries the same dtype rule without the shape check."""
+    return (
+        weight.dtype == torch.int8
+        and scale.dtype == torch.float8_e4m3fn
+        and _shapes_match(weight, scale, NVFP4_SCALE_GROUP_K)
+    )
+def is_preswizzled_mx(weight: torch.Tensor, scale: torch.Tensor) -> bool:
+    """A weight scale already in the SWIZZLE_32_4_4 layout, swizzled once at model load (the
+    deployment contract — the same checkpoint feeds grouped prefill and batched decode with no
+    per-call rearrange). The 5D shape ``(G, row_blocks, cols//4, 2, 256)`` is the marker; the scale is
+    a 1-byte block scale (UE8M0 for MXFP8/MXFP4, E4M3 for NVFP4) against an MX weight (E4M3 or
+    packed E2M1). Format-agnostic — NVFP4 pre-swizzles the same way (the layout cuBLAS wants).
+    6D is the gate|up artifact: the same blocks with the gate-interleave carried in the shape."""
+    return (
+        weight.dtype in (torch.float8_e4m3fn, torch.int8)
+        and scale.ndim in (5, 6)
+        and scale.element_size() == 1
+    )
+def is_mx(weight: torch.Tensor, scale: torch.Tensor) -> bool:
+    """Any microscaled weight/scale pair — MXFP8 (``float8_e4m3fn`` values), MXFP4
+    (``int8``, two E2M1 codes/byte), both UE8M0 group-32, or NVFP4 (packed E2M1 + E4M3
+    group-16); also an already-swizzled MX scale (``is_preswizzled_mx``). The dispatchers route
+    on this into the ``mx_*`` kernels; the op picks the format from the dtypes."""
+    return (
+        is_mxfp8(weight, scale)
+        or is_mxfp4(weight, scale)
+        or is_nvfp4(weight, scale)
+        or is_preswizzled_mx(weight, scale)
+    )
+def is_tensor_wide(block_size, weight: torch.Tensor) -> bool:
+    """True when ``block_size`` selects per-tensor (tensor-dynamic) scaling: ``None`` or
+    equal to the weight's full ``(N, K)`` — one scale block spanning the whole matrix.
+    Handles 2D ``(N, K)`` and 3D ``(E, N, K)`` weights via the last two dims. (2D path
+    only — the grouped/batched dispatchers derive the format from the SCALE shape via
+    ``weight_block_size``.)"""
+    return block_size is None or (
+        block_size[0] == weight.shape[-2] and block_size[1] == weight.shape[-1]
+    )
+def weight_block_size(B: torch.Tensor, Bs: torch.Tensor) -> list[int] | None:
+    """The fp8 weight-quantization block ``[block_n, block_k]``, derived from the scale
+    tensor's shape — the data already says how it was quantized, so no ``block_size``
+    parameter exists to disagree with it. ``None`` = tensor-wide (one scalar per expert:
+    ``Bs`` ``(E,)`` or ``(E, 1, 1)`` spanning the full ``(N, K)``). Expects a 3D
+    ``(E, N, K)`` fp8 ``B``; MX weights never reach here (format keyed by scale dtype)."""
+    num_experts, n_rows, K = B.shape
+    if Bs.numel() == num_experts:
+        return None
+    assert Bs.ndim == 3 and K % Bs.shape[2] == 0, (
+        f"Bs shape {tuple(Bs.shape)} does not tile B {tuple(B.shape)} along K"
+    )
+    # K tiles evenly, so block_k is exact; N (n_rows) may be non-aligned (a partial last block,
+    # n_rows % Bs.shape[1] != 0) — recover block_n from the even K dim (square FP8 blocks). The
+    # routed kernels then reject the non-aligned N via ``require_moe_dims_aligned``; matmul_2d masks it.
+    block_k = K // Bs.shape[2]
+    block_n = n_rows // Bs.shape[1] if n_rows % Bs.shape[1] == 0 else block_k
+    return [block_n, block_k]
+def validate_dense_operands(A: torch.Tensor, B: torch.Tensor) -> None:
+    """Shared (rows, K) x (num_experts, N, K) operand checks for the unpacked formats —
+    the packed-E2M1 ops do their own (K spans two values per stored byte)."""
+    assert A.ndim == 2, f"A must be 2D (rows, K), got ndim={A.ndim}"
+    assert A.is_contiguous(), "A must be contiguous"
+    assert B.ndim == 3, f"B must be 3D (num_experts, N, K), got ndim={B.ndim}"
+    assert B.is_contiguous(), "B must be contiguous"
+    assert A.shape[1] == B.shape[2], (
+        f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
+    )
+def validate_dense_2d_operands(A: torch.Tensor, B: torch.Tensor) -> None:
+    """Shared (rows, K) x (N, K) operand checks for the 2D dense unpacked-fp8 wrappers — matching
+    K, contiguous A, 2D B. The packed-E2M1 (MX) 2D op does its own (K is two values per
+    stored byte, so B is (N, K // 2)).
+    B may be row-major or a transposed view of a row-major tensor. The kernels read both of B's
+    strides, and ``descriptor_box_pruner`` keeps the TMA arms — which require a unit-stride
+    innermost dim — off the transposed form, so a caller can contract over the other axis by
+    passing ``B.t()`` (with its scale grid transposed to match) instead of materializing a copy."""
+    assert A.shape[-1] == B.shape[-1], (
+        f"K mismatch: A has K={A.shape[-1]}, B has K={B.shape[-1]}"
+    )
+    assert A.is_contiguous(), "A must be contiguous"
+    assert B.ndim == 2, f"B must be 2D (N, K), got ndim={B.ndim}"
+    assert B.is_contiguous() or B.t().is_contiguous(), (
+        "B must be contiguous or a transposed view of a contiguous tensor, got strides "
+        f"{tuple(B.stride())} for shape {tuple(B.shape)}"
+    )
+def expert_weight_shape(B: torch.Tensor, gate: bool) -> tuple[int, int, int]:
+    """(num_experts, n_rows, N) of an expert weight stack — under a gate epilogue B is
+    the (E, 2N, K) gate|up stack, so the per-projection width N is half the stored rows."""
+    num_experts, n_rows, _ = B.shape
+    return num_experts, n_rows, (n_rows // 2 if gate else n_rows)
+def routed_rows(A, gather_idx, scatter_idx, expert_start, num_experts) -> int:
+    """S (routed rows) for a grouped launch, carried by the (S,) maps — A's rows are
+    gather SOURCES and under-count S whenever top_k > 1 (gate_up reading raw hidden);
+    only with no maps at all is A itself the expert-sorted (S, K) matrix. Validates the
+    maps and the ``(next_power_of_2(E) + 1,)`` ``expert_start`` schedule."""
+    if gather_idx is not None:
+        S = gather_idx.numel()
+    elif scatter_idx is not None:
+        S = scatter_idx.numel()
+    else:
+        S = A.shape[0]
+    for perm_map in (gather_idx, scatter_idx):
+        assert perm_map is None or (perm_map.numel() == S and perm_map.is_contiguous())
+    assert (
+        expert_start.is_contiguous()
+        and expert_start.numel() == triton.next_power_of_2(num_experts) + 1
+    ), "expert_start must be contiguous (next_power_of_2(num_experts) + 1,)"
+    return S
+def tokens_per_expert_bucket(S: int, num_experts: int) -> int:
+    """log2 bucket of the average routed rows per expert — the grouped kernels' autotune
+    key (raw S would retune per unique token count)."""
+    return int((S + num_experts - 1) // num_experts).bit_length()
+def normalize_global_scale(
+    g: torch.Tensor | None, num_experts: int
+) -> torch.Tensor | None:
+    """An NVFP4 second-level global in one of the kernels' fp32 load layouts: ``(1,)`` per tensor
+    or ``(num_experts,)`` per expert (``(E, 1, 1)`` accepted for the latter). The kernels index it
+    through ``global_scale_stride``, so a shared value broadcasts by stride and nothing is
+    materialized per call. ``None`` passes through.
+    A gate|up stack calibrated per half (modelopt's ``(E, 2)`` ``weight_scale_2``) is merged to one
+    global per expert at load — SwiGLU is linear in the up half, so its global folds into the down
+    projection's — never here: the ops take one global per matrix."""
+    if g is None:
+        return None
+    g = g.float()
+    if g.numel() == 1:
+        return g.reshape(1)
+    assert g.numel() == num_experts, (
+        f"global scale of shape {tuple(g.shape)}: expected a scalar or ({num_experts},)"
+    )
+    return g.reshape(num_experts).contiguous()
+def is_per_expert_global(g: torch.Tensor | None) -> bool:
+    """Whether an NVFP4 second-level global carries one value per expert rather than one for the
+    whole tensor. The activation quant branches on it: a per-expert global has to be applied by
+    each quantized row's own expert, which is what decides the grid that quant runs on."""
+    return g is not None and g.numel() > 1
+def global_scale_stride(g: torch.Tensor | None) -> int:
+    """The expert stride of a ``normalize_global_scale`` layout: 0 for a scalar (or ``None``),
+    which broadcasts it to every expert, else 1."""
+    return 0 if g is None or g.numel() == 1 else 1
+def normalize_per_expert_scale(Bs: torch.Tensor, num_experts: int) -> torch.Tensor:
+    """One per-tensor scale per expert, normalized to ``(num_experts, 1, 1)`` from
+    either that or a bare ``(num_experts,)``."""
+    if Bs.ndim == 1:
+        assert Bs.shape[0] == num_experts, (
+            f"Bs shape {tuple(Bs.shape)} != expected ({num_experts},)"
+        )
+        return Bs.reshape(num_experts, 1, 1)
+    assert Bs.shape == (num_experts, 1, 1), (
+        f"Bs shape {tuple(Bs.shape)} != expected ({num_experts}, 1, 1)"
+    )
+    return Bs
+def mx_scale_family(Bs: torch.Tensor, K: int) -> int:
+    """The group size of an MX/NV weight-scale tensor, in either layout — the wrapper hands ``Bs``
+    as-is and this reads the group off its shape. Row-major (2D/3D): ``K // Bs.shape[-1]``.
+    SWIZZLE_32_4_4 (5D ``(G, row_blocks, cols // 4, 2, 256)``): ``K // (Bs.shape[2] * 4)``. The scale
+    dtype IS the format carrier (E4M3 = NVFP4 group-16, UE8M0 = MX group-32) and the pairing is
+    validated; callers that need the format read it off ``Bs.dtype`` (``== torch.float8_e4m3fn``
+    is NVFP4)."""
+    nvfp4 = Bs.dtype == torch.float8_e4m3fn
+    assert nvfp4 or Bs.dtype in UE8M0_SCALE_DTYPES, (
+        f"Bs must be UE8M0 (float8_e8m0fnu/uint8) or E4M3 (NVFP4), got {Bs.dtype}"
+    )
+    scale_group = K // (Bs.shape[2] * 4) if Bs.ndim == 5 else K // Bs.shape[-1]
+    assert scale_group == (NVFP4_SCALE_GROUP_K if nvfp4 else MX_SCALE_GROUP_K), (
+        f"scale group {scale_group} does not match the scale dtype {Bs.dtype}"
+    )
+    assert K % scale_group == 0, f"K (={K}) must be a multiple of {scale_group}"
+    return scale_group
+def get_supported_act_fns() -> tuple[str, ...]:
+    """The activation names the gate|up epilogue fuses in-kernel; any other activation is applied
+    on the host by the MoE forwards (``act_fn`` as a callable)."""
+    return ("silu", "gelu", "relu")
+ACTIVATION_FORMATS = (None, "bf16", "fp8", "mxfp8", "mxfp4", "nvfp4")
+def check_activation_format(activation_format: str | None, quantize_output: bool) -> None:
+    """Validate the dispatchers' activation kwargs. ``activation_format`` is the activations'
+    quantization, in the same vocabulary transformers' ``FineGrainedConfig`` uses; the weights are
+    never named — their format is read off the weight/scale dtypes (``weight_format``).
+    - ``None`` (default): the weights' own format — W8A8 on FP8/MXFP8 weights, the all-fp4 W4A4
+      chain on MXFP4/NVFP4 weights.
+    - ``"bf16"``: weight-only (W4A16 / W8A16) — the activation stays high precision.
+    - ``"fp8"`` / ``"mxfp8"`` / ``"mxfp4"`` / ``"nvfp4"``: an explicit activation format, e.g.
+      ``"mxfp8"`` on MXFP4 weights is the W4A8 chain. A raw ``A`` is quantized to it by the op; an
+      already-quantized ``A`` arrives with ``As`` (its dtype carries the format).
+    ``quantize_output`` stores the (GLU) output requantized into that same format — bit-identical
+    to quantizing the values offline — so it feeds the next op's activation side directly; the op
+    then returns ``(C, Cs)``. A weight-only activation never requantizes.
+    Per weight family:
+    ==================  ===============================================================
+    weight family       activation_format
+    ==================  ===============================================================
+    block-dynamic FP8   None / "fp8" (E4M3 + per-block scales); quantize_output ok
+    tensor-wide FP8     None / "fp8" (E4M3 + per-token scales); no quantize_output
+    MXFP8 / MXFP4       None / "mxfp8" (E4M3 + UE8M0) / "mxfp4" (packed E2M1 + UE8M0) / "bf16"
+    NVFP4               None / "nvfp4" (packed E2M1 + E4M3 group-16) / "bf16"
+    full-precision      None (nothing to quantize)
+    ==================  ===============================================================
+    (Tensor-wide's whole-row activation scale can't be formed by a tile-local epilogue; the unfused
+    path quantizes on the host between GEMMs.) A format the weight family doesn't support fails
+    loudly at the op."""
+    assert activation_format in ACTIVATION_FORMATS, (
+        f"unknown activation_format {activation_format!r}; expected one of {ACTIVATION_FORMATS}"
+    )
+    assert not (quantize_output and activation_format == "bf16"), (
+        "quantize_output needs quantized activations (activation_format='bf16' is weight-only)"
+    )
+def weight_format(B: torch.Tensor, Bs: torch.Tensor) -> str:
+    """The format the weights are ALREADY in, read off their dtypes — what ``activation_format=None``
+    resolves to. E4M3 scales -> ``"nvfp4"``; else packed E2M1 -> ``"mxfp4"``, MX E4M3 -> ``"mxfp8"``,
+    block/tensor-scaled E4M3 -> ``"fp8"``."""
+    if Bs.dtype == torch.float8_e4m3fn:
+        return "nvfp4"
+    if not is_mx(B, Bs):
+        return "fp8"
+    # packed E2M1 is a WEIGHT dtype fact; is_mxfp4's shape check can't see it through a
+    # pre-swizzled (5D) scale, which would silently downgrade a W4A4 default to W4A8.
+    return "mxfp4" if B.dtype in (torch.int8, torch.uint8) else "mxfp8"
+def resolve_activation_format(activation_format: str | None, B: torch.Tensor, Bs: torch.Tensor) -> str:
+    """The format an MX/NVFP4 GEMM quantizes its activations to. ``None`` resolves to the weights'
+    own format (``weight_format``): NVFP4 weights pin ``"nvfp4"`` (the MMA kind needs matching
+    scale formats on both operands), packed E2M1 -> ``"mxfp4"`` (the all-fp4 W4A4 chain), else
+    ``"mxfp8"``. An explicit name overrides (e.g. ``"mxfp8"`` on MXFP4 weights = the W4A8 chain)
+    and is validated against the family. ``"bf16"`` never reaches here — a weight-only chain is
+    routed to its own kernel by the dispatcher."""
+    if Bs.dtype == torch.float8_e4m3fn:
+        assert activation_format in (None, "nvfp4"), (
+            f"NVFP4 activations are packed E2M1 + E4M3 scales, got {activation_format!r}"
+        )
+        return "nvfp4"
+    assert activation_format in (None, "mxfp8", "mxfp4"), (
+        f"MX activations are E4M3 ('mxfp8') or packed E2M1 ('mxfp4'), got {activation_format!r}"
+    )
+    return activation_format if activation_format is not None else weight_format(B, Bs)
+def resolve_output_dtype(
+    output_dtype: torch.dtype | None,
+    activation: torch.Tensor,
+    act_scale: torch.Tensor | None,
+) -> torch.dtype:
+    """Output element type for a quantized matmul: the explicit ``output_dtype`` if given, else
+    the raw activation dtype (``act_scale`` is None -> ``activation`` is high precision), else
+    ``bfloat16`` (``activation`` is pre-quantized FP8, whose dtype is not a valid output)."""
+    if output_dtype is not None:
+        return output_dtype
+    return activation.dtype if act_scale is None else torch.bfloat16

build/torch-rocm/grouped.py ADDED Viewed

The diff for this file is too large to render. See raw diff

build/torch-rocm/loading/__init__.py ADDED Viewed

	@@ -0,0 +1,17 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""How a K-step's operands get from memory into registers: ``tiles`` loads the activation and
+weight value tiles (pointer / host-TMA / gather arms), ``scales`` loads the block scales (affine or
+SWIZZLE_32_4_4) and quantizes a raw activation tile inline."""

build/torch-rocm/loading/scales.py ADDED Viewed

	@@ -0,0 +1,325 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import triton
+import triton.language as tl
+from ..quant import fp8_act_quant_inline, mx_act_quant_inline
+@triton.jit
+def load_block_fp8_act_tile(
+    a_ptrs,
+    as_ptrs,
+    a_descriptor=0,
+    m_off=0,
+    k_off=0,
+    A_MEMORY_MODE: tl.constexpr = "pointer",
+    as_mask=None,
+    TRANSPOSED: tl.constexpr = False,
+):
+    """Block-FP8 counterpart of ``load_mx_act_tile``: load one activation K-tile as
+    ``(a_fp8, a_scale_f32)`` — the arm folds off ``A_MEMORY_MODE`` and the pointer dtype at
+    compile time. Descriptor mode loads the pre-quantized ``(BM, BK)`` host-TMA box at
+    ``(m_off, k_off)`` (``a_ptrs`` unread); pointer fp8 loads the pre-quantized-offline tile;
+    raw bf16/fp16 pointers quantize inline (``as_ptrs`` then a constexpr-dead placeholder).
+    Either way the per-K-block scales come from ``as_ptrs``. ``TRANSPOSED`` marks a ``(K, M)``
+    tile (the swapped pointer arm) so the inline amax reduces the token axis either way.
+    Unmasked unless ``as_mask`` given: every caller's rows are %-wrapped, expert-advanced, or
+    token-replicated."""
+    if A_MEMORY_MODE != "pointer":  # pre-quantized, host-TMA box [BM, BK]
+        a = a_descriptor.load([m_off, k_off])
+        # as_mask None (batched/grouped: %-wrapped/expert-advanced rows) -> unmasked load;
+        # a bounds mask (the 2D affine-scale path) -> masked. `other` is illegal without mask.
+        if as_mask is None:
+            a_s = tl.load(as_ptrs)
+        else:
+            a_s = tl.load(as_ptrs, mask=as_mask, other=0.0)
+    elif a_ptrs.dtype.element_ty == tl.float8e4nv:  # pre-quantized offline
+        a = tl.load(a_ptrs)
+        if as_mask is None:
+            a_s = tl.load(as_ptrs)
+        else:
+            a_s = tl.load(as_ptrs, mask=as_mask, other=0.0)
+    else:  # raw bf16/fp16 — quantize inline
+        a, a_s = fp8_act_quant_inline(tl.load(a_ptrs).to(tl.float32), TRANSPOSED)
+    return a, a_s
+@triton.jit
+def load_swizzled_scale(
+    desc, blk_idx, k_idx,
+    REP: tl.constexpr, REP_K: tl.constexpr, BLOCK: tl.constexpr, SCALE_COLS: tl.constexpr,
+):
+    """Bulk-load one scale tile from a ``SWIZZLE_32_4_4`` descriptor and un-swizzle it to the
+    ``(BLOCK, SCALE_COLS)`` layout ``tl.dot_scaled`` consumes. ``REP = BLOCK // 128``,
+    ``REP_K = SCALE_COLS // 4``; ``blk_idx``/``k_idx`` are the tile's row-block / K-block ids.
+    Descriptor is over the swizzled scale viewed ``(1, rows//128, cols//4, 2, 256)``."""
+    s = desc.load([0, blk_idx * REP, k_idx * REP_K, 0, 0])
+    return s.reshape(REP, REP_K, 32, 4, 4).trans(0, 3, 2, 1, 4).reshape(BLOCK, SCALE_COLS)
+@triton.jit
+def load_swizzled_scale_tile(
+    descriptor,
+    ptr,
+    group_id,
+    pid,
+    k_idx,
+    rows,
+    K,
+    BLOCK: tl.constexpr,
+    SCALE_COLS: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+):
+    """One swizzled-scale tile ``(BLOCK, SCALE_COLS)`` for a row-tile of ANY operand — batched
+    weight (``group_id = expert``), 2D weight / 2D activation (``group_id = 0``, dense). Scales are
+    SWIZZLE_32_4_4 128-row blocks over ``rows`` (the operand's row count, ``rows//128`` blocks per
+    group).
+    - ``BLOCK`` a multiple of 128 with whole 4-group K bands (``BK % 128 == 0``): the fast path —
+      bulk-load the ``REP = BLOCK//128`` row-blocks via the TMA ``descriptor`` (box
+      ``[1, REP, rep_k, 2, 256]``), un-swizzle, feed ``dot_scaled``. Today's pruners cap every
+      swizzled tile at 128 (REP=1); the REP>1 form is written for a future >128-tile descriptor
+      read, not reachable from the current grids.
+    - Otherwise (sub-128 tile — fp8 ``scalar`` decode / small-M offline; or ``BK<128``): the block
+      layout can't be TMA-sliced, so pointer-GATHER exactly this tile's rows. The swizzle is a fixed
+      permutation: logical (row ``r``, K-group ``col``) → byte
+      ``(blk*cols4 + col//4)*512 + (r%32)*16 + ((r%128)//32)*4 + col%4``. Reads only the needed
+      bytes — no 128-block over-read, no un-swizzle transpose — the row-major fast path's cost with
+      the swizzled layout, so ``scalar`` competes on merit instead of eating a TMA penalty."""
+    # Per-expert 128-row-block count is a CEIL: a non-128-multiple ``rows`` (e.g. N=2880) still
+    # occupies ceil(rows/128) blocks in the buffer (the swizzle builder pads the partial last block;
+    # its tail rows read zero-weight via the TMA OOB clamp, so they don't contribute). ``cdiv == floor``
+    # when ``rows`` is 128-aligned, so this is inert for every aligned shape.
+    nrb = (rows + 127) // 128
+    if BLOCK % 128 == 0 and SCALE_COLS >= 4 and SCALE_COLS % 4 == 0:
+        REP: tl.constexpr = BLOCK // 128
+        # absolute 128-block base = group_id*nrb + pid*REP; load_swizzled_scale multiplies blk by REP.
+        # Non-128 ``rows`` (odd ``nrb``) pins REP=1 (BN=128) in the pruner, so group_id*nrb//REP is exact.
+        blk = (group_id * nrb // REP + pid).to(tl.int32)
+        return load_swizzled_scale(descriptor, blk, k_idx, REP, SCALE_COLS // 4, BLOCK, SCALE_COLS)
+    cols4 = (K // SCALE_GROUP_K + 3) // 4  # cdiv: the buffer pads cols to whole 4-group chunks
+    r = pid * BLOCK + tl.arange(0, BLOCK)
+    lin = r // 128
+    blk = group_id * nrb + lin
+    row = r % 128
+    col = k_idx * SCALE_COLS + tl.arange(0, SCALE_COLS)
+    off = (
+        (blk[:, None] * cols4 + col[None, :] // 4) * 512
+        + (row[:, None] % 32) * 16
+        + (row[:, None] // 32) * 4
+        + col[None, :] % 4
+    )
+    return tl.load(ptr + off)
+@triton.jit
+def load_weight_scale_tile(
+    SWIZZLED_SCALES: tl.constexpr,
+    bs_descriptor,
+    bs_ptr,
+    expert_id,
+    pid_n,
+    k_idx,
+    N,
+    K,
+    stride_bs_e,
+    stride_bs_n,
+    stride_bs_k,
+    BLOCK_SIZE_N: tl.constexpr,
+    SCALE_COLS: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    GATE: tl.constexpr,
+    k_col_off=0,
+):
+    """One batched-decode weight-scale tile ``(n_width, SCALE_COLS)``, hiding the swizzled vs
+    un-swizzled choice behind the ``SWIZZLED_SCALES`` flag — the kernel loop reads scales the same
+    way either layout. ``bs_ptr`` is the un-advanced buffer base; the per-expert offset is applied
+    here (the swizzled path indexes by 128-row block, the un-swizzled by the row-major stride, so
+    it can't be pre-advanced uniformly). Under ``GATE`` the gate|up sub-tiles stack into ``2*BN``.
+    - ``SWIZZLED_SCALES``: SWIZZLE_32_4_4 via ``load_swizzled_scale_tile`` (descriptor bulk at BN=128, or
+      pointer gather below), scales swizzled over the full ``2N`` rows/expert under GATE.
+    - else: affine per-group load off ``(expert, N-tile row, K-group)`` — no in-op swizzle, so an
+      un-swizzled caller pays nothing. ``k_col_off`` shifts that read off the tile grid (the
+      caller's K-tile was clamped back into bounds for a BK that doesn't divide K); it is 0
+      for every aligned caller, and the swizzled layouts don't take it."""
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    if SWIZZLED_SCALES:
+        # gate|up needs no arm: the scale grid is row-interleaved with the weight, so the tile's
+        # gate and up rows are the contiguous n_width span at the same pid_n, i.e. the plain read
+        # over the doubled extent.
+        b_s = load_swizzled_scale_tile(
+            bs_descriptor, bs_ptr, expert_id, pid_n, k_idx,
+            2 * N if GATE else N, K,
+            n_width, SCALE_COLS, SCALE_GROUP_K,
+        )
+    else:
+        # affine per-group load off (expert, N-tile row, K-group) from the un-advanced base.
+        # Clamp the row in-bounds instead of masking (a masked load's other=0.0 blocks WS
+        # lowering): a partial last N-tile reads a clamped scale the epilogue masks off
+        # (N_COLS) — without the clamp the GATE arm's up rows (N + offs_bn) run past 2N.
+        base = bs_ptr + expert_id * stride_bs_e
+        offs_bn = pid_n * n_width + tl.arange(0, n_width)
+        offs_bn = tl.minimum(offs_bn, (2 * N if GATE else N) - 1)
+        offs_sf = k_idx * SCALE_COLS + k_col_off + tl.arange(0, SCALE_COLS)
+        ptrs = base + offs_bn[:, None] * stride_bs_n + offs_sf[None, :] * stride_bs_k
+        b_s = tl.load(ptrs)
+    return b_s
+@triton.jit
+def load_mx_act_tile(
+    a_ptrs,
+    as_ptrs,
+    as_global,  # fp32 NVFP4 act global (None off nvfp4); normalizes the raw tile pre-block-quant
+    as_global_row,  # element offset into ``as_global`` (the expert's entry; 0 for a shared scalar)
+    row_mask,
+    a_descriptor,
+    m_start,
+    ka_off,
+    gather_rows,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    A_MEMORY_MODE: tl.constexpr = "pointer",
+    A_GATHER: tl.constexpr = False,
+    FORMAT: tl.constexpr = "mxfp8",
+):
+    """Load one MX activation K-tile as ``(a_vals, a_scale)`` — the arm is picked
+    off the pointer dtype at compile time: fp8 pointers load pre-quantized E4M3 values +
+    UE8M0 scales (``maybe_act_quant``'s offline arm), uint8 pointers load caller-provided
+    packed-E2M1 values (W4A4 — the ``a_ptrs`` tile spans ``BLOCK_SIZE_K // 2`` bytes) +
+    the same UE8M0 scales, raw bf16/fp16 pointers load and quantize inline onto
+    ``FORMAT``'s grid (``mx_act_quant_inline`` — packed E2M1 under the fp4 formats;
+    ``as_ptrs`` then points at a dead placeholder and is never read). Under NVFP4
+    two-level, ``as_global[as_global_row]`` (the calibrated activation global — per tensor
+    or this expert's) normalizes the raw tile before the block quant — bit-identical to the
+    offline ``nvfp4_act_quant(x, global_scale=g_a)`` pass. ``row_mask`` may be ``None``
+    (unmasked tiles, e.g. the %-wrapped 2D matmul). Callers advance both pointers
+    unconditionally."""
+    if a_ptrs.dtype.element_ty == tl.float8e4nv or a_ptrs.dtype.element_ty == tl.uint8:
+        # pre-quantized (E4M3 offline, or packed E2M1 handed in by the caller); under the
+        # descriptor arm the [BM, BK_bytes] box loads the tile's contiguous sorted rows
+        # (no gather — tail rows past the tensor clamp to zero and are store-masked)
+        if A_MEMORY_MODE != "pointer":
+            if A_GATHER:
+                # sm_100 tma gather4: bulk-load the tile's ARBITRARY source rows
+                a = a_descriptor.gather(gather_rows, ka_off)
+            else:
+                a = a_descriptor.load([m_start, ka_off])
+            if as_ptrs.dtype.element_ty == tl.float8e4nv:
+                # NVFP4 scales: 0.0 encodes as byte 0 — padded rows scale to exact 0
+                a_scale = tl.load(as_ptrs, mask=row_mask[:, None], other=0.0)
+            else:
+                # UE8M0 scales: byte 0 decodes to 2^-127 — padded rows can't make 0*inf
+                a_scale = tl.load(as_ptrs, mask=row_mask[:, None], other=0)
+        elif row_mask is None:
+            a = tl.load(a_ptrs)
+            a_scale = tl.load(as_ptrs)
+        else:
+            a = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0)
+            if as_ptrs.dtype.element_ty == tl.float8e4nv:
+                # NVFP4 scales: 0.0 encodes as byte 0 — padded rows scale to exact 0
+                a_scale = tl.load(as_ptrs, mask=row_mask[:, None], other=0.0)
+            else:
+                # UE8M0 scales: byte 0 decodes to 2^-127 — padded rows can't make 0*inf
+                a_scale = tl.load(as_ptrs, mask=row_mask[:, None], other=0)
+    else:  # raw bf16/fp16 — quantize inline
+        # the descriptor arm mirrors the pre-quantized branch above. Without it every
+        # A_MEMORY_MODE != "pointer" config fell through to `tl.load(a_ptrs)`, and under a
+        # descriptor arm `operand_tile_ptrs` hands back the bare base pointer -- so that load
+        # produced a SCALAR and mx_act_quant_inline's reshape died with "'dtype' object has no
+        # attribute 'numel'". Raw-A is exactly the inline-quant (nvfp4/mxfp4 2D) path, so it
+        # killed every descriptor config there: 321 dead compiles in one tune.
+        if A_MEMORY_MODE != "pointer":
+            if A_GATHER:
+                a_raw = a_descriptor.gather(gather_rows, ka_off).to(tl.float32)
+            else:
+                a_raw = a_descriptor.load([m_start, ka_off]).to(tl.float32)
+        elif row_mask is None:
+            a_raw = tl.load(a_ptrs).to(tl.float32)
+        else:
+            a_raw = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0).to(tl.float32)
+        if as_global is not None:  # NVFP4 two-level: normalize by the calibrated act global
+            a_raw = a_raw / tl.load(as_global + as_global_row).to(tl.float32)
+        a, a_scale = mx_act_quant_inline(
+            a_raw, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K, FORMAT
+        )
+    return a, a_scale
+@triton.jit
+def apply_global_scale(acc, GlobalScale, stride_gs_e, GlobalScaleA, stride_gsa_e, expert_id):
+    """Fold the NVFP4 second-level globals onto the fp32 accumulator: the weight global ``g_b``
+    and the activation global ``g_a``, each per tensor (stride 0, broadcast) or per expert.
+    ``None`` folds the arm out."""
+    if GlobalScale is not None:
+        acc = acc * tl.load(GlobalScale + expert_id * stride_gs_e).to(tl.float32)
+    if GlobalScaleA is not None:
+        acc = acc * tl.load(GlobalScaleA + expert_id * stride_gsa_e).to(tl.float32)
+    return acc
+@triton.jit
+def decode_group_scale(scale):
+    """Decode a group scale to fp32 by its dtype: ``uint8`` = UE8M0 exponent bits
+    (``value = 2^(exp - 127)``, built directly as the fp32 bit pattern), E4M3 = NVFP4's
+    direct fp8 value, fp32 (block-dynamic with float scales) passes through. The dtype
+    branch is a compile-time constant, so only the taken path is emitted (single return —
+    Triton requires all ``return`` statements to share a type)."""
+    if scale.dtype == tl.uint8:
+        scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True)
+    elif scale.dtype == tl.float8e4nv:
+        scale = scale.to(tl.float32)
+    return scale
+@triton.jit
+def mx_2d_scale_ptrs(
+    As,
+    Bs,
+    pid_m,
+    pid_n,
+    M,
+    N,
+    stride_as_m,
+    stride_bs_n,
+    stride_bs_k,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    SCALE_COLS: tl.constexpr,
+    SWIZZLED_SCALES: tl.constexpr,
+    GATE: tl.constexpr = False,
+):
+    """Prologue 2D MX scale-pointer tiles + bounds masks as ``(as_ptrs, bs_ptrs, as_mask,
+    bs_mask)``. Affine arm: per-(row, group) ``as``/``bs`` pointer tiles read off AFFINE
+    row/col offsets (the %-wrapped operand offsets would make the scale load a non-affine
+    gather) with row/col bounds masks. Swizzled arm: the scales are read via the SA/BS
+    descriptors in the loop, so these tiles are dead — return the base scalars + null masks.
+    Single return — the arms have divergent types (base scalars + null masks vs affine pointer
+    tiles + bounds masks), so an early return can't unify them; the constexpr selects one."""
+    if SWIZZLED_SCALES:
+        as_ptrs, bs_ptrs, as_mask, bs_mask = As, Bs, None, None
+    else:
+        offs_am_lin = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+        ROWS: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+        offs_bn_lin = pid_n * ROWS + tl.arange(0, ROWS)
+        offs_sf = tl.arange(0, SCALE_COLS)
+        as_ptrs = As + offs_am_lin[:, None] * stride_as_m + offs_sf[None, :]
+        bs_ptrs = Bs + (offs_bn_lin[:, None] * stride_bs_n + offs_sf[None, :] * stride_bs_k)
+        as_mask = offs_am_lin < M
+        bs_mask = offs_bn_lin < (2 * N if GATE else N)
+    return as_ptrs, bs_ptrs, as_mask, bs_mask

build/torch-rocm/loading/tiles.py ADDED Viewed

	@@ -0,0 +1,585 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import triton
+import triton.language as tl
+from .scales import load_block_fp8_act_tile, load_mx_act_tile, load_swizzled_scale, load_swizzled_scale_tile, load_weight_scale_tile
+# A-sub-tile byte budget for the 1-byte-activation grouped-swizzle depth cap (see
+# swizzle_offsets): the co-scheduled rows' A tile per K-step (depth*BM*BK bytes) must stay
+# L2-hot to reuse, and past this it thrashes and the win collapses. ~512KB is a wave
+# reuse-window, NOT gross L2 (0.4% of the B200's 132MB) — a B200 measurement (2026-07-16), so
+# RE-MEASURE per device. Packed fp4 bypasses this (full grouping wins; see swizzle_offsets).
+SWIZZLE_GROUP_A_BYTES = tl.constexpr(524288)
+@triton.jit
+def swizzle_offsets(
+    M,
+    N,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    WEIGHT_VALUES_PER_BYTE: tl.constexpr = 1,
+    GATE: tl.constexpr = False,
+):
+    """2D-grid tile scheduling shared by the kernels below: grouped-swizzle the
+    ``(pid_m, pid_n)`` program ids for L2 locality on B, then build the operand offset
+    vectors. Returns ``(pid_m, pid_n, offs_am, offs_bn, offs_k)`` — the swizzled ids
+    (reused by the output store) and the ``%``-wrapped row/col offsets plus the K range.
+    ``offs_bn`` is the WEIGHT-side row range: under ``GATE`` the doubled span
+    ``[2*pid_n*BN, 2*(pid_n+1)*BN)`` wrapped at ``2N``. The grid is sized by the OUTPUT tiles.
+    The swizzle keeps the B (weight) column-tile L2-hot while the co-scheduled rows reuse it,
+    so the depth cap is set by the WEIGHT footprint, capped at ``min(num_pid_m, .)``. With
+    1-byte weights the reuse thrashes past ~512KB (a hard cliff): the growing rival is the
+    rows' A sub-tile per K-step (``depth * BM * BK`` bytes), so ``SWIZZLE_GROUP_A_BYTES //
+    (BM * BK)`` — MEASURED on B200 (2026-07-16), bd BK128->32 (cliffs at 64), MX BK256->16,
+    BM64->64 (g*BM*BK ~512KB across BM and BK; BN-independent). Packed-fp4 weights
+    (``WEIGHT_VALUES_PER_BYTE==2``) halve that hot set, so it never thrashes and full grouping
+    wins outright (monotone, no cliff — measured for both W4A4 and W4A8, i.e. weight-driven not
+    activation-driven). Uses an EXPLICIT grouped swizzle (linearize the 2D program ids, then
+    group), NOT ``tl.swizzle2d`` (which degrades with group depth, -3pp+ MFU at depth 32). Same
+    grid launch, same result set — only the program-id -> tile mapping changes."""
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    pid = tl.program_id(axis=1) * num_pid_m + tl.program_id(axis=0)
+    if WEIGHT_VALUES_PER_BYTE == 2:
+        max_group = num_pid_m
+    else:
+        max_group = SWIZZLE_GROUP_A_BYTES // (BLOCK_SIZE_M * BLOCK_SIZE_K)
+    num_pid_in_group = max_group * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * max_group
+    group_size_m = min(num_pid_m - first_pid_m, max_group)
+    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)
+    pid_n = (pid % num_pid_in_group) // group_size_m
+    WEIGHT_ROWS: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_bn = (pid_n * WEIGHT_ROWS + tl.arange(0, WEIGHT_ROWS)) % (2 * N if GATE else N)
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    return pid_m, pid_n, offs_am, offs_bn, offs_k
+@triton.jit
+def operand_tile_descriptor(
+    HostDescriptor,
+    W,
+    N,
+    K,
+    stride_n,
+    stride_k,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    B_MEMORY_MODE: tl.constexpr,
+    GATE: tl.constexpr = False,
+):
+    """Resolve one operand's tile descriptor once per program (weight OR activation — the bd 2D
+    kernel calls it for both): the host-built TMA descriptor as passed, a device-built in-kernel
+    tensormap, or 0 under "pointer" (never read — the constexpr branch folds it out of
+    ``load_weight_tile``). Single return — only the taken branch compiles."""
+    if B_MEMORY_MODE == "host_descriptor":
+        descriptor = HostDescriptor
+    elif B_MEMORY_MODE == "device_descriptor":
+        # under GATE the operand is the (2N, K) gate|up weight and a tile is a contiguous 2*BN
+        # row span of it, so both the tensor extent and the box double
+        descriptor = tl.make_tensor_descriptor(
+            W,
+            shape=((2 * N) if GATE else N, K),
+            strides=(stride_n, stride_k),
+            block_shape=((2 * BLOCK_SIZE_N) if GATE else BLOCK_SIZE_N, BLOCK_SIZE_K),
+        )
+    else:  # pointer
+        descriptor = 0
+    return descriptor
+@triton.jit
+def load_grouped_act_tile(
+    a_ptrs,
+    a_descriptor,
+    m_start,
+    ka_off,
+    row_mask,
+    gather_rows,
+    A_MEMORY_MODE: tl.constexpr,
+    A_GATHER: tl.constexpr = False,
+):
+    """A plain (scale-less) grouped activation K-tile: the masked pointer load, or —
+    under the descriptor arm — sm_100 tma gather4 over the tile's ARBITRARY source rows
+    (gathered launches; padded rows read row 0 and are store-masked) or the ``[BM, BK]``
+    box at the tile's contiguous sorted-row start (no-gather launches; tail rows past
+    the tensor clamp to zero). ``row_mask`` None loads the pointer tile maskless (the
+    %-wrapped 2D fast path). Single return — only the taken arm compiles."""
+    if A_MEMORY_MODE == "pointer":
+        if row_mask is None:
+            a = tl.load(a_ptrs)
+        else:
+            a = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0)
+    elif A_GATHER:
+        a = a_descriptor.gather(gather_rows, ka_off)
+    else:
+        a = a_descriptor.load([m_start, ka_off])
+    return a
+@triton.jit
+def load_grouped_weight_tile(
+    w_ptrs,
+    w_descriptor,
+    row0,
+    n_off,
+    kb_off,
+    BLOCK_SIZE_N: tl.constexpr,
+    KB: tl.constexpr,
+    GATE: tl.constexpr,
+    B_MEMORY_MODE: tl.constexpr,
+    SWAP_AB: tl.constexpr = False,
+):
+    """One K-major (optionally gate|up-stacked) MX weight K-tile for the grouped / batched loop:
+    the explicit-pointer tile flattened to the ``[KB, (2|1)*BN]`` rhs (or, under ``SWAP_AB``, the
+    ``[(2|1)*BN, KB]`` rows-major lhs — the batched-decode orientation), or the ``[(2|1), BN, KB]``
+    descriptor box over the ``(E, 2N|N, K_bytes)`` weight view, reshaped and transposed to the same
+    form (the fused-era TMA arm: natural orientation + per-iteration trans; grouped/2D never swap).
+    Single return — only the taken arm compiles; the caller advances ``w_ptrs`` and passes the box
+    offsets either way."""
+    if B_MEMORY_MODE == "pointer":
+        w = tl.load(w_ptrs)
+    else:
+        w = tl.trans(
+            tl.reshape(
+                w_descriptor.load([row0, n_off, kb_off]),
+                ((2 if GATE else 1) * BLOCK_SIZE_N, KB),
+            )
+        )
+    return w
+@triton.jit
+def load_weight_tile(
+    w_ptrs, w_descriptor, row_off, k_off, B_MEMORY_MODE: tl.constexpr, SWAP_AB: tl.constexpr = False
+):
+    """One weight K-tile. Descriptor modes load the ``(BN, BK)`` box at ``(row_off, k_off)``
+    and, in the natural orientation, transpose it once to the ``(BK, BN)`` K-major rhs the
+    pointer arm builds directly (``SWAP_AB`` keeps the box as-is: the weight rows then sit in
+    the MMA M dim). Pointer mode loads the explicit tile in whatever orientation ``w_ptrs``
+    was built with. Single return — only the taken branch compiles."""
+    if B_MEMORY_MODE == "pointer":
+        w = tl.load(w_ptrs)
+    else:
+        w = w_descriptor.load([row_off, k_off])
+        if not SWAP_AB:
+            w = tl.trans(w)
+    return w
+@triton.jit
+def matmul_weight_ptrs(
+    B,
+    offs_n,
+    offs_k,
+    stride_b_n,
+    stride_b_k,
+    B_MEMORY_MODE: tl.constexpr,
+    SWAP_AB: tl.constexpr = False,
+):
+    """Prologue weight-tile pointers via ``operand_tile_ptrs`` (which folds the
+    descriptor-vs-pointer arm); ``offs_n`` carries the gate|up span."""
+    return operand_tile_ptrs(B, offs_n, offs_k, stride_b_n, stride_b_k, B_MEMORY_MODE, SWAP_AB)
+@triton.jit
+def load_act_mx(
+    a_ptrs, as_ptrs, as_global, as_global_row, value_mask, scale_mask, a_descriptor, m_off, k_off,
+    as_descriptor, as_ptr, gather_rows, stride_as_m, pid_m, k, M, K,
+    A_MEMORY_MODE: tl.constexpr, A_GATHER: tl.constexpr, GROUPED: tl.constexpr,
+    SWIZZLED_SCALES: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr, ACTIVATION_FORMAT: tl.constexpr,
+):
+    """MX activation tile + scale. Offline-quantized A under a swizzled weight -> pre-swizzled
+    SWIZZLE_32_4_4 scale (tcgen05 fast path); grouped -> gathered per-(row, K-group) affine scale off
+    the source-order ``As``; else (2D / decode raw A) -> ``load_mx_act_tile`` inline-quant / in-register.
+    Raw bf16 A stays in-register affine even under a swizzled weight (the dot_scaled mix is fine)."""
+    SCALE_COLS: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K
+    A_OFFLINE: tl.constexpr = (
+        a_ptrs.dtype.element_ty == tl.float8e4nv or a_ptrs.dtype.element_ty == tl.uint8
+    )
+    if SWIZZLED_SCALES and A_OFFLINE:  # pre-swizzled SWIZZLE_32_4_4 scale (tcgen05 fast path)
+        a = load_grouped_act_tile(
+            a_ptrs, a_descriptor, m_off, k_off, value_mask, gather_rows, A_MEMORY_MODE, A_GATHER
+        )
+        a_s = load_swizzled_scale_tile(
+            as_descriptor, as_ptr, 0, pid_m, k, M, K, BLOCK_SIZE_M, SCALE_COLS, SCALE_GROUP_K
+        )
+    elif GROUPED:  # gathered per-(row, K-group) affine scale off the row-major source-order As
+        a = load_grouped_act_tile(
+            a_ptrs, a_descriptor, m_off, k_off, value_mask, gather_rows, A_MEMORY_MODE, A_GATHER
+        )
+        offs_sf = k * SCALE_COLS + tl.arange(0, SCALE_COLS)
+        a_s = tl.load(
+            as_ptrs + gather_rows[:, None] * stride_as_m + offs_sf[None, :],
+            mask=scale_mask[:, None],
+            other=0.0,
+        )
+    else:  # 2D / decode: inline-quant affine or in-register scale
+        a, a_s = load_mx_act_tile(
+            a_ptrs, as_ptrs, as_global, as_global_row, scale_mask, a_descriptor, m_off, k_off, 0,
+            BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K, A_MEMORY_MODE,
+            FORMAT=ACTIVATION_FORMAT,
+        )
+    return a, a_s
+@triton.jit
+def load_act_block_dynamic(
+    a_ptrs, as_ptrs, value_mask, scale_mask, a_descriptor, m_off, k_off, gather_rows, k,
+    A_MEMORY_MODE: tl.constexpr, A_GATHER: tl.constexpr, GROUPED: tl.constexpr,
+    SWAP_AB: tl.constexpr,
+):
+    """block_dynamic activation tile + per-row per-K-block scale. Grouped: gathered value + the scale
+    read contiguously at ``As + k``; else ``load_block_fp8_act_tile`` (swap-aware, inline-quant 2D)."""
+    if GROUPED:  # gathered value + per-row per-K-block scale read contiguously at As + k
+        a = load_grouped_act_tile(
+            a_ptrs, a_descriptor, m_off, k_off, value_mask, gather_rows, A_MEMORY_MODE, A_GATHER
+        )
+        a_s = tl.load(as_ptrs + k, mask=scale_mask, other=0.0)
+    else:
+        a, a_s = load_block_fp8_act_tile(
+            a_ptrs, as_ptrs, a_descriptor, m_off, k_off, A_MEMORY_MODE, scale_mask, SWAP_AB
+        )
+    return a, a_s
+@triton.jit
+def load_act_static(
+    a_ptrs, a_descriptor, m_off, k_off, value_mask, gather_rows, a_s_static,
+    A_MEMORY_MODE: tl.constexpr, A_GATHER: tl.constexpr,
+):
+    """static activation tile. Pre-quantized fp8 A loads as the rows-major MMA lhs; raw bf16/fp16
+    (inline arm, small M, pointer-only) is quantized against the scalar ``a_s_static``. ``a_s`` = the
+    values (the static scale is a scalar folded post-loop)."""
+    if a_ptrs.dtype.element_ty == tl.float8e4nv:  # pre-quantized fp8 A (MMA lhs, rows-major)
+        a = load_grouped_act_tile(
+            a_ptrs, a_descriptor, m_off, k_off, value_mask, gather_rows, A_MEMORY_MODE, A_GATHER
+        )
+    else:  # raw bf16/fp16 (inline arm, M<threshold, pointer-only) — quantize vs the static scale
+        a = (tl.load(a_ptrs).to(tl.float32) / a_s_static).to(tl.float8e4nv)
+    a_s = a
+    return a, a_s
+@triton.jit
+def load_act_plain(
+    a_ptrs, a_descriptor, m_off, k_off, value_mask, gather_rows,
+    A_MEMORY_MODE: tl.constexpr, A_GATHER: tl.constexpr,
+):
+    """tensor / full_precision activation: plain value tile, no scale (returns ``(a, a)`` — the
+    second slot is dead, kept for the uniform (value, scale) shape)."""
+    a = load_grouped_act_tile(
+        a_ptrs, a_descriptor, m_off, k_off, value_mask, gather_rows, A_MEMORY_MODE, A_GATHER
+    )
+    return a, a
+@triton.jit
+def _weight_scale_mx(
+    bs_ptrs, bs_mask, bs_descriptor, bs_ptr, blk_idx, expert_id, pid_n, k, N, K,
+    stride_bs_e, stride_bs_n, stride_bs_k,
+    GROUPED: tl.constexpr, GATE: tl.constexpr, PER_EXPERT: tl.constexpr,
+    SWIZZLED_SCALES: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
+    SCALE_COLS: tl.constexpr, SCALE_GROUP_K: tl.constexpr,
+    k_col_off=0,
+):
+    """MX weight-scale K-tile ``(rows, SCALE_COLS)``. Grouped: pre-swizzled SWIZZLE_32_4_4 via the
+    descriptor (``load_swizzled_scale`` — gate|up is one 2*BN tile off the block-interleaved buffer,
+    GATE folding into REP/width), else an inlined affine per-group read off the un-swizzled 3D ``Bs``
+    (inlined, not ``load_weight_scale_tile``: the nested leaf breaks warp-specialization partitioning
+    in the grouped loop). 2D-GATE / batched decode: the per-(expert, N, K) 3D leaf. 2D dense: the
+    descriptor bulk-load (BN=128) else the bounds-masked affine pointer."""
+    if GROUPED:
+        if SWIZZLED_SCALES:
+            NW: tl.constexpr = (2 if GATE else 1) * BLOCK_SIZE_N
+            # REP counts the tile's whole 128-row blocks, so it derives from the DOUBLED extent:
+            # BN=64 under GATE is one 128-row block, not 2 * (64 // 128) == 0
+            NREP: tl.constexpr = max(1, NW // 128)
+            b_s = load_swizzled_scale(
+                bs_descriptor, blk_idx, k, NREP, SCALE_COLS // 4, NW, SCALE_COLS
+            )
+        else:  # affine per-group read off the un-swizzled 3D Bs (num_experts, n_rows, K//g)
+            NW: tl.constexpr = (2 if GATE else 1) * BLOCK_SIZE_N
+            base = bs_ptrs + expert_id * stride_bs_e
+            # clamp rather than mask (a masked load's other=0.0 blocks WS lowering): a partial
+            # last N-tile reads a clamped scale the epilogue masks off via N_COLS
+            offs_bn = tl.minimum(
+                pid_n * NW + tl.arange(0, NW), ((2 if GATE else 1) * N) - 1
+            )
+            offs_sf = k * SCALE_COLS + k_col_off + tl.arange(0, SCALE_COLS)
+            b_s = tl.load(base + offs_bn[:, None] * stride_bs_n + offs_sf[None, :] * stride_bs_k)
+    elif GATE or PER_EXPERT:  # 2D-GATE (expert 0) / batched decode: per-(expert, N, K) scale leaf
+        b_s = load_weight_scale_tile(
+            SWIZZLED_SCALES, bs_descriptor, bs_ptr, expert_id, pid_n, k, N, K,
+            stride_bs_e, stride_bs_n, stride_bs_k, BLOCK_SIZE_N, SCALE_COLS, SCALE_GROUP_K, GATE,
+            k_col_off=k_col_off,
+        )
+    elif SWIZZLED_SCALES:  # pre-swizzled SWIZZLE_32_4_4 scale — descriptor at BN=128, gather below
+        b_s = load_swizzled_scale_tile(
+            bs_descriptor, bs_ptr, 0, pid_n, k, N, K, BLOCK_SIZE_N, SCALE_COLS, SCALE_GROUP_K,
+        )
+    elif bs_mask is None:  # weight-only 2D: N-clamped in-bounds ptrs, UNMASKED (no other=0.0 constant
+        b_s = tl.load(bs_ptrs)  # -> the K-loop warp-specialization-lowers; tail N masked in epilogue)
+    else:
+        b_s = tl.load(bs_ptrs, mask=bs_mask[:, None], other=0.0)  # 0.0 casts to fp8/uint8
+    return b_s
+@triton.jit
+def _weight_scale_block_dynamic(
+    bs_ptrs, up_s_ptr, bs_mask, k, stride_bs_k,
+    GROUPED: tl.constexpr, GATE: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
+):
+    """block_dynamic per-K-block weight scale. Grouped reads it at ``bs_ptrs + k*stride`` (GATE folds
+    gate on the first ``BN`` columns, up on the rest via ``up_s_ptr``); batched reads the pre-offset
+    maskless pointer; 2D reads the bounds-masked advanced pointer."""
+    if GROUPED:
+        if GATE:  # gate scale on the first BN columns, up on the rest
+            b_s = tl.where(
+                tl.arange(0, 2 * BLOCK_SIZE_N) < BLOCK_SIZE_N,
+                tl.load(bs_ptrs + k * stride_bs_k),
+                tl.load(up_s_ptr + k * stride_bs_k),
+            )
+        else:
+            # One 128-block scale covers the whole BN tile, so the pointer is scalar and the load
+            # is rank-0. Broadcast to the tile's N extent so BOTH grouped arms hand back a rank-1
+            # [.., N] scale: `block_dynamic_dot`'s dot_scaled arm indexes `b_s.shape[0]`, which a
+            # rank-0 value cannot answer, and that killed every BLOCK_SIZE_M>=128 config (the only
+            # ones that take the tcgen05 arm) for the ungated grouped GEMM -- the MoE down proj.
+            # broadcast_to keeps the dtype, so decode_group_scale still sees UE8M0 as uint8.
+            b_s = tl.broadcast_to(tl.load(bs_ptrs + k * stride_bs_k)[None], [BLOCK_SIZE_N])
+    elif bs_mask is None:  # batched: bs_ptrs pre-offset (gate/up folded), maskless decode tile
+        b_s = tl.load(bs_ptrs)
+    else:
+        b_s = tl.load(bs_ptrs, mask=bs_mask, other=0.0)
+    return b_s
+@triton.jit
+def _weight_scale_static(
+    bs_ptrs, up_s_ptr, bs_mask, k, stride_bs_k,
+    GROUPED: tl.constexpr, GATE: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
+):
+    """static weight scale — as block_dynamic, but the grouped non-GATE scale broadcasts to a per-N
+    vector so ``accumulate("static")`` applies it in N (the static act scale is a scalar folded
+    post-loop)."""
+    if GROUPED:
+        if GATE:  # gate on the first BN, up on the rest
+            b_s = tl.where(
+                tl.arange(0, 2 * BLOCK_SIZE_N) < BLOCK_SIZE_N,
+                tl.load(bs_ptrs + k * stride_bs_k),
+                tl.load(up_s_ptr + k * stride_bs_k),
+            )
+        else:
+            b_s = tl.load(bs_ptrs + k * stride_bs_k) + tl.zeros(
+                (BLOCK_SIZE_N,), bs_ptrs.dtype.element_ty
+            )
+    elif bs_mask is None:  # batched: bs_ptrs pre-offset (gate/up folded), maskless decode tile
+        b_s = tl.load(bs_ptrs)
+    else:
+        b_s = tl.load(bs_ptrs, mask=bs_mask, other=0.0)  # affine col index -> mask OOB last tile
+    return b_s
+@triton.jit
+def _weight_value(
+    b_ptrs, b_descriptor, row0, n_off, k_off,
+    GATE: tl.constexpr, GROUPED: tl.constexpr, B_MEMORY_MODE: tl.constexpr,
+    SWAP_AB: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, KB: tl.constexpr,
+):
+    """Format-agnostic weight value tile: the per-expert 3D box (GROUPED) or the plain swap-aware
+    tile. Shared by every ``load_weight_<format>`` — the value load never depends on the scale
+    format. A 2D gate|up weight takes the PLAIN arm: its rows are interleaved, so the tile is a
+    contiguous ``2*BN`` span of a 2D ``(2N, K)`` tensor and needs no 3D box — routing it through
+    the grouped loader issued a 3-offset load against a 2D descriptor ("expected 2 offsets, but
+    got 3"), unreachable only while gate was fenced to pointer mode. Single-exit if/else (an early
+    return would type-check the untaken arm)."""
+    if GROUPED:
+        w = load_grouped_weight_tile(
+            b_ptrs, b_descriptor, row0, n_off, k_off, BLOCK_SIZE_N, KB, GATE, B_MEMORY_MODE, SWAP_AB
+        )
+    else:
+        w = load_weight_tile(b_ptrs, b_descriptor, n_off, k_off, B_MEMORY_MODE, SWAP_AB)
+    return w
+@triton.jit
+def load_weight_mx(
+    b_ptrs, b_descriptor, bs_ptrs, bs_mask, bs_descriptor, bs_ptr, row0, n_off, k_off,
+    blk_idx, expert_id, pid_n, k, N, K, stride_bs_e, stride_bs_n, stride_bs_k,
+    GATE: tl.constexpr, GROUPED: tl.constexpr, PER_EXPERT: tl.constexpr,
+    B_MEMORY_MODE: tl.constexpr, SWAP_AB: tl.constexpr, SWIZZLED_SCALES: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, SCALE_GROUP_K: tl.constexpr,
+    WEIGHT_VALUES_PER_BYTE: tl.constexpr,
+    k_col_off=0,
+):
+    """The MX weight path: value tile + pre-swizzled/affine group scale. ``k_col_off`` shifts the
+    affine scale read for a K-tile the caller clamped back in bounds (0 for aligned callers)."""
+    KB: tl.constexpr = BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE
+    SCALE_COLS: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K
+    w = _weight_value(b_ptrs, b_descriptor, row0, n_off, k_off, GATE, GROUPED, B_MEMORY_MODE, SWAP_AB, BLOCK_SIZE_N, KB)
+    w_s = _weight_scale_mx(
+        bs_ptrs, bs_mask, bs_descriptor, bs_ptr, blk_idx, expert_id, pid_n, k, N, K,
+        stride_bs_e, stride_bs_n, stride_bs_k,
+        GROUPED, GATE, PER_EXPERT, SWIZZLED_SCALES, BLOCK_SIZE_N, SCALE_COLS, SCALE_GROUP_K, k_col_off=k_col_off,
+    )
+    return w, w_s
+@triton.jit
+def load_weight_block_dynamic(
+    b_ptrs, b_descriptor, bs_ptrs, bs_mask, up_s_ptr, row0, n_off, k_off, k, stride_bs_k,
+    GATE: tl.constexpr, GROUPED: tl.constexpr, B_MEMORY_MODE: tl.constexpr, SWAP_AB: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, WEIGHT_VALUES_PER_BYTE: tl.constexpr = 1,
+):
+    """The block_dynamic weight path: value tile + per-K-block scale (gate on the first BN, up on the
+    rest via ``up_s_ptr``)."""
+    KB: tl.constexpr = BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE
+    w = _weight_value(b_ptrs, b_descriptor, row0, n_off, k_off, GATE, GROUPED, B_MEMORY_MODE, SWAP_AB, BLOCK_SIZE_N, KB)
+    w_s = _weight_scale_block_dynamic(bs_ptrs, up_s_ptr, bs_mask, k, stride_bs_k, GROUPED, GATE, BLOCK_SIZE_N)
+    return w, w_s
+@triton.jit
+def load_weight_static(
+    b_ptrs, b_descriptor, bs_ptrs, bs_mask, up_s_ptr, row0, n_off, k_off, k, stride_bs_k,
+    GATE: tl.constexpr, GROUPED: tl.constexpr, B_MEMORY_MODE: tl.constexpr, SWAP_AB: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, WEIGHT_VALUES_PER_BYTE: tl.constexpr = 1,
+):
+    """The static weight path: value tile + per-K-block scale broadcast to a per-N vector (the static
+    act scale is a scalar folded post-loop)."""
+    KB: tl.constexpr = BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE
+    w = _weight_value(b_ptrs, b_descriptor, row0, n_off, k_off, GATE, GROUPED, B_MEMORY_MODE, SWAP_AB, BLOCK_SIZE_N, KB)
+    w_s = _weight_scale_static(bs_ptrs, up_s_ptr, bs_mask, k, stride_bs_k, GROUPED, GATE, BLOCK_SIZE_N)
+    return w, w_s
+@triton.jit
+def load_weight_plain(
+    b_ptrs, b_descriptor, row0, n_off, k_off,
+    GATE: tl.constexpr, GROUPED: tl.constexpr, B_MEMORY_MODE: tl.constexpr, SWAP_AB: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, WEIGHT_VALUES_PER_BYTE: tl.constexpr = 1,
+):
+    """The tensor / full_precision weight path: plain value tile, no block scale (the per-tensor
+    scale, if any, is applied post-loop). Returns ``(w, w)`` so callers keep the uniform (value,
+    scale) shape; the second slot is dead."""
+    KB: tl.constexpr = BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE
+    w = _weight_value(b_ptrs, b_descriptor, row0, n_off, k_off, GATE, GROUPED, B_MEMORY_MODE, SWAP_AB, BLOCK_SIZE_N, KB)
+    return w, w
+@triton.jit
+def advance_ptrs(
+    a_ptrs,
+    as_ptrs,
+    w_ptrs,
+    ws_ptrs,
+    w_up_ptrs,
+    ws_up_ptrs,
+    a_step,
+    as_step,
+    w_step,
+    ws_step,
+    A_MEMORY_MODE: tl.constexpr,
+    W_MEMORY_MODE: tl.constexpr,
+    ADVANCE_AS: tl.constexpr,
+    ADVANCE_WS: tl.constexpr,
+    GATE_STREAMS: tl.constexpr = False,
+):
+    """Advance the shared GEMM operand pointers one K-step, folding the memory-mode / scale-layout /
+    gate|up-stream conditionals out of every loop. The operand set is uniform across the kernels:
+    activation (``a_ptrs`` + affine scale ``as_ptrs``) and weight (``w_ptrs`` + affine scale
+    ``ws_ptrs``), the weight either a single stream or — under ``GATE_STREAMS`` — the gate|up pair
+    (``w_up_ptrs`` + ``ws_up_ptrs``, bumped by the same steps). Value pointers advance only on the
+    pointer arm (a descriptor arm re-derives the box K offset from ``k``); scale pointers advance
+    only when read affine (``ADVANCE_AS`` / ``ADVANCE_WS`` — swizzled / in-leaf / per-tensor scales
+    don't). Pass a dead pointer + step 0 (flag off) for any stream a kernel doesn't carry. Returns
+    the six pointers in argument order."""
+    if A_MEMORY_MODE == "pointer":
+        a_ptrs += a_step
+    if W_MEMORY_MODE == "pointer":
+        w_ptrs += w_step
+        if GATE_STREAMS:
+            w_up_ptrs += w_step
+    if ADVANCE_AS:
+        as_ptrs += as_step
+    if ADVANCE_WS:
+        ws_ptrs += ws_step
+        if GATE_STREAMS:
+            ws_up_ptrs += ws_step
+    return a_ptrs, as_ptrs, w_ptrs, ws_ptrs, w_up_ptrs, ws_up_ptrs
+@triton.jit
+def weight_tile_ptrs(
+    base,
+    offs_n,
+    offs_k,
+    stride_n,
+    stride_k,
+    SWAP_AB: tl.constexpr,
+):
+    """Weight-tile pointers oriented by ``SWAP_AB``. A gate|up tile is a ``2*BN`` row span
+    (``offs_n`` carries it); the epilogue splits the projections."""
+    return oriented_tile_ptrs(base, offs_n, offs_k, stride_n, stride_k, SWAP_AB)
+@triton.jit
+def oriented_tile_ptrs(
+    base, offs_rows, offs_k, stride_rows, stride_k, SWAP_AB: tl.constexpr
+):
+    """Operand-tile pointers oriented by whether these rows sit in the MMA M dim
+    (``SWAP_AB``, from the weight's viewpoint — activation callers pass the flag inverted):
+    rows-major ``[rows, K]`` when they do, else K-major ``[K, rows]``. Only the taken
+    constexpr branch compiles, so the divergent shapes never meet. The per-step K-advance
+    is identical for both layouts, so the caller advances the returned pointer the same
+    way regardless of orientation."""
+    if SWAP_AB:
+        ptrs = base + (offs_rows[:, None] * stride_rows + offs_k[None, :] * stride_k)
+    else:
+        ptrs = base + (offs_k[:, None] * stride_k + offs_rows[None, :] * stride_rows)
+    return ptrs
+@triton.jit
+def operand_tile_ptrs(
+    base,
+    offs_rows,
+    offs_k,
+    stride_rows,
+    stride_k,
+    MEMORY_MODE: tl.constexpr,
+    SWAP_AB: tl.constexpr,
+):
+    """Prologue operand-tile pointer, folding the per-operand memory-mode branch: the explicit
+    oriented ``[rows,K]``/``[K,rows]`` tile on the pointer arm, or ``base`` as a scalar
+    placeholder on a descriptor arm (which reads its box via the descriptor — building the index
+    tensor there would only stay live across the K-loop and spill registers). ``SWAP_AB`` is the
+    orientation from the weight's viewpoint (activation callers pass it inverted). Single
+    return — the arms have divergent types (a ``[rows,K]`` tile vs the scalar base), so an early
+    return can't unify them; the constexpr selects one."""
+    if MEMORY_MODE == "pointer":
+        ptrs = oriented_tile_ptrs(base, offs_rows, offs_k, stride_rows, stride_k, SWAP_AB)
+    else:
+        ptrs = base
+    return ptrs

build/torch-rocm/matmul.py ADDED Viewed

	@@ -0,0 +1,1970 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import torch
+import triton
+import triton.language as tl
+from triton.tools.tensor_descriptor import TensorDescriptor
+from .bayesian_autotuner import bayesian_autotune
+from .compat import add_op_namespace_prefix, FP8_DTYPE, is_sm10x, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, tl_dtype
+from .descriptors import maybe_descriptor, rebind_bd_descriptors, rebind_mx_descriptors, rebind_weight_only_descriptors
+from .formats import check_activation_format, global_scale_stride, normalize_global_scale, e2m1_as_uint8, is_mx, mx_scale_family, resolve_activation_format, resolve_output_dtype, ue8m0_as_uint8, validate_dense_2d_operands, weight_format
+from .swizzle import swizzle_mx_scales, swizzled_scale_descriptor
+from .quant import MX_ACT_QUANT, fp8_act_quant_block_dynamic, fp8_act_quant_tensor_wide, maybe_act_quant, mxfp8_act_quant, nvfp4_act_quant
+from .loading.scales import apply_global_scale, mx_2d_scale_ptrs
+from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot
+from .loading.tiles import (
+    advance_ptrs,
+    load_act_block_dynamic,
+    load_act_mx,
+    load_act_plain,
+    load_act_static,
+    load_weight_block_dynamic,
+    load_weight_mx,
+    load_weight_plain,
+    load_weight_static,
+    matmul_weight_ptrs,
+    operand_tile_descriptor,
+    operand_tile_ptrs,
+    swizzle_offsets,
+)
+from .epilogue import acc_init, add_bias, bias_strides, gemm_epilogue, store_masked, store_masked_oriented
+from .pruners import PATH_ANCHOR_AXES, fp8_dot_warp_pruner, gate_stacked_tmem_trap_pruner, block_dynamic_2d_warp_spec_pruner, block_dynamic_mma_width_pruner, block_within_dim_pruner, compose_pruners, descriptor_box_pruner, descriptor_needs_prequant_pruner, matched_memory_modes_pruner, mx_2d_swap_scope_pruner, mx_config_pruner, scalar_max_m_pruner, smem_pruner, swizzled_out_bm_pruner, swizzled_scale_config_pruner, warp_spec_compile_guard_pruner
+# The 2D-grid kernels' L2-locality swizzle depth is derived per-tile inside
+# ``swizzle_offsets`` (see SWIZZLE_GROUP_A_BYTES there) — no per-kernel constant here.
+# maybe_act_quant crossovers (min rows for offline pre-quant). MX: MEASURED — B200
+# 2D-matmul sweep, graph-timed, H=6144 MXFP8 / H=4096 MXFP4: inline wins only at M=1
+# (33 vs 44us / 20 vs 30us), offline from M=16 for MXFP8 (22 vs 33us), 2-3x by M>=64
+# (MXFP4's M=16 cell marginally favors inline — outweighed); NVFP4 inherits the
+# family gate (not swept — same pack/amax structure as MXFP4). STATIC: inherited
+# estimate, not swept — its inline arm is cheaper elementwise work, so the true
+# crossover is at or above the MX one; the M=1 decode case (the one that matters) is
+# inline either way.
+MX_MATMUL_ACT_PREQUANT_MIN_M = 16
+STATIC_MATMUL_ACT_PREQUANT_MIN_M = 16
+# Block-dynamic has NO gate: offline wins at EVERY M incl. M=1 (24.7 vs 31.0us, isolated
+# arm A/B, graph-timed, H=6144 b128) — the inline arm pays a per-tile fp32 amax+div. The
+# always-offline paths (bd 2D/batched/grouped, fused gate_ups, MX grouped) quantize
+# unconditionally; only the two gates above have a real M=1 inline win.
+@bayesian_autotune(
+    # tune_block_m: BLOCK_SIZE_M is a config axis. tune_block_n: the N tile is DECOUPLED
+    # from the caller's scale granularity (block_n) — a BN=256 tile over 128-wide scale
+    # columns halves activation re-reads; any BN is numerically fine (the offs_bn //
+    # block_n gather spans or splits scale columns), while BK stays pinned to block_k
+    # (activation scale groups are per block_k).
+    #
+    # Two memory regimes, tuner-routed per (N, K, M) key:
+    #  - pointer: wins decode / short-K (big tiles starve here — the loop is load-bound).
+    #  - host-TMA on BOTH operands (A_MEMORY_MODE / B_MEMORY_MODE = descriptor) in the NATURAL
+    #    orientation (no SWAP_AB — the descriptor box transposes once, in load_weight_tile, to the
+    #    pointer arm's K-major rhs): wins wide-N compute prefill, feeding the big tiles the pointer
+    #    arm cannot. WS is load-bearing here: the per-iteration transpose races without it.
+    # Emitting both axes lets the tuner route per regime. Verdicts B200 (sm_100) — re-chart on H100.
+    get_accelerator_autotuning_configs(
+        warp_spec=True,
+        tune_block_m=True,
+        tune_block_n=True,
+        a_memory_modes=("descriptor", "pointer"),
+        b_memory_modes=("descriptor", "pointer"),
+        pre_hook=rebind_bd_descriptors,
+    ),
+    # m_bit_length (log2 M bucket) keys the M tile, mirroring mx_dynamic_matmul_kernel; GATE keys the
+    # gate|up arm separately (its stacked dot is 2*BN wide, a distinct config space). block_n/block_k
+    # (the launch-pinned quant block) gate descriptor-box legality and the K tile — two checkpoints
+    # at one shape with different blocks must not share a winner.
+    ["N", "K", "m_bit_length", "GATE", "block_n", "block_k"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe
+    finite_check_args=("C", "Cs"),
+    # WS compile guard + descriptor TMA-box limits (256/dim) + shared-memory fit (the big
+    # TMA prefill tiles reach the smem ceiling; BM256 BN256 OOMs). The gate|up stack reads via the
+    # pointer arm (a contiguous descriptor box can't span the N-apart gate/up rows), so prune the
+    # B-descriptor configs under GATE. The implicit dot_scaled arm (UE8M0 As, native-M tile)
+    # carries the sm_10x scaled-MMA width cap — GATE at BN=256 would emit a width-512 trap.
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            fp8_dot_warp_pruner(),
+            warp_spec_compile_guard_pruner(),
+            matched_memory_modes_pruner(),
+            descriptor_box_pruner("block_k"),
+            smem_pruner("block_k"),
+            block_dynamic_mma_width_pruner(),
+            block_dynamic_2d_warp_spec_pruner(),
+        )
+    },
+)
+@triton.jit
+def w8a8_block_dynamic_fp8_matmul_kernel(
+    A,  # (M, K) E4M3 activations (pre-quantized once by the wrapper)
+    ADescriptor,  # host TMA descriptor over A (M, K), box (BLOCK_SIZE_M, block_k); read iff A_MEMORY_MODE != "pointer"
+    As,  # (M, K // block_k) per-row, per-K-block activation scales (fp32 or uint8/UE8M0)
+    B,  # (N, K) FP8 weights
+    BDescriptor,  # host TMA descriptor over B (N, K), box (BLOCK_SIZE_N, block_k); read iff B_MEMORY_MODE != "pointer"
+    Bs,  # (N // block_n, K // block_k) weight scales (fp32 or uint8/UE8M0)
+    C,  # (M, N) output
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    # Shape
+    M,
+    N,
+    K,
+    m_bit_length,  # autotune key only (log2 M bucket); unused in body
+    # Strides
+    stride_a_m,
+    stride_a_k,
+    stride_as_m,
+    stride_as_k,  # activation-scale group-axis stride: 1 for a row-major (M, K//block_k) grid,
+    # but a real stride when the caller hands in a transposed view (the weight side has carried
+    # both strides all along; this is the activation side catching up)
+    stride_b_k,
+    stride_b_n,
+    stride_bs_k,
+    stride_bs_n,
+    stride_c_m,
+    stride_c_n,
+    stride_bias_e,
+    stride_bias_n,
+    # Weight-quantization blocks (the caller's block_size); block_k is also the K tile
+    # (the activation scale groups are per block_k)
+    block_n: tl.constexpr,
+    block_k: tl.constexpr,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    WARP_SPEC: tl.constexpr = False,
+    B_MEMORY_MODE: tl.constexpr = "pointer",
+    A_MEMORY_MODE: tl.constexpr = "pointer",
+    SWAP_AB: tl.constexpr = False,
+    # Dense gate|up fusion: B is the (2N, K) gate|up stack, C the [M, N] GLU output (bf16 — the
+    # requant output formats live on the MX kernel). Every arm folds out at compile time when
+    # GATE=False (the plain dense GEMM, unchanged).
+    GATE: tl.constexpr = False,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
+):
+    """Block-scale FP8 GEMM kernel.
+    Computes ``C = A @ B.T``. Activations arrive pre-quantized (one pass in the
+    wrapper — the inline per-N-tile quant would repeat N//BN times per element; see the
+    grouped kernels). 2D grid with swizzle for L2 cache locality on B.
+    Two operand-feed regimes, tuner-routed (see the decorator): explicit pointers, or
+    host-TMA descriptors on both operands. Under the descriptor modes the natural
+    orientation loads the ``(BN, BK)`` weight box and transposes it once (in
+    ``load_weight_tile``) to the same K-major rhs the pointer arm builds — no ``SWAP_AB``
+    (that trans is race-free only under ``WARP_SPEC``). ``SWAP_AB`` remains an
+    independent pointer-arm orientation knob (weight tile in the MMA M dim, activation
+    loaded transposed via strides); it is implemented but not currently emitted.
+    UE8M0 scales (activations quantized to power-of-two, weights UE8M0) on a native-M tile
+    (the MMA M operand is 128-wide: BLOCK_SIZE_M no-swap, BLOCK_SIZE_N swapped) fold the
+    128-group scales into a tcgen05 ``dot_scaled`` MMA — see ``block_dynamic_dot``. fp32
+    scales, or a narrow M tile (decode), keep the plain ``tl.dot`` + software rescale.
+    """
+    USE_DOT_SCALED: tl.constexpr = (As.dtype.element_ty == tl.uint8) and (
+        (BLOCK_SIZE_N if SWAP_AB else BLOCK_SIZE_M) >= 128
+    )
+    pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, block_k, GATE=GATE
+    )
+    # Scale pointers index off AFFINE row/col offsets + a bounds mask; the %-wrapped forms
+    # (offs_am/offs_bn from swizzle_offsets) drive ONLY the pointer arm's operand tiles
+    # (token replication at decode). Keeping scales affine avoids the non-affine gather
+    # the wrap induces (~460us / 18pp at prefill).
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    n_rows = 2 * N if GATE else N
+    offs_am_lin = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_bn_lin = pid_n * n_width + tl.arange(0, n_width)
+    as_ptrs = As + offs_am_lin * stride_as_m
+    as_mask = offs_am_lin < M
+    bs_mask = offs_bn_lin < n_rows
+    a_descriptor = operand_tile_descriptor(
+        ADescriptor, A, M, K, stride_a_m, stride_a_k, BLOCK_SIZE_M, block_k, A_MEMORY_MODE
+    )
+    b_descriptor = operand_tile_descriptor(
+        BDescriptor,
+        B,
+        N,
+        K,
+        stride_b_n,
+        stride_b_k,
+        BLOCK_SIZE_N,
+        block_k,
+        B_MEMORY_MODE, GATE,
+    )
+    # Explicit operand-tile pointers on the pointer arm; a scalar placeholder on a descriptor
+    # arm (which reads its box via the descriptor — the [BM,BK]/[BK,BN] index tensor would stay
+    # live across the K-loop for nothing and spill registers, ~460us / 20pp MFU at wide-N prefill).
+    # Under GATE the weight is the stacked (2N, K) gate|up tile (matmul_weight_ptrs; pointer-only, the
+    # descriptor box can't span the N-apart rows), and its per-weight-row block scale is a 2*BN gather
+    # (gate rows [0,N), up rows offset by N//block_n scale blocks) — the block_dynamic_dot broadcast
+    # then folds it per row exactly as in the dense case.
+    a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, not SWAP_AB)
+    b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, B_MEMORY_MODE, SWAP_AB)
+    # the scale-index gather decouples the tile from the scale grid: a wide tile spans several
+    # scale columns, a narrow one shares a column
+    bs_ptrs = Bs + (offs_bn_lin // block_n) * stride_bs_n
+    accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, SWAP_AB)
+    for k in tl.range(0, tl.cdiv(K, block_k), warp_specialize=WARP_SPEC):
+        a, a_s = load_act_block_dynamic(
+            a_ptrs, as_ptrs, None, as_mask, a_descriptor, pid_m * BLOCK_SIZE_M, k * block_k,
+            0, 0, A_MEMORY_MODE, False, False, SWAP_AB,
+        )
+        b, b_s = load_weight_block_dynamic(
+            b_ptrs, b_descriptor, bs_ptrs, bs_mask, 0, 0, pid_n * n_width, k * block_k, 0, 0,
+            GATE, False, B_MEMORY_MODE, SWAP_AB, BLOCK_SIZE_N, block_k,
+        )
+        accumulator = block_dynamic_dot(
+            accumulator, a, a_s, b, b_s, block_k, SWAP_AB, USE_DOT_SCALED, False
+        )
+        a_ptrs, as_ptrs, b_ptrs, bs_ptrs, _, _ = advance_ptrs(
+            a_ptrs, as_ptrs, b_ptrs, bs_ptrs, b_ptrs, bs_ptrs,
+            block_k * stride_a_k, stride_as_k, block_k * stride_b_k, stride_bs_k,
+            A_MEMORY_MODE, B_MEMORY_MODE, True, True, False,
+        )
+    if GATE:
+        gemm_epilogue(
+            C, None, accumulator, offs_am_lin, pid_n, pid_m, as_mask,
+            stride_c_m, stride_c_n, 0, 0,
+            BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, block_k,
+            ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+            COMPUTE_MODE="dot", SWAP_AB=SWAP_AB, N_COLS=N,
+            Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n,
+        )
+    else:
+        accumulator = add_bias(
+            accumulator, Bias, stride_bias_e, stride_bias_n, 0, pid_n,
+            BLOCK_SIZE_N, N, SWAP_AB,
+        )
+        store_masked_oriented(
+            C,
+            accumulator,
+            pid_m,
+            pid_n,
+            M,
+            N,
+            stride_c_m,
+            stride_c_n,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            SWAP_AB,
+        )
+@bayesian_autotune(
+    # tune_block_m mirrors the block-dynamic kernel (adaptive BM=128 cost it 17% at prefill).
+    get_accelerator_autotuning_configs(
+        tune_block_nk=True, warp_spec=True, tune_block_m=True
+    ),
+    # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide, a distinct config space).
+    ["N", "K", "m_bit_length", "GATE"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe
+    finite_check_args=("C", "Cs"),
+    # block_k is a tuned axis and the loop below is maskless — veto non-dividing BKs;
+    # WS is a pure perf axis here (non-WS is the validated state), compile-guarded.
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            fp8_dot_warp_pruner(),
+            block_within_dim_pruner("K"),
+            warp_spec_compile_guard_pruner(),
+        )
+    },
+)
+@triton.jit
+def w8a8_tensor_dynamic_fp8_matmul_kernel(
+    A,  # (M, K) pre-quantized FP8 activations
+    As,  # (M,) per-token activation scales
+    B,  # (N, K) FP8 weights
+    Bs,  # scalar/(1,) per-tensor weight scale
+    C,  # (M, N) output
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    # Shape
+    M,
+    N,
+    K,
+    m_bit_length,  # autotune key only (log2 M bucket); unused in body
+    # Strides
+    stride_a_m,
+    stride_a_k,
+    stride_as_m,
+    stride_b_k,
+    stride_b_n,
+    stride_c_m,
+    stride_c_n,
+    stride_bias_e,
+    stride_bias_n,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    WARP_SPEC: tl.constexpr = False,
+    # Dense gate|up fusion: B is the (2N, K) gate|up stack, C the [M, N] GLU output. Every arm
+    # folds out at compile time when GATE=False (the plain dense GEMM, unchanged).
+    GATE: tl.constexpr = False,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
+):
+    """Tensor-scale FP8 GEMM kernel.
+    Computes ``C = A @ B.T`` with one activation scale per row and one
+    weight scale for the full matrix.
+    Uses a 2D grid with swizzle for L2 cache locality on B tiles.
+    """
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GATE=GATE
+    )
+    a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, "pointer", True)
+    # Under GATE the weight is the stacked (2N, K) gate|up tile; one per-tensor weight scale
+    # covers both projections, applied (with the per-row act scale) before the GLU split.
+    b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, "pointer")
+    a_s = tl.load(As + offs_am * stride_as_m)
+    b_s = tl.load(Bs)
+    # Accumulate raw dot products, apply scales once after the loop.
+    accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, False)
+    for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC):
+        a, _as = load_act_plain(a_ptrs, a_ptrs, 0, 0, None, 0, "pointer", False)
+        b, _bs = load_weight_plain(
+            b_ptrs, b_ptrs, 0, pid_n * n_width, k * BLOCK_SIZE_K,
+            GATE, False, "pointer", False, BLOCK_SIZE_N, BLOCK_SIZE_K,
+        )
+        accumulator = accumulator + fp8_dot(a, b, False, BLOCK_SIZE_K)
+        a_ptrs, _asp, b_ptrs, _bsp, _, _ = advance_ptrs(
+            a_ptrs, a_ptrs, b_ptrs, b_ptrs, b_ptrs, b_ptrs,
+            BLOCK_SIZE_K * stride_a_k, 0, BLOCK_SIZE_K * stride_b_k, 0,
+            "pointer", "pointer", False, False, False,
+        )
+    accumulator = accumulator * a_s[:, None] * b_s
+    if GATE:
+        out_row = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+        gemm_epilogue(
+            C, None, accumulator, out_row, pid_n, pid_m, out_row < M,
+            stride_c_m, stride_c_n, 0, 0,
+            BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, BLOCK_SIZE_K,
+            ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+            COMPUTE_MODE="dot", N_COLS=N,
+            Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n,
+        )
+    else:
+        accumulator = add_bias(
+            accumulator, Bias, stride_bias_e, stride_bias_n, 0, pid_n,
+            BLOCK_SIZE_N, N, False,
+        )
+        store_masked(
+            C,
+            accumulator,
+            pid_m,
+            pid_n,
+            M,
+            N,
+            stride_c_m,
+            stride_c_n,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+        )
+@bayesian_autotune(
+    # tune_block_m + tune_block_n + WARP_SPEC + host-TMA memory modes mirror the block-dynamic
+    # kernel above (same 2D swizzle loop physics). Both A/B memory axes are emitted and the tuner
+    # routes descriptor vs pointer per (N, K, M) key; the A descriptor is read only on the offline
+    # (pre-quantized fp8 A) arm — descriptor_needs_prequant_pruner keeps the M<threshold inline arm
+    # pointer-only. The N tile is decoupled from the scale granularity.
+    get_accelerator_autotuning_configs(
+        warp_spec=True,
+        tune_block_m=True,
+        tune_block_n=True,
+        a_memory_modes=("descriptor", "pointer"),
+        b_memory_modes=("descriptor", "pointer"),
+        pre_hook=rebind_bd_descriptors,
+    ),
+    # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide, a distinct config space);
+    # block_n/block_k (the launch-pinned quant block) gate descriptor-box legality and the K tile.
+    ["N", "K", "m_bit_length", "GATE", "block_n", "block_k"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe
+    finite_check_args=("C", "Cs"),
+    # The gate|up stack reads via the pointer arm (a contiguous descriptor box can't span the
+    # N-apart gate/up rows), so prune the B-descriptor configs under GATE.
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            fp8_dot_warp_pruner(),
+            warp_spec_compile_guard_pruner(),
+            descriptor_needs_prequant_pruner(),
+            matched_memory_modes_pruner(),
+            descriptor_box_pruner("block_k"),
+            smem_pruner("block_k"),
+        )
+    },
+)
+@triton.jit
+def w8a8_block_static_fp8_matmul_kernel(
+    A,  # (M, K) E4M3 activations (pre-quantized against the static scale by the wrapper)
+    ADescriptor,  # host TMA descriptor over A (M, K), box (BM, block_k); read iff A_MEMORY_MODE != "pointer"
+    As,  # scalar — static per-tensor activation scale (calibration-time)
+    B,  # (N, K) FP8 weights
+    BDescriptor,  # host TMA descriptor over B (N, K), box (BN, block_k); read iff B_MEMORY_MODE != "pointer"
+    Bs,  # (N // block_n, K // block_k) weight scales (fp32 or uint8/UE8M0)
+    C,  # (M, N) output
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    # Shape
+    M,
+    N,
+    K,
+    m_bit_length,  # autotune key only (log2 M bucket); unused in body
+    # Strides
+    stride_a_m,
+    stride_a_k,
+    stride_b_k,
+    stride_b_n,
+    stride_bs_k,
+    stride_bs_n,
+    stride_c_m,
+    stride_c_n,
+    stride_bias_e,
+    stride_bias_n,
+    # Weight-quantization blocks (see the block-dynamic kernel)
+    block_n: tl.constexpr,
+    block_k: tl.constexpr,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    WARP_SPEC: tl.constexpr = False,
+    A_MEMORY_MODE: tl.constexpr = "pointer",
+    B_MEMORY_MODE: tl.constexpr = "pointer",
+    # Dense gate|up fusion: B is the (2N, K) gate|up stack, C the [M, N] GLU output. Every arm
+    # folds out at compile time when GATE=False (the plain dense GEMM, unchanged).
+    GATE: tl.constexpr = False,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
+):
+    """Block-scale FP8 GEMM with static (per-tensor) activation scale.
+    ``A`` arrives pre-quantized (one elementwise ``(A / As).to(fp8)`` pass in the
+    wrapper — an inline division would repeat per N-tile). Per-block weight scales apply
+    per-K-tile during accumulation; the scalar activation scale is applied once at
+    the end.
+    """
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    n_rows = 2 * N if GATE else N
+    pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, block_k, GATE=GATE
+    )
+    a_descriptor = operand_tile_descriptor(
+        ADescriptor, A, M, K, stride_a_m, stride_a_k, BLOCK_SIZE_M, block_k, A_MEMORY_MODE
+    )
+    b_descriptor = operand_tile_descriptor(
+        BDescriptor, B, N, K, stride_b_n, stride_b_k, BLOCK_SIZE_N, block_k, B_MEMORY_MODE, GATE,
+    )
+    a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, True)
+    # Weight-scale index off the AFFINE column offset + a bounds mask (the %-wrapped offs_bn
+    # from swizzle_offsets drives only the pointer operand tile); an affine gather avoids the
+    # non-affine scale read the wrap induces, matching the block-dynamic kernel above.
+    b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, B_MEMORY_MODE)
+    offs_bn_lin = pid_n * n_width + tl.arange(0, n_width)
+    bs_ptrs = Bs + (offs_bn_lin // block_n) * stride_bs_n
+    bs_mask = offs_bn_lin < n_rows
+    a_s_static = tl.load(As)
+    accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, False)
+    for k in tl.range(0, tl.cdiv(K, block_k), warp_specialize=WARP_SPEC):
+        a, _as = load_act_static(
+            a_ptrs, a_descriptor, pid_m * BLOCK_SIZE_M, k * block_k, None, 0, a_s_static,
+            A_MEMORY_MODE, False,
+        )
+        b, b_s = load_weight_static(
+            b_ptrs, b_descriptor, bs_ptrs, bs_mask, 0, 0, pid_n * n_width, k * block_k, 0, 0,
+            GATE, False, B_MEMORY_MODE, False, BLOCK_SIZE_N, block_k,
+        )
+        accumulator = static_dot(accumulator, a, b, b_s, False, block_k, False)
+        a_ptrs, _asp, b_ptrs, bs_ptrs, _, _ = advance_ptrs(
+            a_ptrs, a_ptrs, b_ptrs, bs_ptrs, b_ptrs, bs_ptrs,
+            block_k * stride_a_k, 0, block_k * stride_b_k, stride_bs_k,
+            A_MEMORY_MODE, B_MEMORY_MODE, False, True, False,
+        )
+    accumulator = accumulator * a_s_static
+    if GATE:
+        out_row = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+        gemm_epilogue(
+            C, None, accumulator, out_row, pid_n, pid_m, out_row < M,
+            stride_c_m, stride_c_n, 0, 0,
+            BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, block_k,
+            ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+            COMPUTE_MODE="dot", N_COLS=N,
+            Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n,
+        )
+    else:
+        accumulator = add_bias(
+            accumulator, Bias, stride_bias_e, stride_bias_n, 0, pid_n,
+            BLOCK_SIZE_N, N, False,
+        )
+        store_masked(
+            C,
+            accumulator,
+            pid_m,
+            pid_n,
+            M,
+            N,
+            stride_c_m,
+            stride_c_n,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+        )
+@bayesian_autotune(
+    # tune_block_m: BLOCK_SIZE_M becomes a config axis, so the tuner sizes the M tile
+    # per workload — small at decode, large at prefill.
+    # scalar is in the mode set for M=1 decode (attn projection): it avoids the MMA M->16 pad
+    # that held MXFP8 attn decode 35% over block-dynamic at identical weight bytes; fp4-scalar
+    # is dropped by the pruner (ALU-bound unpack). swap_ab emits the BM=1 swap decode rows,
+    # scoped by mx_2d_swap_scope_pruner to E4M3-scale (NVFP4) decode — its ONLY native M=1
+    # arm (dot_scaled has no sub-M=128 E4M3 form; scalar is unpack-bound: 60-142µs vs swap's
+    # 22-50µs on the GLM attn shape, 2026-08-05). UE8M0 formats keep swap OFF: an 18-cell
+    # forced-swap sweep (cudagraph, M3+dsv4 decode) showed swap losing there (M3 attn −38%).
+    get_accelerator_autotuning_configs(
+        mx=True,
+        tune_block_nk=True,
+        compute_modes=("dot_scaled", "dot", "scalar"),
+        tune_block_m=True,
+        swap_ab=True,
+        # host-TMA on both operands, tuner-picked vs pointers (measured +14-17pp: mxfp8
+        # 54->71%, W4A8 50->67% via mixed A=ptr/B=desc, W4A4 +14%). Mixed A/B modes KEPT
+        # (unlike dense bd): A=pointer + B=descriptor wins the mid-M band.
+        a_memory_modes=("descriptor", "pointer"),
+        b_memory_modes=("descriptor", "pointer"),
+        pre_hook=rebind_mx_descriptors,
+        # WS tuner axis: dot_scaled + WS + TMA compiles + wins at prefill (1918 -> ~1995 @M=4k);
+        # warp_spec_compile_guard_pruner keeps it to num_warps%4==0, BM>=64.
+        warp_spec=True,
+    ),
+    # the MXFP4/MXFP8 split keys itself — the tuner appends every tensor arg's dtype to
+    # its cache key (memory and disk);
+    # m_bit_length (log2 M bucket) keys the M tile — the winner keeps shifting with M well past the
+    # BM ceiling and it is NOT noise: cross-applying configs (N=K=4096) costs +62% at M=128 and +245%
+    # at M=4096 (the thin M=128 tile can't saturate the wide GEMM), so don't collapse the buckets.
+    # ACTIVATION_FORMAT keys the inline act-quant grid: A stays raw bf16 under every
+    # format below the pre-quant M threshold, so the dtype-appended key can't
+    # split W4A8 from W4A4 itself. OUTPUT_FORMAT/SWIZZLED_OUT key the requant epilogue
+    # explicitly (today the C/Cs dtypes split them incidentally; a descriptor Cs has no
+    # dtype token, so don't lean on the append).
+    ["N", "K", "m_bit_length", "ACTIVATION_FORMAT", "SWIZZLED_SCALES", "GATE", "OUTPUT_FORMAT", "SWIZZLED_OUT"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe
+    finite_check_args=("C", "Cs"),
+    # BK-within-K veto (the loop loads are unmasked) + the sm_10x dot_scaled shape guards
+    # + scalar restricted to decode-sized M (a BM=1 GEVM at prefill is TPE poison) + the
+    # gate|up stack reads via the pointer arm (the descriptor's one contiguous box can't span
+    # the N-apart gate/up rows — a (2,N,K) box is a follow-up), so prune descriptor under GATE.
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            gate_stacked_tmem_trap_pruner(),
+            swizzled_out_bm_pruner(),
+            mx_config_pruner("K"),
+            mx_2d_swap_scope_pruner(),
+            scalar_max_m_pruner("M"),
+            smem_pruner(),
+            descriptor_box_pruner("BLOCK_SIZE_K"),
+            swizzled_scale_config_pruner(),
+            warp_spec_compile_guard_pruner(),
+        )
+    },
+)
+@triton.jit
+def mx_dynamic_matmul_kernel(
+    A,  # (M, K) activations: E4M3 (pre-quantized) or raw bf16/fp16 (quantized inline)
+    ADescriptor,  # host TMA descriptor over A (M, K_bytes), box (BM, BK_bytes); read iff A_MEMORY_MODE != "pointer"
+    As,  # activation scales: SWIZZLE_32_4_4 (swizzled arm — bulk via ASDescriptor at BM=128, or this pointer for the BM<128 gather) or (M, K // 32) UE8M0 affine
+    ASDescriptor,  # host TMA descriptor over the SWIZZLE_32_4_4 A scales (BM=128 bulk load); read iff SWIZZLED_SCALES
+    B,  # (N, K) E4M3 (MXFP8) or (N, K // 2) packed E2M1 (MXFP4) weights
+    BDescriptor,  # host TMA descriptor over B (N, K_bytes), box (BN, BK_bytes); read iff B_MEMORY_MODE != "pointer"
+    Bs,  # weight scales: SWIZZLE_32_4_4 (swizzled arm — bulk via BSDescriptor at BN=128, or this pointer for the BN<128 gather) or (N, K // SCALE_GROUP_K) affine
+    BSDescriptor,  # host TMA descriptor over the SWIZZLE_32_4_4 B scales (BN=128 bulk load); read iff SWIZZLED_SCALES
+    C,  # (M, N) output (the GLU intermediate under GATE)
+    Cs,  # (M, N // group) row-major requant output scale — written iff OUTPUT_FORMAT and not SWIZZLED_OUT
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    CSDescriptor,  # SWIZZLE_32_4_4 requant-scale descriptor — written iff SWIZZLED_OUT (dummy else), like AS/BS
+    AsGlobal,  # (1,) fp32 NVFP4 activation global g_a — SOLELY normalizes the inline raw-A quant (A/g_a); read iff not None
+    AsBsGlobal,  # (1,) fp32 NVFP4 WEIGHT global g_b — recovers on the accumulator together with AsGlobal (in-register product); read iff not None
+    CsGlobal,  # (1,) fp32 NVFP4 output global (next proj's provided input_scale); normalizes the requant; read iff not None
+    # Shape
+    M,
+    N,
+    K,
+    m_bit_length,  # autotune key only (log2 M bucket); unused in body
+    # Strides
+    stride_a_m,
+    stride_a_k,
+    stride_as_m,
+    stride_as_k,  # activation-scale group-axis stride, in GROUPS (see the affine advance below):
+    # 1 for the row-major (M, K // SCALE_GROUP_K) grid, a real stride under a transposed view.
+    # Dead on the swizzled arm, which re-derives its block offsets instead of walking pointers.
+    stride_b_k,
+    stride_b_n,
+    stride_bs_k,
+    stride_bs_n,
+    stride_c_m,
+    stride_c_n,
+    stride_cs_m,  # requant output-scale strides (dead unless OUTPUT_FORMAT)
+    stride_cs_n,
+    stride_bias_e,
+    stride_bias_n,
+    stride_as_global,  # expert stride of AsGlobal (0 broadcasts a shared scalar)
+    stride_bs_global,  # expert stride of AsBsGlobal (``global_scale_stride``)
+    stride_cs_global,  # expert stride of CsGlobal (0 broadcasts a shared scalar)
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    COMPUTE_MODE: tl.constexpr,
+    A_MEMORY_MODE: tl.constexpr = "pointer",
+    B_MEMORY_MODE: tl.constexpr = "pointer",
+    ACTIVATION_FORMAT: tl.constexpr = "mxfp8",
+    SWIZZLED_SCALES: tl.constexpr = False,  # scales pre-swizzled (5D weight + matching acts); else affine/inline
+    SWIZZLED_OUT: tl.constexpr = False,  # requant emits Cs in SWIZZLE_32_4_4 (CSDescriptor); single source: wrapper
+    # Dense gate|up fusion: B is the (2N, K) gate|up stack, C the [M, N] GLU output. Every arm
+    # folds out at compile time when GATE=False (the plain dense GEMM, unchanged).
+    GATE: tl.constexpr = False,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    OUTPUT_FORMAT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
+    WARP_SPEC: tl.constexpr = False,  # tuner axis; +4% at prefill (guard: num_warps%4, BM>=64)
+    SWAP_AB: tl.constexpr = False,  # swapped decode orientation (BM=1); scoped by mx_2d_swap_scope_pruner
+):
+    """Unified MXFP4/MXFP8 (W4A8/W8A8; a ``uint8`` ``A`` is packed E2M1 — W4A4) GEMM.
+    ``C = A @ B.T``. ``A``'s dtype picks the activation form (compile-time folded):
+    pre-quantized E4M3 + UE8M0 group-32 scales (one wrapper pass — the inline
+    per-N-tile quant re-ran N//BN times per element, ~2x at prefill) vs raw bf16/fp16
+    quantized inline (small M: the GEMM is weight-bandwidth-bound with idle ALU, so
+    inline is free and a separate quant kernel only adds latency; both forms are
+    bit-exact, same group-32 boundaries). Each operand's format is its dtype (``uint8``
+    = packed E2M1, two values per byte; else E4M3). ``COMPUTE_MODE`` picks the MMA:
+    ``tl.dot_scaled`` (native M=128 scaled MMA) vs fp8 ``tl.dot`` + per-group software
+    rescale (wins at decode; FP4 unpacks E2M1->E4M3 first — lossless). 2D grid with
+    swizzle for L2 reuse.
+    """
+    # uint8 A = packed-E2M1 activations (W4A4 — the dtype IS the activation format)
+    ACT_VALUES_PER_BYTE: tl.constexpr = 2 if A.dtype.element_ty == tl.uint8 else 1
+    WEIGHT_VALUES_PER_BYTE: tl.constexpr = 2 if B.dtype.element_ty == tl.uint8 else 1
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    SCALE_COLS: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K
+    offs_ka = tl.arange(0, BLOCK_SIZE_K // ACT_VALUES_PER_BYTE)
+    offs_kb = tl.arange(0, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE)
+    # packed-fp4 weights (WEIGHT_VALUES_PER_BYTE==2) route swizzle_offsets to full L2 grouping
+    pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, WEIGHT_VALUES_PER_BYTE, GATE=GATE
+    )
+    # Scales read affine off row/col offsets + bounds masks (the %-wrapped operand offsets would
+    # make the scale load a non-affine gather); the SWIZZLED arm reads them via the SA/BS
+    # descriptors in the loop instead, so its pointer tiles are dead (base scalars + null masks).
+    as_ptrs, bs_ptrs, as_mask, bs_mask = mx_2d_scale_ptrs(
+        As, Bs, pid_m, pid_n, M, N, stride_as_m, stride_bs_n, stride_bs_k,
+        BLOCK_SIZE_M, BLOCK_SIZE_N, SCALE_COLS, SWIZZLED_SCALES, GATE,
+    )
+    # Operand tiles: activation + the (GATE-stacked-aware) weight, both single leaf calls (the
+    # descriptor-vs-pointer and gate|up-stack branches fold inside the leaves). The swap arm
+    # loads both naturally — mx_compute reorients at the dot under SWAP_AB, flattening the token tile.
+    if SWAP_AB:
+        tl.static_assert(
+            BLOCK_SIZE_M == 1, "the mx swap arm flattens the single-token tile"
+        )
+    a_ptrs = operand_tile_ptrs(A, offs_am, offs_ka, stride_a_m, stride_a_k, A_MEMORY_MODE, True)
+    b_ptrs = matmul_weight_ptrs(
+        B, offs_bn, offs_kb, stride_b_n, stride_b_k, B_MEMORY_MODE, SWAP_AB
+    )
+    accumulator = acc_init(COMPUTE_MODE if SWAP_AB else "dot", BLOCK_SIZE_M, n_width, SWAP_AB)
+    for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC):
+        a, a_s = load_act_mx(
+            a_ptrs, as_ptrs, AsGlobal, 0, None, as_mask, ADescriptor, pid_m * BLOCK_SIZE_M,
+            k * (BLOCK_SIZE_K // ACT_VALUES_PER_BYTE), ASDescriptor, As, 0, 0, pid_m, k, M, K,
+            A_MEMORY_MODE, False, False, SWIZZLED_SCALES,
+            BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K, ACTIVATION_FORMAT,
+        )
+        b, b_s = load_weight_mx(
+            b_ptrs, BDescriptor, bs_ptrs, bs_mask, BSDescriptor, Bs, 0, pid_n * n_width,
+            k * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE), 0, 0, pid_n, k, N, K,
+            0, stride_bs_n, stride_bs_k,
+            GATE, False, False, B_MEMORY_MODE, SWAP_AB, SWIZZLED_SCALES,
+            BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, WEIGHT_VALUES_PER_BYTE,
+        )
+        accumulator = mx_compute(
+            accumulator, a, a_s, b, b_s, COMPUTE_MODE,
+            n_width, BLOCK_SIZE_K, SCALE_GROUP_K, SWAP_AB,
+        )
+        a_ptrs, as_ptrs, b_ptrs, bs_ptrs, _, _ = advance_ptrs(
+            a_ptrs, as_ptrs, b_ptrs, bs_ptrs, b_ptrs, bs_ptrs,
+            (BLOCK_SIZE_K // ACT_VALUES_PER_BYTE) * stride_a_k,
+            (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_as_k,
+            (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE) * stride_b_k,
+            (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_bs_k,
+            A_MEMORY_MODE, B_MEMORY_MODE,
+            not SWIZZLED_SCALES, (not SWIZZLED_SCALES) and not GATE, False,
+        )
+    # NVFP4 two-level: block e4m3 scales rode through dot_scaled; recover the combined per-tensor
+    # one epilogue for both arms: GATE splits+SwiGLUs (+ optional requant), plain casts+stores.
+    # affine output rows (swizzle's offs_am is %-wrapped — the store scatters to real rows);
+    # N_COLS masks the 2D-dense column tail.
+    out_row = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    gemm_epilogue(
+        C, Cs, accumulator, out_row, pid_n, pid_m, out_row < M,
+        stride_c_m, stride_c_n, stride_cs_m, stride_cs_n,
+        BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, OUTPUT_FORMAT, SCALE_GROUP_K,
+        ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+        COMPUTE_MODE=COMPUTE_MODE, SWAP_AB=SWAP_AB, N_COLS=N, SWIZZLED_OUT=SWIZZLED_OUT,
+        CSDescriptor=CSDescriptor, CsGlobal=CsGlobal, stride_cs_global=stride_cs_global,
+        GlobalScale=AsBsGlobal,
+        stride_global_e=stride_bs_global,
+        GlobalScaleA=AsGlobal,
+        stride_global_a=stride_as_global,
+        global_row=0,
+        Bias=Bias,
+        stride_bias_e=stride_bias_e,
+        stride_bias_n=stride_bias_n,
+    )
+@bayesian_autotune(
+    # weight-only is a plain bf16 dot (the fp4/fp8 weight is upcast per-group in-loop), tuned like the
+    # full-precision kernel — tile + WS — PLUS a host-TMA memory-mode axis on both operands
+    # (tuner-routed vs pointers, like mx_dynamic): the descriptor box feeds the wide-N prefill tiles
+    # the pointer arm starves; pointer wins decode / short-K. Scales stay affine (pointer) — weight-only
+    # never swizzles — so the descriptor covers only the operand values. BK spans any number of
+    # groups (the dequant is in-tile), so no BK==group constraint.
+    get_accelerator_autotuning_configs(
+        tune_block_nk=True,
+        tune_block_m=True,
+        warp_spec=True,
+        compute_modes=("dot", "dot_scaled"),
+        a_memory_modes=("descriptor", "pointer"),
+        b_memory_modes=("descriptor", "pointer"),
+        pre_hook=rebind_weight_only_descriptors,
+    ),
+    ["N", "K", "m_bit_length", "GATE"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe
+    finite_check_args=("C", "Cs"),
+    # WS compile guard + descriptor TMA-box limits (256/dim) + smem fit; the gate|up stack reads via
+    # the pointer arm (a contiguous box can't span the N-apart gate/up rows), so prune B-descriptor
+    # under GATE.
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            block_within_dim_pruner("K"),
+            mx_config_pruner("K", "N"),  # dot_scaled shape gates
+            warp_spec_compile_guard_pruner(),
+            descriptor_box_pruner("BLOCK_SIZE_K"),
+            smem_pruner("BLOCK_SIZE_K"),
+        )
+    },
+)
+@triton.jit
+def mx_weight_only_matmul_2d_kernel(
+    A,  # (M, K) raw BF16/FP16 activations — NOT quantized
+    ADescriptor,  # host TMA descriptor over A (M, K), box (BM, BK); read iff A_MEMORY_MODE != "pointer"
+    B,  # (N, K[/2]) MXFP4/NVFP4 (packed E2M1) / MXFP8 (E4M3) weights; under GATE the (2N, ...) gate|up stack
+    BDescriptor,  # host TMA descriptor over B (N, K_bytes), box (BN, BK_bytes); read iff B_MEMORY_MODE != "pointer"
+    Bs,  # (N, K // SCALE_GROUP_K) group scales — UE8M0 (MX) or E4M3 (NVFP4)
+    BsGlobal,  # scalar fp32 NVFP4 per-tensor global — recovers on the accumulator; read iff not None
+    C,  # (M, N) output (GLU intermediate under GATE)
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    M,
+    N,
+    K,
+    m_bit_length,  # autotune key only (log2 M bucket)
+    stride_a_m,
+    stride_a_k,
+    stride_b_n,
+    stride_b_k,
+    stride_bs_n,
+    stride_bs_k,
+    stride_c_m,
+    stride_c_n,
+    stride_bias_e,
+    stride_bias_n,
+    stride_bs_global,  # expert stride of BsGlobal (``global_scale_stride``)
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    WEIGHT_VALUES_PER_BYTE: tl.constexpr,
+    A_MEMORY_MODE: tl.constexpr = "pointer",
+    B_MEMORY_MODE: tl.constexpr = "pointer",
+    COMPUTE_MODE: tl.constexpr = "dot",
+    WARP_SPEC: tl.constexpr = False,
+    GATE: tl.constexpr = False,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
+):
+    """weight-only dense matmul: ``C = A @ B.T`` with a RAW bf16 activation and an MXFP4/MXFP8 weight
+    upcast to bf16 in-loop (unpack E2M1->E4M3, per-group group-scale multiply), then a plain bf16
+    ``tl.dot``. This is the ``matmul_ogs`` format (fp4/fp8 weight bytes kept in memory, bf16 acts,
+    weight upcast in-MMA) — no activation quant, no ``dot_scaled``. Pointer/affine scales only.
+    ``GATE`` loads the stacked (2N, K) gate|up weight as one ``[BK, 2*BN]`` dot + applies the GLU."""
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    n_rows = 2 * N if GATE else N
+    SCALE_COLS: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K
+    offs_kb = tl.arange(0, BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE)
+    pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, WEIGHT_VALUES_PER_BYTE, GATE=GATE
+    )
+    # Only the weight scale is needed (bf16 acts carry no scale). Build it N-clamped + UNMASKED
+    # (bs_mask=None): valid columns (offs_bn_lin < N) read their real scale; the tail reads a clamped
+    # in-bounds scale that the epilogue masks off (N_COLS / store_masked). Dropping the other=0.0
+    # bounds mask is what lets the K-loop warp-specialization-lower (the masked load's fill constant
+    # is the WS partitioner blocker). GATE routes its stacked gate|up scale through the per-expert
+    # leaf, not these ptrs, so this only governs the plain dense arm.
+    offs_bn_lin = pid_n * n_width + tl.arange(0, n_width)
+    bs_ptrs = Bs + (
+        tl.minimum(offs_bn_lin, n_rows - 1)[:, None] * stride_bs_n
+        + tl.arange(0, SCALE_COLS)[None, :] * stride_bs_k
+    )
+    bs_mask = None
+    a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, A_MEMORY_MODE, True)
+    b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_kb, stride_b_n, stride_b_k, B_MEMORY_MODE)
+    accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, False)
+    for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC):
+        # raw bf16 [BM, BK]: pointer tile or the host-TMA box (no scale — the second slot is dead;
+        # a distinct name, not `_`, which is already bound pre-loop by mx_2d_scale_ptrs)
+        a, a_scale_dead = load_act_plain(
+            a_ptrs, ADescriptor, pid_m * BLOCK_SIZE_M, k * BLOCK_SIZE_K, None, 0, A_MEMORY_MODE, False
+        )
+        b, b_s = load_weight_mx(
+            b_ptrs, BDescriptor, bs_ptrs, bs_mask, None, Bs, 0, pid_n * n_width,
+            k * (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE), 0, 0, pid_n, k, N, K,
+            0, stride_bs_n, stride_bs_k,
+            GATE, False, False, B_MEMORY_MODE, False, False,
+            BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, WEIGHT_VALUES_PER_BYTE,
+        )
+        accumulator = mx_weight_only_compute(
+            accumulator, a, b, b_s, COMPUTE_MODE, BLOCK_SIZE_K, n_width, SCALE_GROUP_K
+        )
+        # descriptor arms read by absolute box offset (m/k above); the pointer arms advance here
+        a_ptrs += BLOCK_SIZE_K * stride_a_k
+        b_ptrs += (BLOCK_SIZE_K // WEIGHT_VALUES_PER_BYTE) * stride_b_k
+        bs_ptrs += SCALE_COLS * stride_bs_k
+    if GATE:
+        out_row = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+        gemm_epilogue(
+            C, None, accumulator, out_row, pid_n, pid_m, out_row < M,
+            stride_c_m, stride_c_n, 0, 0,
+            BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, BLOCK_SIZE_K,
+            ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+            COMPUTE_MODE="dot", N_COLS=N,
+            GlobalScale=BsGlobal, stride_global_e=stride_bs_global,
+            Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n,
+        )
+    else:
+        # explicit (not the epilogue fold): this arm exits through store_masked, which has no
+        # GlobalScale slot
+        accumulator = apply_global_scale(
+            accumulator, BsGlobal, stride_bs_global, None, 0, 0
+        )
+        accumulator = add_bias(
+            accumulator, Bias, stride_bias_e, stride_bias_n, 0, pid_n,
+            BLOCK_SIZE_N, N, False,
+        )
+        store_masked(
+            C, accumulator, pid_m, pid_n, M, N, stride_c_m, stride_c_n,
+            BLOCK_SIZE_M, BLOCK_SIZE_N,
+        )
+@bayesian_autotune(
+    # Same 2D-swizzle loop physics as the tensor kernel — tune the tile, WS a pure perf axis
+    # (compile-guarded), no memory-mode/scale axes (unquantized, no scales anywhere).
+    get_accelerator_autotuning_configs(tune_block_nk=True, warp_spec=True, tune_block_m=True),
+    # GATE keys the gate|up arm separately (its stacked dot is 2*BN wide).
+    ["N", "K", "m_bit_length", "GATE"],
+    n_trials=100,
+    path_anchor_axes=PATH_ANCHOR_AXES,
+    # 2D outputs are fully written (no sentinel/padding contract), so the numerics veto is safe
+    finite_check_args=("C", "Cs"),
+    prune_configs_by={
+        "early_config_prune": compose_pruners(
+            block_within_dim_pruner("K"),
+            warp_spec_compile_guard_pruner(),
+        )
+    },
+)
+@triton.jit
+def full_precision_matmul_2d_kernel(
+    A,  # (M, K) BF16/FP16 activations
+    B,  # (N, K) weights in A's dtype; under GATE the (2N, K) gate|up stack
+    C,  # (M, N) output; under GATE the GLU intermediate
+    Bias,  # (E, N_out) per-expert output bias, N_out = 2N under GATE; read iff not None
+    # Shape
+    M,
+    N,
+    K,
+    m_bit_length,  # autotune key only (log2 M bucket); unused in body
+    # Strides
+    stride_a_m,
+    stride_a_k,
+    stride_b_k,
+    stride_b_n,
+    stride_c_m,
+    stride_c_n,
+    stride_bias_e,
+    stride_bias_n,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    WARP_SPEC: tl.constexpr = False,
+    # Dense gate|up fusion: B is the (2N, K) gate|up stack, C the [M, N] GLU output. Folds out at
+    # compile time when GATE=False (the plain dense GEMM).
+    GATE: tl.constexpr = False,
+    ACT_FN: tl.constexpr = "silu",
+    SWIGLU_ALPHA: tl.constexpr = None,
+    SWIGLU_LIMIT: tl.constexpr = None,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
+):
+    """Full-precision (BF16/FP16) 2D matmul — the single-GEMM sibling of
+    ``full_precision_matmul_grouped/batched``: plain ``tl.dot``, fp32 accumulation, no scales.
+    ``GATE`` loads the stacked (2N, K) gate|up weight as one ``[BK, 2*BN]`` dot and applies the
+    ``ACT_FN``/SwiGLU GLU; ``GATE=False`` is the plain dense GEMM. 2D grid with swizzle for L2 reuse."""
+    n_width: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    pid_m, pid_n, offs_am, offs_bn, offs_k = swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GATE=GATE
+    )
+    a_ptrs = operand_tile_ptrs(A, offs_am, offs_k, stride_a_m, stride_a_k, "pointer", True)
+    b_ptrs = matmul_weight_ptrs(B, offs_bn, offs_k, stride_b_n, stride_b_k, "pointer")
+    accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, False)
+    for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC):
+        a = tl.load(a_ptrs)
+        w, _ = load_weight_plain(
+            b_ptrs, b_ptrs, 0, pid_n * n_width, k * BLOCK_SIZE_K,
+            GATE, False, "pointer", False, BLOCK_SIZE_N, BLOCK_SIZE_K,
+        )
+        accumulator = accumulator + fp8_dot(a, w, False, BLOCK_SIZE_K)
+        a_ptrs += BLOCK_SIZE_K * stride_a_k
+        b_ptrs += BLOCK_SIZE_K * stride_b_k
+    if GATE:
+        out_row = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+        gemm_epilogue(
+            C, None, accumulator, out_row, pid_n, pid_m, out_row < M,
+            stride_c_m, stride_c_n, 0, 0,
+            BLOCK_SIZE_M, BLOCK_SIZE_N, GATE, None, BLOCK_SIZE_K,
+            ACT_FN, SWIGLU_ALPHA, SWIGLU_LIMIT, SIMULATE_UNFUSED, INTERMEDIATE_DTYPE,
+            COMPUTE_MODE="dot", N_COLS=N,
+            Bias=Bias, stride_bias_e=stride_bias_e, stride_bias_n=stride_bias_n,
+        )
+    else:
+        accumulator = add_bias(
+            accumulator, Bias, stride_bias_e, stride_bias_n, 0, pid_n,
+            BLOCK_SIZE_N, N, False,
+        )
+        store_masked(
+            C, accumulator, pid_m, pid_n, M, N, stride_c_m, stride_c_n,
+            BLOCK_SIZE_M, BLOCK_SIZE_N,
+        )
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("w8a8_block_dynamic_fp8_matmul"),
+    mutates_args=(),
+    opaque=True,
+)
+def w8a8_block_dynamic_fp8_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    block_size: list[int],
+    output_dtype: torch.dtype | None = None,
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    bias: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    """Block-scale FP8 matmul: ``C = A @ B.T``; activations quantized offline in one pass.
+    A:  (..., K) raw activations, bf16/fp16/fp32 (quantized to FP8 in one wrapper pass)
+    B:  (N, K) FP8 weights — under ``gate`` the ``(2N, K)`` gate|up stack (gate rows [0,N), up [N,2N))
+    Bs: (N // block_n, K // block_k) per-block weight scales (2N rows under ``gate``)
+    ``gate`` fuses the gate|up projection into one stacked GEMM + SwiGLU, returning the
+    ``[..., N]`` GLU intermediate (``output_dtype``); the fused-requant output formats stay on
+    the MX path. Returns a one-element list (mirrors the MX/grouped/batched op signature).
+    """
+    assert len(block_size) == 2, (
+        f"block_size must be [block_n, block_k], got {block_size}"
+    )
+    block_n, block_k = block_size[0], block_size[1]
+    validate_dense_2d_operands(A, B)
+    rows, K = B.shape
+    # Under gate|up fusion B is the (2N, K) gate|up stack; N is the per-projection output width.
+    N = rows // 2 if gate else rows
+    M = A.numel() // A.shape[-1]
+    assert K % block_k == 0, f"K ({K}) must be divisible by block_k ({block_k})"
+    assert not gate or N % block_n == 0, (
+        f"gate|up fusion needs N ({N}) divisible by block_n ({block_n}) — the up-proj scale block"
+    )
+    assert Bs.ndim == 2, f"Bs must be 2D (rows//block_n, K//block_k), got ndim={Bs.ndim}"
+    assert Bs.shape == (triton.cdiv(rows, block_n), K // block_k), (
+        f"Bs shape {tuple(Bs.shape)} != expected ({triton.cdiv(rows, block_n)}, {K // block_k})"
+    )
+    bs_u8 = ue8m0_as_uint8(Bs)
+    # UE8M0 weight scales are the DeepGEMM-Blackwell format — quantize activations to UE8M0
+    # too so the kernel folds both group scales into the tcgen05 dot_scaled MMA (else fp32).
+    A_q, A_s = fp8_act_quant_block_dynamic(
+        A.view(M, K), block_k, use_ue8m0=bs_u8.dtype == torch.uint8
+    )
+    C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype)
+    # Host TMA descriptors over the (M, K) / (N, K) matrices — the placeholder box is
+    # re-bound per tuned config by rebind_bd_descriptors. Read only by the descriptor
+    # (host-TMA) configs the tuner picks for wide-N prefill; pointer configs never touch
+    # them (the constexpr arm folds out in the load helpers).
+    a_descriptor = TensorDescriptor.from_tensor(A_q, [1, block_k])
+    b_descriptor = maybe_descriptor(B, [1, block_k])
+    def grid(META):
+        return (
+            triton.cdiv(M, META["BLOCK_SIZE_M"]),
+            triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        )
+    with device_context(A.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(w8a8_block_dynamic_fp8_matmul_kernel)[grid](
+            A_q,
+            a_descriptor,
+            A_s,
+            B,
+            b_descriptor,
+            bs_u8,
+            C,
+            bias,
+            M,
+            N,
+            K,
+            int(M).bit_length(),
+            A_q.stride(0),
+            A_q.stride(1),
+            A_s.stride(0),
+            A_s.stride(1),
+            B.stride(1),
+            B.stride(0),
+            bs_u8.stride(1),
+            bs_u8.stride(0),
+            C.stride(-2),
+            C.stride(-1),
+            bias_stride_e,
+            bias_stride_n,
+            block_n=block_n,
+            block_k=block_k,
+            GATE=gate,
+            ACT_FN=act_fn,
+            SWIGLU_ALPHA=swiglu_alpha,
+            SWIGLU_LIMIT=swiglu_limit,
+            SIMULATE_UNFUSED=simulate_unfused,
+            INTERMEDIATE_DTYPE=tl_dtype(resolve_output_dtype(output_dtype, A, None)),
+        )
+    return [C]
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("w8a8_block_static_fp8_matmul"),
+    mutates_args=(),
+    opaque=True,
+)
+def w8a8_block_static_fp8_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor,
+    Bs: torch.Tensor,
+    block_size: list[int],
+    output_dtype: torch.dtype | None = None,
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    bias: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    """Block-scale FP8 matmul with static (per-tensor) activation quantization.
+    A:  (..., K) raw bf16/fp16 activations — pre-quantized against ``As`` in the wrapper
+    B:  (N, K) FP8 weights — under ``gate`` the ``(2N, K)`` gate|up stack
+    As: scalar / (1,) — per-tensor static activation scale
+    Bs: (N // block_n, K // block_k) per-block weight scales (2N rows under ``gate``)
+    ``gate`` fuses the gate|up projection into one stacked GEMM + SwiGLU, returning the
+    ``[..., N]`` GLU intermediate. Returns a one-element list (mirrors the MX/grouped op).
+    """
+    assert len(block_size) == 2, (
+        f"block_size must be [block_n, block_k], got {block_size}"
+    )
+    block_n, block_k = block_size[0], block_size[1]
+    assert B.dtype != torch.int8, (
+        "static activation quant is not supported on the FP4 path"
+    )
+    assert not (block_n == B.size(0) and block_k == B.size(1)), (
+        "static activation quant requires block-wise weights, not tensor-mode"
+    )
+    validate_dense_2d_operands(A, B)
+    assert As.numel() == 1, f"As must be scalar or (1,), got {tuple(As.shape)}"
+    rows, K = B.shape
+    # Under gate|up fusion B is the (2N, K) gate|up stack; N is the per-projection output width.
+    N = rows // 2 if gate else rows
+    M = A.numel() // A.shape[-1]
+    assert Bs.ndim == 2, f"Bs must be 2D (rows//block_n, K//block_k), got ndim={Bs.ndim}"
+    assert K % block_k == 0, f"K ({K}) must be divisible by block_k ({block_k})"
+    assert not gate or N % block_n == 0, (
+        f"gate|up fusion needs N ({N}) divisible by block_n ({block_n}) — the up-proj scale block"
+    )
+    assert Bs.shape == (triton.cdiv(rows, block_n), K // block_k), (
+        f"Bs shape {tuple(Bs.shape)} != expected ({triton.cdiv(rows, block_n)}, {K // block_k})"
+    )
+    def grid(META):
+        return (
+            triton.cdiv(M, META["BLOCK_SIZE_M"]),
+            triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        )
+    bs_u8 = ue8m0_as_uint8(Bs)
+    C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype)
+    As = As.reshape(1).to(torch.float32)
+    # M-gated static pre-quant (bit-exact with the inline arm: same scalar, same cast);
+    # like MX, the inline form is cheap elementwise work — at M=1 a separate kernel is
+    # pure added latency. The kernel picks its arm off A's dtype.
+    A_q, _ = maybe_act_quant(
+        A.view(M, K),
+        lambda x: ((x.to(torch.float32) / As).to(FP8_DTYPE), As),
+        STATIC_MATMUL_ACT_PREQUANT_MIN_M,
+    )
+    # Host-TMA descriptors over (M, K)/(N, K); placeholder boxes re-bound per config by
+    # rebind_bd_descriptors. Read only by the descriptor configs the tuner picks for wide-N
+    # prefill (offline fp8 A); pointer/inline configs never touch them.
+    a_descriptor = TensorDescriptor.from_tensor(A_q, [1, block_k])
+    b_descriptor = maybe_descriptor(B, [1, block_k])
+    with device_context(A.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(w8a8_block_static_fp8_matmul_kernel)[grid](
+            A_q,
+            a_descriptor,
+            As,
+            B,
+            b_descriptor,
+            bs_u8,
+            C,
+            bias,
+            M,
+            N,
+            K,
+            int(M).bit_length(),  # m_bit_length key bucket
+            A_q.stride(0),
+            A_q.stride(1),
+            B.stride(1),
+            B.stride(0),
+            bs_u8.stride(1),
+            bs_u8.stride(0),
+            C.stride(-2),
+            C.stride(-1),
+            bias_stride_e,
+            bias_stride_n,
+            # Meta-parameters (BM and BN come from the config; BK is the caller's
+            # block_k — see the block-dynamic kernel)
+            block_n=block_n,
+            block_k=block_k,
+            GATE=gate,
+            ACT_FN=act_fn,
+            SWIGLU_ALPHA=swiglu_alpha,
+            SWIGLU_LIMIT=swiglu_limit,
+            SIMULATE_UNFUSED=simulate_unfused,
+            INTERMEDIATE_DTYPE=tl_dtype(resolve_output_dtype(output_dtype, A, None)),
+        )
+    return [C]
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("w8a8_tensor_dynamic_fp8_matmul"),
+    mutates_args=(),
+    opaque=True,
+)
+def w8a8_tensor_dynamic_fp8_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    bias: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    """Tensor-scale FP8 matmul: ``C = A @ B.T``; activations quantized offline per row.
+    A:  (..., K) raw activations, bf16/fp16/fp32 (flattened to (M, K)
+        internally) — per-row scales computed via ``fp8_act_quant_tensor_wide(A, K)``.
+    B:  (N, K) FP8 weights — under ``gate`` the ``(2N, K)`` gate|up stack (one per-tensor scale).
+    Bs: scalar, (1,), or (1, 1) — single tensor-scale weight scale.
+    ``gate`` fuses the gate|up projection into one stacked GEMM + SwiGLU, returning the
+    ``[..., N]`` GLU intermediate. Returns a one-element list (mirrors the MX/grouped op).
+    """
+    validate_dense_2d_operands(A, B)
+    rows, K = B.shape
+    # Under gate|up fusion B is the (2N, K) gate|up stack; N is the per-projection output width.
+    N = rows // 2 if gate else rows
+    M = A.numel() // A.shape[-1]
+    assert Bs.numel() == 1, f"Bs must be scalar or (1,), got {tuple(Bs.shape)}"
+    # Per-row scalar activation scale (one per token).
+    qA, As = fp8_act_quant_tensor_wide(A, K)
+    As = As.reshape(M)
+    Bs = Bs.reshape(1)
+    C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype)
+    def grid(META):
+        return (
+            triton.cdiv(M, META["BLOCK_SIZE_M"]),
+            triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        )
+    with device_context(A.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(w8a8_tensor_dynamic_fp8_matmul_kernel)[grid](
+            qA,
+            As,
+            B,
+            Bs,
+            C,
+            bias,
+            M,
+            N,
+            K,
+            int(M).bit_length(),  # m_bit_length key bucket
+            qA.stride(-2),
+            qA.stride(-1),
+            As.stride(0),
+            B.stride(1),
+            B.stride(0),
+            C.stride(-2),
+            C.stride(-1),
+            bias_stride_e,
+            bias_stride_n,
+            GATE=gate,
+            ACT_FN=act_fn,
+            SWIGLU_ALPHA=swiglu_alpha,
+            SWIGLU_LIMIT=swiglu_limit,
+            SIMULATE_UNFUSED=simulate_unfused,
+            INTERMEDIATE_DTYPE=tl_dtype(resolve_output_dtype(output_dtype, A, None)),
+        )
+    return [C]
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("mx_dynamic_matmul"), mutates_args=(), opaque=True
+)
+def mx_dynamic_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor | None,
+    Bs: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+    activation_format: str | None = None,
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    quantize_output: bool = False,
+    a_global_scale: torch.Tensor | None = None,
+    b_global_scale: torch.Tensor | None = None,
+    output_global_scale: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    """MX/NVFP4 matmul ``C = A @ B.T``; activations quantized offline above the
+    ``maybe_act_quant`` M threshold, inline below it. ``activation_format`` sets the
+    activation grid — E4M3 (``"mxfp8"``, the default) or packed E2M1 (``"mxfp4"``,
+    W4A4); NVFP4 weights (E4M3 scales) pin ``"nvfp4"``. Weight format detected from
+    ``B.dtype``: ``int8`` → packed E2M1 (``B`` is ``(N, K//2)``); ``float8_e4m3fn`` →
+    unpacked E4M3. The scale dtype carries the family: UE8M0 = MX group-32, E4M3 =
+    NVFP4 group-16 (``Bs`` is ``(N, K//group)``); tile + dot path are autotuned.
+    A:  (..., K) raw activations, bf16/fp16/fp32;
+        leading dims are flattened to (M, K) and restored on the output
+    """
+    assert B.ndim == 2 and Bs.ndim in (2, 5, 6)  # 5D = pre-swizzled; 6D = the gate|up artifact
+    assert B.dtype in (torch.int8, torch.float8_e4m3fn), (
+        f"B must be int8 (packed E2M1) or float8_e4m3fn (E4M3), got {B.dtype}"
+    )
+    assert A.is_contiguous(), "A must be contiguous"
+    # B may be row-major (the forward weight) or a transposed view of one: the kernel reads
+    # both of B's strides, and `descriptor_box_pruner` keeps the TMA arms — which need a
+    # unit-stride innermost dim — off the transposed form. This is what lets a backward pass
+    # contract over N by handing in `B.t()` instead of materializing a transposed copy.
+    assert B.is_contiguous() or B.t().is_contiguous(), (
+        "B must be contiguous or a transposed view of a contiguous tensor, got strides "
+        f"{tuple(B.stride())} for shape {tuple(B.shape)}"
+    )
+    # cuBLAS block-scaled fast path — INSIDE the op so every caller gets it (dispatchers,
+    # direct op calls, prequantized-As callers) and so autograd keeps working: register_autograd
+    # hangs on this op and saves INPUTS, so the registered dgrad serves a scaled_mm forward
+    # unchanged. Measured 1.4-4.8x over the Triton kernel at every M (B200, 2026-08-27); a None
+    # falls through to the Triton launch exactly as before. Skipped under fake/meta propagation
+    # (the opaque op's fake impl runs this body with launches disabled — the shape-correct C
+    # allocation below is all fake mode needs).
+    from . import compat as _compat
+    if (
+        not gate and not quantize_output and bias is None and not simulate_unfused
+        and not _compat._SKIP_LAUNCHES_MIRROR
+    ):
+        _smm = _torch_scaled_mm_2d(
+            A, B, As, Bs, activation_format, a_global_scale, b_global_scale,
+            resolve_output_dtype(output_dtype, A, None),
+        )
+        if _smm is not None:
+            return [_smm]
+    WEIGHT_VALUES_PER_BYTE = NIBBLES_PER_BYTE if B.dtype == torch.int8 else 1
+    rows, K_b = B.shape
+    K = A.shape[-1]
+    M = A.numel() // K
+    # Under gate|up fusion B is the (2N, K) gate|up stack; N is the per-projection output width.
+    N = rows // 2 if gate else rows
+    assert K == WEIGHT_VALUES_PER_BYTE * K_b, (
+        f"K (={K}) must equal {WEIGHT_VALUES_PER_BYTE} * B.shape[1] (={K_b})"
+    )
+    # Weight scales arrive row-major (rows, K // scale_group) — read affine, no gain swizzling only
+    # one operand — or already SWIZZLE_32_4_4 (5D), swizzled once at load: the deployment contract,
+    # the op never swizzles in the hot path. The format is the scale dtype (E4M3 = NVFP4 group-16,
+    # UE8M0 = MX group-32). SWIZZLED_SCALES governs both operands (a swizzled weight is paired with
+    # swizzled acts). Callers wanting the tcgen05 fast path pre-swizzle the weight (5D).
+    swizzled_scales = Bs.ndim == 5
+    scale_group = mx_scale_family(Bs, K)
+    if not swizzled_scales:
+        assert Bs.shape == (rows, K // scale_group), (
+            f"Bs shape {tuple(Bs.shape)} != ({rows}, {K // scale_group})"
+        )
+    else:
+        # the artifact's row blocks must match THIS weight (one matrix: leading dim 1) — a
+        # wrong-layer artifact with a matching K would otherwise dequantize with garbage scales silently
+        assert Bs.shape[:2] == (1, triton.cdiv(rows, 128)), (
+            f"swizzled Bs carries {tuple(Bs.shape[:2])} (matrices, 128-row blocks), expected "
+            f"{(1, triton.cdiv(rows, 128))} for this ({rows}, K) weight — wrong artifact"
+        )
+    b_u8 = e2m1_as_uint8(B)
+    bs_u8 = ue8m0_as_uint8(Bs)  # caller's layout (5D swizzled / 2D affine); the op never swizzles
+    activation_format = resolve_activation_format(activation_format, B, Bs)
+    # resolve BEFORE the C alloc: with a pre-quantized A (As given) the default output is
+    # bf16, not A's storage dtype (fp8 / packed int8 are not valid output element types)
+    output_dtype = resolve_output_dtype(output_dtype, A, As)
+    # Activation quant is always maybe_act_quant — offline above the M threshold, inline in the
+    # kernel below it (fp4 packs in-register). The offline arm writes swizzled scales directly when
+    # the weight is swizzled (fused, no post-quant pass), else affine. Inline acts (small M) stay
+    # affine in-register even under a swizzled weight — the dot_scaled reads the mixed layout fine.
+    # NVFP4 two-level: a_global_scale (the calibrated activation global) normalizes the quant on both
+    # arms — the offline kernel divides before the block quant, the inline arm via AsGlobal.
+    act_global_scale = normalize_global_scale(a_global_scale, 1)  # g_a; folds in-kernel
+    if act_global_scale is not None:
+        assert activation_format == "nvfp4", "an activation global is NVFP4-only"
+        act_quant = lambda a: MX_ACT_QUANT[activation_format](  # noqa: E731
+            a, swizzled=swizzled_scales, global_scale=act_global_scale
+        )
+    elif swizzled_scales:
+        act_quant = lambda a: MX_ACT_QUANT[activation_format](a, swizzled=True)  # noqa: E731
+    else:
+        act_quant = MX_ACT_QUANT[activation_format]
+    # NVFP4 accumulator correction: the g_a·g_b product folded onto the fp32 accumulator
+    input_global_scale = normalize_global_scale(b_global_scale, 1)  # g_b
+    out_global_scale = normalize_global_scale(output_global_scale, 1)
+    # As given ⇒ A is already quantized (the routed-op parity: a pre-quantized activation + its
+    # scales); else quantize raw A (offline above the M threshold, inline in the kernel below it).
+    if As is not None:
+        assert (As.dtype == torch.float8_e4m3fn) == (Bs.dtype == torch.float8_e4m3fn), (
+            f"activation scales ({As.dtype}) must match the weight scale family ({Bs.dtype})"
+        )
+        assert swizzled_scales or As.ndim != 5, (
+            "un-swizzled weights pair with affine (row-major) activation scales, got 5D As"
+        )
+        a_vals, as_scales = A.view(M, K), As
+        # one SWIZZLED_SCALES flag governs both operands in-kernel: row-major activation
+        # scales under a swizzled weight are swizzled here (the grouped wrapper does the same)
+        if swizzled_scales and As.ndim == 2:
+            as_scales = swizzle_mx_scales(As)
+    else:
+        a_vals, as_scales = maybe_act_quant(A.view(M, K), act_quant, MX_MATMUL_ACT_PREQUANT_MIN_M)
+    A_q = e2m1_as_uint8(a_vals)
+    as_u8 = ue8m0_as_uint8(as_scales)
+    # the act-scale layout is carried by the buffer itself: 5D = swizzled (the offline arm under
+    # a swizzled weight, or pre-swizzled As); the inline arm's raw activation stays 2D (affine)
+    act_swizzled = as_u8.ndim == 5
+    # gate|up output: the GLU intermediate [M, N] (output_dtype), or — under ``quantize_output`` —
+    # requantized (mxfp8: fp8 + (M, N//32) uint8 Cs; fp4: packed (M, N//2) + (M, N//out_g) scale).
+    output_format = activation_format if quantize_output and gate else None
+    requant = output_format is not None
+    out_g = 16 if output_format == "nvfp4" else 32
+    cs_dtype = torch.float8_e4m3fn if output_format == "nvfp4" else torch.uint8
+    # Swizzled in -> swizzled out: a swizzled block requants Cs straight into the down's
+    # SWIZZLE_32_4_4 layout (CSDescriptor), so the fused down reads it on the fast path. 2D dense
+    # output rows are contiguous (no scatter), so it's always swizzle-able. Format-general — the
+    # swizzle is a byte-tiling over the N // out_g scale grid (mxfp8/mxfp4 group-32, nvfp4 group-16).
+    swizzled_out = requant and swizzled_scales
+    if output_format in ("mxfp4", "nvfp4"):
+        C = A.new_empty(A.shape[:-1] + (N // 2,), dtype=torch.int8)
+    elif output_format == "mxfp8":
+        C = A.new_empty(A.shape[:-1] + (N,), dtype=FP8_DTYPE)
+    else:
+        C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype)
+    if swizzled_out:
+        cb_cs = triton.cdiv(N // out_g, 4)
+        cs_ret = torch.empty(1, triton.cdiv(M, 128), cb_cs, 2, 256, device=A.device, dtype=cs_dtype)
+        CSDescriptor = TensorDescriptor.from_tensor(cs_ret, [1, 1, 1, 2, 256])
+        Cs = None  # row-major pointer unread here; CSDescriptor does the store
+    elif requant:
+        cs_ret = torch.empty(M, N // out_g, device=A.device, dtype=cs_dtype)
+        Cs, CSDescriptor = cs_ret, None
+    else:
+        cs_ret, Cs, CSDescriptor = None, None, None  # unread (no OUTPUT_FORMAT)
+    stride_cs_m = cs_ret.stride(0) if (requant and not swizzled_out) else 0
+    stride_cs_n = cs_ret.stride(1) if (requant and not swizzled_out) else 0
+    # Host-TMA descriptors over the packed (M, K_bytes) / (N, K_bytes) matrices — placeholder
+    # box rebound per tuned config by rebind_mx_descriptors. Read only by the descriptor
+    # configs the tuner picks; pointer configs never touch them.
+    a_descriptor = TensorDescriptor.from_tensor(A_q, [1, 32])
+    b_descriptor = TensorDescriptor.from_tensor(b_u8, [1, 32])
+    # One scale pointer per operand (as_u8/bs_u8) — the swizzled buffer on the fast path (also the
+    # base for the BM<128 / BN<128 scalar gather), the affine scale otherwise — plus its descriptor,
+    # built only on the swizzled path (dummy off the operand descriptor when un-swizzled, unread).
+    as_descriptor = swizzled_scale_descriptor(as_u8) if act_swizzled else a_descriptor
+    bs_descriptor = swizzled_scale_descriptor(bs_u8) if swizzled_scales else b_descriptor
+    def grid(META):
+        return (
+            triton.cdiv(M, META["BLOCK_SIZE_M"]),
+            triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        )
+    with device_context(A.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(mx_dynamic_matmul_kernel)[grid](
+            A_q,
+            a_descriptor,
+            as_u8,
+            as_descriptor,
+            b_u8,
+            b_descriptor,
+            bs_u8,
+            bs_descriptor,
+            C,
+            Cs,
+            bias,
+            CSDescriptor,
+            act_global_scale,  # AsGlobal: g_a for the inline-quant arm (A/g_a)
+            input_global_scale,  # AsBsGlobal = g_b (acc; g_a folds in-kernel via AsGlobal)
+            out_global_scale,  # CsGlobal: requant output normalization (next proj's provided input_scale); None folds out
+            M,
+            N,
+            K,
+            int(
+                M
+            ).bit_length(),  # m_bit_length key bucket; int() concretizes M (a SymInt under torch.compile has no .bit_length)
+            A_q.stride(0),
+            A_q.stride(1),
+            as_u8.stride(0),  # As row stride (affine act-scale read; dead on the swizzled arm)
+            as_u8.stride(-1),  # As group-axis stride (1 when row-major; see stride_as_k)
+            b_u8.stride(1),
+            b_u8.stride(0),
+            bs_u8.stride(1),
+            bs_u8.stride(0),
+            C.stride(-2),
+            C.stride(-1),
+            stride_cs_m,
+            stride_cs_n,
+            bias_stride_e,
+            bias_stride_n,
+            global_scale_stride(act_global_scale),
+            global_scale_stride(input_global_scale),
+            global_scale_stride(out_global_scale),
+            SCALE_GROUP_K=scale_group,
+            ACTIVATION_FORMAT=activation_format,
+            SWIZZLED_SCALES=swizzled_scales,
+            SWIZZLED_OUT=swizzled_out,
+            GATE=gate,
+            ACT_FN=act_fn,
+            SWIGLU_ALPHA=swiglu_alpha,
+            SWIGLU_LIMIT=swiglu_limit,
+            OUTPUT_FORMAT=output_format,
+            SIMULATE_UNFUSED=simulate_unfused,
+            INTERMEDIATE_DTYPE=tl_dtype(output_dtype),
+        )
+    # NVFP4's block scales are already normalized by the provided output global (CsGlobal =
+    # output_global_scale) at requant, so the caller pairs cs_ret with that global as the down's As.
+    return [C, cs_ret] if cs_ret is not None else [C]
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("full_precision_matmul_2d"), mutates_args=(), opaque=True
+)
+def full_precision_matmul_2d(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    bias: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    """Full-precision (BF16/FP16) 2D matmul ``C = A @ B.T``; no quantization anywhere. ``gate``
+    fuses the ``(2N, K)`` gate|up projection into one stacked GEMM + SwiGLU, returning the
+    ``[..., N]`` GLU intermediate. Returns a one-element list (mirrors the quantized 2D ops)."""
+    validate_dense_2d_operands(A, B)
+    assert A.dtype == B.dtype and A.dtype in (torch.bfloat16, torch.float16), (
+        f"full-precision path needs matching BF16/FP16 A and B, got {A.dtype} / {B.dtype}"
+    )
+    rows, K = B.shape
+    # Under gate|up fusion B is the (2N, K) gate|up stack; N is the per-projection output width.
+    N = rows // 2 if gate else rows
+    M = A.numel() // A.shape[-1]
+    C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype)
+    def grid(META):
+        return (
+            triton.cdiv(M, META["BLOCK_SIZE_M"]),
+            triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        )
+    with device_context(A.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(full_precision_matmul_2d_kernel)[grid](
+            A,
+            B,
+            C,
+            bias,
+            M,
+            N,
+            K,
+            int(M).bit_length(),  # m_bit_length key bucket
+            A.stride(-2),
+            A.stride(-1),
+            B.stride(1),
+            B.stride(0),
+            C.stride(-2),
+            C.stride(-1),
+            bias_stride_e,
+            bias_stride_n,
+            GATE=gate,
+            ACT_FN=act_fn,
+            SWIGLU_ALPHA=swiglu_alpha,
+            SWIGLU_LIMIT=swiglu_limit,
+            SIMULATE_UNFUSED=simulate_unfused,
+            INTERMEDIATE_DTYPE=tl_dtype(resolve_output_dtype(output_dtype, A, None)),
+        )
+    return [C]
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("mx_weight_only_matmul_2d"),
+    mutates_args=(),
+    opaque=True,
+)
+def mx_weight_only_matmul_2d(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    b_global_scale: torch.Tensor | None = None,
+    bias: torch.Tensor | None = None,
+) -> list[torch.Tensor]:
+    """weight-only dense matmul ``C = A @ B.T``: a RAW bf16/fp16 activation (no quant) against an
+    MXFP4/NVFP4/MXFP8 weight upcast to bf16 in-loop (NVFP4's per-tensor fp32 ``b_global_scale``
+    recovers on the accumulator; ``None`` = single-level) — the ``matmul_ogs`` format (fp4/fp8 weight bytes in
+    memory, bf16 acts, weight upcast in-MMA). Affine (2D) weight scales only. ``gate`` fuses the
+    ``(2N, K)`` gate|up projection + SwiGLU, returning the ``[..., N]`` GLU intermediate. Returns a
+    one-element list (mirrors the other 2D ops). Packed-E2M1 K (B is (rows, K//2)) is handled
+    below, so this op does its own shape checks (not ``validate_dense_2d_operands``)."""
+    assert A.dtype in (torch.bfloat16, torch.float16, torch.float32), (
+        f"weight-only takes a raw bf16/fp16/fp32 activation, got {A.dtype}"
+    )
+    assert Bs.ndim == 2, f"weight-only takes affine (2D) weight scales, got ndim={Bs.ndim}"
+    WEIGHT_VALUES_PER_BYTE = 2 if B.dtype == torch.int8 else 1
+    rows, K_b = B.shape
+    K = A.shape[-1]
+    M = A.numel() // K
+    N = rows // 2 if gate else rows
+    assert K == WEIGHT_VALUES_PER_BYTE * K_b, (
+        f"K (={K}) must equal {WEIGHT_VALUES_PER_BYTE} * B.shape[1] (={K_b})"
+    )
+    scale_group = mx_scale_family(Bs, K)
+    assert Bs.shape == (rows, K // scale_group), (
+        f"Bs shape {tuple(Bs.shape)} != ({rows}, {K // scale_group})"
+    )
+    assert A.is_contiguous(), "A must be contiguous (the host-TMA descriptor arm reads its box)"
+    b_u8 = e2m1_as_uint8(B)
+    bs_u8 = ue8m0_as_uint8(Bs)
+    b_global_scale = normalize_global_scale(b_global_scale, 1)
+    C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype)
+    # Host-TMA descriptors over the (M, K) bf16 A and packed (N, K_bytes) weight — placeholder box
+    # rebound per tuned config by rebind_weight_only_descriptors, read only by the descriptor configs the
+    # tuner picks (pointer configs never touch them). Scales stay affine (no scale descriptor).
+    a_descriptor = TensorDescriptor.from_tensor(A.view(M, K), [1, 32])
+    b_descriptor = TensorDescriptor.from_tensor(b_u8, [1, 32])
+    def grid(META):
+        return (
+            triton.cdiv(M, META["BLOCK_SIZE_M"]),
+            triton.cdiv(N, META["BLOCK_SIZE_N"]),
+        )
+    with device_context(A.device):
+        bias_stride_e, bias_stride_n = bias_strides(bias)
+        compile_time_only_triton_wrap(mx_weight_only_matmul_2d_kernel)[grid](
+            A,
+            a_descriptor,
+            b_u8,
+            b_descriptor,
+            bs_u8,
+            b_global_scale,
+            C,
+            bias,
+            M,
+            N,
+            K,
+            int(M).bit_length(),
+            A.stride(-2),
+            A.stride(-1),
+            b_u8.stride(0),
+            b_u8.stride(1),
+            bs_u8.stride(0),
+            bs_u8.stride(1),
+            C.stride(-2),
+            C.stride(-1),
+            bias_stride_e,
+            bias_stride_n,
+            global_scale_stride(b_global_scale),
+            SCALE_GROUP_K=scale_group,
+            WEIGHT_VALUES_PER_BYTE=WEIGHT_VALUES_PER_BYTE,
+            GATE=gate,
+            ACT_FN=act_fn,
+            SWIGLU_ALPHA=swiglu_alpha,
+            SWIGLU_LIMIT=swiglu_limit,
+            SIMULATE_UNFUSED=simulate_unfused,
+            INTERMEDIATE_DTYPE=tl_dtype(resolve_output_dtype(output_dtype, A, None)),
+        )
+    return [C]
+# ── torch scaled_mm fast path (2D dense MX, Blackwell) ────────────────────────
+#
+# cuBLAS's block-scaled GEMM beats our Triton 2D kernel at EVERY measured M on the
+# quantized-activation MX formats (B200, N=18432 K=6144, 200-trial tunes, 2026-08-27):
+#     mxfp8  M=16 3.42x   M=256 3.17x   M=1024 1.57x   M=4096 1.55x   M=8192 1.42x
+#     nvfp4  M=16 4.83x   M=256 4.15x   M=1024 2.48x   M=4096 1.81x   M=8192 1.80x
+# and its bf16 output is bit-identical to ours at every checked shape (same offline act
+# quant feeds both, fp32 accumulate). Only the M=1 swap-AB decode dispatch stays ahead
+# (18.4us vs ~33), so the route floor is M >= 2.
+#
+# Inductor refuses to lower scaled_mm with SWIZZLE_32_4_4 scales ("does not yet support
+# non-trivial swizzles" — repros/scaled_mm_swizzle_compile.py), so the call lives behind an
+# OPAQUE custom op: dynamo captures the op as a leaf and the real scaled_mm runs at runtime,
+# which serves compiled callers too. F.scaled_mm has no autograd formula while our ops carry
+# registered dgrads, so differentiating callers fall through to the Triton ops. CUDA-graph
+# capture is fine (measured under do_bench_cudagraph).
+_SCALED_MM_MIN_M = int(os.environ.get("FINEGRAINED_SCALED_MM_MIN_M", "2"))
+@torch.library.custom_op(add_op_namespace_prefix("scaled_mm_2d_mx"), mutates_args=())
+def _scaled_mm_2d_op(
+    Aq: torch.Tensor,
+    As: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    g_a: torch.Tensor | None,
+    g_b: torch.Tensor | None,
+    fmt: str,
+) -> torch.Tensor:
+    """``F.scaled_mm`` behind an OPAQUE custom op. Inductor cannot lower scaled_mm with
+    SWIZZLE_32_4_4 scales (repros/scaled_mm_swizzle_compile.py), but it never has to: dynamo
+    captures this op as a leaf and the real call runs at runtime — so compiled callers get the
+    cuBLAS fast path too, not just eager ones."""
+    F = torch.nn.functional
+    ST, SW = F.ScalingType, F.SwizzleType
+    if fmt == "mxfp8":
+        return F.scaled_mm(
+            Aq, B.t(),
+            As.view(torch.float8_e8m0fnu), ST.BlockWise1x32,
+            Bs.view(torch.float8_e8m0fnu), ST.BlockWise1x32,
+            swizzle_a=SW.SWIZZLE_32_4_4, swizzle_b=SW.SWIZZLE_32_4_4,
+            output_dtype=torch.bfloat16,
+        )
+    one = torch.ones((), device=Aq.device, dtype=torch.float32)
+    FP4 = torch.float4_e2m1fn_x2
+    return F.scaled_mm(
+        Aq.view(FP4), B.view(FP4).t(),
+        [As.view(torch.float8_e4m3fn), one if g_a is None else g_a.reshape(()).float()],
+        [ST.BlockWise1x16, ST.TensorWise],
+        [Bs.view(torch.float8_e4m3fn), one if g_b is None else g_b.reshape(()).float()],
+        [ST.BlockWise1x16, ST.TensorWise],
+        swizzle_a=[SW.SWIZZLE_32_4_4, SW.NO_SWIZZLE],
+        swizzle_b=[SW.SWIZZLE_32_4_4, SW.NO_SWIZZLE],
+        output_dtype=torch.bfloat16,
+    )
+@_scaled_mm_2d_op.register_fake
+def _(Aq, As, B, Bs, g_a, g_b, fmt):
+    return Aq.new_empty(Aq.shape[0], B.shape[0], dtype=torch.bfloat16)
+def _smm_reject(reason):
+    if os.environ.get("FINEGRAINED_SCALED_MM_DEBUG"):
+        print(f"[scaled_mm route] rejected: {reason}", flush=True)
+    return None
+def _torch_scaled_mm_2d(A, B, As, Bs, activation_format, a_global_scale, b_global_scale, out_dtype):
+    """The scaled_mm route, or ``None`` to take the Triton ops. Fires only where measured to
+    win AND where the semantics are identical: plain ungated GEMM, quantized activations in
+    the WEIGHT's fmt, SWIZZLE_32_4_4 weight scales (Blackwell scaled_mm rejects row-major
+    outright), bf16 out, eager, no grad."""
+    if os.environ.get("FINEGRAINED_DISABLE_SCALED_MM"):
+        return _smm_reject("env-disabled")
+    if not is_sm10x():
+        return _smm_reject("not-sm10x")
+    if not (hasattr(torch.nn.functional, "scaled_mm") and hasattr(torch.nn.functional, "ScalingType")):
+        return _smm_reject("no-F.scaled_mm")
+    fmt = weight_format(B, Bs)
+    if fmt not in ("mxfp8", "nvfp4") or activation_format == "bf16":
+        return _smm_reject("format/activation_format")
+    if (a_global_scale is not None and a_global_scale.numel() > 1) or (
+        b_global_scale is not None and b_global_scale.numel() > 1
+    ):
+        # scaled_mm's second level is one TensorWise value per operand, so a gate|up stack whose
+        # halves carry their own global has no form here — the Triton kernels fold it by column
+        return _smm_reject("multi-valued global")
+    # None follows the weights; resolve it before comparing (the bench and the integrations pass it)
+    if resolve_activation_format(activation_format, B, Bs) != fmt:
+        return _smm_reject("format-mismatch-after-resolve")
+    if Bs.ndim != 5:  # affine (row-major) scales: scaled_mm rejects them on Blackwell
+        return _smm_reject("affine-Bs")
+    if out_dtype is not torch.bfloat16:
+        return _smm_reject("out-dtype")
+    M = A.numel() // A.shape[-1]
+    group = 32 if fmt == "mxfp8" else 16
+    K = A.shape[-1] if As is None else (A.shape[-1] * (2 if fmt == "nvfp4" else 1))
+    if M < _SCALED_MM_MIN_M or K % group or B.shape[0] % 16:
+        return _smm_reject("shape/minM")
+    if As is not None and As.ndim != 5:
+        return _smm_reject("affine-As")  # pre-quantized activations must arrive with swizzled scales too
+    if fmt == "mxfp8":
+        Aq = A.reshape(M, K) if As is not None else None
+        if Aq is None:
+            Aq, As = mxfp8_act_quant(A.reshape(M, K), swizzled=True)
+    else:
+        Aq = A.reshape(M, K // 2) if As is not None else None
+        if Aq is None:
+            Aq, As = nvfp4_act_quant(A.reshape(M, K), swizzled=True, global_scale=a_global_scale)[:2]
+    out = _scaled_mm_2d_op(Aq, As, B, Bs, a_global_scale, b_global_scale, fmt)
+    return out.reshape(*A.shape[:-1], B.shape[0])
+def matmul_2d(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    As: torch.Tensor | None = None,
+    Bs: torch.Tensor | None = None,
+    *,
+    bias: torch.Tensor | None = None,  # (E, N_out) per-expert output bias; (N_out,) for 2D
+    activation_format: str | None = None,
+    gate: bool = False,
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    quantize_output: bool = False,
+    simulate_unfused: bool = False,
+    output_dtype: torch.dtype | None = None,
+    a_global_scale: torch.Tensor | None = None,
+    b_global_scale: torch.Tensor | None = None,
+    output_global_scale: torch.Tensor | None = None,
+) -> torch.Tensor | list[torch.Tensor]:
+    """Dense (2D) quantized matmul dispatcher (W8A8 FP8, W4A8/W4A4 FP4) — the single-GEMM sibling of
+    ``matmul_grouped``/``matmul_batched``, taking the same ``As``/``Bs`` scale spec,
+    ``activation_format`` / gate|up fusion kwargs, and inferred quant block (no ``block_size`` argument — it
+    falls out of the weight-scale shape, so the data can't disagree with a parameter).
+    ``As`` is the activation block-scale spec, same as the routed ops: ``None`` → the op quantizes raw
+    ``A`` dynamically; a per-tensor scalar → static (calibrated) FP8 activation scale; block scales →
+    pre-quantized ``A``. ``Bs`` is the weight block scale. ``a_global_scale``/``b_global_scale`` are the
+    two-level NVFP4 per-tensor second-level scales (calibrated ``input_scale`` ``g_a`` and
+    ``weight_scale_2`` ``g_b``) — provided, never computed; the op folds ``g_a·g_b`` onto the
+    accumulator (``a_global_scale`` rides raw ``A`` or a pre-quantized ``As`` alike).
+    ``output_global_scale`` (NVFP4 requant only) is the NEXT proj's provided ``input_scale``:
+    the fused requant normalizes the GLU intermediate by it, returning ``(C, Cs)`` the down-proj
+    consumes as ``As=Cs, a_global_scale=output_global_scale``.
+    ``activation_format`` is the activations' quantization: ``None`` (default) = the weights' own
+    format, ``"bf16"`` = weight-only (no activation quant), or an explicit ``"fp8"`` / ``"mxfp8"`` /
+    ``"mxfp4"`` / ``"nvfp4"`` (``"mxfp8"`` on MXFP4 weights = the W4A8 chain; ``check_activation_format``
+    has the per-family matrix). ``gate=True`` reads ``B`` as the ``(2N, K)`` gate|up stack and fuses
+    the ``act_fn`` GLU (``swiglu_alpha``/``swiglu_limit`` its clamped variant), returning the
+    ``[..., N]`` intermediate; ``quantize_output=True`` (MX weights only) stores that intermediate in
+    ``activation_format`` and returns ``(C, Cs)`` for the next op.
+    Routes by weight dtype + inferred block: ``Bs`` None → full-precision; MX/NVFP4 (int8/E4M3 weights
+    with UE8M0 group-32 or E4M3 group-16 ``Bs``) → ``mx_dynamic_matmul``; tensor-wide FP8 (scalar
+    ``Bs``) → ``w8a8_tensor_dynamic_fp8_matmul``; block FP8 → the static variant when ``As`` is a
+    per-tensor scalar, else dynamic.
+    """
+    check_activation_format(activation_format, quantize_output)
+    assert (a_global_scale is None and b_global_scale is None) or (
+        Bs is not None and weight_format(B, Bs) == "nvfp4"
+    ), "two-level globals (a_global_scale / b_global_scale) are NVFP4-only"
+    assert output_global_scale is None or (
+        Bs is not None and quantize_output and resolve_activation_format(activation_format, B, Bs) == "nvfp4"
+    ), (
+        "output_global_scale is the NVFP4 requant second level — it requires quantize_output=True on NVFP4 "
+        "(the epilogue would otherwise normalize by it with nothing downstream to compensate)"
+    )
+    def _unwrap(ret: list[torch.Tensor]) -> torch.Tensor | tuple[torch.Tensor, ...]:
+        # tuple, matching the routed siblings — callers branch on isinstance(out, tuple)
+        return ret[0] if len(ret) == 1 else tuple(ret)
+    if Bs is None:  # unquantized BF16/FP16 weights — plain dot, no scales
+        assert As is None and a_global_scale is None and activation_format in (None, "bf16") and not quantize_output, (
+            "the full-precision path (Bs=None) takes no activation scale or quantized activation format"
+        )
+        return _unwrap(
+            full_precision_matmul_2d(
+                A, B, output_dtype, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, bias=bias)
+        )
+    if is_mx(B, Bs):
+        if activation_format == "bf16":  # weight-only: raw bf16 activation, MX weight upcast in-MMA
+            assert As is None and a_global_scale is None and not quantize_output, (
+                "weight-only (activation_format='bf16') takes a raw bf16 activation, no As/global/requant"
+            )
+            assert output_global_scale is None, (
+                "weight-only has no requant epilogue for output_global_scale to normalize"
+            )
+            return _unwrap(
+                mx_weight_only_matmul_2d(
+                    A, B, Bs, output_dtype, gate, act_fn, swiglu_alpha, swiglu_limit,
+                    simulate_unfused, b_global_scale=b_global_scale, bias=bias)
+            )
+        return _unwrap(
+            mx_dynamic_matmul(
+                A, B, As, Bs, output_dtype, activation_format,
+                gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, quantize_output,
+                a_global_scale, b_global_scale, output_global_scale, bias=bias)
+        )
+    # FP8 activations are always E4M3 with the weight-implied scale granularity, so "fp8" is a
+    # no-op format name (accepted for symmetry with the MoE ops); no other name applies here.
+    assert activation_format in (None, "fp8"), (
+        f"FP8-weight activations are E4M3 ('fp8'), got {activation_format!r}"
+    )
+    assert not quantize_output, "quantize_output (fused-requant gate|up output) is MX-only"
+    # Infer the FP8 quant block from the weight-scale shape (scalar Bs = tensor-wide). The contraction
+    # dim K tiles evenly, so block_k is exact; N may be non-aligned (a partial last block, so
+    # N % Bs.shape[0] != 0 and plain division would under-count) — recover block_n from the even dim,
+    # the square-block convention for these FP8 weights.
+    if Bs.numel() == 1:
+        block_size = None
+    else:
+        block_k = B.shape[1] // Bs.shape[1]
+        block_n = B.shape[0] // Bs.shape[0] if B.shape[0] % Bs.shape[0] == 0 else block_k
+        block_size = [block_n, block_k]
+    if block_size is None:  # tensor-wide (per-tensor) scale
+        assert As is None, "tensor-wide FP8 quantizes A dynamically — no As"
+        return _unwrap(
+            w8a8_tensor_dynamic_fp8_matmul(
+                A, B, Bs, output_dtype, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, bias=bias)
+        )
+    # Block-wise FP8: a per-tensor scalar As is the static (calibrated) activation scale; else dynamic.
+    if As is not None:
+        assert As.numel() == 1, "block-wise FP8 As is a per-tensor static scale (scalar)"
+        return _unwrap(
+            w8a8_block_static_fp8_matmul(
+                A, B, As, Bs, block_size, output_dtype,
+                gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, bias=bias)
+        )
+    return _unwrap(
+        w8a8_block_dynamic_fp8_matmul(
+            A, B, Bs, block_size, output_dtype,
+            gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, bias=bias)
+    )

build/torch-rocm/metadata.json ADDED Viewed

	@@ -0,0 +1,48 @@

+{
+  "name": "finegrained-kernels",
+  "id": "_finegrained_kernels_rocm_1348465",
+  "version": 0,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "rocm"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "43nBPHxybam5ChWZmsz1Ji+pjJ6K6zF5lzwbqPcMwE4=",
+      "_ops.py": "Ssu+B/vlZZOXRh+hKJNWiv3V59Z7J6Wf0w/+9c7uT+M=",
+      "backward.py": "ri9lKPhb9SmkYBvUS8KvuORiW5PcKViNamFNVWxb69g=",
+      "batched.py": "cSkV7Mnv6kK+k5MANm8f6ZFu3Jdbh0tAz9u9t4e+R3c=",
+      "bayesian_autotuner.py": "DPpIATTls8aY7n12oKggpuAOvC3RlugYl1Wid1t5J+U=",
+      "compat.py": "yj7b7kCAOrq3j1Y3WxSb/YKoVcPRFfkeXmMj7dkqQoc=",
+      "descriptors.py": "4acRMazGPFlbydNt1Hbo8RpAqb3fOtnOQBv5yb4onqw=",
+      "epilogue.py": "tGALTI/VMgO3Mz0sQ6H6wuRZpO4UKqM1/4h7d81SnVM=",
+      "finegrained_kernels/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "formats.py": "V6F0PXdNhXQn/UBz0/frbvcp5Lgterh053XDrvQ/Cpw=",
+      "grouped.py": "PXhZU9Rn6UIIiQtR6IUmI060Mank5GMuooTWapRvya8=",
+      "loading/__init__.py": "so7kYmlVoiASfoOzemdU5fBTiP8yA85gVPNQExIq/KU=",
+      "loading/scales.py": "8HJpflmImEXYJquA6hVLStUUEecvExAxBEYVb3qPPWE=",
+      "loading/tiles.py": "hvFIDikI1k0+vV1QOjqdl1bovI/JEElfNwLWg7VC3ZM=",
+      "matmul.py": "qLV1KsPyCeAq3O0swnk38GUOT+lS12WxUM8bQZ9DF9A=",
+      "mma.py": "1fg2an7Dc8Vd6qpSmq+Q0781A6sNilOJHU0PqEbouT8=",
+      "moe.py": "rpiYMM+kzTtNuxXFUrQARLO9kLHcEgKR8rs/dZ0wn/g=",
+      "norm.py": "VvE6GUE0Oztf7PDcviLCcB4fhnXt7/zgdE+aApv6DBM=",
+      "pruners.py": "v03bXbvsZp1znVIj/eiR7LAWctU66DJlJr51wn8Pheg=",
+      "quant.py": "PZhzgOHSDH8lLV6zvrv040B4g9SgFE6Eq6bc8XxfPY8=",
+      "scheduling.py": "UGgQ5uVdm4cWPnH+4w6Ffou/PSX3iUarMeRo2E04BhY=",
+      "swizzle.py": "ueIJ9FsO8QZhVt3nMFlA+8b+U9/ZHabHFlhufsONV04="
+    }
+  },
+  "provenance": {
+    "kernel-builder": {
+      "version": "0.17.0-dev0",
+      "sha": "81580bb92577f2f7228661ca2a221fc052375709",
+      "dirty": false
+    },
+    "kernel": {
+      "sha": "1348465af00f2d4c596d37bdb262ae0875dc1d77",
+      "dirty": false
+    }
+  }
+}

build/torch-rocm/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHTDCCBtKgAwIBAgIUKwy7DRLjNl12oimua1LRYxRBXBQwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTE2MTU0MjA3WhcNMjYwOTE2MTU1MjA3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEkX51E0716VRnV+N7XqZPQSuT+8+WFqVZ8jyLnA6PxNlNCIgr4O1c3VI4np8PMZ+9k0EozfVf3XbIiPsqY9vU/aOCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUo1O1wLYeWm7XlvpitaZXgdwuluUwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoMTM0ODQ2NWFmMDBmMmQ0YzU5NmQzN2JkYjI2MmFlMDg3NWRjMWQ3NzATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoMTM0ODQ2NWFmMDBmMmQ0YzU5NmQzN2JkYjI2MmFlMDg3NWRjMWQ3NzAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoMTM0ODQ2NWFmMDBmMmQ0YzU5NmQzN2JkYjI2MmFlMDg3NWRjMWQ3NzAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDEzNDg0NjVhZjAwZjJkNGM1OTZkMzdiZGIyNjJhZTA4NzVkYzFkNzcwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzUxMTY2NDkyOTcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoKrh6swAAAQDAEgwRgIhANX/7k2eKwwKjy8ZxKOpmVkg1xAd5OZZtFP+JoYLm5xdAiEA541PEWCxlJ8vBCctUblYJ8CwahUQrQHAB1az4ZQBRvAwCgYIKoZIzj0EAwMDaAAwZQIxAOrFZmrON23FouSKXOvH3VhA3MDrJ0zVLduijdnHKbEB6XAwuqVcNzLd8e+SystARQIwOdpWPy0GsRlCrzjovhZ/iOd3XOZtYhCHbwuhADV5U8LyW0FHy/Jl3g8mCcNkM+hY"}, "tlogEntries":[{"logIndex":"2863678766", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1789573327", "inclusionPromise":{"signedEntryTimestamp":"MEUCIQDOUMA7NduNwOul83R0X3Ocp2ybE707LUZ6Z7MxIDEDHwIgY3TKJTPvt660IBYW6CxkGv0W7iyJWsMQ6DE9TAVfu2A="}, "inclusionProof":{"logIndex":"2741774504", "rootHash":"h7JnsSI02oktCuUf2hPr37WQeR1av+o8FGUwr34KmHo=", "treeSize":"2741774526", "hashes":["1O4OieGYomUIdapwcW0fU0+2+02pFm1aEo0ti4sg08s=", "mcj/XMUVpOVz+S5LUKHconnoT5H4349XaBWB7RFeOkI=", "Wdf+QLOD4triMw+JPnM0auon4nMf4ww9lucCEay94fU=", "6OX+tq7oFddP7/kDlt87XDsQ0WlCJnt1SdJUYyCdR+8=", "B5h5p0vseS2PGxobaWllh1w2JdneoJtoZSCxox8DQMc=", "66RI6MJezQM4018f5MMtQln23ah4XGjWHARyTP4OGH4=", "Ojxn+u3l9okV6fn1qhDIlj0xB1NlpdMwwfQZrYDTTDA=", "dNQ0Q6e05QOPvF+9aMH7hKeNMC+sdDn/+tAhAdOJnpQ=", "cuse2fF4Cl9QdYLPPxCWd9gSEI01WyuvnM3KB1Ah78Y=", "g+kZXC8JlqtkZn+3BtMPZOIjPv/1QgiVvlH7tpuF5rU=", "z7hfW7IyWq58tcWTZuli3Brmb7xA7392X7mpl47fnCc=", "/DBUF30Sdk4LDt4XUgHCGSr10fT7nDkSitYbygYwfNg=", "vC+yZlSpf9tVD+3MdNtHCe71eNNHHi6gplgltL/tmDc=", "gGfwuhE97xnhmxu6exffPtS05ABS4/ldyg66sFcxSRE=", "jhevWFOJ2cWU6SvzMnAe08zHGcVquRlP7gAed39+wI4=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2741774526\nh7JnsSI02oktCuUf2hPr37WQeR1av+o8FGUwr34KmHo=\n\n— rekor.sigstore.dev wNI9ajBGAiEA1L58wqYgcdSBeA6sGCi5zk9nSn7U/bA/v0obbkuIovICIQDm2hhhrVGvjLnHrfJ4aGYqS45uNAgpXJMtOVd90NuZzQ==\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJkZmI0MDlkYzcyN2RjMWM1M2I1MGJjZGZjMzU0NmE4NmZjNWNjOWVlY2UzOTVhZmJkZmEzZWJhOWIwNmRjZDY5In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNjRFZZNFd2Ylhob3pHcDFFSTZSdWxNUTJOZ295M2w1OWpxUWxtcVpXem93SWdKaUYxQ0RhWUl0Tm5icC84WDdMd05kQW03bVVycW93ZWJaL2RMall1dzRNPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFVSRU5EUW5STFowRjNTVUpCWjBsVlMzZDVOMFJTVEdwT2JERXliMmx0ZFdFeFRGSlplRkpDV0VKUmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVUlRKTlZGVXdUV3BCTTFkb1kwNU5hbGwzVDFSRk1rMVVWVEZOYWtFelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZyV0RVeFJUQTNNVFpXVW01V0swNDNXSEZhVUZGVGRWUXJPQ3RYUm5GV1dqaHFlVXdLYmtFMlVIaE9iRTVEU1dkeU5FOHhZek5XU1RSdWNEaFFUVm9yT1dzd1JXOTZabFptTTFoaVNXbFFjM0ZaT1haVkwyRlBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZ2TVU4eENuZE1XV1ZYYlRkWWJIWndhWFJoV2xoblpIZDFiSFZWZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOU5WRTB3VDBSUk1rNVhSbTFOUkVKdFRXMVJNRmw2VlRWT2JWRjZDazR5U210WmFra3lUVzFHYkUxRVp6Tk9WMUpxVFZkUk0wNTZRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDFVVFRCUFJGRXlUbGRHYlUxRVFtMU5iVkV3V1hwVk5VNXRVWHBPTWtwcldXcEpNazF0Um13S1RVUm5NMDVYVW1wTlYxRXpUbnBCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDFVVFRCUFJGRXlUbGRHYlUxRVFtMU5iVkV3V1hwVk5VNXRVWG9LVGpKS2ExbHFTVEpOYlVac1RVUm5NMDVYVW1wTlYxRXpUbnBCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUkZlazVFWnpBS1RtcFdhRnBxUVhkYWFrcHJUa2ROTVU5VVdtdE5lbVJwV2tkSmVVNXFTbWhhVkVFMFRucFdhMWw2Um10T2VtTjNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWxWNFRWUlpNazVFYTNsUFZHTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOUxjbWcyYzNkQlFVRlJSQXBCUldkM1VtZEphRUZPV0M4M2F6SmxTM2QzUzJwNU9GcDRTMDl3YlZaclp6RjRRV1ExVDFwYWRFWlFLMHB2V1V4dE5YaGtRV2xGUVRVME1WQkZWME40Q214S09IWkNRMk4wVldKc1dVbzRRM2RoYUZWUmNsRklRVUl4WVhvMFdsRkNVblpCZDBObldVbExiMXBKZW1vd1JVRjNUVVJoUVVGM1dsRkplRUZQY2tZS1dtMXlUMDR5TTBadmRWTkxXRTkyU0ROV2FFRXpUVVJ5U2pCNlZreGtkV2xxWkc1SVMySkZRalpZUVhkMWNWWmpUbnBNWkRobEsxTjVjM1JCVWxGSmR3cFBaSEJYVUhrd1IzTlNiRU55ZW1wdmRtaGFMMmxQWkROWVQxcDBXV2hEU0dKM2RXaEJSRlkxVlRoTWVWY3dSa2g1TDBwc00yYzRiVU5qVG10TksyaFpDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyDADAgEAMIICvwYJKoZIhvcNAQcCoIICsDCCAqwCAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgEEWXb1aAU9uWYKBG9k5yRgiCru6MaWJqEJ7eotv6QLsCFAVQ7WjJb8AfpOlPrlg4vundjMeXGA8yMDI2MDkxNjE1NDIwN1owAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdowggHWAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwOTE2MTU0MjA3WjAvBgkqhkiG9w0BCQQxIgQgnhWF5DW3tegD61eDeAX0KIpasZoCwslhIsaRZtJyxHwwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGYwZAIwQ8JSNadEl6jXNYx0SryoCfLxVtpQES91e45sgot7605M9yXk+nFgOfZ5qRSu1fJdAjB3nmagQ7GQmID30bAxY5/nRpO4VFWhFXqvRv5R0UcT17jG0LE4x4TOt9vkd1x9JHE="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"37QJ3HJ9wcU7ULzfw1Rqhvxcye7OOVr736PrqbBtzWk="}, "signature":"MEUCIQCcDVY4WvbXhozGp1EI6RulMQ2Ngoy3l59jqQlmqZWzowIgJiF1CDaYItNnbp/8X7LwNdAm7mUrqowebZ/dLjYuw4M="}}

build/torch-rocm/mma.py ADDED Viewed

	@@ -0,0 +1,367 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import triton
+import triton.language as tl
+from .quant import e2m1_cols_to_e4m3, e2m1_cols_to_f32, e2m1_to_bf16, e2m1_to_e4m3, e2m1_to_f32
+from .loading.scales import decode_group_scale
+@triton.jit
+def mx_dot_scaled(
+    acc,
+    a,
+    a_scale,
+    w,
+    w_scale,
+    SWAP_AB: tl.constexpr = False,
+    BLOCK_SIZE_K: tl.constexpr = 0,
+    SCALE_GROUP_K: tl.constexpr = 0,
+):
+    """MX 'dot_scaled' step: the tcgen05 scaled MMA folds the group scales into the tensor core;
+    each operand's format is its tile's dtype (``uint8`` = packed E2M1, E4M3, or bf16 for the
+    weight-only unscaled lhs). fp4 on BOTH operands lowers to the native ``kind::mxf4`` MMA (2x the
+    fp8 rate; native iff the M operand is 128 — same gate as mxf8f6f4). Plain: ``a`` [BM, BK] x
+    ``w`` [BK, BN] (the caller pre-shapes ``w``/``w_scale``, e.g. ``tl.trans(gu)``). ``SWAP_AB``
+    (decode): the weight ``w`` [BN, BK] is the lhs (output rows in M) and the single [BK] token is
+    padded to the N=16 rhs (col 0 real — 16 is Triton's tcgen05-selection gate, N=8 drops to the
+    bf16-upcast fallback, bare-1 was 1.83x); ``acc`` is the persistent ``[BN, MMA_N_ATOM]``
+    accumulator the caller takes column 0 of (never a fresh per-step init, which trips the sm_100
+    accumulator-init pass). A packed ``a`` stays packed (the E4M3-scaled mxf4nvf4 kind is fp4 x fp4
+    only) and the token's group scale broadcasts to the rhs columns."""
+    a_fmt: tl.constexpr = (
+        "e2m1" if a.dtype == tl.uint8
+        else ("e4m3" if a.dtype == tl.float8e4nv else "bf16")
+    )
+    w_fmt: tl.constexpr = "e2m1" if w.dtype == tl.uint8 else "e4m3"
+    if SWAP_AB:
+        rhs = swap_pad_rhs(a, BLOCK_SIZE_K // 2 if a.dtype == tl.uint8 else BLOCK_SIZE_K)
+        if a_scale.dtype == tl.uint8:  # UE8M0 broadcast via the zero-add idiom
+            asc = tl.trans(
+                a_scale[:, None]
+                + tl.zeros((BLOCK_SIZE_K // SCALE_GROUP_K, MMA_N_ATOM), tl.uint8)
+            )
+        else:  # E4M3 (NVFP4) — no fp8 arithmetic; materialize the broadcast directly
+            asc = tl.trans(
+                tl.broadcast_to(
+                    a_scale[:, None], (BLOCK_SIZE_K // SCALE_GROUP_K, MMA_N_ATOM)
+                )
+            )
+        acc = tl.dot_scaled(w, w_scale, w_fmt, rhs, asc, a_fmt, acc)
+    else:
+        acc = tl.dot_scaled(a, a_scale, a_fmt, w, w_scale, w_fmt, acc)
+    return acc
+@triton.jit
+def mx_dot_rescale(
+    acc,
+    a,
+    a_scale,
+    w,
+    w_scale,
+    SWAP_AB: tl.constexpr = False,
+    BLOCK_SIZE_K: tl.constexpr = 0,
+):
+    """MX 'dot' step (BK == one scale group): fp8 ``tl.dot`` on E4M3-decoded operands + per-group
+    software rescale (both UE8M0 scales decoded here), accumulating into ``acc``. Plain: ``a``
+    [BM, BK] x ``w`` [BK, BN] with the weight's per-column scale transposed onto the product (the
+    batched gate_up kernel passes the stacked 2*BN tile — per-column independence keeps that
+    bit-exact). ``SWAP_AB`` (decode): the weight ``w`` [ROWS, BK] is the lhs (E2M1 column-unpacked,
+    K order low nibble first) and the [BK] token padded to the N=16 atom is the rhs — the
+    well-shaped fp8 MMA at M=1 (M quantizes to 64/128, N only to 8, so weight rows fill the big
+    atom); the weight's per-output-row scale broadcasts down the acc columns and the token's single
+    group scale is a scalar; ``acc`` is the persistent ``[ROWS, MMA_N_ATOM]`` accumulator."""
+    aq = e2m1_cols_to_e4m3(a) if a.dtype == tl.uint8 else a
+    a_s = decode_group_scale(a_scale)
+    w_s = decode_group_scale(w_scale)
+    if SWAP_AB:
+        wq = e2m1_cols_to_e4m3(w) if w.dtype == tl.uint8 else w
+        acc = acc + tl.dot(wq, swap_pad_rhs(aq, BLOCK_SIZE_K)) * w_s * a_s
+    else:
+        wq = e2m1_to_e4m3(w) if w.dtype == tl.uint8 else w
+        acc = acc + tl.dot(aq, wq) * a_s * tl.trans(w_s)
+    return acc
+@triton.jit
+def mx_weight_upcast(w, w_scale, BLOCK_SIZE_K: tl.constexpr, N: tl.constexpr, SCALE_GROUP_K: tl.constexpr,
+                     OUT_DTYPE: tl.constexpr = tl.bfloat16):
+    """Upcast one MXFP4/MXFP8 weight K-tile ``[BK, N]`` to ``OUT_DTYPE`` for the weight-only path: unpack
+    E2M1 -> bf16 directly (fp8 passes through a cast) and apply each ``SCALE_GROUP_K``-row K-group's
+    UE8M0 group scale. ``w_scale`` is ``[N, BK // SCALE_GROUP_K]`` (per-N-row, per-group); transposed
+    and broadcast across each group's rows to ``[BK, N]``. Unlike the ``dot`` arm (BK == group, scale
+    folded onto the [M,N] product) this dequantizes IN the tile, so BK spans any number of groups — a
+    full-BK bf16 ``tl.dot`` (the matmul_ogs format: fp4 weight, bf16 acts).
+    The scale is applied as a plain bf16 multiply, NOT a hand-rolled exponent-add: the UE8M0 scale is
+    a power of two so a ``bf16 * 2^k`` is already an exact exponent shift (a cheap FMA), and folding it
+    into the bits as an int add + zero-mask ``tl.where`` measured ~13% SLOWER (gate_up 4.68->5.29ms) —
+    the mask + uint16<->int32 bitcast chain costs more than the multiply it replaces."""
+    ng: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K
+    wq = e2m1_to_bf16(w) if w.dtype == tl.uint8 else w.to(tl.bfloat16)  # [BK, N] bf16 (direct, no E4M3)
+    # Reshape the (real, contiguous) weight to [ng, g, N] and broadcast the group scale [ng, 1, N]
+    # across each group's g rows in the multiply, then reshape back to [BK, N]. (Reshaping a
+    # broadcasted stride-0 tensor doesn't lower, so broadcast happens IN the op, not before it.)
+    ws = tl.trans(decode_group_scale(w_scale)).to(tl.bfloat16)[:, None, :]  # [ng, 1, N]
+    w3 = wq.reshape(ng, SCALE_GROUP_K, N) * ws
+    # dequant + scale happen in bf16 (exact: E2M1/E4M3 codes and power-of-two UE8M0 scales all fit),
+    # then widen if the activation is wider — bf16->fp32 is lossless and keeps the multiply cheap.
+    return w3.reshape(BLOCK_SIZE_K, N).to(OUT_DTYPE)
+@triton.jit
+def mx_weight_only_compute(
+    acc,
+    a,
+    w,
+    w_scale,
+    COMPUTE_MODE: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    N_WIDTH: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    SWAP_AB: tl.constexpr = False,
+):
+    """Weight-only (W4A16/W8A16) MMA step: raw bf16/fp16 activations against an MXFP4/MXFP8 weight
+    tile, dispatched on ``COMPUTE_MODE``. ``"dot_scaled"`` hands the packed weight and its UE8M0
+    group scales straight to the tcgen05 scaled MMA with an UNSCALED activation operand (the
+    ``lhs_scale=None`` / bf16-format form), so the tensor core does the fp4 decode and the group
+    rescale; ``"dot"`` upcasts the tile to the activation dtype in-loop and runs a plain
+    ``tl.dot`` (the matmul_ogs Hopper format). Both are correct everywhere and the tuner picks per
+    workload — the two differ by arm: measured on gpt-oss K=2880, ``dot_scaled`` wins the down
+    projection by 18.5% while the stacked gate|up tile prefers ``dot`` by 7.3``%``. ``SWAP_AB``
+    is the decode form: weight output rows lead the tile and the FMA reduce replaces the MMA
+    (both dot arms keep the token in M). Single return — only the taken branch compiles."""
+    if SWAP_AB:
+        acc = mx_weight_only_scalar_swapped(
+            acc, tl.reshape(a, (BLOCK_SIZE_K,)), w, w_scale,
+            N_WIDTH, BLOCK_SIZE_K, SCALE_GROUP_K,
+        )
+    elif COMPUTE_MODE == "dot_scaled":
+        acc = mx_dot_scaled(acc, a, None, w, w_scale)
+    else:
+        acc = acc + tl.dot(
+            a, mx_weight_upcast(w, w_scale, BLOCK_SIZE_K, N_WIDTH, SCALE_GROUP_K, a.dtype)
+        )
+    return acc
+@triton.jit
+def mx_weight_only_scalar_swapped(
+    acc,
+    a,
+    w,
+    w_scale,
+    ROWS_W: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+):
+    """Swapped weight-only reduce: CUDA-core FMA GEVM against RAW bf16 activations (no act
+    scale — the weight-only counterpart of ``mx_scalar_reduce_swapped``). Weight
+    ``w`` is output-rows-major ``[ROWS_W, BK]``, ``a`` the ``[BK]`` token; the UE8M0 group
+    scale factors out of the inner sum, so it costs one multiply per group instead of per
+    element. Returns ``acc + [1, ROWS_W]``.
+    This is the decode structure for W4A16: with no MMA there is no M->16 pad to waste, and
+    the loop runs at the dequantized weight stream's own rate (gpt-oss gate_up 32.0 -> 13.7us,
+    down 19.8 -> 9.1us, against a 12.3us/6us dequant-inclusive load floor)."""
+    NG: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K
+    wq = e2m1_cols_to_f32(w) if w.dtype == tl.uint8 else w.to(tl.float32)  # [ROWS_W, BK]
+    prod = wq * a.to(tl.float32)[None, :]
+    grp = tl.sum(tl.reshape(prod, (ROWS_W, NG, SCALE_GROUP_K)), axis=2)  # [ROWS_W, NG]
+    return acc + tl.reshape(
+        tl.sum(grp * decode_group_scale(w_scale), axis=1), (1, ROWS_W)
+    )
+@triton.jit
+def mx_scalar_reduce(
+    acc,
+    a,
+    a_scale,
+    w,
+    w_scale,
+    ROWS_W: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    SWAP_AB: tl.constexpr = False,
+):
+    """MX 'scalar' step: CUDA-core FMA GEMV, unpacking MXFP4 to fp32 and dequantizing activation +
+    weight by their group scales, reducing over K into ``acc``. No tensor core (so no M->16 MMA
+    pad) — wins the memory-bound decode GEMV (M=1). The gate_up kernels pass the stacked gate|up
+    tile (ROWS_W = 2*BN). The UE8M0 scale is constant within each ``SCALE_GROUP_K`` group, so it
+    factors out of the inner sum: reduce the raw products per group, then apply ONE combined
+    (act x weight) scale per group — ``SCALE_GROUP_K``x fewer scale-muls (~18% faster, bit-identical
+    to the expanded form). Plain: ``a`` [BM, BK] against ``w`` [BK, ROWS_W]. ``SWAP_AB`` (decode):
+    ``w`` output-rows-major [ROWS_W, BK] against the [BK] token, MXFP4 unpacked along columns (K),
+    no transposes. Both return ``acc + [1, ROWS_W]``."""
+    NG: tl.constexpr = BLOCK_SIZE_K // SCALE_GROUP_K
+    aq = e2m1_cols_to_f32(a) if a.dtype == tl.uint8 else a.to(tl.float32)
+    if SWAP_AB:
+        wq = e2m1_cols_to_f32(w) if w.dtype == tl.uint8 else w.to(tl.float32)  # [ROWS_W, BK]
+        grp = tl.sum(tl.reshape(aq[None, :] * wq, (ROWS_W, NG, SCALE_GROUP_K)), axis=2)
+        scale = decode_group_scale(a_scale)[None, :] * decode_group_scale(w_scale)  # [ROWS_W, NG]
+        acc = acc + tl.reshape(tl.sum(grp * scale, axis=1), (1, ROWS_W))
+    else:
+        wq = e2m1_to_f32(w) if w.dtype == tl.uint8 else w.to(tl.float32)  # [BK, ROWS_W]
+        grp = tl.sum(tl.reshape(tl.trans(aq) * wq, (NG, SCALE_GROUP_K, ROWS_W)), axis=1)
+        scale = tl.trans(decode_group_scale(a_scale)) * tl.trans(decode_group_scale(w_scale))
+        acc = acc + tl.sum(grp * scale, axis=0)[None, :]
+    return acc
+@triton.jit
+def mx_compute(
+    acc,
+    a,
+    a_scale,
+    w,
+    w_scale,
+    COMPUTE_MODE: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    SWAP_AB: tl.constexpr = False,
+):
+    """Single-projection MMA step, dispatched on ``COMPUTE_MODE`` (scaled-MMA on the raw weight,
+    fp8 ``tl.dot`` + per-group rescale, or the scalar reduce); ``SWAP_AB`` selects each leaf's
+    decode form — weight output rows in the MMA M dim, the single token flattened to the [BK] rhs
+    (a packed-E2M1 token flattens to its BYTE length; dot_scaled consumes it packed, dot/scalar
+    column-unpack it losslessly). The acc shapes diverge across modes and orientations, but only
+    the taken constexpr branch compiles, so the single return never has to unify them.
+    ``BLOCK_SIZE_N`` is the weight tile's row count — the gate_up kernels pass ``2*BN`` with the
+    stacked gate|up tile (split back via ``split_gate_up``): one load and one MMA for both
+    projections keeps the native microscaled-MMA M=128 operand at BN=64, doubling the CTAs on the
+    parallelism-starved decode grid (dsv4 gate_up 1.34x, bit-exact)."""
+    if SWAP_AB:
+        if a.dtype == tl.uint8:
+            a1 = tl.reshape(a, (BLOCK_SIZE_K // 2,))
+        else:
+            a1 = tl.reshape(a, (BLOCK_SIZE_K,))
+        as1 = tl.reshape(a_scale, (BLOCK_SIZE_K // SCALE_GROUP_K,))
+    else:
+        a1 = a
+        as1 = a_scale
+    if COMPUTE_MODE == "dot_scaled":
+        acc = mx_dot_scaled(acc, a1, as1, w, w_scale, SWAP_AB, BLOCK_SIZE_K, SCALE_GROUP_K)
+    elif COMPUTE_MODE == "dot":
+        acc = mx_dot_rescale(acc, a1, as1, w, w_scale, SWAP_AB, BLOCK_SIZE_K)
+    else:  # scalar
+        acc = mx_scalar_reduce(acc, a1, as1, w, w_scale, BLOCK_SIZE_N, BLOCK_SIZE_K, SCALE_GROUP_K, SWAP_AB)
+    return acc
+# ── swap-AB decode compute: M=1 batched GEMV with output rows in the MMA M dim ───────
+#
+# The batched (decode) kernels are structurally M=1, where the sm_100 scaled MMA pads M→128.
+# Putting the WEIGHT's output rows in M (fully used) and the single token in N cuts that pad to
+# the N-atom (16) — a ~1.5× decode win for fp4 (dot_scaled), neutral-to-worse for fp8 (scalar
+# stays memory-bound and wins, so the tuner keeps it). Weight is loaded output-rows-major
+# ``[BN, BK]`` for every mode, so the kernel does ONE load; each helper returns ``[1, BN]``.
+# The sm_100 MMA's minimum N tile (16). In the swap path the single decode token sits in the MMA's
+# N dim, so it must be padded up to this width (col 0 = the token, cols 1..15 = zero). It is NOT a
+# block size — BLOCK_SIZE_M stays 1 under swap; this is the token's *padded N extent*, fixed by the
+# hardware. The plain int is the source of truth: host-side consumers (the autotune pruners, which
+# filter configs in ordinary Python) compare against it directly, while kernels read the
+# tl.constexpr view derived from it — the only module-global form a @triton.jit fn can read.
+MMA_N_ATOM_WIDTH = 16
+MMA_N_ATOM = tl.constexpr(MMA_N_ATOM_WIDTH)
+@triton.jit
+def swap_pad_rhs(a, BLOCK_SIZE_K: tl.constexpr):
+    """Pad the ``[BLOCK_SIZE_K]`` M=1 token to the ``[BLOCK_SIZE_K, MMA_N_ATOM]`` swap-AB MMA rhs —
+    only column 0 is the real token (16 is measured, not a hardware floor — see acc_init). Used
+    by the M=1 batched / fused-MoE fp8 ``tl.dot`` swap paths (weight output rows in the MMA M dim);
+    the caller takes column 0 of the ``[BN, MMA_N_ATOM]`` result after the K-loop."""
+    return tl.where(
+        tl.arange(0, MMA_N_ATOM)[None, :] == 0,
+        a[:, None],
+        tl.zeros((BLOCK_SIZE_K, MMA_N_ATOM), a.dtype),
+    )
+@triton.jit
+def fp8_dot(a, b, SWAP_AB: tl.constexpr, BLOCK_SIZE_K: tl.constexpr):
+    """Swap-aware plain ``tl.dot`` for the fp8 (block/tensor-dynamic) paths — no scaling; the caller
+    applies its own per-block/per-tensor scales. Swap: weight ``b`` ``[N, BK]`` × the single token
+    padded to the N=16 atom → ``[N, MMA_N_ATOM]`` (col 0 real). No-swap: token ``a`` ``[M, BK]`` ×
+    weight ``b`` ``[BK, N]`` → ``[M, N]``. ``BLOCK_SIZE_K`` is the contraction tile (the down
+    projection passes its intermediate tile). Single return: only the taken branch compiles."""
+    if SWAP_AB:
+        out = tl.dot(b, swap_pad_rhs(tl.reshape(a, (BLOCK_SIZE_K,)), BLOCK_SIZE_K))
+    else:
+        out = tl.dot(a, b)
+    return out
+@triton.jit
+def block_dynamic_dot(
+    acc, a, a_s, b, b_s,
+    block_k: tl.constexpr, SWAP_AB: tl.constexpr, USE_DOT_SCALED: tl.constexpr,
+    FAKE_BATCH: tl.constexpr = False,
+):
+    """Accumulate one block-dynamic (1x128/128x128) fp8 K-tile into ``acc``, oriented by
+    ``SWAP_AB``. ``USE_DOT_SCALED`` (UE8M0 scales on a native-M tile): fold the group scales
+    into the tcgen05 MMA — the tile's single 128-group scale broadcasts in-register to the
+    ``block_k // 32`` group-32 columns ``dot_scaled`` consumes, identical to a 128-group
+    rescale but with no 4x scale memory and no software multiply. Else: plain fp8 ``tl.dot``
+    + per-group software rescale (``decode_group_scale`` is a no-op on fp32 scales, decodes
+    UE8M0). ``FAKE_BATCH`` (single-token decode): ``fp8_dot`` pads the lone token to the MMA N
+    atom and both scales broadcast down the weight-row (M) dim. Single-exit if/else so only the
+    taken arm type-checks (a trailing fall-through arm would be checked even when an earlier
+    branch is taken)."""
+    if USE_DOT_SCALED:
+        reps: tl.constexpr = block_k // 32
+        a_sg = a_s[:, None].broadcast_to(a_s.shape[0], reps)
+        b_sg = b_s[:, None].broadcast_to(b_s.shape[0], reps)
+        if SWAP_AB:
+            acc = tl.dot_scaled(b, b_sg, "e4m3", a, a_sg, "e4m3", acc)
+        else:
+            acc = tl.dot_scaled(a, a_sg, "e4m3", b, b_sg, "e4m3", acc)
+    else:
+        # plain fp8 tl.dot + per-group decoded scales, oriented by SWAP_AB (weight rows in the MMA
+        # M dim under swap). decode_group_scale: fp32 passthrough, UE8M0 -> 2^(e-127).
+        a_sd = decode_group_scale(a_s)
+        b_sd = decode_group_scale(b_s)
+        if FAKE_BATCH:
+            d = fp8_dot(a, b, SWAP_AB, block_k)
+            if SWAP_AB:  # [BN, N-atom]: weight-row scale down the M=BN dim, act scalar broadcasts
+                acc = acc + d * a_sd[:, None] * b_sd[:, None]
+            else:  # no-swap [BM, BN] GEVM: act scale down M, weight-N scale across N (upstream form)
+                acc = acc + d * a_sd[:, None] * b_sd[None, :]
+        elif SWAP_AB:
+            acc = acc + tl.dot(b, a) * b_sd[:, None] * a_sd[None, :]
+        else:
+            acc = acc + tl.dot(a, b) * a_sd[:, None] * b_sd[None, :]
+    return acc
+@triton.jit
+def static_dot(acc, a, b, b_s, SWAP_AB: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, FAKE_BATCH: tl.constexpr):
+    """static format K-step: plain (swap-aware) fp8 dot + per-K-block weight rescale. FAKE_BATCH
+    (single-token decode) routes the rescale down the weight-row (M) dim; else it broadcasts across
+    the N columns. The per-tensor activation scale is applied post-loop."""
+    b_sd = decode_group_scale(b_s)
+    d = fp8_dot(a, b, SWAP_AB, BLOCK_SIZE_K)
+    if FAKE_BATCH:
+        acc = acc + d * b_sd[:, None]
+    else:
+        acc = acc + d * b_sd[None, :]
+    return acc

build/torch-rocm/moe.py ADDED Viewed

	@@ -0,0 +1,864 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""MoE forwards — thin orchestrations over the base ``matmul_grouped`` / ``matmul_batched`` ops.
+The base ops carry the gate|up fusion (SwiGLU + FP8/MX requant) and the gather/scatter row
+maps, so both the fused and unfused MoE forwards are pure sequencing here — no MoE-specific
+kernels live in this module:
+  fused:   gate_up (``gate=True`` + ``quantize_output=True``) -> down -> ``weighted_reduce``. The
+           SwiGLU + intermediate requant happen inside the gate_up kernel epilogue. A model's
+           ``post_expert_norm`` runs on the down's routed rows, before the reduce.
+  unfused: gate_up (plain GEMM) -> host ``apply_glu`` -> down (plain GEMM) -> ``weighted_reduce``.
+           The activation + requant happen between two plain GEMMs; the GEMMs self-quantize their
+           raw inputs (``As=None``). Same math as the fused path, split across kernels.
+grouped (prefill) shares one on-device routing pass (``compute_grouped_scheduling``): gate_up
+gathers hidden by routed row and leaves its output expert-ordered; down reads it in place and
+scatters to routed rows. batched (decode) dispatches per token: ``gather_idx`` reads each routed
+row from the unexpanded hidden in-kernel (no copy), and EP-sentinel rows (``id >= num_experts``)
+are left uninit by the GEMM and skipped in ``weighted_reduce``. ``moe_fused_*`` / ``moe_unfused_*`` are format-neutral:
+the base ops dispatch on the weight dtypes / scale layout (block-dynamic FP8, MXFP4/MXFP8,
+NVFP4), and every forward takes ``activation_format`` naming the block's activation quantization
+(``None`` = the weights' own format, ``"bf16"`` = weight-only, or an explicit format such as
+``"mxfp8"`` on MXFP4 weights for W4A8)."""
+import functools
+from collections.abc import Callable
+import torch
+import torch.nn.functional as F
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait
+from .grouped import matmul_grouped
+from .batched import GATE_UNSTACK_MAX_S, matmul_batched
+from .bayesian_autotuner import bayesian_autotune
+from .compat import MX_SCALE_GROUP_K, NVFP4_SCALE_GROUP_K, compile_time_only_triton_wrap, decode_pdl, device_context
+from .formats import get_supported_act_fns, is_mx, is_mxfp4, weight_format
+from .norm import norm_column_factor, rms_inv_rows, rms_norm_rows
+from .quant import _launch_act_quant
+from .scheduling import compute_grouped_scheduling
+from .epilogue import fused_glu
+@bayesian_autotune(
+    [
+        triton.Config({"BLOCK_H": block_h}, num_warps=warps)
+        for block_h in (256, 512, 1024, 2048)
+        for warps in (4, 8)
+    ],
+    # the H tile width trades off against grid occupancy: at few groups (decode) narrow
+    # tiles spread more H-blocks across SMs, at many groups (prefill) wide tiles amortize
+    # the per-row weight load — so key on H and the group-count bucket.
+    ["H", "num_groups_bit_length", "NORM"],
+    n_trials=8,
+)
+@triton.jit
+def weighted_reduce_kernel(
+    Rows,  # (num_groups * NUM_TOP_K, H) — rows to reduce, group-major
+    Out,  # (num_groups, H) — one reduced row per group
+    Ids,  # (num_groups, NUM_TOP_K) — per-row id; a row is skipped when its id >= NUM_EXPERTS
+    Weights,  # (num_groups * NUM_TOP_K,) — per-row scale
+    NormWeight,  # (H,) fused post-expert norm weight; None (with NormInv) folds the arm out
+    NormInv,  # (num_groups * NUM_TOP_K,) fp32 per-row rsqrt from rms_inv_rows
+    H,
+    stride_rows_m,
+    stride_rows_h,
+    stride_o_m,
+    stride_o_h,
+    stride_ids_m,
+    stride_ids_k,
+    num_groups_bit_length,  # autotune key only (log2 group-count bucket); unused in body
+    NUM_TOP_K: tl.constexpr,
+    NUM_EXPERTS: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+    NORM: tl.constexpr = None,  # a get_supported_norms() name folded into the reduce
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    PDL: tl.constexpr = False,
+):
+    """Per group ``g``, the weighted sum of its ``NUM_TOP_K`` rows into ``Out[g]``:
+    ``sum_k Weights[g*NUM_TOP_K + k] * Rows[g*NUM_TOP_K + k]``, skipping rows whose id is
+    ``>= NUM_EXPERTS`` (out-of-range rows are never written upstream and contribute 0).
+    fp32 accumulate; ~2.8x a generic ``view(g, k, H).sum(1)``. ``SIMULATE_UNFUSED`` rounds
+    each weighted row to ``Out``'s dtype before summing, matching a reference that weights
+    in that dtype; production leaves the accumulation in fp32.
+    ``NORM`` folds a per-expert output norm into the pass: each row is scaled by its
+    ``NormInv`` (``rms_inv_rows``, one pass ahead of this one) and the norm's column factor
+    before the routing weight, so the normalized rows are never written or re-read."""
+    if PDL:
+        gdc_wait()
+    g = tl.program_id(0)
+    offs_h = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask = offs_h < H
+    acc = tl.zeros((BLOCK_H,), tl.float32)
+    if NORM is not None:  # the tile's column factor is row-independent — load it once
+        factor = norm_column_factor(NormWeight, offs_h, mask, NORM)
+    for k in tl.static_range(NUM_TOP_K):
+        flat = g * NUM_TOP_K + k
+        valid = tl.load(Ids + g * stride_ids_m + k * stride_ids_k) < NUM_EXPERTS
+        weight = tl.load(Weights + flat)
+        row = tl.load(
+            Rows + flat * stride_rows_m + offs_h * stride_rows_h,
+            mask=mask & valid,
+            other=0.0,
+        ).to(tl.float32)
+        if NORM is not None:
+            # masked like the value load: a sentinel row's rsqrt comes off uninitialized memory,
+            # and 0 * NaN would poison the sum
+            row = row * factor * tl.load(NormInv + flat, mask=valid, other=0.0)
+            if SIMULATE_UNFUSED:  # the reference materializes the normalized row in Out's dtype
+                row = row.to(Out.dtype.element_ty).to(tl.float32)
+        contrib = weight * row
+        if SIMULATE_UNFUSED:
+            contrib = contrib.to(Out.dtype.element_ty).to(tl.float32)
+        acc += contrib
+    if PDL:
+        gdc_launch_dependents()
+    tl.store(
+        Out + g * stride_o_m + offs_h * stride_o_h,
+        acc.to(Out.dtype.element_ty),
+        mask=mask,
+    )
+def weighted_reduce(
+    rows: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    num_experts: int,
+    simulate_unfused: bool = False,
+    norm: str | None = None,
+    norm_weight: torch.Tensor | None = None,
+    norm_inv: torch.Tensor | None = None,
+) -> torch.Tensor:
+    """Routing-weighted top-k reduce — the bookend of the fused-MoE chain. Folds each token's
+    ``num_top_k`` expert-output rows (``rows``, group-major, scaled by ``top_k_weights``, with
+    EP-sentinel rows ``id >= num_experts`` skipped) from the routed-row layout back to
+    ``(num_tokens, H)``. See ``weighted_reduce_kernel``. ``norm`` (a ``get_supported_norms()``
+    name, with ``norm_weight`` and the ``rms_inv_rows`` ``norm_inv``) folds a per-expert output
+    norm into this pass rather than normalizing the rows in one of their own."""
+    assert (norm is None) == (norm_inv is None), "a fused norm needs its per-row rsqrt"
+    num_tokens, num_top_k = top_k_index.shape
+    H = rows.size(1)
+    reduced = torch.empty(num_tokens, H, device=rows.device, dtype=rows.dtype)
+    with device_context(rows.device):
+        compile_time_only_triton_wrap(weighted_reduce_kernel)[
+            lambda meta: (num_tokens, triton.cdiv(H, meta["BLOCK_H"]))
+        ](
+            rows,
+            reduced,
+            top_k_index,
+            top_k_weights,
+            norm_weight,
+            norm_inv,
+            H,
+            rows.stride(0),
+            rows.stride(1),
+            reduced.stride(0),
+            reduced.stride(1),
+            top_k_index.stride(0),
+            top_k_index.stride(1),
+            num_groups_bit_length=int(num_tokens).bit_length(),
+            NUM_TOP_K=num_top_k,
+            NUM_EXPERTS=num_experts,
+            NORM=norm,
+            SIMULATE_UNFUSED=simulate_unfused,
+            PDL=decode_pdl(),
+            launch_pdl=decode_pdl(),
+        )
+    return reduced
+def _validate_moe(gate_up_proj, gate_up_proj_scale, down_proj, down_proj_scale):
+    """gate_up and down must share the format (both MX or both block-dynamic FP8 — the
+    intermediate handed between them carries one quant format). Returns whether the format
+    is MX (the fused dispatchers branch on it); the fp8 quantization block is derived from
+    the scale shapes (``weight_block_size``), never passed. Scales are the pure block scales
+    (per the decoupled API — per-tensor globals ride as separate ``*_global_scale`` args);
+    the format predicates read the block scale's dtype/grouping."""
+    gate_up_is_mx = is_mx(gate_up_proj, gate_up_proj_scale)
+    if gate_up_is_mx != is_mx(down_proj, down_proj_scale):
+        raise ValueError(
+            "gate_up_proj and down_proj must use the same format (both MX or both block-dynamic FP8)."
+        )
+    if is_mxfp4(gate_up_proj, gate_up_proj_scale) != is_mxfp4(down_proj, down_proj_scale):
+        raise ValueError("gate_up_proj and down_proj must use the same MX format (both MXFP4 or both MXFP8).")
+    return gate_up_is_mx
+def _gather_idx(top_k_index: torch.Tensor) -> torch.Tensor:
+    """The batched routed-row gather: routed row ``s`` (``= t*K + k``) reads token ``s // num_top_k``
+    of the unexpanded hidden. ``matmul_batched`` applies it in-kernel, so no ``(S, H)`` copy.
+    The map depends only on the SHAPE, so it is cached per (tokens, top_k, device) — its two
+    elementwise launches (~3µs) were the largest non-GEMM cost of the fp8 decode chain."""
+    return _gather_idx_cached(
+        top_k_index.shape[0], top_k_index.shape[1], top_k_index.device
+    )
+@functools.lru_cache(maxsize=64)
+def _gather_idx_cached(
+    num_tokens: int, num_top_k: int, device: torch.device
+) -> torch.Tensor:
+    return (
+        torch.arange(num_tokens * num_top_k, device=device, dtype=torch.int32)
+        // num_top_k
+    )
+def _torch_weighted_reduce(down_out, top_k_index, top_k_weights, num_experts):
+    """Naive (unfused) routing-weighted top-k reduce in plain torch — NOT the fused
+    ``weighted_reduce`` kernel. Materializes the (bf16) weighted contribs, masks EP-sentinel rows
+    (``id >= num_experts``, left uninit in ``down_out``) to 0, and torch-sums to ``(num_tokens, H)``
+    (fp32 accumulate, activation-dtype out). This is the independent reference the fused
+    ``weighted_reduce`` is checked against; the fused path's ``simulate_unfused`` reproduces its
+    bf16-contrib rounding."""
+    num_tokens, num_top_k = top_k_index.shape
+    dropped = (top_k_index.reshape(-1) >= num_experts).reshape(-1, 1)
+    # masked in place on the PRODUCT: a sentinel row's down_out is uninitialized, so zeroing the
+    # weight instead would leave 0 * NaN == NaN.
+    contrib = down_out * top_k_weights.reshape(-1, 1)
+    contrib.masked_fill_(dropped, 0)
+    return contrib.view(num_tokens, num_top_k, down_out.size(1)).sum(dim=1)
+def _host_glu(gate_up_out, act_fn, swiglu_alpha, swiglu_limit, gate):
+    """The activation on the host, for the arms the epilogue does not fuse: a callable takes the
+    raw GEMM output (interleaved gate|up columns when ``gate``) and returns the intermediate —
+    the caller owns its semantics, so any activation runs without a kernel change; a name from
+    ``get_supported_act_fns()`` runs the torch GLU (``gate``) or the bare activation."""
+    if callable(act_fn):
+        return act_fn(gate_up_out)
+    if act_fn not in get_supported_act_fns():
+        raise ValueError(f"act_fn must be one of {get_supported_act_fns()} or a callable, got {act_fn!r}")
+    if gate:
+        return fused_glu(gate_up_out, act_fn, swiglu_alpha, swiglu_limit)
+    return {"silu": F.silu, "gelu": F.gelu, "relu": F.relu}[act_fn](gate_up_out)
+def _fused_glu(act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, gate) -> dict:
+    """The gate_up GEMM's fused-GLU kwargs when the activation is a fusable name; ``{}`` (a plain
+    GEMM) sends a callable — or an unfusable name, rejected in ``_host_glu`` — to the host."""
+    if not isinstance(act_fn, str) or act_fn not in get_supported_act_fns():
+        return {}
+    return dict(
+        gate=gate,
+        act_fn=act_fn,
+        swiglu_alpha=swiglu_alpha,
+        swiglu_limit=swiglu_limit,
+        simulate_unfused=simulate_unfused,
+    )
+# ── Fused (gate_up epilogue owns SwiGLU + intermediate requant) ──────────────
+def _block_format(gate_up_proj, gate_up_proj_scale, down_proj, down_proj_scale, activation_format):
+    """The MoE block's activation format, resolved to a name: validates the weight pairing; an
+    explicit ``activation_format`` (``"bf16"`` weight-only, or a format) is respected as-is,
+    ``None`` follows the weights (fp8 / mxfp8 / mxfp4 / nvfp4 — mxfp4 weights default to mxfp4
+    activations, the all-fp4 W4A4 chain; unquantized BF16/FP16 weights carry no scales, so their
+    activations stay ``"bf16"``, the full-precision path)."""
+    _validate_moe(gate_up_proj, gate_up_proj_scale, down_proj, down_proj_scale)
+    if activation_format is not None:
+        return activation_format
+    if gate_up_proj_scale is None:
+        return "bf16"
+    return weight_format(gate_up_proj, gate_up_proj_scale)
+def _post_expert_norm(down_out, post_expert_norm, weight, eps):
+    """A model's per-expert output norm applied to the down projection's routed rows — one row
+    per (token, expert) application, which is what such a norm is defined over. ``None`` passes
+    through; a ``get_supported_norms()`` name runs ``rms_norm_rows`` against ``weight``; anything
+    else is a host callable, like ``act_fn``'s."""
+    if post_expert_norm is None:
+        return down_out
+    if isinstance(post_expert_norm, str):
+        return rms_norm_rows(down_out, weight, eps, post_expert_norm)
+    return post_expert_norm(down_out)
+def _fused_post_expert_norm(down_out, post_expert_norm, weight, eps):
+    """``(rows, weighted_reduce kwargs)`` for a fused forward: a named norm rides INTO the reduce
+    — one pass for the row's ``rsqrt``, then the reduce it already runs applies that and the
+    column factor — so the normalized rows are never materialized. A callable has to run first."""
+    if isinstance(post_expert_norm, str):
+        return down_out, {
+            "norm": post_expert_norm,
+            "norm_weight": weight,
+            "norm_inv": rms_inv_rows(down_out, weight, eps, post_expert_norm),
+        }
+    return _post_expert_norm(down_out, post_expert_norm, weight, eps), {}
+def moe_fused_grouped(
+    hidden_states: torch.Tensor,  # (T, H)
+    top_k_index: torch.Tensor,  # (T, K) int
+    top_k_weights: torch.Tensor,  # (T, K)
+    gate_up_proj: torch.Tensor,  # (E, 2I, H)
+    down_proj: torch.Tensor,  # (E, H, I)
+    gate_up_proj_scale_inv: torch.Tensor,
+    down_proj_scale_inv: torch.Tensor,
+    gate_up_proj_weight_global_scale: torch.Tensor | None = None,
+    down_proj_weight_global_scale: torch.Tensor | None = None,
+    gate_up_proj_input_global_scale: torch.Tensor | None = None,
+    down_proj_input_global_scale: torch.Tensor | None = None,
+    post_expert_norm=None,  # the model's per-expert output norm on the routed rows: a
+    # get_supported_norms() name (fused into the reduce) or a host callable
+    post_expert_norm_weight: torch.Tensor | None = None,  # (H,) weight of a named norm
+    post_expert_norm_eps: float = 1e-6,
+    gate_up_proj_bias: torch.Tensor | None = None,  # (E, 2I) pre-activation bias
+    down_proj_bias: torch.Tensor | None = None,  # (E, H)
+    act_fn: str | Callable = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    activation_format: str | None = None,
+    gate: bool = True,
+) -> torch.Tensor:
+    """Fused grouped MoE (prefill): gather gate_up + SiLU + requant epilogue → quantized
+    expert-ordered intermediate → grouped down → routing-weighted top-k reduce. Returns
+    ``(num_tokens, hidden_dim)``. The base ops dispatch on the weight dtypes / scale
+    layout (block-dynamic FP8, MXFP8/MXFP4, NVFP4); ``activation_format`` names the activation
+    quantization for the whole block — activations and the fused intermediate requant
+    carry it (``"mxfp4"``/``"nvfp4"`` run all-fp4 W4A4 chains, ``"bf16"`` is weight-only);
+    ``None`` follows the weight format, and the ops validate the pairing. The
+    ``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's
+    calibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own,
+    requants the intermediate against the down's, and the down consumes that as its activation
+    global. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each step through
+    the activation dtype so the output matches the unfused reference to reduce order. ``act_fn`` is a
+    ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
+    callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
+    projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
+    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine."""
+    fmt = _block_format(
+        gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
+    )
+    num_top_k = top_k_index.size(-1)
+    glu = _fused_glu(act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, gate)
+    NUM_EXPERTS = gate_up_proj.size(0)
+    expert_start, gather_idx, scatter_idx = compute_grouped_scheduling(
+        top_k_index, NUM_EXPERTS, num_top_k
+    )
+    # Phase 1: gate_up + SiLU + requant in the block format -> expert-ordered quantized
+    # intermediate (the op quantizes the raw hidden itself and owns the expand-vs-gather
+    # regime policy — this forward is pure sequencing). scatter_idx=None: the down reads
+    # the intermediate in place. (C, Cs) under a requant format; a bare Tensor otherwise.
+    gate_up_out = matmul_grouped(
+        hidden_states,
+        gate_up_proj,
+        Bs=gate_up_proj_scale_inv,
+        a_global_scale=gate_up_proj_input_global_scale,
+        b_global_scale=gate_up_proj_weight_global_scale,
+        # the intermediate requant normalizes against the DOWN's calibrated input global,
+        # which the down then consumes as its activation global — the two-level handoff
+        output_global_scale=down_proj_input_global_scale,
+        expert_start=expert_start,
+        **glu,
+        bias=gate_up_proj_bias,
+        # fmt is the resolved format; "bf16" (weight-only) leaves the GLU intermediate bf16,
+        # no requant.
+        activation_format=fmt,
+        quantize_output=bool(glu) and fmt != "bf16",
+        output_dtype=hidden_states.dtype,
+        gather_idx=gather_idx,
+    )
+    inter, inter_scale = (
+        gate_up_out if isinstance(gate_up_out, tuple) else (gate_up_out, None)
+    )
+    if not glu:
+        inter = _host_glu(inter, act_fn, swiglu_alpha, swiglu_limit, gate)
+    # Phase 2: grouped down over the expert-ordered pre-quantized intermediate (its dtypes
+    # carry the format; gather_idx=None), scattering to routed rows (scatter_idx).
+    down_out = matmul_grouped(
+        inter,
+        down_proj,
+        As=inter_scale,
+        Bs=down_proj_scale_inv,
+        a_global_scale=down_proj_input_global_scale,
+        b_global_scale=down_proj_weight_global_scale,
+        expert_start=expert_start,
+        bias=down_proj_bias,
+        # weight-only: the intermediate is bf16 (As None) — the down goes weight-only too; a
+        # quantized intermediate carries its format in its own dtypes.
+        activation_format=fmt if inter_scale is None else None,
+        output_dtype=hidden_states.dtype,
+        scatter_idx=scatter_idx,
+    )
+    # Phase 3: routing-weighted top-k reduce -> (num_tokens, hidden_dim). simulate_unfused
+    # rounds each weighted contrib to the activation dtype before summing, matching the
+    # unfused path's torch reduce (which materializes bf16 contribs); production
+    # accumulates in fp32.
+    rows, norm_kwargs = _fused_post_expert_norm(
+        down_out, post_expert_norm, post_expert_norm_weight, post_expert_norm_eps
+    )
+    return weighted_reduce(
+        rows, top_k_index, top_k_weights, NUM_EXPERTS, simulate_unfused, **norm_kwargs
+    )
+def moe_fused_batched(
+    hidden_states: torch.Tensor,  # (T, H)
+    top_k_index: torch.Tensor,  # (T, K) int
+    top_k_weights: torch.Tensor,  # (T, K)
+    gate_up_proj: torch.Tensor,  # (E, 2I, H)
+    down_proj: torch.Tensor,  # (E, H, I)
+    gate_up_proj_scale_inv: torch.Tensor,
+    down_proj_scale_inv: torch.Tensor,
+    gate_up_proj_weight_global_scale: torch.Tensor | None = None,
+    down_proj_weight_global_scale: torch.Tensor | None = None,
+    gate_up_proj_input_global_scale: torch.Tensor | None = None,
+    down_proj_input_global_scale: torch.Tensor | None = None,
+    post_expert_norm=None,  # the model's per-expert output norm on the routed rows: a
+    # get_supported_norms() name (fused into the reduce) or a host callable
+    post_expert_norm_weight: torch.Tensor | None = None,  # (H,) weight of a named norm
+    post_expert_norm_eps: float = 1e-6,
+    gate_up_proj_bias: torch.Tensor | None = None,  # (E, 2I) pre-activation bias
+    down_proj_bias: torch.Tensor | None = None,  # (E, H)
+    act_fn: str | Callable = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    simulate_unfused: bool = False,
+    activation_format: str | None = None,
+    gate: bool = True,
+) -> torch.Tensor:
+    """Fused batched MoE (decode): gate_up + SiLU + requant epilogue → per-row quantized
+    intermediate → batched down → routing-weighted top-k reduce. Returns
+    ``(num_tokens, hidden_dim)``. The base ops dispatch on the weight dtypes / scale
+    layout (block-dynamic FP8, MXFP8/MXFP4, NVFP4 — decode runs the software/swap arms
+    below the native mxf4nvf4 M=128 staging); ``activation_format`` names the activation
+    quantization for the whole block — activations and the fused intermediate requant
+    carry it (``"mxfp4"`` runs the all-fp4 W4A4 chain, ``"bf16"`` is weight-only); ``None``
+    follows the weight format, and the ops validate the pairing. The
+    ``*_proj_input_global_scale`` pair is the NVFP4 activation second level (a checkpoint's
+    calibrated per-projection ``input_scale``): the gate_up quantizes hidden against its own,
+    requants the intermediate against the down's, and the down consumes that as its activation
+    global. ``None`` is dynamic quant everywhere else. ``simulate_unfused`` (testing) rounds each
+    step through the activation dtype so the output matches the unfused reference to
+    reduce order. ``act_fn`` is a
+    ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
+    callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
+    projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
+    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine."""
+    fmt = _block_format(
+        gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
+    )
+    glu = _fused_glu(act_fn, swiglu_alpha, swiglu_limit, simulate_unfused, gate)
+    NUM_EXPERTS = gate_up_proj.size(0)
+    expert_ids = top_k_index.reshape(-1)
+    gather_idx = _gather_idx(top_k_index)
+    # Phase 1: gate_up + SiLU + requant in the block format -> per-row quantized
+    # intermediate (the op quantizes the raw activations). gather_idx reads each routed
+    # row from the unexpanded hidden in-kernel (no copy).
+    # (C, Cs) under a requant format; a bare Tensor on the full-precision path
+    gate_up_out = matmul_batched(
+        hidden_states,
+        gate_up_proj,
+        Bs=gate_up_proj_scale_inv,
+        a_global_scale=gate_up_proj_input_global_scale,
+        b_global_scale=gate_up_proj_weight_global_scale,
+        # the two-level handoff, as in the grouped sibling
+        output_global_scale=down_proj_input_global_scale,
+        expert_ids=expert_ids,
+        **glu,
+        bias=gate_up_proj_bias,
+        # Decode (batched): "bf16" (weight-only) leaves the intermediate bf16, no requant.
+        # Block-FP8: INSIDE the unstacked decode band the requant fuses into the GLU kernel
+        # (``fused_glu(quant_group=...)`` — one launch, hands the down a ready fp8+scales intermediate and
+        # kills its offline act quant); ABOVE the band the stacked epilogue's requant pins
+        # the gate|up tile to the whole block scale and halves the grid, so the bf16 handoff
+        # (down inline-quants) stays the win there.
+        activation_format=fmt,
+        quantize_output=(
+            bool(glu) and fmt != "bf16" and (fmt != "fp8" or expert_ids.numel() <= GATE_UNSTACK_MAX_S)
+        ),
+        output_dtype=hidden_states.dtype,
+        gather_idx=gather_idx,
+    )
+    inter, inter_scale = (
+        gate_up_out if isinstance(gate_up_out, tuple) else (gate_up_out, None)
+    )
+    if not glu:
+        inter = _host_glu(inter, act_fn, swiglu_alpha, swiglu_limit, gate)
+    # Phase 2: batched down over the intermediate (its dtypes carry the format; already
+    # routed-order, no gather).
+    down_out = matmul_batched(
+        inter,
+        down_proj,
+        As=inter_scale,
+        Bs=down_proj_scale_inv,
+        a_global_scale=down_proj_input_global_scale,
+        b_global_scale=down_proj_weight_global_scale,
+        expert_ids=expert_ids,
+        bias=down_proj_bias,
+        # weight-only / block-FP8: the intermediate is bf16 (As is None), so the down carries the
+        # format and quantizes it, mirroring the unfused sibling; a quantized MX intermediate
+        # carries its format in its own dtypes.
+        activation_format=fmt if inter_scale is None or fmt == "fp8" else None,
+        output_dtype=hidden_states.dtype,
+    )
+    # Phase 3: routing-weighted top-k reduce -> (num_tokens, hidden_dim). simulate_unfused
+    # rounds each weighted contrib to the activation dtype before summing, matching the
+    # unfused path's torch reduce (which materializes bf16 contribs); production
+    # accumulates in fp32.
+    rows, norm_kwargs = _fused_post_expert_norm(
+        down_out, post_expert_norm, post_expert_norm_weight, post_expert_norm_eps
+    )
+    return weighted_reduce(
+        rows, top_k_index, top_k_weights, NUM_EXPERTS, simulate_unfused, **norm_kwargs
+    )
+# ── Unfused (plain GEMMs + host GLU) ──────────────────────────────────────────
+def moe_unfused_grouped(
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    gate_up_proj: torch.Tensor,
+    down_proj: torch.Tensor,
+    gate_up_proj_scale_inv: torch.Tensor,
+    down_proj_scale_inv: torch.Tensor,
+    gate_up_proj_weight_global_scale: torch.Tensor | None = None,
+    down_proj_weight_global_scale: torch.Tensor | None = None,
+    gate_up_proj_input_global_scale: torch.Tensor | None = None,
+    down_proj_input_global_scale: torch.Tensor | None = None,
+    post_expert_norm=None,  # the model's per-expert output norm on the routed rows: a
+    # get_supported_norms() name (fused into the reduce) or a host callable
+    post_expert_norm_weight: torch.Tensor | None = None,  # (H,) weight of a named norm
+    post_expert_norm_eps: float = 1e-6,
+    gate_up_proj_bias: torch.Tensor | None = None,  # (E, 2I) pre-activation bias
+    down_proj_bias: torch.Tensor | None = None,  # (E, H)
+    act_fn: str | Callable = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    activation_format: str | None = None,
+    gate: bool = True,
+) -> torch.Tensor:
+    """Unfused grouped MoE: gate_up (plain grouped GEMM, gather hidden) → host ``apply_glu`` →
+    down (plain grouped GEMM, scatter to routed rows) → routing-weighted reduce. Same math as
+    ``moe_fused_grouped`` but the SwiGLU + intermediate quant happen between two plain GEMMs
+    rather than inside the gate_up epilogue; each GEMM quantizes its raw input in
+    ``activation_format`` (``None`` follows the weight format, mirroring the fused forward — mxfp4
+    weights run the all-fp4 W4A4 chain). All formats route through the shared ``matmul_grouped``. The NVFP4 activation globals thread the same way as the fused sibling:
+    each GEMM quantizes its raw input against its ``*_input_global_scale``. ``act_fn`` is a
+    ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
+    callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
+    projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
+    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine."""
+    fmt = _block_format(
+        gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
+    )
+    num_top_k = top_k_index.size(-1)
+    NUM_EXPERTS = gate_up_proj.size(0)
+    expert_start, gather_idx, scatter_idx = compute_grouped_scheduling(
+        top_k_index, NUM_EXPERTS, num_top_k
+    )
+    # gate_up as a plain GEMM (no gate epilogue) over gathered hidden -> expert-ordered (S, 2I).
+    gate_up_out = matmul_grouped(
+        hidden_states,
+        gate_up_proj,
+        Bs=gate_up_proj_scale_inv,
+        a_global_scale=gate_up_proj_input_global_scale,
+        b_global_scale=gate_up_proj_weight_global_scale,
+        expert_start=expert_start,
+        bias=gate_up_proj_bias,
+        activation_format=fmt,
+        output_dtype=hidden_states.dtype,
+        gather_idx=gather_idx,
+    )
+    inter = _host_glu(gate_up_out, act_fn, swiglu_alpha, swiglu_limit, gate)
+    # down over the expert-ordered intermediate (quantized in the same format), scattering
+    # to routed rows.
+    down_out = matmul_grouped(
+        inter,
+        down_proj,
+        Bs=down_proj_scale_inv,
+        a_global_scale=down_proj_input_global_scale,
+        b_global_scale=down_proj_weight_global_scale,
+        expert_start=expert_start,
+        bias=down_proj_bias,
+        activation_format=fmt,
+        output_dtype=hidden_states.dtype,
+        scatter_idx=scatter_idx,
+    )
+    return _torch_weighted_reduce(
+        _post_expert_norm(down_out, post_expert_norm, post_expert_norm_weight, post_expert_norm_eps),
+        top_k_index, top_k_weights, NUM_EXPERTS,
+    )
+def moe_torch_grouped(
+    hidden_states: torch.Tensor,  # (T, H)
+    top_k_index: torch.Tensor,  # (T, K) int
+    top_k_weights: torch.Tensor,  # (T, K)
+    gate_up_proj: torch.Tensor,  # (E, 2I, H) E4M3
+    down_proj: torch.Tensor,  # (E, H, I) E4M3
+    gate_up_proj_scale_inv: torch.Tensor,  # gate_up scale through torchao's triton_mx_block_rearrange_per_group_3d (NOT swizzle_mx_scales)
+    down_proj_scale_inv: torch.Tensor,  # down scale through the same torchao rearrange
+    gate_up_proj_weight_global_scale: torch.Tensor | None = None,
+    down_proj_weight_global_scale: torch.Tensor | None = None,
+    gate_up_proj_input_global_scale: torch.Tensor | None = None,
+    down_proj_input_global_scale: torch.Tensor | None = None,
+    post_expert_norm=None,  # the model's per-expert output norm on the routed rows: a
+    # get_supported_norms() name (fused into the reduce) or a host callable
+    post_expert_norm_weight: torch.Tensor | None = None,  # (H,) weight of a named norm
+    post_expert_norm_eps: float = 1e-6,
+    gate_up_proj_bias: torch.Tensor | None = None,  # (E, 2I) pre-activation bias
+    down_proj_bias: torch.Tensor | None = None,  # (E, H)
+    act_fn: str = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    activation_format: str | None = None,
+) -> torch.Tensor:
+    """Torch-only MX grouped MoE — the fair cuBLAS baseline for ``moe_fused_grouped`` /
+    ``moe_unfused_grouped`` on the PUBLIC ``torch.nn.functional.scaled_grouped_mm``. Same weights,
+    scales, and routing as our forwards; the only difference is the machinery torch forces:
+    - Routing by **sort**, not our on-device gather/scatter: stable-argsort the ``T*K`` routed slots
+      by expert into contiguous groups (cumulative ``offs``).
+    - Two ``scaled_grouped_mm`` calls (per-format ``ScalingType``: group-32 ``BlockWise1x32`` for
+      mxfp8/mxfp4, group-16 ``BlockWise1x16`` for nvfp4; fp4 operands viewed as ``e2m1_x2``).
+    - Our Triton MX act-quant (so torch is timed on the same fast quant), the shared host ``apply_glu``,
+      and the shared ``_torch_weighted_reduce``. All three MX formats.
+    WEIGHT scales arrive already SWIZZLE_32_4_4-blocked by **torchao's**
+    ``triton_mx_block_rearrange_per_group_3d`` (done once offline — a real deployment doesn't
+    reblock a fixed weight every forward); this is scaled_grouped_mm's own layout, NOT the
+    ``swizzle_mx_scales`` artifact the other four forwards consume. The timed loop only blocks the
+    ACTIVATION scale (which changes each call). The format is read off the dtypes (the block
+    preserves them: E4M3 scale = NVFP4, uint8 = MX; packed-E2M1 weight = int8) since the blocked
+    shape no longer matches the group-shape detectors."""
+    assert gate_up_proj.dtype in (torch.int8, torch.float8_e4m3fn), (
+        "torch grouped baseline is MX-only (packed E2M1 or E4M3 weights)"
+    )
+    assert gate_up_proj_scale_inv.ndim not in (5, 6) and down_proj_scale_inv.ndim not in (5, 6), (
+        "the torch baseline consumes torchao's triton_mx_block_rearrange_per_group_3d layout, "
+        "not the swizzle_mx_scales artifact — a byte-compatible wrong layout would misread silently"
+    )
+    assert activation_format != "bf16", (
+        "the torch baseline always quantizes activations (scaled_grouped_mm has no bf16-act x "
+        "MX-weight form) — activation_format='bf16' (W4A16/W8A16) is not representable here"
+    )
+    import torch.nn.functional as F
+    from torch.nn.functional import ScalingType, SwizzleType
+    # torchao >= 0.18 required with cutlass-dsl >= 4.6 (0.17 imports a helper path 4.6
+    # removed; fixed upstream in pytorch/ao).
+    # torchao's per-group blocked-scale builder for the per-forward ACTIVATION scale (graph-capturable
+    # @triton_op, same S+128·E static padding + SWIZZLE_32_4_4 layout scaled_grouped_mm consumes). The
+    # weight scale is already blocked (offline); only the act scale is blocked here.
+    from torchao.prototype.moe_training.kernels.mxfp8 import (
+        triton_mx_block_rearrange_2d_M_groups,
+    )
+    nvfp4 = gate_up_proj_scale_inv.dtype == torch.float8_e4m3fn
+    packed = gate_up_proj.dtype == torch.int8  # fp4 formats pack e2m1
+    family = "nvfp4" if nvfp4 else "mxfp4" if packed else "mxfp8"
+    act_format = family if activation_format is None else activation_format
+    scale_group = NVFP4_SCALE_GROUP_K if nvfp4 else MX_SCALE_GROUP_K
+    scale_dtype = (
+        torch.float8_e4m3fn if nvfp4 else torch.uint8
+    )  # our act-quant/swizzle carry uint8
+    # scaled_grouped_mm dispatches on the scale dtype — view the uint8 MX scales as e8m0 for it
+    f_dtype = torch.float8_e4m3fn if nvfp4 else torch.float8_e8m0fnu
+    SWZ = SwizzleType.SWIZZLE_32_4_4
+    BW = ScalingType.BlockWise1x16 if nvfp4 else ScalingType.BlockWise1x32
+    FP4 = getattr(torch, "float4_e2m1fn_x2", None)
+    E = gate_up_proj.shape[0]
+    # NVFP4's tcgen05 MMA kind requires TWO-level scaling: the e4m3 per-16 block scale AND a
+    # per-tensor global fp32 scale — unlike our kernels there is no "no global" form, so a
+    # missing calibrated global rides as identity 1.0 (dynamic quant). Weight globals arrive
+    # as the separate *_global_scale args; a calibrated *_input_global_scale normalizes that
+    # GEMM's activation quant and rides as its TensorWise scale. MX formats are single-level.
+    tensorwise = ScalingType.TensorWise
+    def _tensorwise_global(g, n):  # (n,) fp32 TensorWise operand, identity when uncalibrated
+        if g is None:
+            return torch.ones(n, device=hidden_states.device, dtype=torch.float32)
+        if g.numel() == 1:
+            return g.reshape(-1).float().expand(n).contiguous()
+        assert g.shape == (n,), (
+            f"torch's scaled_grouped_mm takes ONE TensorWise global per expert, got "
+            f"{tuple(g.shape)} — a per-expert activation global has no slot in this baseline "
+            "(the Triton ops take one)"
+        )
+        return g.float()
+    top_k = top_k_index.shape[1]
+    out_dtype = hidden_states.dtype
+    # route: stable-sort routed slots by expert into contiguous groups (torch has no gather/scatter fuse)
+    flat_e = top_k_index.reshape(-1)
+    order = torch.argsort(flat_e, stable=True)
+    counts = torch.histc(flat_e.float(), bins=E, min=0, max=E - 1).to(torch.int32)
+    offs = counts.cumsum(0).to(torch.int32)
+    tok = (order // top_k).to(torch.long)  # source token of each sorted slot
+    slot_e = flat_e[order].to(torch.long)  # expert of each sorted slot, for the per-expert biases
+    def pk(t):  # view a packed-e2m1 operand as torch's fp4 dtype for scaled_grouped_mm
+        return t.view(FP4) if packed else t
+    def aswz(a_s):  # (S, K//G) -> per-group blocked layout, one launch
+        return triton_mx_block_rearrange_2d_M_groups(a_s.view(torch.uint8), offs).view(
+            f_dtype
+        )
+    def wswz(
+        w_s,
+    ):  # weight scale is pre-blocked offline (SWIZZLE_32_4_4) — pass through, no per-call kernel
+        return w_s.view(f_dtype)
+    def grouped_mm(a, w_q, w_s, w_g=None, a_g=None):
+        assert a_g is None or nvfp4, "an activation global is NVFP4-only"  # match the ops
+        assert a_g is None or a_g.numel() == 1, (
+            "this baseline quantizes the routed rows in one pass, so the activation global is "
+            "per tensor; per-expert activation globals are a Triton-op path"
+        )
+        # our Triton MX act-quant (format-taking launcher) — torch is timed on the same fast quant
+        aq, a_s = _launch_act_quant(
+            a, act_format, scale_group, scale_dtype, global_scale=a_g
+        )
+        sa, ra = aswz(a_s), BW
+        sb, rb = wswz(w_s), BW
+        if nvfp4:  # two-level: block e4m3 + the per-tensor/per-expert fp32 globals
+            sa, ra = [sa, _tensorwise_global(a_g, 1)], [BW, tensorwise]
+            sb, rb = [sb, _tensorwise_global(w_g, E)], [BW, tensorwise]
+        return F.scaled_grouped_mm(
+            pk(aq),
+            pk(w_q).transpose(-2, -1),
+            sa,
+            ra,
+            sb,
+            rb,
+            swizzle_a=SWZ,
+            swizzle_b=SWZ,
+            offs=offs,
+            output_dtype=out_dtype,
+        )
+    gate_up = grouped_mm(
+        hidden_states[tok],
+        gate_up_proj,
+        gate_up_proj_scale_inv,
+        gate_up_proj_weight_global_scale,
+        gate_up_proj_input_global_scale,
+    )
+    # torch has no fused bias, so this baseline adds both host-side; rows are expert-sorted here,
+    # so the per-expert bias indexes by the sorted expert ids
+    if gate_up_proj_bias is not None:
+        gate_up = gate_up + gate_up_proj_bias[slot_e]
+    inter = fused_glu(gate_up, act_fn, swiglu_alpha, swiglu_limit)
+    down_out = grouped_mm(
+        inter, down_proj, down_proj_scale_inv,
+        down_proj_weight_global_scale, down_proj_input_global_scale,
+    )
+    # One weighted scatter-reduce: down_out is expert-sorted, so index_add_ over the source-token
+    # map fuses unroute + routing-weight + top-k sum into (T, H) directly — no separate unsort pass.
+    out = torch.zeros_like(hidden_states)
+    w = top_k_weights.reshape(-1)[order].unsqueeze(-1).to(out.dtype)
+    if down_proj_bias is not None:
+        down_out = down_out + down_proj_bias[slot_e]
+    down_out = _post_expert_norm(
+        down_out, post_expert_norm, post_expert_norm_weight, post_expert_norm_eps
+    )
+    return out.index_add_(0, tok, down_out * w)
+def moe_unfused_batched(
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    gate_up_proj: torch.Tensor,
+    down_proj: torch.Tensor,
+    gate_up_proj_scale_inv: torch.Tensor,
+    down_proj_scale_inv: torch.Tensor,
+    gate_up_proj_weight_global_scale: torch.Tensor | None = None,
+    down_proj_weight_global_scale: torch.Tensor | None = None,
+    gate_up_proj_input_global_scale: torch.Tensor | None = None,
+    down_proj_input_global_scale: torch.Tensor | None = None,
+    post_expert_norm=None,  # the model's per-expert output norm on the routed rows: a
+    # get_supported_norms() name (fused into the reduce) or a host callable
+    post_expert_norm_weight: torch.Tensor | None = None,  # (H,) weight of a named norm
+    post_expert_norm_eps: float = 1e-6,
+    gate_up_proj_bias: torch.Tensor | None = None,  # (E, 2I) pre-activation bias
+    down_proj_bias: torch.Tensor | None = None,  # (E, H)
+    act_fn: str | Callable = "silu",
+    swiglu_alpha: float | None = None,
+    swiglu_limit: float | None = None,
+    activation_format: str | None = None,
+    gate: bool = True,
+) -> torch.Tensor:
+    """Unfused batched MoE: gate_up (plain batched GEMM, gather hidden) → host ``apply_glu`` →
+    down (plain batched GEMM) → routing-weighted reduce. Same math as ``moe_fused_batched`` but
+    the SwiGLU + intermediate quant happen between two plain GEMMs; each GEMM quantizes its raw
+    input in ``activation_format`` (``None`` follows the weight format, ``"bf16"`` is weight-only). All
+    formats route through the shared ``matmul_batched``. The NVFP4 activation globals thread the same way as the fused sibling:
+    each GEMM quantizes its raw input against its ``*_input_global_scale``. ``act_fn`` is a
+    ``get_supported_act_fns()`` name (fused into the gate_up epilogue where the forward fuses) or any
+    callable applied on the host to the raw gate_up output; ``gate=False`` runs an ungated
+    projection. Scales are affine or pre-swizzled (``SWIZZLE_32_4_4``, self-describing 5-D)
+    — the swizzled layout is for the dot_scaled arm; weight-only chains take affine."""
+    fmt = _block_format(
+        gate_up_proj, gate_up_proj_scale_inv, down_proj, down_proj_scale_inv, activation_format
+    )
+    NUM_EXPERTS = gate_up_proj.size(0)
+    expert_ids = top_k_index.reshape(-1)
+    gather_idx = _gather_idx(top_k_index)
+    # gate_up as a plain GEMM (no gate epilogue) over gathered hidden -> (S, 2I).
+    gate_up_out = matmul_batched(
+        hidden_states,
+        gate_up_proj,
+        Bs=gate_up_proj_scale_inv,
+        a_global_scale=gate_up_proj_input_global_scale,
+        b_global_scale=gate_up_proj_weight_global_scale,
+        expert_ids=expert_ids,
+        bias=gate_up_proj_bias,
+        activation_format=fmt,
+        output_dtype=hidden_states.dtype,
+        gather_idx=gather_idx,
+    )
+    inter = _host_glu(gate_up_out, act_fn, swiglu_alpha, swiglu_limit, gate)
+    # down over the intermediate (quantized in the same format), routed-order output.
+    down_out = matmul_batched(
+        inter,
+        down_proj,
+        Bs=down_proj_scale_inv,
+        a_global_scale=down_proj_input_global_scale,
+        b_global_scale=down_proj_weight_global_scale,
+        expert_ids=expert_ids,
+        bias=down_proj_bias,
+        activation_format=fmt,
+        output_dtype=hidden_states.dtype,
+    )
+    return _torch_weighted_reduce(
+        _post_expert_norm(down_out, post_expert_norm, post_expert_norm_weight, post_expert_norm_eps),
+        top_k_index, top_k_weights, NUM_EXPERTS,
+    )

build/torch-rocm/norm.py ADDED Viewed

	@@ -0,0 +1,156 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Row-wise RMS normalization, in the three forms models write it.
+Every form is ``x * factor(h) * rsqrt(mean(y**2) + eps)`` over the row: the per-column
+``factor`` is the norm weight (``"rms_norm"``) or ``1 + weight`` (the zero-centered
+parameterization), and ``y`` is the raw row except under ``"input_scaled_rms_norm"``, which
+scales BEFORE normalizing so the row's mean square is taken on the scaled values.
+The MoE chain never materializes the normalized rows: ``rms_inv_rows`` computes the per-row
+``rsqrt`` in one pass and ``weighted_reduce`` folds it — with the column factor — into the
+reduce it already runs. ``rms_norm_rows`` is the standalone form for the paths that have no
+reduce to fold into."""
+import torch
+import triton
+import triton.language as tl
+from .bayesian_autotuner import bayesian_autotune
+from .compat import compile_time_only_triton_wrap, device_context
+# the fused names, by how a model composes the weight and the normalization
+RMS_NORMS = ("rms_norm", "centered_rms_norm", "input_scaled_rms_norm")
+def get_supported_norms() -> tuple[str, ...]:
+    """The norm names the kernels fuse (``post_expert_norm``), each ``x * factor * rsqrt(mean +
+    eps)`` over the row: ``"rms_norm"`` scales the normalized row by ``weight``,
+    ``"centered_rms_norm"`` by ``1 + weight`` (the zero-centered parameterization), and
+    ``"input_scaled_rms_norm"`` scales by ``1 + weight`` BEFORE normalizing, so the mean square
+    is taken on the scaled row. Anything else stays a host callable."""
+    return RMS_NORMS
+@triton.jit
+def norm_column_factor(W, offs_h, mask, NORM: tl.constexpr):
+    """The per-column multiplier of a fused norm: the weight, or ``1 + weight`` under the
+    zero-centered parameterizations."""
+    w = tl.load(W + offs_h, mask=mask, other=0.0).to(tl.float32)
+    if NORM == "rms_norm":
+        factor = w
+    else:
+        factor = 1.0 + w
+    return factor
+@bayesian_autotune(
+    [
+        triton.Config({"BLOCK_H": block_h}, num_warps=warps)
+        for block_h in (256, 512, 1024, 2048)
+        for warps in (4, 8)
+    ],
+    # one program per row, so the tile width trades the row's loop length against occupancy
+    ["H", "NORM", "WRITE_ROWS"],
+    n_trials=8,
+)
+@triton.jit
+def _rms_norm_kernel(
+    X,  # (S, H) rows to normalize
+    W,  # (H,) norm weight
+    Out,  # (S, H) normalized rows; None writes only Inv (the fused chain's one-pass arm)
+    Inv,  # (S,) fp32 per-row rsqrt(mean + eps); None skips it
+    H,
+    eps,
+    stride_x_m,
+    stride_x_h,
+    stride_o_m,
+    stride_o_h,
+    NORM: tl.constexpr,
+    WRITE_ROWS: tl.constexpr,  # Out is written (else the pass stops at Inv)
+    BLOCK_H: tl.constexpr,
+):
+    """One program per row: accumulate the row's mean square in fp32 over ``BLOCK_H`` tiles, then
+    ``Inv[row] = rsqrt(mean + eps)`` and, under ``WRITE_ROWS``, a second pass writing
+    ``x * factor * inv`` (the row is L2-resident by then). ``input_scaled_rms_norm`` squares the
+    SCALED values, the other forms the raw ones — the one place the three differ."""
+    row = tl.program_id(0)
+    x_row = X + row * stride_x_m
+    square_sum = tl.zeros((), tl.float32)
+    for h in range(0, tl.cdiv(H, BLOCK_H)):
+        offs_h = h * BLOCK_H + tl.arange(0, BLOCK_H)
+        mask = offs_h < H
+        x = tl.load(x_row + offs_h * stride_x_h, mask=mask, other=0.0).to(tl.float32)
+        if NORM == "input_scaled_rms_norm":
+            x = x * norm_column_factor(W, offs_h, mask, NORM)
+        square_sum += tl.sum(x * x, 0)
+    inv = tl.rsqrt(square_sum / H + eps)
+    if Inv is not None:
+        tl.store(Inv + row, inv)
+    if WRITE_ROWS:
+        out_row = Out + row * stride_o_m
+        for h in range(0, tl.cdiv(H, BLOCK_H)):
+            offs_h = h * BLOCK_H + tl.arange(0, BLOCK_H)
+            mask = offs_h < H
+            x = tl.load(x_row + offs_h * stride_x_h, mask=mask, other=0.0).to(tl.float32)
+            out = x * norm_column_factor(W, offs_h, mask, NORM) * inv
+            tl.store(out_row + offs_h * stride_o_h, out.to(Out.dtype.element_ty), mask=mask)
+def _launch_rms_norm(x, weight, eps, norm, write_rows):
+    """One launch of ``_rms_norm_kernel`` over ``x``'s rows, returning the normalized rows or the
+    per-row ``rsqrt`` — whichever was asked for is the only one allocated."""
+    assert norm in RMS_NORMS, f"norm must be one of {RMS_NORMS} or a callable, got {norm!r}"
+    assert x.ndim == 2, f"rows must be 2D (rows, H), got ndim={x.ndim}"
+    assert weight.ndim == 1 and weight.shape[0] == x.shape[1], (
+        f"norm weight {tuple(weight.shape)} does not match the row width ({x.shape[1]},)"
+    )
+    S, H = x.shape
+    out = torch.empty_like(x) if write_rows else None
+    inv = None if write_rows else torch.empty(S, device=x.device, dtype=torch.float32)
+    with device_context(x.device):
+        compile_time_only_triton_wrap(_rms_norm_kernel)[(S,)](
+            x,
+            weight,
+            out,
+            inv,
+            H,
+            eps,
+            x.stride(0),
+            x.stride(1),
+            out.stride(0) if write_rows else 0,
+            out.stride(1) if write_rows else 0,
+            NORM=norm,
+            WRITE_ROWS=write_rows,
+        )
+    return out if write_rows else inv
+def rms_norm_rows(
+    x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6, norm: str = "rms_norm"
+) -> torch.Tensor:
+    """RMS-normalize ``(S, H)`` rows in one launch, in the ``get_supported_norms()`` form named by
+    ``norm`` (fp32 accumulate, input dtype out). The MoE chain folds this into
+    ``weighted_reduce`` instead; this is the standalone form."""
+    return _launch_rms_norm(x, weight, eps, norm, write_rows=True)
+def rms_inv_rows(
+    x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6, norm: str = "rms_norm"
+) -> torch.Tensor:
+    """The ``(S,)`` fp32 ``rsqrt(mean + eps)`` of each row under ``norm`` — the half of
+    ``rms_norm_rows`` a consumer that already reads the rows needs, so the rows are never
+    written back."""
+    return _launch_rms_norm(x, weight, eps, norm, write_rows=False)

build/torch-rocm/pruners.py ADDED Viewed

	@@ -0,0 +1,1213 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+import triton
+from .compat import get_active_device_type, is_sm10x, is_sm90, sm_count, sm_shared_memory_limit
+from .mma import MMA_N_ATOM_WIDTH
+# ── config pruners ────────────────────────────────────────────────────────────
+# Every guard exists for one of four reasons; the map (pruner -> rule -> attached to):
+#
+#   SILENTLY-WRONG configs (must be pruned, the tuner would time and pick them):
+#     block_within_dim_pruner   BK not dividing K over-reads rows (maskless K-loops) —
+#                               tensor/weight-only/full-precision 2D+batched+grouped, and the
+#                               first stage of mx_config_pruner, whose BN|N stage is scoped
+#                               to the mx-dynamic AFFINE arm (its grouped scale gather has no
+#                               wrap; swizzled masks its N tail). The weight-only batched and
+#                               grouped kernels carry their own ungated BN|N stages.
+#     scalar "scalar" -> BM=1   the scalar GEVM broadcasts wrong at BM>1 (config builder)
+#     descriptor_needs_prequant_pruner  raw-A inline-quant arm reads a_ptrs, dead under a
+#                               descriptor mode — static 2D
+#     weight_only_swap_scope_pruner  SWAP_AB iff COMPUTE_MODE == "scalar" — the swapped dot
+#                               arms are unimplemented and a non-swapped scalar has no
+#                               accumulator shape — weight-only batched
+#     mma_trap_ok               swapped dot_scaled BK<128 (sticky 716), single-trip
+#                               accumulator-init miscompile, width > 256 (sticky trap),
+#                               packed-fp4 GATE BN<64 silent zeros — mx 2D/batched/grouped;
+#                               raises when empty (never falls back to a miscompile)
+#     gate_stacked_tmem_trap_pruner  GATE 2*BN=512 acc overflows tensor memory undetected
+#                               at exactly BN=256/BM=64 (checker gap -> sticky 716) —
+#                               weight-only grouped; raises when empty
+#     block_dynamic_mma_width_pruner  the same width trap on the bd implicit dot_scaled
+#                               arm (GATE 2*BN) — bd 2D only (grouped/batched bd pin their
+#                               N tile to the quant block, within the cap); UE8M0-As-gated
+#     scale_subblock_pruner (whole-block clause)  a tile wider than the quant block reads
+#                               one scale across two blocks; raises when the grid has no
+#                               fitting tile — bd/static batched
+#   COMPILER BUGS / unsupported combos on this triton+arch (benign inf, pruned to save
+#   compiles and to keep can't-win configs from poisoning the TPE densities):
+#     warp_spec_compile_guard_pruner   WS compiles iff (BN if swapped else BM)>=64 &
+#                               num_warps%4==0 — bd/tensor/mx/weight-only/fp 2D+grouped
+#     gated_pointer_weight_warp_spec_pruner  GATE+pointer-weight WS lowering failures —
+#                               weight-only grouped
+#     affine_scale_warp_spec_pruner    affine scale gather's other=0.0 can't take
+#                               ttg.partition under WS — mx-dynamic grouped
+#     global_scale_warp_spec_pruner    the NVFP4 global load can't take ttg.partition
+#                               under WS — weight-only grouped, BsGlobal-gated
+#     nvfp4_native_ok (mx_config_pruner)  E4M3-scale dot_scaled PassManager-fails off the
+#                               native staging: M operand < 128, num_warps outside {4, 8},
+#                               or a weight-only kernel (no As — never native)
+#     descriptor_box_pruner     TMA box > 256/dim is illegal — all descriptor-axis kernels
+#     smem_pruner               empirical smem model (B200-sampled, keep-side error) —
+#                               bd/static/mx/weight-only 2D; tensor/mx/weight-only/fp grouped;
+#                               mx batched
+#   RACE guards (Triton 3.7.1 pipeliner, per-loop-structure flake maps):
+#     block_dynamic_grouped_matmul_pruner   WS-only at BM>=64, non-WS below (disjoint) —
+#                               bd + static grouped (plain and GATE arms, same loop)
+#     block_dynamic_2d_warp_spec_pruner  bd 2D WS compile-fail laws: w16+WS at any tile,
+#                               and the implicit dot_scaled arm (UE8M0 As, native-M tile)
+#                               + WS — bd 2D only (COMPILER BUG class, listed here next to
+#                               its sibling)
+#   TPE-POISON fences (valid configs that can't win in a regime):
+#     scalar_max_m_pruner       scalar above M=64 (prefill GEVM) — mx 2D
+#     mx_config_pruner poison fences   the dot arm where the
+#                               native MMA exists (A/B'd 2026-07-14; weight-only kernels
+#                               keep dot), mma_poison_ok: no-swap BK<128 where 128 divides K
+#                               (basin claim) + swapped rows<128 (bf16 fallback never wins)
+#     scale_subblock_pruner     sub-quant-block BN only when the grid starves — bd/static
+#                               batched
+#     mx_2d_swap_scope_pruner   2D mx SWAP_AB rows only for E4M3-scale decode (BM=1,
+#                               M<=16) — swap loses everywhere else (18-cell sweep)
+#     packed_schedule_scope_pruner  PACKED_SCHEDULE rows only at BM>=128 (register relief
+#                               is the whole win) — full-precision grouped
+#     swizzled_scales_bm_pruner grouped mx swizzled arm pinned to BM=BN=128 (slab layout);
+#                               affine arm keeps every BM/BN
+#     swizzled_scale_config_pruner  swizzled-arm BK%128, BN<=128, GATE block pairs —
+#                               mx 2D/batched
+#     matched_memory_modes_pruner   mixed pointer/descriptor never wins a DENSE 2D GEMM —
+#                               bd/static 2D (grouped keeps mixed: it wins there)
+# Every pruner is a ``config_filter``: a per-config ``ok(config, args)`` predicate, an
+# optional ``when(args)`` regime gate, and an empty-survivor policy (advisory guards fall
+# back to the untouched grid; the BK-within-K correctness veto raises via ``on_empty`` —
+# silence there means wrong results). Ordered below: shared infrastructure, then the four
+# categories above.
+# sm_10x caps a single scaled MMA (dot_scaled) at N=256; Triton miscompiles wider ones
+# into a sticky illegal-address device trap. Plain-dot MMAs have NO such cap: a forced
+# width-512 GATE sweep on the tensor grouped kernel (2026-07-14) ran bit-exact wherever
+# it fit shared memory and failed only as benign launch-time smem overflows.
+SM10X_SCALED_MMA_MAX_N = 256
+# Branch axes that RELOCATE the tile optimum (compute unit / operand orientation): the
+# tuner's guaranteed max-tile anchors group by these (``path_anchor_axes`` — a declaration,
+# so the tuner itself stays independent of configuration details). Scheduling axes
+# (WARP_SPEC, PACKED_SCHEDULE, memory modes) share the tile basin shape and stay out —
+# per-basin anchors over the full branch cross starved a 100-trial budget (see
+# _basin_anchor_indices).
+PATH_ANCHOR_AXES = ("COMPUTE_MODE", "SWAP_AB")
+def config_dim(c, all_args, name):
+    """A tile dimension for config ``c`` — from its tuned meta, else the launch args.
+    No default: a pruner reasoning about a dimension the kernel doesn't have is a wiring
+    bug, and a made-up value would silently mis-prune."""
+    v = c.kwargs.get(name, all_args.get(name))
+    if v is None:
+        raise ValueError(
+            f"pruner needs {name} (autotune config meta or launch arg); none found"
+        )
+    return v
+def config_filter(ok, when=None, on_empty=None):
+    """Wrap a per-config predicate into an ``early_config_prune``: keeps the configs where
+    ``ok(config, all_args)`` (``all_args`` = launch args + autotune kwargs). Every pruner
+    below is one of these — the predicate is always named ``ok``, so the factories differ
+    only in their rule, not their plumbing. ``when(all_args)`` gates the whole filter —
+    hardware/launch-regime checks live there, per-config logic in ``ok``. When nothing
+    survives, the untouched grid is returned (guards are advisory) unless ``on_empty(configs,
+    all_args)`` is given — correctness vetoes pass one that raises, because falling back to
+    configs that compute wrong results must be loud."""
+    def prune(configs, named_args, **kwargs):
+        all_args = {**named_args, **kwargs}
+        if when is not None and not when(all_args):
+            return configs
+        kept = [c for c in configs if ok(c, all_args)]
+        if kept:
+            return kept
+        return configs if on_empty is None else on_empty(configs, all_args)
+    return prune
+def compose_pruners(*pruners):
+    """Chain ``early_config_prune`` callbacks left to right (each sees the previous
+    survivors)."""
+    def prune(configs, named_args, **kwargs):
+        for p in pruners:
+            configs = p(configs, named_args, **kwargs)
+        return configs
+    return prune
+def block_within_dim_pruner(dim_arg: str, block_key: str = "BLOCK_SIZE_K", when=None):
+    """``early_config_prune`` dropping configs whose ``block_key`` tile does not divide the
+    launch dim named by ``dim_arg``: the K-loops load unmasked and the batched/grouped
+    N-tiles store row-masked only, so a non-dividing tile's last trip reads or writes past
+    the row — silently wrong results the tuner would happily time and pick (a BN=256
+    winner at N=128 corrupted 40/64 rows before the N veto existed). A dim smaller than
+    every grid tile is a hard error. ``when`` gates the veto to a launch regime (e.g. the
+    grouped swizzled arm masks its N-tail — %-wrapped value load + ``N_COLS`` store mask — so
+    BN needn't divide N there). Used standalone by the tensor-dynamic kernels and as the first
+    stages of ``mx_config_pruner``."""
+    def ok(c, args):
+        block = c.kwargs.get(block_key, 0)
+        return block == 0 or args[dim_arg] % block == 0
+    def raise_no_dividing_block(configs, args):
+        min_block = min(c.kwargs.get(block_key, 0) for c in configs)
+        raise ValueError(
+            f"{dim_arg}={args[dim_arg]} is not a multiple of any {block_key} in the "
+            f"autotune grid; the unmasked tile would run past the row. Pad the problem "
+            f"along {dim_arg} (smallest grid tile: {min_block})."
+        )
+    return config_filter(ok, when=when, on_empty=raise_no_dividing_block)
+def block_fits_dim_pruner(dim_arg: str, block_key: str = "BLOCK_SIZE_K"):
+    """``early_config_prune`` for a loop whose tail tile SLIDES back in bounds instead of
+    running past the row (the weight-only batched K-loop): the tile need not divide the dim,
+    it only has to fit inside it, since the clamp would otherwise walk backwards off the
+    buffer. Everywhere the tail is unmasked, use ``block_within_dim_pruner`` — this one keeps
+    non-dividing tiles, and on a maskless loop those read past the row."""
+    def ok(c, args):
+        block = c.kwargs.get(block_key, 0)
+        return block == 0 or block <= args[dim_arg]
+    def raise_no_fitting_block(configs, args):
+        min_block = min(c.kwargs.get(block_key, 0) for c in configs)
+        raise ValueError(
+            f"{dim_arg}={args[dim_arg]} is smaller than every {block_key} in the autotune "
+            f"grid (smallest: {min_block}); the sliding tail tile has nowhere to sit."
+        )
+    return config_filter(ok, on_empty=raise_no_fitting_block)
+def require_moe_dims_aligned(N: int, K: int, block_n: int, block_k: int) -> None:
+    """Routed MoE (grouped/batched) GEMMs tile expert dims by the quant block — N-tiles store
+    column-unmasked and K loads unmasked — so ``N`` and ``K`` must each be a multiple of the block.
+    Non-aligned dims are a dense capability: use ``matmul_2d`` (its single GEMM masks the N/K tail).
+    Raised early with a clear message rather than letting a tile run past the row."""
+    if N % block_n != 0 or K % block_k != 0:
+        raise ValueError(
+            f"routed MoE GEMM needs N ({N}) and K ({K}) each a multiple of the quant block "
+            f"(block_n={block_n}, block_k={block_k}); non-aligned dims are supported only by the "
+            f"dense matmul_2d op."
+        )
+def warp_spec_compile_guard_pruner():
+    """``early_config_prune`` dropping ``warp_specialize`` configs that can never compile: WS needs
+    ``num_warps % 4 == 0`` (its async producer/consumer partitions) and enough M work — the MMA's M
+    operand, which ``SWAP_AB`` moves to ``BLOCK_SIZE_N`` (so BM=1 decode can WS when swapped), so
+    ``(BN if swapped else BM) >= 64``. ``dot_scaled`` + WS + TMA + gate all compile and win here
+    (measured 1995 TFLOPS on a clean context) — the earlier "dot_scaled/gate never WS-compiles"
+    verdicts were a ``num_warps=2`` / trap-contaminated-GPU misdiagnosis. Non-WS configs all stay —
+    WS is purely a perf axis. CUDA-only."""
+    def ok(c, args):
+        if not c.kwargs.get("WARP_SPEC"):
+            return True
+        # WS needs num_warps a multiple of 4 (its async producer/consumer partitions) and enough M
+        # work — the MMA's M operand, which SWAP_AB moves to BLOCK_SIZE_N (so BM=1 decode can still
+        # WS when swapped). dot_scaled + WS + TMA DOES compile here (measured 1995 on a clean
+        # context) — the old "dot_scaled never WS-compiles" was a num_warps=2 misdiagnosis on a
+        # trap-contaminated GPU. Callers needing dot_scaled+WS off (block-dynamic swap arm) fence it.
+        mma_m = "BLOCK_SIZE_N" if c.kwargs.get("SWAP_AB") else "BLOCK_SIZE_M"
+        return config_dim(c, args, mma_m) >= 64 and c.num_warps % 4 == 0
+    return config_filter(ok, when=lambda args: get_active_device_type() == "cuda")
+def gated_pointer_weight_warp_spec_pruner():
+    """``early_config_prune`` dropping the ``warp_specialize`` configs that cannot lower on the
+    weight-only grouped MX kernel's GATE arm, where the stacked gate|up tile doubles the weight
+    extent to ``2 * BLOCK_SIZE_N``. With a pointer-mode weight (its masked load carries the
+    ``other=0.0`` constant the WS partitioner cannot tag) the automatic-warp-specialization pass
+    fails in ``TritonGPUOptimizePartitionWarps`` with ``RuntimeError: PassManager::run failed``,
+    but only where the partitioner also has to split a descriptor-mode A, 16 warps, or a
+    sub-128 M tile — plain pointer/pointer at BM=128 and w4/w8 compiles and is often the winner.
+    Region measured by tuning the GPT-OSS W4A16 problem (N=K=2880) with a fresh Triton cache and
+    ``FINEGRAINED_AUTOTUNE_LOG``: of 129 WS configs, the 12 that failed are exactly the ones this
+    predicate names, and it drops none of the 117 that compiled. Every failure sat in the GATE
+    shape — the ungated shape's 63 WS configs all compiled. Unfenced these are wasted compiles per
+    new shape (the tuner memoizes them only after paying once). The ``BM < 128`` clause rests on a
+    single observation, so it is the clause to revisit if a later probe contradicts it. CUDA-only;
+    a no-op for ungated launches.
+    NOT covered: the later ``dot_scaled`` arm adds 4 UNGATED failures (descriptor-A +
+    pointer-B at num_warps>=8), outside this predicate's GATE scope. Left unfenced on
+    purpose — the region separates only on a GATE x COMPUTE_MODE x A-mode x warps
+    interaction, and every predicate wide enough to catch it also dropped working configs
+    across the five probes. The tuner memoizes the 4 after paying once per shape, and this
+    fence was re-checked against the two-arm grid: 0/207 working configs dropped."""
+    def ok(c, args):
+        if not c.kwargs.get("WARP_SPEC"):
+            return True
+        if c.kwargs.get("B_MEMORY_MODE", "pointer") != "pointer":
+            return True
+        return not (
+            c.kwargs.get("A_MEMORY_MODE", "pointer") != "pointer"
+            or c.num_warps >= 16
+            or config_dim(c, args, "BLOCK_SIZE_M") < 128
+        )
+    return config_filter(
+        ok, when=lambda args: args.get("GATE", False) and get_active_device_type() == "cuda"
+    )
+def affine_scale_warp_spec_pruner():
+    """``early_config_prune`` dropping ``warp_specialize`` on the grouped MX AFFINE (row-major,
+    non-``SWIZZLED_SCALES``) scale path. Its per-(gathered-row, K-group) 2D pointer-gather scale
+    load carries an ``other=0.0`` constant the WS partitioner cannot tag — the automatic-warp-
+    specialization pass fails with ``'arith.constant' op does not have expected attribute
+    ttg.partition``. The swizzled arm (gate + non-gate) loads the scale through a TMA descriptor
+    (no such constant) and WS-compiles + wins, so the fused / deployment path keeps WS. CUDA-only;
+    a no-op where ``SWIZZLED_SCALES`` isn't a kernel arg (defaults to allowing WS)."""
+    def ok(c, args):
+        if not c.kwargs.get("WARP_SPEC"):
+            return True
+        return args.get("SWIZZLED_SCALES", True)
+    return config_filter(ok, when=lambda args: get_active_device_type() == "cuda")
+def mx_config_pruner(k_arg: str, n_arg: str | None = None, block_within_k: bool = True):
+    """``early_config_prune`` for the MX kernels (2D, batched, grouped — their tile is always
+    tuned): a BK-within-K veto plus sm_10x MMA-shape guards (no-ops elsewhere and for scalar
+    configs). GATE-aware: under the GATE constexpr (read from the launch args) the kernel
+    computes gate|up as one stacked 2*BN extent, so the swapped dot_scaled lhs has ``2*BN``
+    rows and the no-swap combined dot is ``2*BN`` wide; a plain launch's counts are just
+    ``BN``.
+    ``block_within_k=False`` drops the BK veto for a caller whose K-loop masks or slides its
+    tail (the weight-only batched kernel); the MMA-shape guards still apply.
+    - ``BLOCK_SIZE_K`` not dividing the launch's contraction dim (``k_arg`` names it) →
+      dropped: the K-loop loads are unmasked, so any non-dividing BK's last trip reads
+      past the row — silently wrong results the tuner would happily time and pick (bit us when
+      BK=512 met a K=256 test problem).
+    - MXFP4 ``scalar`` configs are KEPT. They used to be dropped on sm_10x: fp4 scalar decode
+      was ALU-bound in the E2M1 unpack and measured 1.8x slower than dot_scaled. That verdict
+      was against the INTEGER unpack. Decoding E2M1 through the sm_100 CVT unit
+      (``_e2m1_pair_to_e4m3``) inverts it — on the dsv4 W4A8 decode, scalar+swap beats the
+      dot_scaled crown 20.7 vs 41.8us (gate_up) and 11.3 vs 22.0us (down), bit-identical —
+      because at BM=1 dot_scaled cannot reach the native microscaled MMA either (it needs an
+      M=128 operand) and emulates through the same upcast, while scalar skips the MMA
+      entirely. Re-fence only against a fresh measurement.
+    - Swapped ``dot_scaled`` rows < 128 → dropped: the native microscaled MMA gates on the M
+      operand being exactly 128, so smaller rows run the bf16-upcast fallback and never win —
+      the same poison mechanism (an earlier dsv4 gate_up tune shipped 63µs missing the 43µs
+      swap winner).
+    - No-swap ``dot_scaled`` width > ``SM10X_SCALED_MMA_MAX_N`` → dropped: Triton
+      miscompiles wider scaled MMAs into a sticky illegal-address device trap that poisons
+      the context mid-autotune.
+    Packed-E2M1 activations (a ``uint8`` activation tile, W4A4 — caller-packed or
+    quantized inline from a raw tile) are consumed by EVERY arm: dot_scaled natively
+    (the ``kind::mxf4`` MMA was probed bit-exact: BK down to 64, single-trip, and
+    BN=256 all clean; BM<128 falls back to the correct-but-slow f16 kind), and the
+    dot/scalar/swap arms column-unpack them to E4M3 (lossless) first — no arm is
+    structurally packed-incompatible, so W4A4 needs no shape gate of its own.
+    The shape gates above are scoped to ``dot_scaled`` — they are native scaled-MMA bug
+    gates. The ``dot`` arm (BK structurally the UE8M0 group, 32) is CORRECT everywhere
+    probed (forced-config sweep 2026-07-14, GATE and plain, MXFP4/MXFP8, incl. width-512
+    stacked dots) but is fenced on sm_10x as TPE POISON: an A/B of fresh 100-trial tunes
+    (mx grouped, S=4096 E=32 N=1024 K=2048 MXFP8) shipped dot winners 2-3x SLOWER with the
+    dot rows in the grid (gate 0.297ms vs 0.095ms, plain 0.147ms vs 0.079ms) — the ~2x
+    larger grid dilutes the trial budget and the early-sampled dot basin hijacks the TPE
+    densities, so the native scaled-MMA winner is never found. dot never wins on sm_10x
+    (native dot_scaled owns prefill, scalar owns M=1 decode); other targets keep the arm.
+    Revisit if a dot win materializes (e.g. the batched decode BM=16 observation).
+    The fence applies only where the kernel carries an ``As`` activation-scale operand, so
+    the weight-only (W4A16/W8A16) kernels keep their dot rows while still getting the shape
+    gates. Their ``dot`` is a weight upcast + plain ``tl.dot`` over a free BK — with raw
+    activations there is no native microscaled MMA to lose to, so both arms upcast and this
+    one WINS the stacked gate|up tile (gpt-oss K=2880: dot 4645.7us vs dot_scaled 5011.3us,
+    while the down projection goes the other way, 2348.1us vs 2882.6us).
+    The gates are two classes with opposite empty-survivor policies, composed trap-first:
+    ``mma_trap_ok`` (miscompiles: swapped BK<128, single-trip, width, fp4 GATE zeros) raises
+    when nothing survives — falling back to configs that compute wrong results must be loud —
+    while the poison fences (dot/scalar arms, ``mma_poison_ok``) stay advisory, so emptying
+    one of them re-admits CORRECT-but-slow configs, never a miscompile. A contraction dim
+    smaller than every grid BK is a hard error: any config would over-read past the row and
+    return silently wrong results."""
+    def scaled_mma_available(args):
+        # dot_scaled rows exist for this launch iff the grid's smallest BK (64)
+        # divides the contraction dim with at least two trips; off the 64 grid no
+        # row survives the shape gates, so the software arms are the only correct
+        # ones and their poison fences lift.
+        return args[k_arg] % 64 == 0 and args[k_arg] > 64
+    def mma_trap_ok(c, args):
+        # MISCOMPILE gates only — every rule here names a config that runs and returns
+        # wrong results or traps the device; the composed filter raises when they empty.
+        if c.kwargs.get("COMPUTE_MODE") != "dot_scaled":
+            return True
+        # Swapped dot_scaled below BK=128 is a Triton sm_10x lowering bug — a STICKY
+        # misaligned-address device trap — on the UE8M0 (group-32) scale family. Probed
+        # per-arm 2026-07-15 (fresh processes): batched W4A4 swap BK=64 traps at K=2880
+        # AND K=2944 alike (K-alignment is NOT the variable), swap BK=128 runs clean and
+        # correct, no-swap BK=64 is compute-sanitizer-clean, and the grouped kernel (no
+        # swap axis) is clean at BK=64 in every memory mode. The E4M3 (NVFP4 group-16)
+        # variant does NOT trap: 40 raw-grid sweep configs + 6 forced fresh-process runs
+        # at K∈{256,2048,2944} all bit-correct (2026-08-05) — and swap BN=64 GATE is the
+        # native M=128 staging for gate decode, so it stays admitted. All probing on
+        # B200 — re-chart on H100 or the target device before trusting the fence there.
+        if (
+            c.kwargs.get("SWAP_AB")
+            and config_dim(c, args, "BLOCK_SIZE_K") < 128
+            and getattr(args.get("Bs"), "dtype", None) != torch.float8_e4m3fn
+        ):
+            return False
+        # Single-trip dot_scaled (BK >= contraction dim) trips the sm_10x
+        # accumulator-init miscompile (uninitialized TMEM alloc must be mutable).
+        # Bites only small K (e.g. a K=512 gate_up with BK=512): silently wrong results
+        # the tuner would happily time and pick (surfaced as a 35% MXFP4 fused-MoE error).
+        if config_dim(c, args, "BLOCK_SIZE_K") >= args[k_arg]:
+            return False
+        # tcgen05 microscaled-fp4 MMA silently no-ops at the GATE stacked N=2*BN=64 shape (BN<64):
+        # the IR emits a correct tc_gen5_mma_scaled (128x64 acc, 64x2 scale) but it never writes the
+        # accumulator, which reads back its zero-init → EXACTLY-ZERO output. N-independent (probed
+        # zero at both non-aligned N=320 and aligned N=512; BN=64 / stacked N=128 is correct
+        # everywhere), so this is a LATENT poison config: any packed-fp4 gate|up whose tuner crowned
+        # BN=32 (favoured at small N) would silently return zeros — larger shapes pass only because
+        # BN>=64 wins. (BN=32 is recoverable only by emitting gate/up as two separate BN dots
+        # instead of one stacked 2*BN dot; not worth it — a sub-64 gated tile rarely wins over BN=64.)
+        if (
+            args.get("GATE")
+            and config_dim(c, args, "BLOCK_SIZE_N") < 64
+            and getattr(args.get("B"), "dtype", None) == torch.uint8
+        ):
+            return False
+        # no-swap width cap: Triton miscompiles a scaled MMA wider than 256 into a sticky
+        # illegal-address device trap; the stacked 2*BN extent under GATE counts in full
+        rows = (2 if args.get("GATE") else 1) * config_dim(c, args, "BLOCK_SIZE_N")
+        return c.kwargs.get("SWAP_AB") or rows <= SM10X_SCALED_MMA_MAX_N
+    def raise_all_mma_trapped(configs, args):
+        raise ValueError(
+            f"every autotune config hits an sm_10x scaled-MMA miscompile for "
+            f"{k_arg}={args[k_arg]} (GATE={bool(args.get('GATE'))}); no correct config "
+            f"exists in this grid — extend the grid (a no-swap BK strictly inside the "
+            f"contraction dim) rather than lifting the fence: the pruned configs return "
+            f"wrong results or trap the device."
+        )
+    def mma_poison_ok(c, args):
+        # CORRECT-but-can't-win fences (advisory: emptying re-admits correct configs)
+        if c.kwargs.get("COMPUTE_MODE") != "dot_scaled":
+            return True
+        # no-swap BK<128 rows are CORRECT but admitted only where BK=128 can't divide the
+        # contraction dim (e.g. gpt-oss K=2880). "The BK>=128 basin owns the shape" is a
+        # basin claim, not an A/B — the fence only trims trial-budget dilution, so lifting
+        # it costs tune time, not results.
+        if (
+            not c.kwargs.get("SWAP_AB")
+            and config_dim(c, args, "BLOCK_SIZE_K") < 128
+            and args[k_arg] % 128 == 0
+        ):
+            return False
+        # swapped rows < 128 run the bf16-upcast fallback (the native microscaled MMA
+        # gates on the M operand being exactly 128) — correct, never wins, TPE poison
+        rows = (2 if args.get("GATE") else 1) * config_dim(c, args, "BLOCK_SIZE_N")
+        return not c.kwargs.get("SWAP_AB") or rows >= 128
+    def dot_arm_ok(c, args):
+        return c.kwargs.get("COMPUTE_MODE") != "dot"
+    def acts_are_scaled(args):
+        # The dot-arm poison verdict holds only where dot_scaled can reach the NATIVE
+        # microscaled MMA, which needs both operands microscaled — i.e. the kernel carries an
+        # activation-scale operand. The weight-only (W4A16/W8A16) kernels have no ``As``: their
+        # activations are raw bf16/fp16, so no native low-precision MMA exists, both arms upcast,
+        # and their "dot" (weight upcast + plain tl.dot) genuinely competes. Keyed on the operand,
+        # not A's dtype — the MX-dynamic inline-quant arm takes a raw bf16 A and still reaches the
+        # native MMA.
+        return "As" in args
+    def nvfp4_native_ok(c, args):
+        # mxf4nvf4 (E4M3 scales) has NO fallback below the native staging — the decompose
+        # path materializes the scale through an integer-only DenseElementsAttr and
+        # PassManager fails outright (charted BM {16,32,64} x BK {128,256}, 2026-07-15;
+        # warps/operand law 2026-08-12). Native staging needs the M operand >= 128 AND
+        # num_warps in {4, 8} (2 and 16 both fail across every BN/BK/stages cell) AND
+        # both operands microscaled — the weight-only kernels (no ``As``: raw bf16 acts,
+        # in-register inline quant) never stage it, so their E4M3-scale dot_scaled fails
+        # at every tile (their dot/scalar arms compete instead). dot_scaled-only
+        # (dot/scalar decode E4M3 scales in software at any M);
+        # the MMA's M operand is BLOCK_SIZE_M, or the weight rows when SWAP_AB puts the
+        # token in N.
+        if c.kwargs.get("COMPUTE_MODE") != "dot_scaled":
+            return True
+        if not acts_are_scaled(args) or not 4 <= c.num_warps <= 8:
+            return False
+        if c.kwargs.get("SWAP_AB"):
+            return (2 if args.get("GATE") else 1) * config_dim(c, args, "BLOCK_SIZE_N") >= 128
+        return config_dim(c, args, "BLOCK_SIZE_M") >= 128
+    def scales_are_e4m3(args):
+        return getattr(args.get("Bs"), "dtype", None) == torch.float8_e4m3fn
+    stages = [block_within_dim_pruner(k_arg)] if block_within_k else []
+    if n_arg is not None:
+        # BN | N is needed only on the mx-dynamic AFFINE arm: its weight-scale gather has no
+        # column wrap. The swizzled arm masks its N-tail (grouped/batched: %-wrapped value load
+        # + N_COLS store mask; 2D: swizzle-offset wrap), and the weight-only kernels (no
+        # SWIZZLED_SCALES axis) clamp their scale pointers in-bounds and mask the tail in the
+        # epilogue (N_COLS), so both keep non-dividing BN.
+        stages.append(
+            block_within_dim_pruner(
+                n_arg, "BLOCK_SIZE_N", when=lambda args: not args.get("SWIZZLED_SCALES", True)
+            )
+        )
+    # trap gates run FIRST (on the full grid, raising when empty); the poison fences after
+    # them are advisory, so an emptying poison fence falls back to trap-clean survivors —
+    # never to a miscompile (the old single mixed filter fell back to the untouched grid,
+    # re-admitting the single-trip/width traps it existed to remove).
+    return compose_pruners(
+        *stages,
+        config_filter(nvfp4_native_ok, when=scales_are_e4m3),
+        config_filter(
+            mma_trap_ok, when=lambda args: is_sm10x(), on_empty=raise_all_mma_trapped
+        ),
+        config_filter(
+            dot_arm_ok,
+            when=lambda args: is_sm10x()
+            and scaled_mma_available(args)
+            and acts_are_scaled(args),
+        ),
+        config_filter(mma_poison_ok, when=lambda args: is_sm10x()),
+    )
+def dot_scaled_staging_pruner():
+    """``early_config_prune`` dropping ``dot_scaled`` where the MMA's M operand degenerates to a
+    single row: ``BLOCK_SIZE_M == 1`` with no ``SWAP_AB`` to move the operand to ``BLOCK_SIZE_N``.
+    The scaled MMA has no 1-row staging and the config dies in ``PassManager::run failed`` at every
+    warps/stages cell — the decode arms that DO serve BM=1 are scalar, or dot_scaled swapped (which
+    puts the weight rows in M). Measured on the batched mx kernel: the last 6 dead compiles per
+    tune after the swapped-scalar act fence.
+    Two staging laws, both of which ``mx_config_pruner`` already enforces for NVFP4/E4M3 and
+    neither of which was applied to the UE8M0 half:
+    - ``num_warps`` outside {4, 8}: the scaled MMA has no staging for it (w16 s6 measured failing).
+    - a 1-row M operand, per above.
+    Measured on the batched mx kernel: 128 dead compiles per tune on HEAD, 2 after both clauses."""
+    def ok(c, args):
+        if c.kwargs.get("COMPUTE_MODE") != "dot_scaled":
+            return True
+        if c.num_warps not in (4, 8):
+            return False
+        if c.kwargs.get("SWAP_AB"):
+            return True
+        return config_dim(c, args, "BLOCK_SIZE_M") > 1
+    return config_filter(ok)
+def block_dynamic_2d_warp_spec_pruner():
+    """``early_config_prune`` for the block-dynamic 2D kernel's two WARP_SPEC compile-fail
+    laws (Triton 3.7.1 PassManager, TritonGPUPipeline; charted 2026-08-13 at the
+    DeepSeek-V4-Flash ue8m0 checkpoint shapes, forced-config matrices):
+      - WS + num_warps=16 fails at EVERY tile/scale-dtype (w8 compiles everywhere). The mx
+        2D kernel WS-compiles w16 fine; the bd GROUPED kernel does NOT on its pointer arm
+        (same law, see block_dynamic_grouped_matmul_pruner) — so this tracks the loop's
+        structure and memory mode, not a global WS law.
+      - the implicit ``dot_scaled`` arm (``USE_DOT_SCALED``: UE8M0/uint8 ``As`` on a
+        native-M tile — (BN if SWAP_AB else BM) >= 128) + WS fails at every warps/stages/
+        memory-mode cell; BM=64 and non-WS all compile. fp32-scale launches never emit the
+        scaled MMA and keep the (w<=8) WS grid.
+    Mostly 2D-only: the grouped bd loop WS-compiles from its descriptor arm and its own
+    race guard REQUIRES WS at BM >= 64, but it shares the dot_scaled+WS failure on the
+    POINTER arm (see block_dynamic_grouped_matmul_pruner). Compile soundness is
+    per-loop-structure AND per-memory-mode — chart both before reusing a law."""
+    def ok(c, args):
+        if not c.kwargs.get("WARP_SPEC"):
+            return True
+        if c.num_warps > 8:
+            return False
+        if getattr(args.get("As"), "dtype", None) == torch.uint8:
+            native_m = "BLOCK_SIZE_N" if c.kwargs.get("SWAP_AB") else "BLOCK_SIZE_M"
+            return config_dim(c, args, native_m) < 128
+        return True
+    return config_filter(ok, when=lambda args: is_sm10x())
+def block_dynamic_mma_width_pruner():
+    """``early_config_prune`` for the block-dynamic kernels' implicit ``dot_scaled`` arm
+    (UE8M0 scales on a native-M tile — ``USE_DOT_SCALED`` in the kernel): the sm_10x
+    width cap ``mx_config_pruner``'s trap gate applies to the MX kernels, which the bd
+    family never got. Under GATE the stacked gate|up dot is ``2 * BLOCK_SIZE_N`` wide, so
+    a BN=256 GATE config emits a width-512 scaled MMA — Triton miscompiles those into a
+    sticky illegal-address device trap the tuner trips mid-tune. Scoped to launches whose
+    ``As`` is UE8M0 (uint8) on sm_10x; fp32-scale launches never emit the scaled MMA.
+    Configs below the native-M tile keep the plain ``tl.dot`` + software rescale and pass
+    untouched; the swapped N side is the token tile (within the cap on every grid)."""
+    def ok(c, args):
+        native_m = "BLOCK_SIZE_N" if c.kwargs.get("SWAP_AB") else "BLOCK_SIZE_M"
+        if config_dim(c, args, native_m) < 128 or c.kwargs.get("SWAP_AB"):
+            return True
+        width = (2 if args.get("GATE") else 1) * config_dim(c, args, "BLOCK_SIZE_N")
+        return width <= SM10X_SCALED_MMA_MAX_N
+    def raise_all_too_wide(configs, args):
+        raise ValueError(
+            "every autotune config emits a scaled MMA wider than "
+            f"{SM10X_SCALED_MMA_MAX_N} (GATE={bool(args.get('GATE'))}) — sm_10x "
+            "miscompiles those into a sticky device trap; extend the grid below the cap."
+        )
+    return config_filter(
+        ok,
+        when=lambda args: is_sm10x()
+        and getattr(args.get("As"), "dtype", None) == torch.uint8,
+        on_empty=raise_all_too_wide,
+    )
+def fp8_dot_warp_pruner():
+    """SILENTLY-WRONG fence: the fp8 ``tl.dot`` paths lose accuracy above ``num_warps=2`` at both
+    ends of the M-tile range, on sm_90. Two clauses, one phenomenon.
+    (1) The single-token swapped GEVM. ``swap_pad_rhs`` pads the lone token to the ``MMA_N_ATOM``
+        rhs and the epilogue takes column 0 after the K-loop (``acc_init``'s ``[BN, MMA_N_ATOM]``
+        accumulator). Above 2 warps that column comes back wrong.
+    (2) Large M tiles, ``BLOCK_SIZE_M >= 64``, in the grouped and 2D kernels.
+    Measured on H100 (sm_90), triton 3.6 AND 3.8, by forcing every admitted config against the
+    torch oracle and scoring error in bf16-ULP units (the floor — what perfect bf16 output
+    storage costs — is 0.49 mean / 1.00 max; the torch reference itself sits exactly there)::
+        BLOCK_SIZE_M 16 / 32, any num_warps   0.49 mean / 1.0 max ULP   bit-optimal
+        BLOCK_SIZE_M 64 / 128, num_warps 8    0.58 mean / 52.1 max ULP  WRONG   (grouped)
+        BLOCK_SIZE_M 64 / 128, num_warps >2   280 of 578 configs WRONG          (2D)
+        SWAP_AB + BLOCK_SIZE_M=1, num_warps >2                         WRONG   (batched)
+    The tuner benches by SPEED, so these get crowned silently: block-static crowned a wrong one
+    outright, while block-dynamic / tensor-dynamic (batched) and the grouped kernel carried
+    30-90 wrong configs each behind a GREEN test, because the winner happened to be correct.
+    In the grouped kernel ``BLOCK_SIZE_M >= 64`` is admitted only at ``num_warps=8``, so those two
+    axes are confounded there and the data cannot separate them; the 2D kernel does separate them
+    (BM=64 at num_warps=2 is clean), which is why the rule is written on num_warps.
+    Scoped by a POSITIVE ``is_sm90()`` test, not by ``not is_sm10x()``: the latter is also true on
+    XPU, ROCm and driverless build boxes, which do not run the NVIDIA codegen this was probed
+    against, and restricting them on sm_90 evidence would be a guess. sm_10x is likewise untouched
+    — it is where this repo's tuning and validation were done and its tcgen05 MMA lowers these
+    paths differently. Other CUDA arches (sm_89, sm_120) are unmeasured and so unfenced; each
+    needs its own probe. The sm_90 cost is real — the grouped kernel is the prefill path, where
+    the large-M tiles are the throughput — and is accepted: the alternative is a crowned config
+    that returns wrong results."""
+    def ok(c, args):
+        bm = c.kwargs.get("BLOCK_SIZE_M", 1)
+        if c.num_warps <= 2:
+            return True
+        if bm >= 64:  # (2) large-M tile
+            return False
+        # (1) swapped single-token GEVM; native-M tiles need no pad/column-extract
+        return not (c.kwargs.get("SWAP_AB") and bm < MMA_N_ATOM_WIDTH)
+    def raise_all_wrong(configs, args):
+        raise ValueError(
+            "every autotune config is an fp8 tl.dot at num_warps>2 with a BLOCK_SIZE_M that "
+            "returns WRONG results off sm_10x; no correct config exists in this grid — extend it "
+            "with a num_warps=2 or a 16/32-row M tile rather than lifting the fence."
+        )
+    return config_filter(ok, when=lambda args: is_sm90(), on_empty=raise_all_wrong)
+def scale_subblock_pruner(min_ctas_per_sm: int = 4):
+    """``early_config_prune`` for kernels whose compute tile (``BLOCK_SIZE_N``) may subdivide the
+    quant block (``BLOCK_N``). One scale covers the whole block, so a narrower tile just reads its
+    own block — and it multiplies the N grid, the lever while the whole-block grid is too thin to
+    hide the loop's latency (a gate|up fusion halves that grid again by folding both into one program).
+    Sub-block tiles are admitted ONLY there: below ``min_ctas_per_sm`` whole-block programs per SM.
+    "Fills the device" (one program per SM) is NOT enough at decode — the loop is latency-bound, so
+    a 192-448-program grid still gains from more, narrower tiles: forced BN=64 vs the BN=128 crown,
+    bit-identical, DSV4 fp8 decode gate 24.9 -> 23.4us and down 13.4 -> 12.2, DSV3 down 22.3 -> 21.7
+    but DSV3 gate 0.88x (2026-08-29) — a tuner call, so the fence only trims prefill, where a
+    thousands-of-programs grid never wants the 3x wider config space (measurably worse winners
+    when it carried it).
+    A fused requant always pins the tile to the whole block: its output scale is one amax per
+    (row, ``BLOCK_N``) group, which a partial tile can't compute. Raises when nothing
+    survives: a tile wider than the quant block reads ONE scale across two blocks — silently
+    wrong results — so a checkpoint whose ``BLOCK_N`` is off the tile grid must fail loudly,
+    not fall back to the full grid."""
+    def ok(c, args):
+        block_n = args.get("BLOCK_N")
+        if block_n is None:
+            return True
+        bn = config_dim(c, args, "BLOCK_SIZE_N")
+        starved = args["S"] * triton.cdiv(args["N"], block_n) < min_ctas_per_sm * sm_count(
+            args["B"].device.index
+        )
+        if args.get("OUTPUT_FORMAT") is not None or not starved:
+            return bn == block_n
+        return bn <= block_n and args["N"] % bn == 0
+    def raise_no_block_tile(configs, args):
+        raise ValueError(
+            f"no BLOCK_SIZE_N in the autotune grid fits the weight quant block "
+            f"(BLOCK_N={args.get('BLOCK_N')}); a wider tile would read one scale across "
+            f"two quant blocks and return wrong results — extend the grid or requantize "
+            f"the checkpoint to a gridded block size."
+        )
+    return config_filter(ok, on_empty=raise_no_block_tile)
+def packed_schedule_scope_pruner(min_bm: int = 128):
+    """``early_config_prune`` scoping ``PACKED_SCHEDULE`` rows to big-BM configs — the only
+    regime the packed resolve wins (register relief unlocking pressure-walled tiles: DSV3
+    bf16 forced pairs +10.7% at BM128 vs −12.5% at BM64/BN256 and a wash at BM32,
+    2026-08-06; both tuned crowns chose BM128). Everywhere else the inline resolve's free
+    ALU beats the packed arm's per-tile dependent loads, so those rows are tune-budget
+    dilution (the axis would double the grid; this keeps it +25%)."""
+    def ok(c, args):
+        if not c.kwargs.get("PACKED_SCHEDULE"):
+            return True
+        return config_dim(c, args, "BLOCK_SIZE_M") >= min_bm
+    return config_filter(ok)
+def weight_only_warp_spec_matched_mode_pruner():
+    """``early_config_prune`` for the weight-only grouped kernel's WARP_SPEC lowering wall
+    (Triton 3.7.1 ``TritonGPUOptimizePartitionWarps`` -> PassManager::run failed; the pass
+    cannot partition this loop when BOTH operand loads share a memory mode). Charted
+    2026-08-26 (48-cell forced-config matrix, BM 64 and 128 IDENTICAL — BM-independent):
+        CM          A/B modes   w4    w8    w16
+        dot         ptr/ptr     ok    ok    FAIL
+        dot         desc/desc   ok    ok    ok      (was charted FAIL/FAIL — refuted, see _TRAPPING_NUM_WARPS)
+        dot_scaled  ptr/ptr     FAIL  FAIL  FAIL
+        dot_scaled  desc/desc   ok    ok    ok      (w8 was charted FAIL — refuted)
+        MIXED modes (ptr/desc, desc/ptr): 24/24 ok at every warp count.
+        BN=32 + WS: FAIL in every cell (the tile law, fenced here).
+    The non-monotone warp dependence (w16 rescues desc/desc dot but breaks ptr/ptr dot) marks
+    this as the pass's internal partition feasibility, not kernel source — the source is
+    identical across passing and failing cells. Two root-cause attempts on the sibling family
+    are refuted and recorded (block_dynamic_grouped_matmul_pruner); the kernel-side remedy is
+    to stop emitting configs the pass cannot lower. Cost in the wild: 19 dead compiles per
+    weight-only tune, and under inductor ONE failing config kills the whole torch.compile
+    cell instead of scoring inf — this fence is what recovers those cells."""
+    # desc/desc rows REMOVED 2026-08-29: forced through the tuner's own Config objects (pre_hook
+    # intact) at the charted GPT-OSS shape, (dot, desc/desc) w4 AND w8 and (dot_scaled, desc/desc) w8
+    # all compile and run bit-identical — dot+WS+BK=128 is the gate_up's best config (2128 vs 2235us).
+    # The original matrix was charted with hand-built Configs, which skip the descriptor pre_hooks
+    # and fail every descriptor cell for the wrong reason. The pointer rows stand unrefuted.
+    _TRAPPING_NUM_WARPS = {
+        ("dot", "pointer"): {16},
+        ("dot_scaled", "pointer"): {4, 8, 16},
+    }
+    def ok(c, args):
+        if not c.kwargs.get("WARP_SPEC"):
+            return True
+        # BN=32 + WS fails to lower in every memory mode / warp count / compute mode probed
+        # (9/9 cells, GPT-OSS N=K=2880, 2026-08-29) — the tile law the matrix never encoded
+        if config_dim(c, args, "BLOCK_SIZE_N") == 32:
+            return False
+        a, b = c.kwargs.get("A_MEMORY_MODE"), c.kwargs.get("B_MEMORY_MODE")
+        if a != b:  # mixed modes lower everywhere
+            return True
+        return c.num_warps not in _TRAPPING_NUM_WARPS.get((c.kwargs.get("COMPUTE_MODE"), a), set())
+    return config_filter(ok, when=lambda args: is_sm10x())
+def weight_only_swap_scope_pruner():
+    """``early_config_prune`` binding the weight-only kernel's ``SWAP_AB`` axis to its one
+    implemented arm: the CUDA-core reduce. Swap exists here to drop the MMA entirely — with
+    raw bf16 activations there is no native fp4 x bf16 MMA, so both dot arms pay an M->16 pad
+    and a full-tile upcast, and a swapped ``dot_scaled`` measured no better than the row-major
+    one (35.2 vs 38.0us, gpt-oss gate_up decode). So ``SWAP_AB`` iff ``COMPUTE_MODE`` is
+    ``"scalar"``: the swapped-dot rows would compile into an arm the compute helper does not
+    implement, and a non-swapped scalar row has no accumulator shape."""
+    def ok(c, args):
+        return bool(c.kwargs.get("SWAP_AB")) == (c.kwargs.get("COMPUTE_MODE") == "scalar")
+    return config_filter(ok)
+def mx_2d_swap_scope_pruner(max_m: int = 16):
+    """``early_config_prune`` scoping the 2D mx kernel's ``SWAP_AB`` rows to the ONE regime
+    they exist for: E4M3-scale (NVFP4) single-token decode. NVFP4 has no other native M=1
+    arm — ``dot_scaled`` has no sub-M=128 form for E4M3 scales (PassManager) and scalar is
+    E2M1-unpack-bound (measured 60-142µs vs swap's 22-50µs on the GLM attn shape,
+    2026-08-05). Everywhere else swap is a proven loser (an 18-cell forced-swap sweep:
+    UE8M0 formats' scalar/dot arms win, M3 attn swap −38%), so its rows are dropped:
+    UE8M0-scale launches, M above the decode band, and any non-BM=1 row (the mx swap
+    compute flattens the token tile — BM=1 is structural, enforced by a static_assert).
+    Also drops UNSWAPPED ``scalar``: the arm exists only in its swapped form here, so those
+    configs cannot build an accumulator and fail to compile."""
+    def ok(c, args):
+        swap = bool(c.kwargs.get("SWAP_AB"))
+        # Under GATE, scalar exists only as the swapped arm: unswapped it has no accumulator
+        # shape (``acc_init`` is handed "dot" when not swapped) and every such config dies in
+        # CompilationError (3 dead compiles per gated mxfp4 tune). NOT true ungated — unswapped
+        # scalar is a legitimate (and sometimes winning) W4A4 arm there, and fencing it globally
+        # crowned a wrong config at N=320.
+        if c.kwargs.get("COMPUTE_MODE") == "scalar" and not swap and args.get("GATE"):
+            return False
+        if not swap:
+            return True
+        return (
+            config_dim(c, args, "BLOCK_SIZE_M") == 1
+            and getattr(args.get("Bs"), "dtype", None) == torch.float8_e4m3fn
+            and args["M"] <= max_m
+        )
+    return config_filter(ok)
+def swizzled_out_bm_pruner():
+    """``early_config_prune`` pinning the SWIZZLE_32_4_4 requant OUTPUT arm (``SWIZZLED_OUT``) to
+    ``BLOCK_SIZE_M == BLOCK_SIZE_N == 128`` — the output-side twin of
+    ``swizzled_scales_bm_pruner``.
+    ``_epilogue_requant_mx`` writes the group scales straight into the down proj's swizzled
+    layout as ``q_s.reshape(1, 4, 32, REP_K_CS, 4)`` with ``REP_K_CS = (BN // group) // 4``.
+    ``q_s`` is ``[BM, BN // group]``, so the reshape balances only at ``BM == 128``
+    (``BM * (BN//group) == 128 * (BN//group)``); every other BM raises "reshape() cannot change
+    total number of elements". BN is pinned with it because the descriptor's block shape is the
+    tuned ``REP_K_CS``. The code comment there already said "BM/BN pinned 128" — this makes the
+    config space agree, instead of handing the tuner configs that cannot compile.
+    Measured 2026-08-26 on the gated fp4-requant cells: 288 dead compiles per tune, i.e. every
+    non-128 config of the swizzled-out key (``SWIZZLED_OUT`` is in the tune key, so the key is
+    entirely swizzled). Un-swizzled requant keeps every BM/BN — it stores Cs row-major."""
+    def ok(c, args):
+        return (
+            config_dim(c, args, "BLOCK_SIZE_M") == 128
+            and config_dim(c, args, "BLOCK_SIZE_N") == 128
+        )
+    return config_filter(ok, when=lambda args: args.get("SWIZZLED_OUT"))
+def swizzled_scales_bm_pruner():
+    """``early_config_prune`` pinning the grouped MX pre-swizzled (``SWIZZLED_SCALES``) arm to
+    ``BLOCK_SIZE_M == 128`` and a whole-128-block ``BLOCK_SIZE_N``. BM: the offline act-quant lays each expert's scale
+    slab out 128-padded, and the kernel's per-tile scale-block index (``pid_m``) only lines up
+    when the M tile is exactly 128 (``build_tile_layout`` pads experts on the same granularity).
+    BN: the scale is read as whole 128-row SWIZZLE_32_4_4 blocks off the descriptor (under GATE,
+    the gate 128-block + the up 128-block stacked into the 2*BN tile); a sub-128 BN would need a
+    partial-block read the descriptor can't express, while any whole multiple is a multi-block bulk
+    load the pre_hook sizes per config — BN=256 wins the down GEMM (665 -> 576us at DSV4 prefill,
+    bit-identical; GATE at 256 is the scaled-MMA width cap's veto, ``mx_config_pruner``).
+    Both pins are swizzled-slab properties, so the un-swizzled (affine) arm keeps every BM and
+    BN: ``build_tile_layout`` tiles any BM with a row mask, and the affine act-scale read is
+    indexed by the gathered row, not ``pid_m``."""
+    def ok(c, args):
+        return (
+            config_dim(c, args, "BLOCK_SIZE_M") == 128
+            and config_dim(c, args, "BLOCK_SIZE_N") % 128 == 0
+        )
+    return config_filter(ok, when=lambda args: args.get("SWIZZLED_SCALES"))
+def swizzled_scale_config_pruner(allow_gate_subblock=False):
+    """Drop pre-swizzled-scale configs the SWIZZLE_32_4_4 descriptor load can't serve, gated on
+    ``SWIZZLED_SCALES`` (the un-swizzled arm loads scales per group and takes any tile):
+    - ``BLOCK_SIZE_K < 4 * SCALE_GROUP_K``: the un-swizzle reads whole 4-group K bands
+      (``REP_K = (BK // SCALE_GROUP_K) // 4``), so a smaller BK collapses ``REP_K`` to 0 and the
+      reshape traps — 128 for group-32 (MX), 64 for group-16 (NVFP4). BK must also tile whole
+      bands (``BK % (4 * group) == 0``).
+    - ``BLOCK_SIZE_N > 128`` not a whole-block multiple: the descriptor arm bulk-loads
+      ``BN // 128`` swizzled blocks (each kernel's pre_hook sizes the box per config), a sub-128 BN
+      reads a scalar slice of one block; a BN between is neither, and a multi-block BN whose
+      last tile passes the 128-padded N extent zero-fills (see the fit bound below). BN=256 is
+      a legal two-block load where it fits (the GATE tile has always read one) and wins at
+      prefill — grouped down 665 -> 576us, dense 2D 703 -> 552us at M=49152, bit-identical; the
+      scaled-MMA width cap is ``mx_config_pruner``'s.
+    - under ``GATE``, ``BLOCK_SIZE_N != 128``: the gate|up scale is interleaved as whole 128-row
+      block pairs [g0,u0,g1,u1,...], read as one 2*BN tile; a sub-128 BN can't index a block pair
+      off the descriptor. ``allow_gate_subblock`` (batched decode) admits BN in (32, 64): its
+      scale leaf pointer-GATHERs the sub-128 gate|up tile out of the interleaved buffer, which is
+      the raw path's BN=32 gate|up tile without un-swizzling (the whole swizzled-vs-raw decode
+      gap). Packed fp4 is admitted too: the historical sub-tile fp4->bf16 lowering trap
+      (``xVals.size() % 4``) no longer reproduces — mxfp4 AND nvfp4 forced at BN∈{32,64},
+      BK=128, scalar, K=2048 all run bit-correct (2026-08-05); a resurfacing trap is a loud
+      compile failure (forgiven inf), never silent."""
+    def ok(c, args):
+        band = 4 * args["SCALE_GROUP_K"]  # one whole swizzled K band (4 scale groups)
+        if config_dim(c, args, "BLOCK_SIZE_K") % band != 0:
+            return False
+        bn = config_dim(c, args, "BLOCK_SIZE_N")
+        if args.get("GATE"):
+            return bn == 128 or (allow_gate_subblock and bn in (32, 64))
+        if bn > 128:
+            # BN=256 pipelines OOM smem at num_stages >= 5 (253 KB at s5, 304 KB at s6 for the
+            # fp8 x fp4 and fp8 x fp8 dot_scaled tiles, BK=128; s4 fits) — below smem_pruner's
+            # keep-side bound, so without this the TPE burns whole 200-trial budgets on dead
+            # compiles (observed 2026-08-29).
+            if c.num_stages >= 5:
+                return False
+            # A multi-block tile must FIT the swizzled grid: the bulk load reads BN // 128 whole
+            # blocks with no row mask, and the grid is padded to 128 rows — a tile whose blocks
+            # pass the padded extent TMA-zero-fills the scales: silent garbage, never a trap
+            # (mxfp4_swizzled N=128 forced at BN=256 fails parity on triton 3.7.1 and 3.8 alike;
+            # the BN=256 prefill wins were measured at N % 256 == 0 shapes).
+            n = args.get("N")
+            if n is not None and -(-n // bn) * bn > -(-n // 128) * 128:
+                return False
+        return bn <= 128 or bn % 128 == 0
+    def raise_no_swizzled_tile(configs, args):
+        raise ValueError(
+            "no autotune config can serve pre-swizzled scales for this launch (the "
+            "SWIZZLE_32_4_4 read needs BLOCK_SIZE_K % 128 == 0 and a <=128-row or 128-multiple N tile; "
+            f"GATE={bool(args.get('GATE'))}) — the contraction dim likely has no "
+            "128-dividing tile; pass affine (row-major) scales for this shape."
+        )
+    return config_filter(
+        ok, when=lambda args: args.get("SWIZZLED_SCALES"), on_empty=raise_no_swizzled_tile
+    )
+def block_dynamic_grouped_matmul_pruner():
+    """``early_config_prune`` for the block-dynamic grouped kernel: the Triton 3.7.1
+    pipeliner-race guard, sized to its four-load-stream single-dot K-loop. The GATE arm
+    shares the guard: post-activation-prequant it dots the stacked gate|up as the same
+    single-dot loop (one wider weight tile, not a second dot/load stream) and was
+    flake-verified race-free cross-process. Measured on sm_100, 15-fresh-process flake
+    runs per cell, big and tiny shapes:
+    - ``warp_specialize`` compiles iff ``num_warps % 4 == 0`` (its async partitions) and
+      ``BLOCK_SIZE_M >= 64``, and is race-free everywhere it compiles — but on the POINTER
+      arm it hits a lowering wall: "'arith.constant' op does not have expected attribute
+      ttg.partition" -> PassManager::run failed (TritonGPUPipeline). The descriptor arm
+      always compiles. The obvious suspect — the masked activation load's ``other=0.0``,
+      which is the known cause of the analogous mx-grouped fence — was PROBED and is NOT
+      it: dropping the mask entirely leaves every pointer cell below still failing, so the
+      trigger is structural to the pointer arm, not that constant. Do not re-try the
+      clamp-not-mask rewrite on this basis. Two INDEPENDENT legs, forced-config matrices,
+      2026-08-21 B200:
+        * ``num_warps >= 16`` fails at BM 64 AND 128, with fp32 scales as well as UE8M0 —
+          a property of this loop under WS, not of the scaled MMA.
+        * uint8/UE8M0 ``As`` on a native-M ``BLOCK_SIZE_M >= 128`` tile (the implicit
+          ``dot_scaled`` arm) fails at every warps/stages cell: pointer 0/6, descriptor
+          4/6 + 2 smem.
+      Combined with the race guard below, the descriptor arm is the only BM=128 route for
+      UE8M0 grouped — which is what makes the native tcgen05 tile reachable at all here.
+    - The default pipeliner RACES at ``BLOCK_SIZE_M >= 64`` (3/15 wrong at BM64/w8) and is
+      clean at BM16/32.
+    The regions are disjoint, so per ``BLOCK_SIZE_M``: BM >= 64 keeps only the compilable
+    WS configs, BM < 64 keeps only non-WS. CUDA-only — the race is a CUDA pipeliner
+    artifact and the WS axis isn't emitted elsewhere.
+    Also attached to the STATIC grouped kernel as an INHERITED guard: its three-load-stream
+    K-loop (scalar act scale preloaded, no per-row ``as`` stream) was never flake-mapped —
+    race soundness is per-loop-structure, so this map may over-fence it (perf cost) or, if
+    the 3-stream loop races elsewhere, under-fence it. Re-map with multi-process flake runs
+    before trusting the static attachment."""
+    def ok(c, args):
+        bm = config_dim(c, args, "BLOCK_SIZE_M")
+        if bm < 64:
+            return not c.kwargs.get("WARP_SPEC")
+        if not (c.kwargs.get("WARP_SPEC") and c.num_warps % 4 == 0):
+            return False
+        if c.kwargs.get("A_MEMORY_MODE") == "pointer":
+            if c.num_warps >= 16:
+                return False
+            if getattr(args.get("As"), "dtype", None) == torch.uint8 and bm >= 128:
+                return False
+        return True
+    return config_filter(ok, when=lambda args: get_active_device_type() == "cuda")
+def scalar_max_m_pruner(m_arg: str, max_m: int = 64):
+    """``early_config_prune`` dropping ``scalar`` configs when the launch's row count
+    (``m_arg``) exceeds ``max_m``: scalar is a BM=1 GEVM — sensible for decode-sized M,
+    hopeless at prefill, and hopeless-but-benched configs poison the TPE's per-dimension
+    densities (measured: with scalar in the M=8192 grid the 2D MX attn prefill tune
+    landed 0.48x vs hub; without it, 2.06x)."""
+    def ok(c, args):
+        return c.kwargs.get("COMPUTE_MODE") != "scalar"
+    return config_filter(ok, when=lambda args: args[m_arg] > max_m)
+def matched_memory_modes_pruner():
+    """Drop configs that mix pointer and descriptor operand loads — for a DENSE 2D GEMM
+    only (contiguous A, whole M in one tile). Measured across M=1..8192 on the block-scale
+    2D matmul: mixed only ever TIES the matched combos (<1us at small M) and loses at large
+    M, so pruning it costs no win and halves the descriptor search. Do NOT use this on the
+    GROUPED kernels: there M is distributed across experts (S/E small), A is token-GATHERED,
+    and A=pointer + B=descriptor (mixed) is the dominant winner — those keep both axes free.
+    Non-empty (the matched configs remain)."""
+    def ok(c, args):
+        a_ptr = c.kwargs.get("A_MEMORY_MODE", "pointer") == "pointer"
+        b_ptr = c.kwargs.get("B_MEMORY_MODE", "pointer") == "pointer"
+        return a_ptr == b_ptr
+    return config_filter(ok)
+def descriptor_needs_prequant_pruner():
+    """Keep only pointer-mode configs when ``A`` is raw (bf16/fp16 — the ``maybe_act_quant``
+    M<threshold inline-quant arm). A descriptor operand load needs a pre-quantized fp8 ``A``: the
+    inline arm reads ``a_ptrs``, which ``operand_tile_ptrs`` leaves a dead placeholder under a
+    descriptor mode, so a descriptor config there would compute garbage yet bench fast (a broken
+    config the tuner would crown). Gated on ``A``'s dtype so the offline (fp8) path keeps both
+    memory axes; a no-op for callers that always pre-quantize (``A`` already fp8)."""
+    def ok(c, args):
+        return (
+            c.kwargs.get("A_MEMORY_MODE", "pointer") == "pointer"
+            and c.kwargs.get("B_MEMORY_MODE", "pointer") == "pointer"
+        )
+    return config_filter(
+        ok, when=lambda args: args["A"].dtype not in (torch.float8_e4m3fn, torch.int8, torch.uint8)
+    )
+def descriptor_box_pruner(k_dim="BLOCK_SIZE_K"):
+    """Drop descriptor configs whose TMA box exceeds the tensormap's 256-per-dim limit:
+    the weight inner dim is the K tile in BYTES over values-per-byte (packed E2M1
+    halves it; E4M3 is one byte per value), so BK=512 e4m3 boxes are illegal while
+    BK=512 packed ones fit. Pointer configs pass untouched. ``k_dim`` names the K-tile
+    axis — ``BLOCK_SIZE_K`` for the tune_block_nk kernels, ``block_k`` for the block-scale
+    kernels whose K tile is the caller's fixed scale block."""
+    def ok(c, args):
+        # tiles may be tuned (config kwargs) or pinned launch args — config_dim reads both
+        bk = config_dim(c, args, k_dim)
+        if c.kwargs.get("B_MEMORY_MODE", "pointer") != "pointer":
+            weight_vpb = (
+                2 if getattr(args.get("B"), "dtype", None) == torch.uint8 else 1
+            )
+            if bk // weight_vpb > 256 or config_dim(c, args, "BLOCK_SIZE_N") > 256:
+                return False
+            # A TMA descriptor needs a unit-stride innermost dim; a transposed operand
+            # (column-major view, as a backward pass hands in) has none, so only the
+            # pointer arm can serve it. The kernel itself is stride-general.
+            b = args.get("B")
+            if b is not None and getattr(b, "ndim", 0) == 2 and b.stride(-1) != 1:
+                return False
+        if c.kwargs.get("A_MEMORY_MODE", "pointer") != "pointer":
+            # descriptor A, both arms tuner-routed per the tokens-per-expert key:
+            # contiguous boxes (no gather) win at dense routing (-7.1/-7.4%) and lose at
+            # skew; tma gather4 (gathered rows, 1-row box, sm_100 only) is MONOTONIC in
+            # tokens/expert — 2x loss at 32/expert, parity at 1024, -5.3/-4.7% at 4096
+            # (32k tokens, 2026-07-16). Only the arch is fenced.
+            act_vpb = 2 if getattr(args.get("A"), "dtype", None) == torch.uint8 else 1
+            if args.get("GatherIdx") is not None and not is_sm10x():
+                return False
+            if bk // act_vpb > 256:
+                return False
+        return True
+    return config_filter(ok)
+def global_scale_warp_spec_pruner():
+    """Drop WARP_SPEC configs when the launch carries an NVFP4 global (``BsGlobal`` not None):
+    the per-expert global load cannot take ``ttg.partition`` under warp specialization — the
+    same Triton lowering gap as the affine scale gather's ``other=0.0`` — so every such config
+    is a benign compile-failure inf. Fenced to save the compiles; single-level launches
+    (``BsGlobal`` None) keep their WS rows."""
+    def ok(c, args):
+        return not (c.kwargs.get("WARP_SPEC") and args.get("BsGlobal") is not None)
+    return config_filter(ok)
+def gate_stacked_tmem_trap_pruner():
+    """Miscompile gate: drop GATE configs at exactly ``BLOCK_SIZE_N == 256, BLOCK_SIZE_M == 64``
+    — the stacked gate|up accumulator (``2*BN = 512`` columns) exceeds sm_100 tensor memory, and
+    at BM=64 Triton's resource checker misses the overflow, so the config runs and TRAPS the
+    device (sticky CUDA 716 "misaligned address"; one such trial poisons the whole tune and
+    every launch after it). Probed per-arm in fresh processes (2026-08-10, B200, weight-only
+    grouped, K=2880 and K=2048 alike — N-divisibility is NOT the variable): BM=64 traps in every
+    memory mode, WS on and off, BK 64 and 128; BM=128 fails CLEANLY (OutOfResources: tensor
+    memory 544 > 512) and self-prunes; BM=32 and BN<=128 run correct. The checker gap is the
+    only reason this needs a fence at all — drop precisely the probed trap point and let every
+    neighbor keep self-reporting.
+    ON sm_90 the whole ``2*BN = 512`` stacked tile goes, not just the BM=64 point. The
+    self-pruning above is a TENSOR-MEMORY overflow, and sm_90 has no tensor memory to overflow:
+    nothing raises OutOfResources, so the config runs and returns garbage instead. Measured on
+    H100, triton 3.8, nvfp4 gate|silu with fused requant through the 2D kernel — sampling the
+    admitted grid, ``BLOCK_SIZE_N`` splits cleanly::
+        BN 32 / 64 / 128   60 correct, 0 wrong
+        BN 256             0 correct, 9 wrong   (BM=128; the crowned config varies per tune,
+                                                 so the scenario test passes or fails by luck)
+    The mechanism is the stacked accumulator width, which is format-independent, so this fences
+    every GATE config at BN=256 rather than only the nvfp4 cell that exposed it. The BM=64 trap
+    point stays fenced on every arch (it was probed on B200); the wider sm_90 clause is scoped by
+    a positive arch test so XPU/ROCm, which never ran this codegen, keep their full grid."""
+    def ok(c, args):
+        if not args.get("GATE") or config_dim(c, args, "BLOCK_SIZE_N") != 256:
+            return True
+        if config_dim(c, args, "BLOCK_SIZE_M") == 64:
+            return False  # the probed sm_100 trap point — fenced on every arch
+        return not is_sm90()  # sm_90 has no tensor memory, so nothing self-reports the overflow
+    def raise_all_trap(configs, args):
+        raise ValueError(
+            "every config in the grid is the (GATE, BN=256, BM=64) tensor-memory trap point; "
+            "falling back would run configs that trap the device (sticky CUDA 716)"
+        )
+    return config_filter(ok, on_empty=raise_all_trap)
+def smem_pruner(k_dim="BLOCK_SIZE_K", min_sets: float = 0.5):
+    """``early_config_prune`` dropping configs whose shared memory certainly cannot fit,
+    with the bound picked by the operand dtypes (sampled from ``metadata.shared`` across
+    every kernel family, 22 cells):
+    - both operands >= 2-byte floats (the unquantized loops): ``num_stages`` full
+      buffers of both tiles (exact) plus the fixed barrier terms at their sampled
+      MINIMA (+16 base; +28 warp-spec, observed up to 48 varying with warps; +32
+      host-descriptor) — minima keep the bound below the true usage, so a config is
+      only ever dropped when its operand buffers alone cannot fit.
+    - quantized ``dot_scaled`` (the only quantized arm whose tiles can reach the limit):
+      the per-allocation law, every term named by its TTGIR ``local_alloc`` and exact to
+      the allocator's alignment/packing slack (<= 16 B observed; barriers pack into
+      padding, so no fixed constant is added — the slack is keep-side) —
+      ``num_stages - 1`` buffers of operands+scales at the native M=128 staging (weights
+      full-width even when fp4-packed; GATE shifts the A tile to a full ``num_stages``),
+      or the bf16-upcast fallback below the gate (one single-buffered upcast tile,
+      packed weights, no A-scale buffers).
+    - quantized ``dot``/``scalar``: the raw-operand floor (their 32-wide K tiles cannot
+      reach the limit).
+    ``min_sets`` scales the quantized ``dot_scaled`` bound: the default HALF a set is the
+    family-wide keep-side envelope (the 2D nvfp4 arm sub-tiles a request into half a set);
+    the GROUPED kernel never compiles below TWO full sets — fresh-tune logs at the DSV4 W4A8
+    gate_up key (244 trials, 2026-08-29): 0 of 82 compiling configs had 2*set > smem, while
+    120 of 162 dead ones did (every BK=512 request, 59/59), so ``min_sets=2`` there halves the
+    dead compiles a 100-trial tune pays without dropping a single compiling config.
+    Both arms are EMPIRICAL models of this Triton's allocator, fit to and verified
+    against ``metadata.shared`` samples (not derived from a spec) — chosen so the error
+    direction under the sampled behavior is always toward KEEPING a config, which then
+    merely compiles and self-prunes as inf. All samples taken on B200 (sm_100) —
+    re-verify on a Triton bump AND on H100 or the target device (smem capacity and
+    allocator behavior are per-arch). Non-empty fallback."""
+    def ok(c, args):
+        a, b = args["A"], args["B"]
+        tiles = 2 if args.get("GATE") else 1
+        bk = config_dim(c, args, k_dim)
+        unquantized = a.element_size() >= 2 and b.element_size() >= 2
+        stages = c.num_stages if unquantized else c.num_stages - 1
+        packed = 2 if b.dtype == torch.uint8 else 1
+        bm = config_dim(c, args, "BLOCK_SIZE_M")
+        bn = config_dim(c, args, "BLOCK_SIZE_N")
+        if unquantized:
+            # sampled minima of the fixed allocations beside the operand buffers
+            # (the barrier terms vary by tens of bytes with warps; minima only ever
+            # err toward keeping a config)
+            fixed = (
+                16
+                + (28 if c.kwargs.get("WARP_SPEC") else 0)
+                + (32 if c.kwargs.get("B_MEMORY_MODE", "pointer") != "pointer" else 0)
+            )
+            need = fixed + stages * bk * (
+                a.element_size() * bm + b.element_size() * tiles * bn
+            )
+        elif c.kwargs.get("COMPUTE_MODE") == "dot_scaled":
+            # KEEP-SIDE bound = HALF of one buffer set: operands stored PACKED (fp4 = 2
+            # values/byte) plus both group scale tiles (BK/group cols; group 32 UE8M0, 16
+            # E4M3). Requested-stages models over-estimate — the pipeliner CLAMPS stages
+            # (an s3 request can compile with an s2 footprint) and the 2D nvfp4 arm can
+            # compile half a set — so any tighter bound drops configs the compiler would
+            # have rescued. Proven est<=actual on 1,099 metadata.shared samples across the
+            # family (see OPTIMIZATION_LOG: smem-keep-side-refit).
+            group = 16 if getattr(args.get("Bs"), "dtype", None) == torch.float8_e4m3fn else 32
+            # A's smem bytes follow its LAUNCH dtype: packed E2M1 (int8/uint8) = 2 values/byte,
+            # E4M3 = 1 B/value, raw bf16/fp16 (weight-only) = element_size B/value
+            a_dtype = a.dtype
+            if a_dtype in (torch.int8, torch.uint8):
+                a_bytes = (bk // 2) * bm
+            else:
+                a_bytes = bk * bm * a.element_size()
+            # HALF a set: the sampled lower envelope. The 2D kernel's nvfp4 arm can compile an
+            # s2 request into half a buffer set (sub-tile pipelining: BN=64/128/256 x BK=512
+            # all landed at set/2 + <=8 B), so the full set is not a lower bound; the half set
+            # is, over all 1099 samples (largest remaining margin 68 B on a 4 KB decode tile,
+            # covered by the halving). A config whose HALF set exceeds the limit cannot compile
+            # in any observed form — everything else is kept and self-reports as a benign inf.
+            # The family-wide keep-side model counts fp4 weights PACKED; the grouped pipeline stages
+            # them at full E4M3 width (BN=256, BK=128, s5 OOMs at 253 KB = 5 x 50 KB sets, which only
+            # a full-width B tile adds up to), so the whole-set bound (min_sets >= 1) uses full width.
+            b_bytes = bk * tiles * bn if min_sets >= 1 else (bk // packed) * tiles * bn
+            need = int(
+                (a_bytes + b_bytes + (bk // group) * bm + (bk // group) * tiles * bn) * min_sets
+            )
+        else:
+            # dot (BK = the 32-group) / scalar: tiles too small to reach the limit —
+            # the raw-operand floor suffices
+            need = stages * (bk * bm + (bk // packed) * tiles * bn)
+        return need <= sm_shared_memory_limit()
+    return config_filter(ok, when=lambda args: get_active_device_type() == "cuda")

build/torch-rocm/quant.py ADDED Viewed

	@@ -0,0 +1,955 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+import triton
+import triton.language as tl
+from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait
+from .bayesian_autotuner import bayesian_autotune
+from .formats import global_scale_stride, is_per_expert_global
+from .compat import add_op_namespace_prefix, FP8_DTYPE, MX_SCALE_GROUP_K, NVFP4_SCALE_GROUP_K, compile_time_only_triton_op, compile_time_only_triton_wrap, decode_pdl, device_context, is_sm10x
+from .swizzle import swizzle_store_block
+from .scheduling import build_tile_layout, resolve_tile_inline
+# ── Triton-side helpers (inlined by ``@triton.jit`` callers) ──────────────────
+@triton.jit
+def fp8_act_quant_inline(
+    a_raw, TRANSPOSED: tl.constexpr = False, UE8M0: tl.constexpr = False
+):
+    """Inline FP8 (E4M3) activation quant for the W8A8 block-scale path.
+    Per-token amax → fp32 scale ``amax/448`` (floored at 1e-12 against zero rows)
+    → cast values to FP8. Returns ``(a_fp8, a_s)`` with ``a_s`` shaped ``(M,)``;
+    ``TRANSPOSED`` marks a ``(K, M)`` tile (the swapped descriptor arm), where the
+    token axis is 0 instead of 1.
+    ``UE8M0`` ceils each scale up to a power of two and returns the E8M0 exponent
+    byte (uint8) instead of the fp32 scale — the group format the tcgen05
+    ``dot_scaled`` MMA consumes. The ceil mirrors DeepGEMM's ``ceil_to_ue8m0``
+    (add ``0x7FFFFF``, clear the mantissa) so our scales match its checkpoints.
+    """
+    if TRANSPOSED:
+        a_s = tl.max(tl.abs(a_raw), axis=0) / 448.0
+    else:
+        a_s = tl.max(tl.abs(a_raw), axis=1) / 448.0
+    denom = tl.maximum(a_s, 1e-12)
+    if UE8M0:
+        bits = (denom.to(tl.int32, bitcast=True) + 0x7FFFFF) & ~0x7FFFFF
+        denom = bits.to(tl.float32, bitcast=True)
+    if TRANSPOSED:
+        a_fp8 = (a_raw / denom[None, :]).to(tl.float8e4nv)
+    else:
+        a_fp8 = (a_raw / denom[:, None]).to(tl.float8e4nv)
+    if UE8M0:
+        a_scale = ((bits >> 23) & 0xFF).to(tl.uint8)
+    else:
+        a_scale = a_s
+    return a_fp8, a_scale
+# cvt.e2m1x2.f32 (hardware FP4 pack) exists only on sm_100 (Blackwell). Resolved once at
+# import as a compile-time constexpr for the jit helper below; the ALU fallback compiles
+# everywhere else. ``is_sm10x`` is driverless-safe, so no import-time guard is needed.
+_E2M1_HW_CVT = tl.constexpr(is_sm10x())
+@triton.jit
+def _quant_e2m1_packed(v, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_K: tl.constexpr):
+    """Pack signed, pre-scaled values on the E2M1 grid to ``(M, K//2)`` uint8 (first value
+    of each pair in the low nibble). On sm_100 the Blackwell hardware convert
+    ``cvt.rn.satfinite.e2m1x2.f32`` (two f32 → one packed byte, first operand → HIGH nibble)
+    does it in one instruction; elsewhere a ``>=``-threshold bucketize builds the code. The
+    two agree except at exact E2M1 midpoints (0.25, 0.75, …), where the hardware rounds to
+    nearest-even and the ALU form rounds half-up. Only the taken arm compiles."""
+    if _E2M1_HW_CVT:
+        lo, hi = tl.split(tl.reshape(v, (BLOCK_SIZE_M, BLOCK_SIZE_K // 2, 2)))
+        packed = tl.inline_asm_elementwise(
+            "{ .reg .b8 t8; cvt.rn.satfinite.e2m1x2.f32 t8, $1, $2; cvt.u16.u8 $0, t8; }",
+            "=h,f,f",
+            [hi, lo],
+            dtype=tl.uint16,
+            is_pure=True,
+            pack=1,
+        )
+        values = (packed & 0xFF).to(tl.uint8)
+    else:
+        av = tl.abs(v)
+        code = (
+            (av >= 0.25).to(tl.int32)
+            + (av >= 0.75).to(tl.int32)
+            + (av >= 1.25).to(tl.int32)
+            + (av >= 1.75).to(tl.int32)
+            + (av >= 2.5).to(tl.int32)
+            + (av >= 3.5).to(tl.int32)
+            + (av >= 5.0).to(tl.int32)
+        ) | ((v < 0).to(tl.int32) << 3)
+        lo, hi = tl.split(tl.reshape(code, (BLOCK_SIZE_M, BLOCK_SIZE_K // 2, 2)))
+        values = (lo | (hi << 4)).to(tl.uint8)
+    return values
+@triton.jit
+def mx_act_quant_inline(
+    a_raw,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    FORMAT: tl.constexpr = "mxfp8",
+):
+    """Inline MX activation quant, one helper for both value grids. Per-row, per-K-group
+    amax → UE8M0 scale (ceil to the next power of two via the exponent-bump trick, the
+    divisor being the grid's largest magnitude) → values onto the format's grid:
+    - ``"mxfp8"``: cast to E4M3 — returns ``((M, K) fp8, (M, K // SCALE_GROUP_K) uint8)``.
+    - ``"mxfp4"``: round to E2M1 and pack nibble pairs (``_quant_e2m1_packed`` — hardware
+      ``cvt.e2m1x2`` on sm_100, else a ``>=``-threshold bucketize; they agree off exact
+      midpoints) — returns ``((M, K//2) uint8, (M, K // SCALE_GROUP_K) uint8)``.
+    - ``"nvfp4"``: E4M3 scale (amax/6, not a power of two), values divide by the DECODED
+      scale before the E2M1 grid — returns ``((M, K//2) uint8, (M, K // SCALE_GROUP_K) E4M3)``.
+    Only the taken format arm compiles."""
+    a_groups = tl.reshape(
+        a_raw, (BLOCK_SIZE_M, BLOCK_SIZE_K // SCALE_GROUP_K, SCALE_GROUP_K)
+    )
+    amax = tl.max(tl.abs(a_groups), axis=2)
+    if FORMAT == "nvfp4":
+        # E4M3 scale (amax/6 rounded to E4M3, NOT a power of two); values divide by the
+        # DECODED scale before hitting the E2M1 grid — the standard NVFP4 two-step
+        scales = (amax / 6.0).to(tl.float8e4nv)
+        decoded = tl.maximum(scales.to(tl.float32), 1.1754944e-38)
+        v = tl.reshape(a_groups / decoded[:, :, None], (BLOCK_SIZE_M, BLOCK_SIZE_K))
+        values = _quant_e2m1_packed(v, BLOCK_SIZE_M, BLOCK_SIZE_K)
+    elif FORMAT == "mxfp4":
+        bits = (amax / 6.0).to(tl.int32, bitcast=True)
+        # ceil_to_ue8m0: bump exponent by 1 when mantissa is non-zero.
+        exp_ceil = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32)
+        exp_ceil = tl.minimum(tl.maximum(exp_ceil, 1), 254)
+        exp_ceil = tl.where(amax == 0, 127, exp_ceil)
+        scales = exp_ceil.to(tl.uint8)
+        a_s_pow2 = (exp_ceil << 23).to(tl.float32, bitcast=True)
+        v = tl.reshape(a_groups / a_s_pow2[:, :, None], (BLOCK_SIZE_M, BLOCK_SIZE_K))
+        values = _quant_e2m1_packed(v, BLOCK_SIZE_M, BLOCK_SIZE_K)
+    else:
+        bits = (amax / 448.0).to(tl.int32, bitcast=True)
+        # ceil_to_ue8m0: bump exponent by 1 when mantissa is non-zero.
+        exp_ceil = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32)
+        exp_ceil = tl.minimum(tl.maximum(exp_ceil, 1), 254)
+        scales = exp_ceil.to(tl.uint8)
+        a_s_pow2 = (exp_ceil << 23).to(tl.float32, bitcast=True)
+        values = tl.reshape(
+            a_groups / tl.maximum(a_s_pow2[:, :, None], 1e-12),
+            (BLOCK_SIZE_M, BLOCK_SIZE_K),
+        ).to(tl.float8e4nv)
+    return values, scales
+@triton.jit
+def _e2m1_code_to_e4m3_bits(code):
+    """One E2M1 4-bit code -> the E4M3 byte holding the same value, in pure integer
+    ops. Every E2M1 magnitude ``{0, .5, 1, 1.5, 2, 3, 4, 6}`` is exact in E4M3, and
+    above the 0.5 subnormal the mapping is affine in the code: ``bits = (mag + 12) << 2``
+    (exponent re-bias +6, mantissa bit lands at bit 2). No float math, no converts —
+    callers bitcast the byte to ``float8e4nv``."""
+    code = code.to(tl.int32)
+    mag = code & 7
+    bits = tl.where(mag == 0, 0, tl.where(mag == 1, 0x30, (mag + 12) << 2))
+    return bits | ((code >> 3) << 7)
+@triton.jit
+def _e2m1_code_to_bf16_bits(code):
+    """One E2M1 4-bit code -> the bf16 bit pattern (uint16) holding the same value, pure int ops.
+    The 8 magnitudes ``{0,.5,1,1.5,2,3,4,6}`` are exact in bf16; above the 0.5 subnormal the mapping
+    is affine in the code: ``bits = 0x3F80 + (mag - 2) * 0x40`` (each step doubles/1.5x's cleanly on
+    the bf16 exponent/mantissa). Callers bitcast the uint16 to bf16 — no float convert, and no E4M3
+    intermediate. bf16 is the exact intermediate for every activation dtype, not the final one: the
+    caller widens the dequantized tile if the activation is wider."""
+    code = code.to(tl.int32)
+    mag = code & 7
+    bits = tl.where(mag == 0, 0, tl.where(mag == 1, 0x3F00, 0x3F80 + (mag - 2) * 0x40))
+    return bits | ((code >> 3) << 15)  # sign at bit 15
+@triton.jit
+def _e2m1_pair_to_f16(packed):
+    """One packed E2M1 byte -> the two fp16 values it holds, as a uint32 (low half = low nibble),
+    through the sm_100 CVT unit: ONE instruction per pair against the bit-twiddle form's ~10
+    integer ops. Every E2M1 magnitude is exact in fp16, so this is a value-identical decode, and
+    the direction of the unpack is irrelevant to the cost — the arithmetic is the whole bill
+    (measured on the gpt-oss W4A16 decode weight stream: 39.6us bit-twiddle, either axis, vs
+    11.6us here, against an 8.0us load-only floor). PTX has no 8-bit asm operand, so the byte
+    rides in a 16-bit register and is narrowed inside the block."""
+    return tl.inline_asm_elementwise(
+        "{ .reg .b8 t; cvt.u8.u16 t, $1; cvt.rn.f16x2.e2m1x2 $0, t; }",
+        "=r,h",
+        [packed.to(tl.uint16)],
+        dtype=tl.uint32,
+        is_pure=True,
+        pack=1,
+    )
+@triton.jit
+def e2m1_to_bf16(b_packed):
+    """Unpack packed MXFP4 (E2M1, two nibbles/byte along K) straight to bf16, doubling the K (row)
+    dim: ``(R, C) uint8 -> (2R, C) bf16``. Direct (no E4M3 round-trip), and the low/high nibbles
+    interleave along K with a SINGLE 3D axis-swap (``join`` -> ``[R, C, 2]`` -> swap last two ->
+    reshape) instead of the three 2D transposes ``e2m1_to_e4m3`` needs — the weight-only unpack hot path
+    (fp4 weight upcast in-loop, the caller widening to the activation dtype). K order low-nibble-first, matching
+    ``e2m1_to_e4m3``: ``[byte0_lo, byte0_hi, byte1_lo, ...]``. On sm_100 the codes decode through
+    the hardware CVT (``_e2m1_pair_to_f16``, fp16 being exact for the grid); elsewhere the integer
+    bit-build. Only the taken arm compiles."""
+    if _E2M1_HW_CVT:
+        pair = _e2m1_pair_to_f16(b_packed)
+        lo = (pair & 0xFFFF).to(tl.uint16)
+        hi = (pair >> 16).to(tl.uint16)
+        paired = tl.trans(tl.join(lo, hi), 0, 2, 1)
+        f16 = tl.reshape(
+            paired, (2 * b_packed.shape[0], b_packed.shape[1])
+        ).to(tl.float16, bitcast=True)
+        return f16.to(tl.bfloat16)
+    lo = _e2m1_code_to_bf16_bits(b_packed & 0xF)  # [R, C] int32 (bf16 bits)
+    hi = _e2m1_code_to_bf16_bits(b_packed >> 4)   # [R, C]
+    paired = tl.join(lo, hi)                       # [R, C, 2]  (pair in last dim)
+    paired = tl.trans(paired, 0, 2, 1)             # [R, 2, C]  (pair adjacent in row dim)
+    out = tl.reshape(paired, (2 * b_packed.shape[0], b_packed.shape[1]))  # [2R, C]
+    return out.to(tl.uint16).to(tl.bfloat16, bitcast=True)
+@triton.jit
+def _e2m1_pair_to_e4m3(packed):
+    """One packed E2M1 byte -> its two values as an E4M3 pair (uint16, low byte = low nibble),
+    through the CVT unit in two instructions — fp4 to fp16 to fp8 — against the bit-build's ~8
+    integer ops per value. Both hops are exact: every E2M1 magnitude is representable in fp16
+    and in E4M3, so the pair is value-identical to ``_e2m1_code_to_e4m3_bits``."""
+    return tl.inline_asm_elementwise(
+        "{ .reg .b8 t; .reg .b32 f; cvt.u8.u16 t, $1; cvt.rn.f16x2.e2m1x2 f, t; "
+        "cvt.rn.satfinite.e4m3x2.f16x2 $0, f; }",
+        "=h,h",
+        [packed.to(tl.uint16)],
+        dtype=tl.uint16,
+        is_pure=True,
+        pack=1,
+    )
+@triton.jit
+def e2m1_cols_to_bf16(packed):
+    """Column-unpack packed E2M1 to bf16: ``(..., C) uint8 -> (..., 2C)``, low nibble first — the
+    column-axis counterpart of the row-doubling ``e2m1_to_bf16``, for the swapped weight-only
+    tile (output rows leading, K contiguous). Same sm_100 hardware decode; the pair lands in
+    adjacent columns, so ``tl.interleave`` places it (no axis swap)."""
+    if _E2M1_HW_CVT:
+        pair = _e2m1_pair_to_f16(packed)
+        lo = (pair & 0xFFFF).to(tl.uint16)
+        hi = (pair >> 16).to(tl.uint16)
+        f16 = tl.interleave(lo, hi).to(tl.float16, bitcast=True)
+        return f16.to(tl.bfloat16)
+    bits = tl.interleave(
+        _e2m1_code_to_bf16_bits(packed & 0xF), _e2m1_code_to_bf16_bits(packed >> 4)
+    )
+    return bits.to(tl.uint16).to(tl.bfloat16, bitcast=True)
+@triton.jit
+def e2m1_cols_to_f32(packed):
+    """Column-unpack packed E2M1 straight to fp32, for the CUDA-core reduces: ``(..., C) uint8 ->
+    (..., 2C)``. The narrow-float hop its siblings owe their MMA callers (bf16 for ``tl.dot``, E4M3
+    for the scaled MMA) is pure work for a reduce that widens to fp32 anyway — measured 12.6 ->
+    15.4us on the gpt-oss decode tile for the bf16 round trip alone. Bit-identical to going via
+    either: every E2M1 value is exact in f16, bf16, E4M3 and fp32."""
+    if _E2M1_HW_CVT:
+        pair = _e2m1_pair_to_f16(packed)
+        lo = (pair & 0xFFFF).to(tl.uint16)
+        hi = (pair >> 16).to(tl.uint16)
+        f16 = tl.interleave(lo, hi).to(tl.float16, bitcast=True)
+        return f16.to(tl.float32)
+    bits = tl.interleave(
+        _e2m1_code_to_bf16_bits(packed & 0xF), _e2m1_code_to_bf16_bits(packed >> 4)
+    )
+    return bits.to(tl.uint16).to(tl.bfloat16, bitcast=True).to(tl.float32)
+@triton.jit
+def e2m1_cols_to_e4m3(packed):
+    """Column-unpack packed E2M1 (two nibbles per byte along the last dim, low nibble
+    first) to E4M3: ``(..., C) uint8 -> (..., 2C)`` — the column-axis counterpart of the
+    row-doubling ``e2m1_to_e4m3``; lossless (every E2M1 value is exact in E4M3). sm_100
+    decodes through the CVT unit (``_e2m1_pair_to_e4m3``), elsewhere the integer bit-build;
+    either way the bytes are bitcast once. Only the taken arm compiles."""
+    if _E2M1_HW_CVT:
+        pair = _e2m1_pair_to_e4m3(packed)
+        bits = tl.interleave(pair & 0xFF, pair >> 8)
+    else:
+        bits = tl.interleave(
+            _e2m1_code_to_e4m3_bits(packed & 0xF), _e2m1_code_to_e4m3_bits(packed >> 4)
+        )
+    return bits.to(tl.uint8).to(tl.float8e4nv, bitcast=True)
+@triton.jit
+def e2m1_to_f32(b_packed):
+    """Row-doubling counterpart of ``e2m1_cols_to_f32``: ``(R, C) uint8 -> (2R, C)`` fp32, for the
+    non-swapped CUDA-core reduce. Skips the E4M3 hop ``e2m1_to_e4m3`` owes its ``tl.dot`` callers;
+    bit-identical, since every E2M1 value is exact in E4M3, f16 and fp32."""
+    if _E2M1_HW_CVT:
+        pair = _e2m1_pair_to_f16(b_packed)
+        lo = (pair & 0xFFFF).to(tl.uint16)
+        hi = (pair >> 16).to(tl.uint16)
+        unpacked = tl.trans(tl.interleave(tl.trans(lo), tl.trans(hi)))
+        return unpacked.to(tl.uint16).to(tl.float16, bitcast=True).to(tl.float32)
+    lo = _e2m1_code_to_bf16_bits(b_packed & 0xF)
+    hi = _e2m1_code_to_bf16_bits(b_packed >> 4)
+    unpacked = tl.trans(tl.interleave(tl.trans(lo), tl.trans(hi)))
+    return unpacked.to(tl.uint16).to(tl.bfloat16, bitcast=True).to(tl.float32)
+@triton.jit
+def e2m1_to_e4m3(b_packed):
+    """Unpack packed MXFP4 (E2M1, two nibbles/byte along K) to E4M3, doubling the K
+    (row) dim: ``(R, C) uint8 -> (2R, C) E4M3``. E2M1's 8 magnitudes are all exact in
+    E4M3, so this is lossless — it lets the FP8 ``tl.dot`` path stand in for
+    ``tl.dot_scaled`` at decode (avoiding its M->128 pad). K order is the low nibble
+    first: ``[byte0_lo, byte0_hi, byte1_lo, ...]``. sm_100 decodes through the CVT unit
+    (``_e2m1_pair_to_e4m3``); elsewhere the integer bit-build. Only the taken arm compiles."""
+    if _E2M1_HW_CVT:
+        pair = _e2m1_pair_to_e4m3(b_packed)
+        lo, hi = pair & 0xFF, pair >> 8
+    else:
+        lo = _e2m1_code_to_e4m3_bits(b_packed & 0xF)
+        hi = _e2m1_code_to_e4m3_bits(b_packed >> 4)
+    # interleave along the K (row) dim via trans -> interleave-last-dim -> trans back
+    unpacked = tl.trans(tl.interleave(tl.trans(lo), tl.trans(hi)))
+    return unpacked.to(tl.uint8).to(tl.float8e4nv, bitcast=True)
+def _quant_block_k_pruner(configs, named_args, **kwargs):
+    """Keep configs whose BLOCK_K divides K (the quant grid is K // BLOCK_K programs per row;
+    K is always a multiple of 32, so the BLOCK_K=32 configs guarantee a non-empty list). On the
+    SWIZZLED path, additionally require BLOCK_K a multiple of ``4 * SCALE_GROUP_K`` so a whole
+    SWIZZLE_32_4_4 col-block (4 scale groups) lands inside one K-tile. BLOCK_T is free on the dense
+    grid (the per-element store handles any tile height) but pinned to 128 on the GROUPED grid,
+    where one program == one 128-row expert-sorted tile (``build_tile_layout``'s pad granularity)."""
+    args = {**named_args, **kwargs}
+    k, g = args["K"], args["SCALE_GROUP_K"]
+    swizzled, grouped = args.get("SWIZZLED", False), args.get("GROUPED", True)
+    return [
+        c
+        for c in configs
+        if k % c.kwargs["BLOCK_K"] == 0
+        and (not swizzled or c.kwargs["BLOCK_K"] % (4 * g) == 0)
+        and (not grouped or c.kwargs["BLOCK_T"] == 128)
+    ]
+@triton.jit
+def store_mx_act_scales(
+    S,
+    SOut,
+    out_row,  # per-lane output row (any int vector: contiguous grid rows or scattered dests)
+    s,
+    row_mask,
+    kb,
+    T,
+    K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    SWIZZLED: tl.constexpr,
+    MASK: tl.constexpr,  # "rows" (per-lane), "blocks" (whole 128-row blocks past T), "none"
+):
+    """Store one act-quant scale tile. Row-major (``S``, per ``out_row``) or straight into the
+    SWIZZLE_32_4_4 buffer (``SOut``) as PER-ELEMENT ptr arithmetic (no 128-row reshape): s[t, g]
+    lands at (block*cb_total + kb*REP_K + rep)*512 + r32*16 + outer4*4 + c4, where block =
+    out_row//128, r32 = out_row%32, outer4 = (out_row%128)//32, and the col-block splits g into
+    rep = g//4, c4 = g%4. Byte-identical to the old reshape form at BLOCK_T=128 (grouped grid)
+    but valid for any out_row vector, so the dense grid, the grouped-sorted grid, and the
+    scatter (quantize-once) stores all share it. Masking is the caller\'s store policy: the
+    grouped-sorted grid writes its padding rows too ("none" — quantized zeros the GEMM never
+    reads), the dense grid masks whole tiles past the real row-blocks ("blocks"), scattered
+    and row-major stores mask per lane ("rows")."""
+    groups: tl.constexpr = BLOCK_K // SCALE_GROUP_K
+    if SWIZZLED:
+        cb_total = K // SCALE_GROUP_K // 4
+        lg = tl.arange(0, groups)
+        rep = lg // 4
+        block = out_row // 128
+        off = (
+            (block[:, None] * cb_total + kb * (groups // 4) + rep[None, :]) * 512
+            + (out_row % 32)[:, None] * 16
+            + ((out_row % 128) // 32)[:, None] * 4
+            + (lg % 4)[None, :]
+        )
+        if MASK == "rows":
+            tl.store(SOut + off, s, mask=row_mask[:, None])
+        elif MASK == "blocks":
+            tl.store(SOut + off, s, mask=block[:, None] < tl.cdiv(T, 128))
+        else:
+            tl.store(SOut + off, s)
+    else:
+        sg = kb * groups + tl.arange(0, groups)
+        tl.store(
+            S + out_row.to(tl.int64)[:, None] * (K // SCALE_GROUP_K) + sg[None, :],
+            s,
+            mask=row_mask[:, None],
+        )
+@bayesian_autotune(
+    [
+        triton.Config({"BLOCK_K": bk, "BLOCK_T": bt}, num_warps=w)
+        for bk in (32, 64, 128, 256)
+        for bt in (8, 16, 32, 64, 128)
+        for w in (2, 4, 8)
+    ],
+    # t_bucket (log2 of the token count) is in the key: at small T the tile is the only
+    # parallelism lever while at prefill scale it isn't — same bucketing as the grouped
+    # kernels (raw T would retune per unique token count). SWIZZLED keys the swizzled-scale
+    # store separately (a disjoint config basin), GROUPED the grid — its 128-row expert tiles
+    # pin BLOCK_T, so its crowns cannot be shared with the dense grid's. FORMAT keys the value
+    # dtype/packing (E4M3 vs packed E2M1, and SCALE_GROUP_K 32 vs 16) — a dtype-blind key hands
+    # packed MXFP4 the E4M3 config and mistunes it.
+    ["K", "t_bucket", "SWIZZLED", "GROUPED", "FORMAT"],
+    n_trials=100,
+    prune_configs_by={"early_config_prune": _quant_block_k_pruner},
+)
+@triton.jit
+def _mx_act_quant_kernel(
+    X,
+    Y,
+    S,  # (T, K // SCALE_GROUP_K) row-major scales; None on the swizzled path
+    SOut,  # flat SWIZZLE_32_4_4 scale buffer (1, n_tiles, cb, 2, 256); None on the plain path
+    GatherIdx,  # (S,) int32 sorted position -> source row of X; read only when GROUPED and not None
+    ExpertStart,  # (NUM_EXPERTS_POW2 + 1,) int32 cumulative sorted-row starts; read iff GROUPED
+    GlobalScale,  # fp32 NVFP4 second-level global, per tensor or per expert; None ⇒ single-level (arm folds out)
+    GlobalIdx,  # (T,) row -> expert map, for a per-expert global on the DENSE grid (the batched op's own expert_ids); None on the grouped grid, whose tiles carry an expert
+    stride_global,  # expert stride of GlobalScale (0 broadcasts a scalar)
+    last_global,  # index of GlobalScale's last entry — clamps GlobalIdx, whose EP sentinel rows name an out-of-range expert (they are masked out of the GEMM, so any in-range global serves them)
+    stride_x_t,
+    stride_x_k,
+    T,
+    t_bucket,  # autotune key only (log2 token-count bucket); unused in body
+    K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    # dynamo's triton wrapper appends the tuner's config kwargs after the call kwargs
+    # and requires signature order — the tuned axes stay LAST
+    FORMAT: tl.constexpr = "mxfp8",
+    SWIZZLED: tl.constexpr = False,
+    GROUPED: tl.constexpr = True,  # grid: expert-sorted tiles (True) vs plain dense rows (False)
+    NUM_EXPERTS_POW2: tl.constexpr = 1,  # always passed explicitly; see the dense launch
+    BLOCK_K: tl.constexpr = 32,
+    BLOCK_T: tl.constexpr = 32,
+    PDL: tl.constexpr = False,
+):
+    """One-pass activation quant, one launch per format (``mx_act_quant_inline`` does
+    the math, so the offline and inline forms are bit-identical by construction): E4M3 +
+    UE8M0 ("mxfp8"), packed E2M1 + UE8M0 ("mxfp4"), or packed E2M1 + E4M3 group-16
+    ("nvfp4"). Group boundaries are identical across forms (SCALE_GROUP_K | BLOCK_K | K).
+    Arbitrary input strides.
+    Two independent axes. ``GROUPED`` picks the GRID: dense ``(cdiv(T, BLOCK_T), K // BLOCK_K)``
+    over plain row tiles (the one-row-per-program form starved memory at 1.5-1.8 TB/s on the
+    packed formats; the row tile coalesces), or ``(n_m_tiles, K // BLOCK_K)`` over the
+    expert-sorted, 128-padded tile layout (``build_tile_layout``, so BLOCK_T is pinned there).
+    A grouped program gathers its tile's source rows through ``GatherIdx`` (padding masked) and
+    scatters the VALUES back to source row order — the GEMM still TMA-gathers them, and duplicate
+    sorted->source writes store identical bytes — and its tile carries the expert a per-expert
+    NVFP4 global is indexed by, which is why only the dense grid needs a row map for that.
+    ``SWIZZLED`` picks the SCALE STORE: straight into the SWIZZLE_32_4_4 layout the grouped GEMM
+    reads affine (the inverse of ``load_swizzled_scale``), so no post-quant gather/swizzle pass,
+    or row-major ``S``. Padding-row scales are quantized zeros (harmless: the GEMM masks those
+    rows' values to 0)."""
+    if PDL:
+        gdc_wait()
+    kb = tl.program_id(1)
+    if SWIZZLED:
+        pid_m = tl.program_id(0)
+        # scale output-row position (== source row on the dense path); the swizzled block index is
+        # so // 128, valid for ANY BLOCK_T (no 128-row pin) so the dense grid autotunes BLOCK_T.
+        so = pid_m * BLOCK_T + tl.arange(0, BLOCK_T)
+    if GROUPED:
+        exp_start, freqs, tile_start_excl, total_m_tiles, e_offs = build_tile_layout(
+            ExpertStart, NUM_EXPERTS_POW2, BLOCK_T
+        )
+        expert_id, sorted_idx, row_mask = resolve_tile_inline(
+            tl.program_id(0), exp_start, freqs, tile_start_excl, e_offs, BLOCK_T
+        )
+        if GatherIdx is not None:
+            in_row = tl.load(GatherIdx + sorted_idx, mask=row_mask, other=0).to(tl.int64)
+        else:
+            in_row = sorted_idx.to(tl.int64)
+    else:
+        expert_id = 0
+        rows = tl.program_id(0) * BLOCK_T + tl.arange(0, BLOCK_T)
+        in_row = rows.to(tl.int64)
+        row_mask = rows < T
+    offs = kb * BLOCK_K + tl.arange(0, BLOCK_K)
+    x = tl.load(
+        X + in_row[:, None] * stride_x_t + offs[None, :] * stride_x_k,
+        mask=row_mask[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    if GlobalScale is not None:
+        # NVFP4 two-level: normalize by the calibrated global before the block quant — block
+        # scales are then amax/6 of x/g (the canonical two-step); the GEMM folds g back onto the
+        # accumulator (g_a·g_b). A per-expert global is indexed by the tile's OWN expert, which
+        # the grouped grid resolved for free — the rows of a tile belong to one expert; the dense
+        # grid has no such layout and takes the caller's row map. None folds the arm out.
+        if GlobalIdx is not None:
+            e = tl.minimum(tl.load(GlobalIdx + in_row, mask=row_mask, other=0), last_global)
+            x = x / tl.load(GlobalScale + e * stride_global).to(tl.float32)[:, None]
+        else:
+            x = x / tl.load(GlobalScale + expert_id * stride_global).to(tl.float32)
+    y, s = mx_act_quant_inline(x, BLOCK_T, BLOCK_K, SCALE_GROUP_K, FORMAT)
+    width: tl.constexpr = BLOCK_K // 2 if FORMAT != "mxfp8" else BLOCK_K
+    y_row: tl.constexpr = K // (BLOCK_K // width)  # per-row element count of Y
+    yo = kb * width + tl.arange(0, width)
+    # values -> source row order (the swizzled grid scatters via the gathered in_row)
+    if PDL:
+        gdc_launch_dependents()
+    tl.store(Y + in_row[:, None] * y_row + yo[None, :], y, mask=row_mask[:, None])
+    if SWIZZLED:
+        store_mx_act_scales(
+            S, SOut, so, s, row_mask, kb, T, K, SCALE_GROUP_K, BLOCK_K, True,
+            "none" if GROUPED else "blocks",
+        )
+    else:
+        store_mx_act_scales(
+            S, SOut, in_row, s, row_mask, kb, T, K, SCALE_GROUP_K, BLOCK_K, False, "rows"
+        )
+def mx_act_quant_grouped(
+    x: torch.Tensor,
+    activation_format: str,
+    scale_group: int,
+    scale_dtype: torch.dtype,
+    gather_idx: torch.Tensor | None,
+    expert_start: torch.Tensor,
+    global_scale: torch.Tensor | None = None,
+    swizzled: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor, int]:
+    """Offline MX act-quant on the GROUPED grid of ``_mx_act_quant_kernel``: one program per
+    expert-sorted tile, so each tile knows its own expert and a per-expert NVFP4 global needs no
+    per-row index. Returns ``(values, scales, num_m_tiles)`` — the values in source order (E4M3,
+    or packed E2M1 uint8 for the fp4 formats) and the scales either as the ``(1, num_m_tiles, cb,
+    2, 256)`` SWIZZLE_32_4_4 tensor the GEMM reads through a descriptor (``swizzled``) or
+    row-major ``(T, K // scale_group)`` for the affine arm. Only the ``expert_start[-1]``
+    scheduled rows are laid out; expert padding is per 128 (the ``BLOCK_T`` pin).
+    ``n_m_tiles`` is a STATIC host-side upper bound (``S//128 + E``) on the padded tile count —
+    never ``.item()`` (a CPU-GPU sync / cudagraph break). The buffer/grid over-allocate to it; the
+    kernel writes only the real tiles (``build_tile_layout`` in-kernel), extra tile-programs mask
+    out, and the GEMM reads only ``pid_m < total_m_tiles``. ``sum(ceil(freq/128)) <= S//128 + E``."""
+    T, K = x.shape
+    E = expert_start.numel() - 1
+    S = gather_idx.numel() if gather_idx is not None else T
+    n_m_tiles = S // 128 + E
+    packed = activation_format != "mxfp8"
+    y = torch.empty(
+        T, K // 2 if packed else K, device=x.device, dtype=torch.uint8 if packed else FP8_DTYPE
+    )
+    if swizzled:
+        cb = triton.cdiv(K // scale_group, 4)
+        scales = torch.empty(1, n_m_tiles, cb, 2, 256, device=x.device, dtype=scale_dtype)
+    else:
+        scales = torch.empty(T, K // scale_group, device=x.device, dtype=scale_dtype)
+    with device_context(x.device):
+        compile_time_only_triton_wrap(_mx_act_quant_kernel)[
+            lambda META: (n_m_tiles, K // META["BLOCK_K"])
+        ](
+            x,
+            y,
+            None if swizzled else scales,  # S: row-major scales, affine arm only
+            scales if swizzled else None,  # SOut: SWIZZLE_32_4_4 buffer, swizzled arm only
+            gather_idx,  # None = no gather (the is-not-None guard folds the load out)
+            expert_start,
+            global_scale,  # fp32 NVFP4 two-level global; None ⇒ single-level (arm folds out)
+            None,  # GlobalIdx: the grouped tiles carry their own expert
+            global_scale_stride(global_scale),
+            0,  # last_global: unread without a row map
+            x.stride(0),
+            x.stride(1),
+            T,
+            T.bit_length(),
+            K=K,
+            SCALE_GROUP_K=scale_group,
+            FORMAT=activation_format,
+            SWIZZLED=swizzled,
+            GROUPED=True,
+            NUM_EXPERTS_POW2=E,
+            PDL=decode_pdl(),
+            launch_pdl=decode_pdl(),
+        )
+    return y, scales, n_m_tiles
+@triton.jit
+def _swizzle_grouped_scales_kernel(
+    SRC,  # (rows, cols) row-major pre-quantized group scales (uint8 / e8m0 / e4m3, 1 byte)
+    DST,  # flat SWIZZLE_32_4_4 output buffer (1, n_tiles, NCB, 2, 256)
+    GatherIdx,  # (S,) int sorted position -> source row of SRC; read only when not None
+    ExpertStart,  # (NUM_EXPERTS_POW2 + 1,) cumulative sorted-row starts, S sentinel
+    COLS,
+    NCB,  # number of 4-wide column blocks (cols // 4)
+    stride_src_m,
+    NUM_EXPERTS_POW2: tl.constexpr,
+):
+    """Gather + expert-sorted 128-pad + swizzle a PRE-QUANTIZED grouped scale into the
+    SWIZZLE_32_4_4 layout, in one launch — the padded tile layout is derived in-kernel from
+    ``ExpertStart`` (``build_tile_layout``), so no torch index build. One 128x4 block per
+    (M-tile, col-block): inverse of ``load_swizzled_scale``. Padding rows write quantized-zero
+    scales (the GEMM masks those rows)."""
+    exp_start, freqs, tile_start_excl, total_m_tiles, e_offs = build_tile_layout(
+        ExpertStart, NUM_EXPERTS_POW2, 128
+    )
+    pid_m = tl.program_id(0)
+    cb = tl.program_id(1)
+    _, sorted_idx, row_mask = resolve_tile_inline(
+        pid_m, exp_start, freqs, tile_start_excl, e_offs, 128
+    )
+    if GatherIdx is not None:
+        src = tl.load(GatherIdx + sorted_idx, mask=row_mask, other=0)
+    else:
+        src = sorted_idx
+    cj = cb * 4 + tl.arange(0, 4)
+    s = tl.load(
+        SRC + src[:, None] * stride_src_m + cj[None, :],
+        mask=row_mask[:, None] & (cj[None, :] < COLS),
+        other=0,
+    )
+    swizzle_store_block(DST, s, pid_m, cb, NCB)
+def swizzle_grouped_mx_scales(
+    scale: torch.Tensor,
+    expert_start: torch.Tensor,
+    gather_idx: torch.Tensor | None = None,
+    pad_bm: int = 128,
+) -> tuple[torch.Tensor, int]:
+    """Fused gather + expert-sorted 128-pad + swizzle of a pre-quantized grouped scale (a
+    down projection's externally supplied ``As``), one triton launch — the layout comes from
+    ``expert_start`` in-kernel (no torch index build). Returns ``(swizzled_scale, num_m_tiles)``
+    as the ``(1, num_m_tiles, cols // 4, 2, 256)`` tensor; the caller builds the GEMM's scale
+    descriptor from it. ``gather_idx`` None = activations already expert-sorted.
+    ``n_m_tiles`` is a STATIC host-side upper bound (``S//pad_bm + E``) — never ``.item()`` (a
+    CPU-GPU sync / cudagraph break); the buffer/grid over-allocate and the extra tile-programs
+    mask out (``sum(ceil(freq/pad_bm)) <= S//pad_bm + E``)."""
+    E = expert_start.numel() - 1
+    S = gather_idx.numel() if gather_idx is not None else scale.shape[0]
+    n_m_tiles = S // pad_bm + E
+    cols = scale.shape[1]
+    cb = triton.cdiv(cols, 4)
+    src = scale.view(torch.uint8)  # byte-level; the binder rejects e8m0/e4m3
+    out = torch.empty(1, n_m_tiles, cb, 2, 256, device=scale.device, dtype=torch.uint8)
+    with device_context(scale.device):
+        compile_time_only_triton_wrap(_swizzle_grouped_scales_kernel)[(n_m_tiles, cb)](
+            src,
+            out,  # flat SWIZZLE_32_4_4 buffer (pointer store; no descriptor)
+            gather_idx,  # None = no gather (the is-not-None guard folds the load out)
+            expert_start,
+            cols,
+            cb,
+            src.stride(0),
+            NUM_EXPERTS_POW2=E,
+        )
+    return out.view(scale.dtype), n_m_tiles
+def maybe_act_quant(x, act_quant, min_m):
+    """Row-count-gated offline activation pre-quant. Apply ``act_quant`` (a one-pass
+    quant kernel, e.g. ``mxfp8_act_quant``) when the GEMM consuming ``x`` is
+    compute-bound (``rows >= min_m``) — the inline form re-quantizes per N-tile there.
+    ``min_m`` is the consumer kernel's crossover, defined next to its wrapper with its
+    provenance (measured sweep or inherited estimate). Below the threshold return ``x``
+    raw: the weight-bandwidth-bound GEMM quantizes its
+    one thin tile inline for free (the UE8M0 inline quant is exponent-only), and a
+    separate quant kernel is pure added latency (M=1 graph attn decode measured
+    0.66-0.85x offline). Bit-exact either way. Returns ``(a, a_scale)``; the consumer
+    kernel picks its arm off ``a``'s dtype at compile time (fp8 = pre-quantized, raw
+    bf16/fp16 = quantize inline), so in the inline arm ``a_scale`` is a constexpr-dead
+    placeholder."""
+    if x.shape[0] >= min_m:
+        return act_quant(x)
+    return x, x
+def mxfp8_act_quant(x: torch.Tensor, swizzled: bool = False) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize ``(T, K)`` activations to MX once (E4M3 values + UE8M0 group-32 uint8 scales)
+    instead of inline per weight-tile — the fused gate_up re-ran the inline quant per N-tile
+    (16x redundant amax/convert ALU + 2x act bytes), which held it at ~380 TFLOPS while the
+    pre-quantized down kernel ran at ~1080. One pass costs ~50µs at 8k tokens. Bit-exact with
+    the inline form (same group boundaries). ``swizzled=True`` emits the scale directly in
+    SWIZZLE_32_4_4 for the tcgen05 fast path (same dense grid, per-element store)."""
+    return _launch_act_quant(x, "mxfp8", MX_SCALE_GROUP_K, torch.uint8, swizzled)
+def mxfp4_act_quant(x: torch.Tensor, swizzled: bool = False) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize ``(T, K)`` activations to MXFP4 in one kernel pass: packed-E2M1 values
+    (``(T, K//2)`` int8, two codes per byte, first value in the low nibble) + UE8M0
+    group-32 uint8 scales (``(T, K//32)``, amax/6 ceil'd to a power of two). Bit-identical
+    to the fused epilogues' inline form (shared ``mx_act_quant_inline`` arm). Feeds the
+    W4A4 arm of the MX matmul ops; quantizing activations to fp4 at runtime is an
+    accuracy call the caller owns — the ops never do it implicitly. ``swizzled=True`` emits the
+    scale directly in SWIZZLE_32_4_4 for the tcgen05 fast path."""
+    return _launch_act_quant(x, "mxfp4", MX_SCALE_GROUP_K, torch.uint8, swizzled)
+def nvfp4_act_quant(
+    x: torch.Tensor,
+    swizzled: bool = False,
+    global_scale: torch.Tensor | None = None,
+    expert_index: torch.Tensor | None = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize ``(T, K)`` activations to NVFP4 in one kernel pass: packed-E2M1 values
+    (``(T, K//2)`` int8) + E4M3 group-16 block scales (``(T, K//16)`` — ``amax/6`` rounded to
+    E4M3, NOT a power of two; values divide by the DECODED scale before the E2M1 grid,
+    the standard two-step). ``global_scale`` is the CALIBRATED per-tensor second level
+    (fp32, the checkpoint's ``input_scale``): values are normalized by it before
+    the block quant, so block scales stay in e4m3 range regardless of the activation's
+    dynamic range — the canonical two-level format. The GEMM folds ``g_a·g_b`` back onto
+    the accumulator (pass ``As=scales`` with ``a_global_scale=global_scale``). Per-expert globals
+    need each row's expert, which ``expert_index`` gives where the caller already holds that map
+    (the batched op's ``expert_ids``, sentinels and all) — it is read only when the globals are
+    per expert, and ``mx_act_quant_grouped`` needs none at all, its tiles carrying an expert.
+    ``None`` = single-level (``g_a = 1``). ``swizzled=True`` emits the scale directly in
+    SWIZZLE_32_4_4 for the tcgen05 fast path."""
+    return _launch_act_quant(
+        x, "nvfp4", NVFP4_SCALE_GROUP_K, torch.float8_e4m3fn, swizzled, global_scale, expert_index
+    )
+def nvfp4_quantize_two_level(
+    weight: torch.Tensor, swizzled: bool = False
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Canonical two-level NVFP4 quant of a ``(N, K)`` (or ``(K,)``-last) tensor. Returns
+    ``(packed_e2m1 int8, e4m3 group-16 block scales, fp32 per-tensor global)`` — the block scale
+    is the op's ``Bs`` and the global its ``b_global_scale`` (the decoupled API keeps the two levels
+    as separate arguments).
+    The second level is a per-tensor fp32 global = ``amax / (6 · 448)`` — the smallest global that
+    keeps every e4m3 block scale in range (block ``amax/6`` after dividing by the global stays
+    ``≤ 448``). Two-level quant IS single-level quant of the globally-normalized tensor, so the block
+    values + scales come straight from ``nvfp4_act_quant(weight / global)``. The kernels multiply
+    the folded ``g_a · g_b`` onto the accumulator; the e4m3 block scales ride ``dot_scaled`` as
+    usual (activations are single-level ⇒ ``g_a = 1``)."""
+    global_scale = (weight.abs().amax() / (6.0 * 448.0)).clamp(min=1e-30).float()
+    # normalize in fp32: bf16 / 0-dim-fp32 stays bf16 under type promotion, and the extra
+    # rounding before the block quant flips E2M1 codes vs the exact normalized tensor
+    packed, block = nvfp4_act_quant((weight.float() / global_scale).contiguous(), swizzled)
+    return packed.view(torch.int8), block, global_scale.reshape(1)
+# offline act-quant pass per activation format (keys = ``resolve_activation_format`` results)
+MX_ACT_QUANT = {
+    "mxfp8": mxfp8_act_quant,
+    "mxfp4": mxfp4_act_quant,
+    "nvfp4": nvfp4_act_quant,
+}
+def _launch_act_quant(
+    x, activation_format, scale_group, scale_dtype, swizzled=False, global_scale=None,
+    expert_index=None,
+):
+    """One-pass activation quant for every format (``mxfp8`` = E4M3 values, else packed E2M1) and
+    both scale layouts. ``swizzled=True`` writes the scale straight into the SWIZZLE_32_4_4 buffer
+    ``(1, cdiv(T, 128), cb, 2, 256)`` (per-element ptr store, dense autotuned ``BLOCK_T`` — same grid
+    as the affine path, just the store address flips); ``swizzled=False`` writes row-major
+    ``(T, K // scale_group)``. ``global_scale`` (NVFP4 two-level, ``(1,)`` fp32) normalizes the
+    values before the block quant. Returns ``(values, scales)``."""
+    T, K = x.shape
+    packed = activation_format != "mxfp8"
+    per_expert = is_per_expert_global(global_scale)
+    if packed:
+        assert K % (2 * scale_group) == 0, (
+            f"K (={K}) must be a multiple of {2 * scale_group} to pack E2M1 pairs"
+        )
+    values = torch.empty(
+        T, K // 2 if packed else K, device=x.device,
+        dtype=torch.uint8 if packed else FP8_DTYPE,
+    )
+    if swizzled:
+        cb = triton.cdiv(K // scale_group, 4)
+        scales = torch.empty(
+            1, triton.cdiv(T, 128), cb, 2, 256, device=x.device, dtype=scale_dtype
+        )
+    else:
+        scales = torch.empty(T, K // scale_group, device=x.device, dtype=scale_dtype)
+    with device_context(x.device):
+        compile_time_only_triton_wrap(_mx_act_quant_kernel)[
+            lambda META: (triton.cdiv(T, META["BLOCK_T"]), K // META["BLOCK_K"])
+        ](
+            x,
+            values,
+            None if swizzled else scales,  # S: row-major scales, plain arm only
+            scales if swizzled else None,  # SOut: SWIZZLE_32_4_4 buffer, swizzled arm only
+            None,  # GatherIdx: the dense grid reads rows directly
+            None,  # ExpertStart: no expert-sorted tiles on the dense grid
+            global_scale,  # fp32 NVFP4 two-level global; None ⇒ single-level (arm folds out)
+            expert_index if per_expert else None,  # row -> expert map, read iff the globals are per expert
+            global_scale_stride(global_scale),
+            (global_scale.numel() - 1) if per_expert else 0,
+            x.stride(0),
+            x.stride(1),
+            T,
+            T.bit_length(),
+            K=K,
+            SCALE_GROUP_K=scale_group,
+            FORMAT=activation_format,
+            SWIZZLED=swizzled,
+            GROUPED=False,
+            # explicit even though the dense grid never reads it: inductor's wrap_triton path
+            # drops constexprs left to their default, and the arity mismatch that produces kills
+            # the whole torch.compile launch ("launcher() missing 1 required positional argument")
+            NUM_EXPERTS_POW2=1,
+            PDL=decode_pdl(),
+            launch_pdl=decode_pdl(),
+        )
+    return (values.view(torch.int8) if packed else values), scales
+@bayesian_autotune(
+    [
+        triton.Config({"BLOCK_T": bt}, num_warps=w)
+        for bt in (16, 32, 64, 128)
+        for w in (1, 2, 4, 8)
+    ],
+    ["K", "BLOCK_K", "t_bucket"],
+    n_trials=100,
+)
+@triton.jit
+def _fp8_act_quant_block_dynamic_kernel(
+    X,
+    Y,
+    S,
+    stride_x_t,
+    stride_x_k,
+    T,
+    t_bucket,  # autotune key only (log2 token-count bucket); unused in body
+    K: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+    BLOCK_T: tl.constexpr,
+    UE8M0: tl.constexpr = False,
+    PDL: tl.constexpr = False,
+):
+    """One-pass block-FP8 activation quant: rows → E4M3 + one ``amax/448`` scale per
+    ``BLOCK_K`` span (fp32, or a UE8M0 exponent byte under ``UE8M0``). Grid
+    ``(cdiv(T, BLOCK_T), K // BLOCK_K)`` — each program quantizes a ``[BLOCK_T, BLOCK_K]``
+    tile (``BLOCK_T`` rows over one K-block, each row its own scale). The one-row-per-program
+    form starved memory (128-byte transactions across ~T*K/BK tiny programs); the row tile
+    gives the loads something to coalesce. The span equals the consumer's ``BLOCK_SIZE_K``,
+    so results are bit-exact with the inline quant. Arbitrary input strides (no host-side
+    copy); ``BLOCK_K`` is fixed by the scale layout, ``BLOCK_T`` and warps are tuned."""
+    if PDL:
+        gdc_wait()
+    pid_t = tl.program_id(0)
+    kb = tl.program_id(1)
+    rows = (pid_t * BLOCK_T + tl.arange(0, BLOCK_T)).to(tl.int64)
+    offs = kb * BLOCK_K + tl.arange(0, BLOCK_K)
+    row_mask = rows < T
+    x = tl.load(
+        X + rows[:, None] * stride_x_t + offs[None, :] * stride_x_k,
+        mask=row_mask[:, None],
+        other=0.0,
+    ).to(tl.float32)
+    y, s = fp8_act_quant_inline(x, UE8M0=UE8M0)
+    tl.store(Y + rows[:, None] * K + offs[None, :], y, mask=row_mask[:, None])
+    if PDL:
+        gdc_launch_dependents()
+    tl.store(S + rows * (K // BLOCK_K) + kb, s, mask=row_mask)
+def fp8_act_quant_block_dynamic(
+    x: torch.Tensor, block_k: int, use_ue8m0: bool = False
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize ``(T, K)`` activations to block-FP8 once (E4M3 + per-``block_k`` scales)
+    instead of inline per weight-tile — same rationale and layout as ``mxfp8_act_quant`` (a
+    GEMM re-reads its activation once per N-tile). Bit-exact with the inline form. Scales are
+    fp32 (``amax/448``) or, under ``use_ue8m0``, UE8M0 exponent bytes (power-of-two scales)
+    for the tcgen05 ``dot_scaled`` path — the DeepGEMM-Blackwell format."""
+    T, K = x.shape
+    y = torch.empty(T, K, device=x.device, dtype=FP8_DTYPE)
+    s_dtype = torch.uint8 if use_ue8m0 else torch.float32
+    s = torch.empty(T, K // block_k, device=x.device, dtype=s_dtype)
+    def grid(META):
+        return (triton.cdiv(T, META["BLOCK_T"]), K // block_k)
+    with device_context(x.device):
+        compile_time_only_triton_wrap(_fp8_act_quant_block_dynamic_kernel)[grid](
+            x, y, s, x.stride(0), x.stride(1), T, T.bit_length(),
+            K=K, BLOCK_K=block_k, UE8M0=use_ue8m0,
+            PDL=decode_pdl(),
+            launch_pdl=decode_pdl(),
+        )
+    return y, s
+# ── fp8_act_quant_tensor_wide kernel (used by tensor-mode FP8 wrappers) ───────────────────
+@triton.jit
+def _fp8_act_quant_kernel(
+    x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr, PADDED_BLOCK: tl.constexpr,
+    PDL: tl.constexpr = False,
+):
+    if PDL:
+        gdc_wait()
+    # ``tl.arange`` needs a power-of-2 length, so iterate over PADDED_BLOCK (the next
+    # power of 2) and mask the tail — lets block_size be non-power-of-2 (e.g. a full
+    # row K=14336 in tensor-mode). Masked lanes load 0, which can't affect ``amax``.
+    pid = tl.program_id(axis=0)
+    cols = tl.arange(0, PADDED_BLOCK)
+    mask = cols < BLOCK_SIZE
+    offs = pid * BLOCK_SIZE + cols
+    x = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+    s = tl.max(tl.abs(x)) / 448.0  # float8_e4m3fn max
+    y = (x / tl.maximum(s, 1e-12)).to(y_ptr.dtype.element_ty)
+    tl.store(y_ptr + offs, y, mask=mask)
+    if PDL:
+        gdc_launch_dependents()
+    tl.store(s_ptr + pid, s)
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("fp8_act_quant_tensor_wide"), mutates_args=(), opaque=True
+)
+def fp8_act_quant_tensor_wide(
+    x: torch.Tensor, block_size: int = 128
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert x.is_contiguous()
+    assert x.shape[-1] % block_size == 0
+    y = torch.empty_like(x, dtype=FP8_DTYPE)
+    grid = (triton.cdiv(x.numel(), block_size),)
+    s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
+    with device_context(x.device):
+        compile_time_only_triton_wrap(_fp8_act_quant_kernel)[grid](
+            x,
+            y,
+            s,
+            BLOCK_SIZE=block_size,
+            PADDED_BLOCK=triton.next_power_of_2(block_size),
+            PDL=decode_pdl(),
+            launch_pdl=decode_pdl(),
+        )
+    return y, s

build/torch-rocm/scheduling.py ADDED Viewed

	@@ -0,0 +1,403 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import torch
+import triton
+import triton.language as tl
+from .compat import add_op_namespace_prefix, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, is_sm10x
+@triton.jit
+def build_tile_layout(
+    ExpertStart, NUM_EXPERTS: tl.constexpr, BLOCK_SIZE_M: tl.constexpr
+):
+    """Load ``expert_start`` once and derive the per-BM tile layout vectors (kept in
+    registers for the whole persistent loop): per-expert first sorted row, token count,
+    exclusive tile-start cumsum, and the total M-tile count. ``ExpertStart`` is
+    ``(NUM_EXPERTS + 1,)`` with a trailing ``S`` sentinel (``expert_start[E] == S``)."""
+    e_offs = tl.arange(0, NUM_EXPERTS)
+    exp_start = tl.load(ExpertStart + e_offs)
+    exp_end = tl.load(ExpertStart + e_offs + 1)
+    freqs = exp_end - exp_start
+    tiles_per_e = (freqs + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M
+    tile_start_excl = (
+        tl.cumsum(tiles_per_e, 0) - tiles_per_e
+    )  # first tile index of expert e
+    total_m_tiles = tl.sum(tiles_per_e, 0)
+    return exp_start, freqs, tile_start_excl, total_m_tiles, e_offs
+@triton.jit
+def resolve_tile_inline(
+    pid_m, exp_start, freqs, tile_start_excl, e_offs, BLOCK_SIZE_M: tl.constexpr
+):
+    """Map an M-tile id to its owning expert + the tile's sorted row range, from the
+    register-resident layout (no global loads). Returns ``(expert_id, sorted_indices,
+    row_mask)``."""
+    # Bucketize via the exclusive tile cumsum: #experts whose tile-start <= pid_m, minus 1.
+    expert_id = tl.sum((tile_start_excl <= pid_m).to(tl.int32), 0) - 1
+    sel = (
+        e_offs == expert_id
+    )  # scalar-index the E-vectors via mask-sum (no dynamic index)
+    e_start = tl.sum(tl.where(sel, exp_start, 0), 0)
+    e_tile_start = tl.sum(tl.where(sel, tile_start_excl, 0), 0)
+    freq = tl.sum(tl.where(sel, freqs, 0), 0)
+    within = pid_m - e_tile_start
+    m_start = e_start + within * BLOCK_SIZE_M
+    offs = tl.arange(0, BLOCK_SIZE_M)
+    row_mask = offs < freq - within * BLOCK_SIZE_M
+    sorted_indices = tl.max_contiguous(m_start + offs, BLOCK_SIZE_M)
+    return expert_id, sorted_indices, row_mask
+# Flat-slot tile per program for the O(S) routing kernels (count + scatter). These are small
+# latency-bound atomic kernels that want many programs: a sweep over {256..4096} x prefill shapes
+# put 256 best (or within ~1%) for both, with 1024 up to ~1.5x slower. The grid derives from it
+# so the two can't drift. Power of 2.
+_ROUTING_BLOCK_SIZE = 256
+@triton.jit
+def _exclusive_offsets_kernel(
+    ExpertFreq, ExpertStart, Counters, NUM_EXPERTS: tl.constexpr
+):
+    """Exclusive cumsum of per-expert token counts → ``expert_start`` (leading 0, trailing
+    S), and zero the scatter counters — one launch."""
+    offs = tl.arange(0, NUM_EXPERTS)
+    incl = tl.cumsum(tl.load(ExpertFreq + offs), 0)
+    tl.store(ExpertStart, 0)
+    tl.store(ExpertStart + 1 + offs, incl)
+    tl.store(Counters + offs, tl.zeros([NUM_EXPERTS], tl.int32))
+@triton.jit
+def _scatter_kernel(
+    ExpertIds,
+    Perm,
+    PermToken,
+    ExpertStart,
+    Counters,
+    S,
+    NUM_TOP_K: tl.constexpr,
+    NUM_EXPERTS: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Counting-sort scatter: each flat slot atomically claims the next slot of its expert
+    (``expert_start[e] + counter[e]++``). O(S), replaces an O(S·logS) argsort. Within-expert
+    order is arbitrary (atomic race) — fine, the per-token reduce is order-invariant. Slots whose
+    expert is non-local (EP sentinel id ``>= NUM_EXPERTS``) are skipped — matches ``_count_kernel``,
+    and avoids the atomic/store landing at an out-of-range (invalid) global address."""
+    offs = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    expert_id = tl.load(ExpertIds + offs, mask=offs < S, other=NUM_EXPERTS)
+    valid = expert_id < NUM_EXPERTS
+    dest = tl.load(ExpertStart + expert_id, mask=valid, other=0) + tl.atomic_add(
+        Counters + expert_id, 1, mask=valid, sem="relaxed"
+    )
+    tl.store(Perm + dest, offs, mask=valid)
+    tl.store(PermToken + dest, offs // NUM_TOP_K, mask=valid)
+@triton.jit
+def _count_kernel(
+    ExpertIds, ExpertFreq, S, NUM_EXPERTS: tl.constexpr, BLOCK_SIZE: tl.constexpr
+):
+    """Per-expert token count via atomics — replaces ``torch.histc`` (no float cast), fixed
+    ``(NUM_EXPERTS,)`` output stays CUDA-graph friendly. ``ExpertFreq`` is pre-zeroed."""
+    offs = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offs < S
+    expert_id = tl.load(ExpertIds + offs, mask=mask, other=NUM_EXPERTS)
+    tl.atomic_add(
+        ExpertFreq + expert_id, 1, mask=mask & (expert_id < NUM_EXPERTS), sem="relaxed"
+    )
+def compute_grouped_scheduling(
+    expert_ids: torch.Tensor, num_experts: int, num_top_k: int
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    """On-device routing: expert-sorted index (no copy of the activations) via two Triton
+    launches — exclusive offsets + an atomic counting-sort scatter (replaces host ``argsort``).
+    Run it once per layer and pass the results to every grouped GEMM of that layer. Returns
+    ``(expert_start, gather_idx, scatter_idx)``:
+    - ``expert_start`` — ``(E+1,)`` cumulative sorted-row starts padded with S; the tiling
+      schedule the kernels build their register-resident tile layout from.
+    - ``gather_idx`` — each sorted position's source row of hidden (``perm // num_top_k``,
+      many-to-one for top_k > 1: the gather that reads hidden without replication). Pass as the
+      GEMM's input map (``None`` = ``A`` already expert-sorted, e.g. the down projection).
+    - ``scatter_idx`` — each sorted position's token-major routed destination row ``(t*K + j)``,
+      the ``perm = torch.sort(expert_ids)`` indices (kernels un-permute by SCATTERING at store
+      time, never materializing ``inv_perm``). Pass as the output map (``None`` = leave the output
+      expert-sorted, e.g. the gate_up projection's intermediate).
+    E must be a power of 2 (the scheduling kernels hold the per-expert vectors in one
+    ``tl.arange`` block)."""
+    # the scheduling kernels hold the (E,) frequency/offset vectors in one tl.arange
+    # block, which requires a power of 2 — fail here with a clear message instead of a
+    # Triton compile error from an internal kernel
+    assert num_experts & (num_experts - 1) == 0, (
+        f"num_experts ({num_experts}) must be a power of 2"
+    )
+    gather_idx, scatter_idx, expert_start = _compute_grouped_scheduling(
+        expert_ids, num_experts, num_top_k
+    )
+    return expert_start, gather_idx, scatter_idx
+@compile_time_only_triton_op(
+    add_op_namespace_prefix("compute_grouped_scheduling"), mutates_args=(), opaque=True
+)
+def _compute_grouped_scheduling(
+    expert_ids: torch.Tensor, num_experts: int, num_top_k: int
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    device = expert_ids.device
+    expert_ids = expert_ids.int().contiguous()  # routing kernels index with unit stride
+    num_routed_tokens = expert_ids.numel()  # S = num_tokens * num_top_k
+    expert_freq = torch.zeros(num_experts, dtype=torch.int32, device=device)
+    expert_start = torch.empty(num_experts + 1, dtype=torch.int32, device=device)
+    counters = torch.empty(num_experts, dtype=torch.int32, device=device)
+    perm = torch.empty(num_routed_tokens, dtype=torch.int32, device=device)
+    perm_token = torch.empty(num_routed_tokens, dtype=torch.int32, device=device)
+    with device_context(device):
+        compile_time_only_triton_wrap(_count_kernel)[
+            (triton.cdiv(num_routed_tokens, _ROUTING_BLOCK_SIZE),)
+        ](
+            expert_ids,
+            expert_freq,
+            num_routed_tokens,
+            NUM_EXPERTS=num_experts,
+            BLOCK_SIZE=_ROUTING_BLOCK_SIZE,
+        )
+        compile_time_only_triton_wrap(_exclusive_offsets_kernel)[(1,)](
+            expert_freq,
+            expert_start,
+            counters,
+            NUM_EXPERTS=num_experts,
+        )
+        compile_time_only_triton_wrap(_scatter_kernel)[
+            (triton.cdiv(num_routed_tokens, _ROUTING_BLOCK_SIZE),)
+        ](
+            expert_ids,
+            perm,
+            perm_token,
+            expert_start,
+            counters,
+            num_routed_tokens,
+            NUM_TOP_K=num_top_k,
+            NUM_EXPERTS=num_experts,
+            BLOCK_SIZE=_ROUTING_BLOCK_SIZE,
+        )
+    return perm_token, perm, expert_start
+# The packed schedule exists for ONE tile size: BM=128, the only region the resolve trade
+# wins (packed_schedule_scope_pruner) — so the table is a single section: header slot 0 =
+# total_m_tiles, then (expert_id, m_start, row_end) per tile.
+PACKED_SCHEDULE_BM = 128
+def packed_schedule_capacity(num_routed_tokens: int, num_experts: int) -> int:
+    """Worst-case M-tile count: every expert pads its own last tile, so
+    ``cdiv(S, BM) + E`` bounds ``sum_e(cdiv(freq_e, BM))``. Fixed shape — graph-capturable."""
+    return (num_routed_tokens + PACKED_SCHEDULE_BM - 1) // PACKED_SCHEDULE_BM + num_experts
+def build_packed_schedule(
+    expert_start: torch.Tensor, num_routed_tokens: int, num_experts: int
+) -> torch.Tensor:
+    """Allocate + fill the packed tile schedule (one launch, fixed shape — graph-safe)."""
+    schedule = torch.empty(
+        1 + 3 * packed_schedule_capacity(num_routed_tokens, num_experts),
+        dtype=torch.int32, device=expert_start.device,
+    )
+    pow2 = triton.next_power_of_2(num_experts)
+    with device_context(expert_start.device):
+        compile_time_only_triton_wrap(_packed_schedule_kernel)[(pow2,)](
+            expert_start, schedule, NUM_EXPERTS=pow2
+        )
+    return schedule
+@triton.jit
+def _packed_schedule_kernel(ExpertStart, Schedule, NUM_EXPERTS: tl.constexpr):
+    """Build the packed tile schedule in one launch: one program per expert writes its
+    tiles' ``(expert_id, m_start, row_end)`` triples (header slot 0 = total tile count,
+    written by expert 0 from the full cumsum). Replaces the kernels' O(E)
+    register-resident tile layout with O(1) loads — the packed-schedule tuner axis
+    (register relief; wins at the pressure-walled BM=128 configs)."""
+    BM: tl.constexpr = 128  # PACKED_SCHEDULE_BM
+    e = tl.program_id(0)
+    e_start = tl.load(ExpertStart + e)
+    e_end = tl.load(ExpertStart + e + 1)
+    offs = tl.arange(0, NUM_EXPERTS)
+    starts = tl.load(ExpertStart + offs)
+    ends = tl.load(ExpertStart + offs + 1)
+    tiles_per = (ends - starts + BM - 1) // BM
+    before = tl.sum(tl.where(offs < e, tiles_per, 0), 0)
+    if e == 0:
+        tl.store(Schedule, tl.sum(tiles_per, 0))  # header: total_m_tiles
+    n_tiles = (e_end - e_start + BM - 1) // BM
+    for t in range(n_tiles):
+        base = 1 + (before + t) * 3
+        tl.store(Schedule + base, e)
+        tl.store(Schedule + base + 1, e_start + t * BM)
+        tl.store(Schedule + base + 2, e_end)
+@triton.jit
+def load_packed_schedule(Schedule, BLOCK_SIZE_M: tl.constexpr):
+    """The packed counterpart of ``build_tile_layout``: read the header (total_m_tiles).
+    The schedule is built for BM=128 only — the sole region the packed resolve wins (see
+    ``packed_schedule_scope_pruner``); the static_assert is the self-fence."""
+    tl.static_assert(BLOCK_SIZE_M == 128, "the packed schedule is built for BM=128")
+    return tl.load(Schedule)
+@triton.jit
+def prefetch_packed_entry(Schedule, pid_m, total_m_tiles):
+    """Issue a future M-tile's ``(expert_id, m_start, row_end)`` loads (masked past the
+    schedule end). Consumed by ``resolve_grouped_tile_packed`` one persistent iteration
+    later, so the entry's dependent-load latency hides under the current tile's K-loop
+    instead of serializing into every tile (~300ns/tile, the packed arm's whole overhead)."""
+    m = pid_m < total_m_tiles
+    entry = Schedule + 1 + pid_m * 3
+    return (tl.load(entry, mask=m, other=0),
+            tl.load(entry + 1, mask=m, other=0),
+            tl.load(entry + 2, mask=m, other=0))
+@triton.jit
+def resolve_grouped_tile_packed(
+    tile_id,
+    num_n_tiles,
+    prefetched,
+    GatherIdx,
+    ScatterIdx,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    GATE: tl.constexpr = False,
+):
+    """``resolve_grouped_tile`` with the M-tile resolved from the packed schedule —
+    consuming the triple ``prefetch_packed_entry`` issued last iteration (zero dependent
+    loads on this tile's critical path).
+    Same return contract; bit-identical tiles."""
+    pid_n = tile_id % num_n_tiles
+    expert_id, tile_row0, row_end = prefetched
+    offs_global_m = tl.max_contiguous(tile_row0 + tl.arange(0, BLOCK_SIZE_M), BLOCK_SIZE_M)
+    row_mask = offs_global_m < row_end
+    if GatherIdx is not None:
+        in_row = tl.load(GatherIdx + offs_global_m, mask=row_mask, other=0)
+    else:
+        in_row = offs_global_m
+    if ScatterIdx is not None:
+        out_row = tl.load(ScatterIdx + offs_global_m, mask=row_mask, other=0)
+    else:
+        out_row = offs_global_m
+    expert_id64 = expert_id.to(tl.int64)
+    WEIGHT_ROWS: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    offs_bn = pid_n * WEIGHT_ROWS + tl.arange(0, WEIGHT_ROWS)
+    row0 = expert_id64.to(tl.int32)
+    n_off = pid_n * WEIGHT_ROWS
+    m_start = tl.min(in_row).to(tl.int32)
+    return pid_n, expert_id, expert_id64, in_row, out_row, row_mask, offs_bn, row0, n_off, m_start
+@triton.jit
+def resolve_grouped_tile(
+    tile_id,
+    num_n_tiles,
+    exp_start,
+    freqs,
+    tile_start_excl,
+    e_offs,
+    GatherIdx,
+    ScatterIdx,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    GATE: tl.constexpr = False,
+):
+    """One persistent grouped tile: split the flat ``tile_id`` into (M-tile, N-tile), map
+    the M-tile to its expert + rows via ``resolve_tile_inline`` (on the register-resident
+    layout ``build_tile_layout`` builds once per program, passed in), and apply the virtual
+    sort — rows load from ``in_row`` and store to ``out_row``, mapped by ``GatherIdx`` /
+    ``ScatterIdx`` when present (``None`` = expert-sorted, the position itself; the ``None``
+    check folds at trace time, so no separate has-gather/has-scatter flag is needed).
+    Returns ``(pid_n, expert_id, expert_id64, in_row, out_row, row_mask, offs_bn, row0, n_off,
+    m_start)`` — both expert-id widths (``expert_id`` int32 for TMA descriptor row indices,
+    ``expert_id64`` int64 for byte-offset pointer arithmetic that overflows int32 at full expert
+    counts), plus the tile's derived load coordinates every grouped K-loop needs: ``row0`` (the
+    weight descriptor's expert-slab row base, ×2 under GATE for the gate|up stack), ``n_off`` (the
+    N-tile byte offset), and ``m_start`` (the contiguous sorted-row box start — ``min(in_row)``,
+    exact because a no-gather tile's rows are contiguous). Shared by the base grouped GEMMs and the
+    fused kernels; callers ``_``-ignore whichever fields they don't use."""
+    pid_m = tile_id // num_n_tiles
+    pid_n = tile_id % num_n_tiles
+    expert_id, offs_global_m, row_mask = resolve_tile_inline(
+        pid_m, exp_start, freqs, tile_start_excl, e_offs, BLOCK_SIZE_M
+    )
+    if GatherIdx is not None:
+        in_row = tl.load(GatherIdx + offs_global_m, mask=row_mask, other=0)
+    else:
+        in_row = offs_global_m
+    if ScatterIdx is not None:
+        out_row = tl.load(ScatterIdx + offs_global_m, mask=row_mask, other=0)
+    else:
+        out_row = offs_global_m
+    expert_id64 = expert_id.to(tl.int64)
+    WEIGHT_ROWS: tl.constexpr = 2 * BLOCK_SIZE_N if GATE else BLOCK_SIZE_N
+    offs_bn = pid_n * WEIGHT_ROWS + tl.arange(0, WEIGHT_ROWS)
+    row0 = expert_id64.to(tl.int32)
+    n_off = pid_n * WEIGHT_ROWS
+    m_start = tl.min(in_row).to(tl.int32)
+    return pid_n, expert_id, expert_id64, in_row, out_row, row_mask, offs_bn, row0, n_off, m_start
+def expand_gather_below_parity(vals, scales, gather_idx, num_experts):
+    """Trade an in-kernel activation gather for one contiguous copy of the (packed) rows +
+    scales, when the copy beats what gathering can offer: on sm_10x the gathered descriptor
+    arm is tma gather4, MONOTONIC in tokens/expert (2x loss at 32/expert, parity ~1024 — the
+    measured law in ``smem_pruner``'s gather clause), so expand below parity; every other
+    arch has no gather TMA and forfeits descriptor-A entirely under a gather, so any large-S
+    gather expands. Small S (decode) keeps the gather — nothing to amortize, pointer-A wins.
+    Call AFTER the activation quant: expanding packed rows costs a fraction of raw bf16 and
+    leaves the quant at (num_tokens, K) instead of top_k times that.
+    Returns ``(vals, scales, gather_idx)`` with ``gather_idx=None`` when expanded."""
+    if gather_idx is None:
+        return vals, scales, gather_idx
+    if expand_regime(gather_idx.shape[0], num_experts):
+        g = gather_idx.to(torch.long)
+        vals = vals[g].contiguous()
+        if scales is not None and scales.ndim == 2:  # row-major per-row scales follow the rows
+            scales = scales[g].contiguous()
+        gather_idx = None
+    return vals, scales, gather_idx
+def expand_regime(S_rows: int, num_experts: int) -> bool:
+    """Where one contiguous copy of the routed rows beats in-kernel gathering: prefill row
+    counts (S >= 4096 amortizes the copy) below the sm_10x tma-gather4 parity (~1024
+    tokens/expert — gather4's 1-row box is MONOTONIC in tokens/expert; every other arch has
+    no gather TMA at all, so any large-S gather expands).
+    Raw bf16 activations (the full-precision grouped path) were A/B'd against keeping the
+    in-kernel gather, 200-trial tunes each side, 2026-08-28: DeepSeek-V4 3455 -> 3344us (-3.2%),
+    DeepSeek-V3 7267 -> 7290 (wash), GLM-5.2 6039 -> 6439 (+6.6%). The copy (150-370us) is
+    offset by what the gather costs the GEMM, and which side wins is shape-specific within the
+    tune's roll variance — so the law is NOT split on ``act_itemsize``; it stays as measured.
+    ``FINEGRAINED_FORCE_GATHER=1`` is the A/B knob (never expand)."""
+    if os.environ.get("FINEGRAINED_FORCE_GATHER"):
+        return False
+    return S_rows >= 4096 and (not is_sm10x() or S_rows < 1024 * num_experts)

build/torch-rocm/swizzle.py ADDED Viewed

	@@ -0,0 +1,170 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+import triton
+import triton.language as tl
+from triton.tools.tensor_descriptor import TensorDescriptor
+from .compat import compile_time_only_triton_wrap, device_context
+@triton.jit
+def swizzle_store_block(DST, s, blk, cb, NCB):
+    """Pack one row-major ``(128, 4)`` scale block ``s`` into its SWIZZLE_32_4_4 ``(32, 16)`` block
+    and store it at flat offset ``(blk * NCB + cb) * 512`` — the inverse of the un-swizzle in
+    ``load_swizzled_scale``. Shared by every scale-swizzle kernel below."""
+    sw = s.reshape(4, 32, 4).trans(1, 0, 2).reshape(32, 16)
+    r = tl.arange(0, 32)
+    c = tl.arange(0, 16)
+    tl.store(DST + (blk * NCB + cb) * 512 + r[:, None] * 16 + c[None, :], sw)
+@triton.jit
+def _swizzle_scales_kernel(
+    SRC,  # (rows, cols) row-major block scales (uint8 / e8m0 / e4m3 — 1 byte)
+    DST,  # flat SWIZZLE_32_4_4 output, (n_row_blocks * n_col_blocks * 512,)
+    GatherIdx,  # (padded_rows,) output-row -> source row (-1 = pad); read only when not None
+    ROWS,
+    COLS,
+    NCB,  # number of 4-wide column blocks
+    stride_src_m,
+):
+    """One 128x4 SWIZZLE_32_4_4 block per (row-block, col-block) program: gather+pad+swizzle in
+    a single launch (replaces the torch view/permute/transpose chain — several kernels + two
+    full-tensor copies). Per block: ``(128, 4) -> (4, 32, 4) -> trans(1,0,2) -> (32, 16)`` — the
+    exact reorder the torch packer did, so bit-identical."""
+    rb = tl.program_id(0)
+    cb = tl.program_id(1)
+    ri = rb * 128 + tl.arange(0, 128)
+    if GatherIdx is not None:
+        src = tl.load(GatherIdx + ri, mask=ri < ROWS, other=-1)
+        valid = src >= 0
+        src = tl.where(valid, src, 0)
+    else:
+        src = ri
+        valid = ri < ROWS
+    cj = cb * 4 + tl.arange(0, 4)
+    s = tl.load(
+        SRC + src[:, None] * stride_src_m + cj[None, :],
+        mask=valid[:, None] & (cj[None, :] < COLS),
+        other=0,
+    )
+    swizzle_store_block(DST, s, rb, cb, NCB)
+def unswizzle_mx_scales(
+    swizzled: torch.Tensor, rows: int, cols: int, *, num_experts: int | None = None
+) -> torch.Tensor:
+    """Inverse of ``swizzle_mx_scales``: a SWIZZLE_32_4_4 artifact back to the plain row-major
+    ``(rows, cols)`` grid — or ``(num_experts, rows, cols)`` for an expert stack — with the
+    (128, 4) padding trimmed. Pure reshape/permute on the byte view (no kernel); it runs at
+    save/dequantize time, never on a forward, so the torch-level cost is irrelevant.
+    ``rows``/``cols`` are the LOGICAL dims of ONE matrix (for a gate|up slab, the full ``2N``);
+    both are trimmed back out of the (128, 4) padding. ``num_experts`` selects the stacked
+    ``(E, rows, cols)`` result and must match the artifact's leading dim."""
+    nrb, ncb = -(-rows // 128), -(-cols // 4)
+    n_matrices = num_experts if num_experts is not None else 1
+    assert swizzled.shape == (n_matrices, nrb, ncb, 2, 256), (
+        f"swizzled artifact is {tuple(swizzled.shape)}, expected {(n_matrices, nrb, ncb, 2, 256)} "
+        f"for {n_matrices}x({rows}, {cols})"
+    )
+    plain = (
+        swizzled.reshape(n_matrices, nrb * ncb, 32, 4, 4)
+        .transpose(2, 3)
+        .reshape(n_matrices, nrb, ncb, 128, 4)
+        .permute(0, 1, 3, 2, 4)
+        .reshape(n_matrices, nrb * 128, ncb * 4)[:, :rows, :cols]
+    )
+    return (plain if num_experts is not None else plain[0]).contiguous()
+def swizzled_scale_descriptor(scale_u8: torch.Tensor) -> TensorDescriptor:
+    """TMA descriptor over a SWIZZLE_32_4_4 artifact ``(G, row_blocks, cols // 4, 2, 256)`` addressed
+    as ``(1, G * row_blocks, cols // 4, 2, 256)``: the kernels index 128-row blocks flat
+    (``group * row_blocks + block``), so the expert axis folds into the block axis."""
+    assert scale_u8.ndim == 5 and scale_u8.is_contiguous(), (
+        f"swizzled scale must be a contiguous 5D artifact, got {tuple(scale_u8.shape)}"
+    )
+    groups, row_blocks, cols4 = scale_u8.shape[:3]
+    blocks = groups * row_blocks
+    return TensorDescriptor(
+        scale_u8, [1, blocks, cols4, 2, 256], [blocks * cols4 * 512, cols4 * 512, 512, 256, 1], [1, 1, 1, 2, 256]
+    )
+def _swizzle_to_blocks(
+    scale: torch.Tensor, gather_idx: torch.Tensor | None
+) -> torch.Tensor:
+    """One ``(rows, cols)`` scale matrix -> its ``(row_blocks, cols//4, 2, 256)``
+    SWIZZLE_32_4_4 block stack (single kernel launch)."""
+    cols = scale.shape[1]
+    rows = gather_idx.shape[0] if gather_idx is not None else scale.shape[0]
+    nrb = triton.cdiv(rows, 128)
+    ncb = triton.cdiv(cols, 4)
+    # the reorder is byte-level; view as uint8 so the triton binder accepts e8m0/e4m3 scales
+    src = scale.view(torch.uint8)
+    out = torch.empty(nrb * ncb * 512, device=scale.device, dtype=torch.uint8)
+    with device_context(scale.device):
+        compile_time_only_triton_wrap(_swizzle_scales_kernel)[(nrb, ncb)](
+            src,
+            out,
+            gather_idx,  # None = no gather (the is-not-None guard folds the load out)
+            rows,
+            cols,
+            ncb,
+            src.stride(0),
+        )
+    return out.view(scale.dtype).reshape(nrb, ncb, 2, 256)
+def swizzle_mx_scales(
+    scale: torch.Tensor, gather_idx: torch.Tensor | None = None
+) -> torch.Tensor:
+    """Reorder a block-scale tensor into the ``SWIZZLE_32_4_4`` layout the Blackwell tcgen05
+    scaled-MMA consumes, one triton launch per matrix (``_swizzle_scales_kernel`` — no torch
+    permute/transpose copies). The scale values are unchanged — this only rearranges them from
+    plain row-major into the swizzled order the tensor core reads directly (the same layout
+    cuBLAS/CUTLASS require for MXFP8/NVFP4 ``scaled_mm``). Plain row-major forces a gather that
+    caps the scaled dot below the fp8/fp4 peak; the swizzle removes it. Apply to weight scales
+    once at model load; the ops accept the result directly.
+    ``scale`` is the per-block scale grid — UE8M0 (``float8_e8m0fnu``, group-32 MXFP8/MXFP4) or
+    E4M3 (``float8_e4m3fn``, group-16 NVFP4); the reorder is dtype-agnostic (moves bytes):
+    - 2D ``(rows, K // group)``: one matrix (a dense weight's scale).
+    - 3D ``(E, rows, K // group)``: an expert stack, each expert swizzled independently so its
+      blocks stay ``ceil(rows/128)``-aligned (the reader indexes expert ``e`` at
+      ``e * ceil(rows/128)``; byte-identical to a flat swizzle when ``rows % 128 == 0``).
+    ``gather_idx`` (2D only): a 1-D ``(padded_rows,)`` index mapping each output (sorted) row to
+    its source row in ``scale`` (``-1`` = padding → zero row), folded into the kernel's load — the
+    routed/expert-sorted, per-tile-padded layout a grouped GEMM reads affine per BM=128 tile.
+    ``padded_rows`` must be a multiple of 128.
+    ``rows``/``cols`` are zero-padded to (128, 4) multiples; returns the 5D
+    ``(G, row_blocks, ceil(cols/4), 2, 256)`` artifact the ops read — ``G`` the expert count (1 for
+    a matrix), so an expert stack shards/gathers on its leading dim. Bit-identical to CUTLASS's
+    packer (verified)."""
+    assert gather_idx is None or gather_idx.shape[0] % 128 == 0, (
+        f"gather_idx rows must be 128-padded, got {None if gather_idx is None else gather_idx.shape[0]}"
+    )
+    if scale.ndim == 3:
+        assert gather_idx is None, "gather_idx applies to a single 2D scale, not an expert stack"
+        return torch.stack([_swizzle_to_blocks(e, None) for e in scale])
+    assert scale.ndim == 2, (
+            f"expected a 2D (rows, K//group) or 3D (E, rows, K//group) scale, got {tuple(scale.shape)}"
+    )
+    return _swizzle_to_blocks(scale, gather_idx).unsqueeze(0)