Kernels:

kernels-community
/

finegrained-kernels

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on 19 days ago

Commit

eaedc44

verified ·

1 Parent(s): 811eb42

Uploaded using `kernel-builder`.

Browse files

Files changed (12) hide show

build/torch-rocm/_ops.py +1 -1
build/torch-rocm/batched.py +79 -66
build/torch-rocm/bayesian_autotuner.py +6 -5
build/torch-rocm/compat.py +8 -0
build/torch-rocm/grouped.py +49 -27
build/torch-rocm/loading/tiles.py +12 -9
build/torch-rocm/matmul.py +71 -20
build/torch-rocm/metadata.json +12 -12
build/torch-rocm/metadata.json.sigstore +1 -1
build/torch-rocm/moe.py +53 -14
build/torch-rocm/pruners.py +116 -40
build/torch-rocm/quant.py +116 -4

build/torch-rocm/_ops.py CHANGED Viewed

@@ -22,7 +22,7 @@ def get_backend() -> str:
 def _find_ops_name() -> str:
     kernel_name = "finegrained_kernels"
-    unique_id = "1348465"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

 def _find_ops_name() -> str:
     kernel_name = "finegrained_kernels"
+    unique_id = "b0c6080"
     backend = get_backend()
     return f"_{kernel_name}_{backend}_{unique_id}"

build/torch-rocm/batched.py CHANGED Viewed

@@ -27,7 +27,7 @@ from .compat import add_op_namespace_prefix, FP8_DTYPE, MX_SCALE_GROUP_K, NIBBLE
 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 (
@@ -44,7 +44,7 @@ from .loading.tiles import (
     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
@@ -94,30 +94,6 @@ def expert_setup(
     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
@@ -281,7 +257,7 @@ def w8a8_block_dynamic_fp8_matmul_batched_kernel(
 @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
@@ -300,6 +276,7 @@ def w8a8_block_static_fp8_matmul_batched_kernel(
     stride_b_e,
     stride_b_k,
     stride_b_n,
     stride_bs_e,
     stride_bs_k,
     stride_bs_n,
@@ -334,7 +311,6 @@ def w8a8_block_static_fp8_matmul_batched_kernel(
     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(
@@ -355,6 +331,10 @@ def w8a8_block_static_fp8_matmul_batched_kernel(
     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)
@@ -368,7 +348,9 @@ def w8a8_block_static_fp8_matmul_batched_kernel(
     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,
@@ -412,10 +394,10 @@ def w8a8_block_static_fp8_matmul_batched_kernel(
 @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
@@ -428,6 +410,7 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel(
     stride_a_m,
     stride_a_k,
     stride_as_m,
     stride_b_e,
     stride_b_k,
     stride_b_n,
@@ -443,6 +426,13 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel(
     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.
@@ -452,7 +442,10 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel(
     ``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(
@@ -473,18 +466,20 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel(
     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(
@@ -493,14 +488,21 @@ def w8a8_tensor_dynamic_fp8_matmul_batched_kernel(
             "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
@@ -1170,7 +1172,8 @@ def w8a8_block_static_fp8_matmul_batched(
     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)
@@ -1197,11 +1200,8 @@ def w8a8_block_static_fp8_matmul_batched(
         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)
@@ -1235,6 +1235,7 @@ def w8a8_block_static_fp8_matmul_batched(
             B.stride(0),
             B.stride(2),
             B.stride(1),
             bs_u8.stride(0),
             bs_u8.stride(2),
             bs_u8.stride(1),
@@ -1274,6 +1275,11 @@ def w8a8_tensor_dynamic_fp8_matmul_batched(
     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,
@@ -1285,8 +1291,9 @@ def w8a8_tensor_dynamic_fp8_matmul_batched(
     (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)
@@ -1294,16 +1301,15 @@ def w8a8_tensor_dynamic_fp8_matmul_batched(
     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):
@@ -1328,7 +1334,8 @@ def w8a8_tensor_dynamic_fp8_matmul_batched(
             K,
             qA.stride(0),
             qA.stride(1),
-            As.stride(0),
             B.stride(0),
             B.stride(2),
             B.stride(1),
@@ -1339,6 +1346,11 @@ def w8a8_tensor_dynamic_fp8_matmul_batched(
             bias_stride_n,
             expert_ids.stride(0),
             num_experts=num_experts,
             PDL=decode_pdl(),
             launch_pdl=decode_pdl(),
         )
@@ -1819,12 +1831,15 @@ def matmul_batched(
         "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,
@@ -1912,14 +1927,12 @@ def matmul_batched(
             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(

 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 fp8_act_quant_block_dynamic, MX_ACT_QUANT, static_expert_act_operands, tensor_wide_act_operands
 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 (
     oriented_tile_ptrs,
     weight_tile_ptrs,
 )
+from .epilogue import acc_init, 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
     return batch_id, pid_n, expert_id, A, B, C, Bs, in_row, out_row
 @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
 @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,  # calibrated (static) activation scale: one value, or one per expert
     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
     stride_b_e,
     stride_b_k,
     stride_b_n,
+    stride_as_e,  # 0 = one calibrated scale shared by every expert
     stride_bs_e,
     stride_bs_k,
     stride_bs_n,
     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."""
     if PDL:
         gdc_wait()
     batch_id, pid_n, expert_id, A, B, C, Bs, in_row, out_row = expert_setup(
     if expert_id >= num_experts:
         return
+    # this batch's expert scale: stride 0 means one calibrated scale for every expert. It feeds
+    # the inline quant arm and folds back onto the accumulator post-loop.
+    a_s_static = tl.load(As + expert_id.to(tl.int32) * stride_as_e)
     # 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)
     acc = acc_init("dot", BLOCK_SIZE_M, n_width, SWAP_AB)
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        # raw A quantizes in register against this expert's scale; a pre-quantized E4M3 token
+        # (fake-batch replicated) takes the other arm and ignores it
+        a, _ = load_act_static(a_ptrs, 0, 0, 0, None, 0, a_s_static, "pointer", False)
         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,
 @triton.jit
 def w8a8_tensor_dynamic_fp8_matmul_batched_kernel(
     A,  # (S, K) pre-quantized FP8 activations
+    As,  # per-token activation scales (S,), or a calibrated (static) one: shared, or per expert
+    B,  # (num_experts, N, K) FP8 weights; under GATE the (num_experts, 2N, K) gate|up stack
     Bs,  # (num_experts, 1, 1) per-tensor weight scales
+    C,  # (S, N) output; under GATE the bf16 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
     stride_a_m,
     stride_a_k,
     stride_as_m,
+    stride_as_e,  # 0 = the scale is per token; 1 = one calibrated scale per expert
     stride_b_e,
     stride_b_k,
     stride_b_n,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_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,
+    SIMULATE_UNFUSED: tl.constexpr = False,
+    INTERMEDIATE_DTYPE: tl.constexpr = tl.bfloat16,
     PDL: tl.constexpr = False,
 ):
     """Tensor-scale batched FP8 expert matmul kernel.
     ``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.
+    ``GATE`` fuses the gate|up projection (``B`` the interleaved (2N, K) stack, one per-tensor
+    scale over both halves) into one tile + SwiGLU, emitting the bf16 intermediate; ``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(
     if expert_id >= num_experts:
         return
+    # under GATE the gate|up rows are interleaved, so the tile is one 2*BN span
+    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)
     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 + expert_id * stride_as_e)
+    accumulator = acc_init("dot", BLOCK_SIZE_M, n_width, SWAP_AB)
     for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a, _ = load_act_static(a_ptrs, 0, 0, 0, None, 0, a_s, "pointer", False)
         b, _ = 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, b, SWAP_AB, BLOCK_SIZE_K)
         a_ptrs, _, b_ptrs, _, _, _ = advance_ptrs(
             "pointer", "pointer", False, False, False,
         )
+    # per-tensor scales fold on the raw accumulator; `gemm_epilogue` owns the finalize.
+    # Finalizing here too was invisible under no-swap (a pass-through) and collapsed the
+    # swapped tile twice, so every SWAP_AB config failed to compile and the tuner quietly
+    # dropped the whole swap arm at this shape.
+    accumulator = accumulator * a_s * b_s
     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, 1,
+        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,
+    )
 # The MXFP4/MXFP8 (and packed-activation) splits key themselves — the tuner appends every tensor
     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,) for one calibrated scale, or (num_experts,) for a MoE that
+        calibrates each expert separately
     Bs: (num_experts, N // block_n, K // block_k) per-block weight scales (2N under gate)
     """
     validate_dense_operands(A, B)
         f"the fused 'fp8' requant needs square quant blocks, got {block_size}"
     )
+    A_q, As, as_stride = static_expert_act_operands(A, As, num_experts)
     bs_u8 = ue8m0_as_uint8(Bs)
     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)
             B.stride(0),
             B.stride(2),
             B.stride(1),
+            as_stride,
             bs_u8.stride(0),
             bs_u8.stride(2),
             bs_u8.stride(1),
     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,
     output_dtype: torch.dtype | None = None,
     gather_idx: torch.Tensor | None = None,
     scatter_idx: torch.Tensor | None = None,
     (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; under ``gate`` the (num_experts, 2N, K) stack
+    As: (rows,) per-token scales alongside a pre-quantized A, or — on a raw A — the calibrated
+        (static) activation scale: one value, or one per expert
     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, rows, _ = B.shape
+    # under gate|up fusion B is the (E, 2N, K) stack; N is the per-projection output width
+    N = rows // 2 if gate else rows
     # Normalize Bs to (num_experts, 1, 1)
     Bs = normalize_per_expert_scale(Bs, num_experts)
     bs_u8 = ue8m0_as_uint8(Bs)
+    qA, As, as_stride_m, as_stride_e = tensor_wide_act_operands(A, As, num_experts)
     C = A.new_empty(S, N, dtype=output_dtype)
     def grid(META):
             K,
             qA.stride(0),
             qA.stride(1),
+            as_stride_m,
+            as_stride_e,
             B.stride(0),
             B.stride(2),
             B.stride(1),
             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,
             PDL=decode_pdl(),
             launch_pdl=decode_pdl(),
         )
         "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)"
     )
+    # a calibrated (static) scale on a raw A (see `tensor_wide_act_operands`): block-scale weights
+    # have a dedicated static kernel, per-tensor ones read the same As on the tensor-wide arm
+    static_act = As is not None and As.ndim <= 1 and As.numel() in (1, B.shape[0]) and A.dtype != FP8_DTYPE
+    if static_act:
+        assert Bs is not None and not is_mx(B, Bs), (
+            "a calibrated (static) activation scale is an FP8 form — MX activations carry a scale "
+            "per group, derived per call"
         )
+    if static_act and weight_block_size(B, Bs) is not None:
         out = w8a8_block_static_fp8_matmul_batched(
             A,
             B,
             bias=bias,
         )
     elif (block_size := weight_block_size(B, Bs)) is None:
         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, gate, act_fn, swiglu_alpha, swiglu_limit, simulate_unfused,
+            output_dtype, gather_idx, scatter_idx, bias=bias,
         )
     else:
         out = w8a8_block_dynamic_fp8_matmul_batched(

build/torch-rocm/bayesian_autotuner.py CHANGED Viewed

@@ -101,11 +101,12 @@ class BayesianAutotuner(Autotuner):
         # 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

         # 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.
+        # an explicit argument wins over the env override, which wins over the default; `or`
+        # would drop a deliberate `max_failures=0`, the strictest setting there is
+        if max_failures is None:
+            env_max_failures = os.environ.get("FINEGRAINED_AUTOTUNE_MAX_FAILURES")
+            max_failures = env_max_failures if env_max_failures is not None else self.n_trials
+        self.max_failures = int(max_failures)
         self.n_startup_trials = n_startup_trials
         # top fraction of measured configs the TPE treats as "good"
         self.gamma = gamma

build/torch-rocm/compat.py CHANGED Viewed

@@ -40,6 +40,14 @@ def decode_pdl() -> bool:
     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.

     no fusion). Deployment decode runs eager launches under cudagraphs, where PDL applies."""
     return DECODE_PDL and not torch.compiler.is_compiling()
+# The scaled_grouped_mm scaling enums the torch MoE baseline hands its scale operands. Older
+# torch has neither them nor the op, and only that one baseline reads them, so a missing pair is
+# ``None`` here and the baseline refuses on it rather than the whole package failing to import.
+try:
+    from torch.nn.functional import ScalingType, SwizzleType
+except ImportError:
+    ScalingType = SwizzleType = None
 # 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.

build/torch-rocm/grouped.py CHANGED Viewed

@@ -24,10 +24,10 @@ from .bayesian_autotuner import bayesian_autotune
 from .compat import add_op_namespace_prefix, FP8_DTYPE, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, sm_count, tl_dtype
 from .descriptors import build_grouped_operand_descriptors, rebind_grouped_descriptors, rebind_grouped_mx_descriptors
 from .formats import check_activation_format, global_scale_stride, is_per_expert_global, normalize_global_scale, e2m1_as_uint8, expert_weight_shape, is_mx, mx_scale_family, normalize_per_expert_scale, resolve_activation_format, resolve_output_dtype, routed_rows, tokens_per_expert_bucket, ue8m0_as_uint8, validate_dense_operands, weight_block_size, weight_format
-from .quant import MX_ACT_QUANT, fp8_act_quant_block_dynamic, fp8_act_quant_tensor_wide, mx_act_quant_grouped, swizzle_grouped_mx_scales
 from .swizzle import swizzled_scale_descriptor
 from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot
-from .scheduling import build_tile_layout, expand_gather_below_parity, build_packed_schedule, load_packed_schedule, prefetch_packed_entry, resolve_grouped_tile, resolve_grouped_tile_packed
 from .loading.tiles import (
     load_act_block_dynamic,
     load_act_mx,
@@ -41,7 +41,7 @@ from .loading.tiles import (
     weight_tile_ptrs,
 )
 from .epilogue import acc_init, bias_strides, gemm_epilogue
-from .pruners import PATH_ANCHOR_AXES, fp8_dot_warp_pruner, global_scale_warp_spec_pruner, weight_only_warp_spec_matched_mode_pruner, packed_schedule_scope_pruner, affine_scale_warp_spec_pruner, block_dynamic_grouped_matmul_pruner, block_fits_dim_pruner, block_within_dim_pruner, compose_pruners, descriptor_box_pruner, gate_stacked_tmem_trap_pruner, gated_pointer_weight_warp_spec_pruner, mx_config_pruner, require_moe_dims_aligned, smem_pruner, swizzled_out_bm_pruner, swizzled_scale_config_pruner, swizzled_scales_bm_pruner, warp_spec_compile_guard_pruner
 @bayesian_autotune(
@@ -69,6 +69,7 @@ from .pruners import PATH_ANCHOR_AXES, fp8_dot_warp_pruner, global_scale_warp_sp
             fp8_dot_warp_pruner(),
             packed_schedule_scope_pruner(),
             block_dynamic_grouped_matmul_pruner(),
             descriptor_box_pruner(),
         )
     },
@@ -274,6 +275,7 @@ def w8a8_block_dynamic_fp8_matmul_grouped_kernel(
             packed_schedule_scope_pruner(),
             block_dynamic_grouped_matmul_pruner(),
             descriptor_box_pruner(),
         )
     },
 )
@@ -281,7 +283,7 @@ def w8a8_block_dynamic_fp8_matmul_grouped_kernel(
 def w8a8_block_static_fp8_matmul_grouped_kernel(
     A,  # (num_tokens, K) E4M3 activations (pre-quantized against the static scale by the wrapper)
     ADescriptor,  # host TMA descriptor over A (rows, K), box (BM, BK); read iff A_MEMORY_MODE != "pointer"
-    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
     BDescriptor,  # host TMA descriptor over B viewed (E, 2N|N, K), box (1, (2|1)*BN, BK); read iff B_MEMORY_MODE != "pointer"
     Bs,  # (num_experts, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales (2N under GATE)
@@ -302,6 +304,7 @@ def w8a8_block_static_fp8_matmul_grouped_kernel(
     stride_b_e,
     stride_b_k,
     stride_b_n,
     stride_bs_e,
     stride_bs_k,
     stride_bs_n,
@@ -342,7 +345,6 @@ def w8a8_block_static_fp8_matmul_grouped_kernel(
     weight scales apply per-K-tile (plain ``tl.dot`` + software rescale, ``accumulate("static")``),
     and the scalar activation scale multiplies the accumulator once after the loop. GATE=False is
     the plain grouped GEMM (down projection), bit-identical."""
-    a_s_static = tl.load(As)  # per-tensor static activation scale, applied post-loop
     start_pid = tl.program_id(axis=0)
     if PACKED_SCHEDULE:
         total_m_tiles = load_packed_schedule(Schedule, BLOCK_SIZE_M)
@@ -391,10 +393,13 @@ def w8a8_block_static_fp8_matmul_grouped_kernel(
         gate_s_ptr = Bs + expert_id64 * stride_bs_e + ((2 if GATE else 1) * pid_n) * stride_bs_n
         up_s_ptr = gate_s_ptr + stride_bs_n
         acc = acc_init("dot", BLOCK_SIZE_M, (2 if GATE else 1) * BLOCK_SIZE_N, False)
         for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC):
             a, a_dead = load_act_static(
-                a_ptrs, ADescriptor, m_start, k * BLOCK_SIZE_K, row_mask, in_row, 0.0,
                 A_MEMORY_MODE, GatherIdx is not None,
             )
             w, w_s = load_weight_static(
@@ -466,6 +471,7 @@ def w8a8_block_static_fp8_matmul_grouped_kernel(
             warp_spec_compile_guard_pruner(),
             descriptor_box_pruner(),
             smem_pruner(),
         )
     },
 )
@@ -473,7 +479,7 @@ def w8a8_block_static_fp8_matmul_grouped_kernel(
 def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel(
     A,  # (num_tokens, K) pre-quantized FP8 activations, any row order
     ADescriptor,  # host TMA descriptor over A (rows, K), box (BM, BK); read iff A_MEMORY_MODE != "pointer"
-    As,  # (S,) per-token activation scales
     B,  # (num_experts, N, K) FP8 weights; under GATE the (num_experts, 2N, K) gate|up stack
     BDescriptor,  # host TMA descriptor over B viewed (E, 2N|N, K), box (1, (2|1)*BN, BK); read iff B_MEMORY_MODE != "pointer"
     Bs,  # (num_experts, 1, 1) per-tensor weight scales (one scalar covers the gate|up stack)
@@ -492,6 +498,7 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel(
     stride_a_m,
     stride_a_k,
     stride_as_m,
     stride_b_e,
     stride_b_k,
     stride_b_n,
@@ -580,13 +587,17 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel(
             stride_b_k,
             False,
         )
-        a_s = tl.load(As + in_row * stride_as_m, mask=row_mask, other=0.0)
         b_s = tl.load(Bs + expert_id64 * stride_bs_e)
         acc = acc_init("dot", BLOCK_SIZE_M, (2 if GATE else 1) * BLOCK_SIZE_N, False)
         for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC):
-            a, _as = load_act_plain(
-                a_ptrs, ADescriptor, m_start, k * BLOCK_SIZE_K, row_mask, in_row,
                 A_MEMORY_MODE, GatherIdx is not None,
             )
             w, _ws = load_weight_plain(
@@ -665,6 +676,7 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel(
             smem_pruner(min_sets=2),  # two full buffer sets: the grouped pipeline's floor (see smem_pruner)
             warp_spec_compile_guard_pruner(),
             affine_scale_warp_spec_pruner(),
         )
     },
 )
@@ -916,7 +928,7 @@ def mx_dynamic_matmul_grouped_kernel(
     path_anchor_axes=PATH_ANCHOR_AXES,
     prune_configs_by={
         "early_config_prune": compose_pruners(
-            weight_only_warp_spec_matched_mode_pruner(),
             packed_schedule_scope_pruner(),
             # the tail tile slides back in bounds (see the K-loop) — BK only has to fit
             block_fits_dim_pruner("K"),
@@ -1436,7 +1448,8 @@ def w8a8_block_static_fp8_matmul_grouped(
     A:  (S, 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)
@@ -1463,11 +1476,8 @@ def w8a8_block_static_fp8_matmul_grouped(
     )
     output_dtype = resolve_output_dtype(output_dtype, A, None)
-    As = As.reshape(1).to(torch.float32)
     bs_u8 = ue8m0_as_uint8(Bs)
-    # Pre-quantize the raw activations against the calibrated scalar (offline — MoE always
-    # pre-quants; 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)
         # UE8M0 model (ue8m0 weights) -> UE8M0 intermediate scales; the epilogue infers the format
@@ -1510,6 +1520,7 @@ def w8a8_block_static_fp8_matmul_grouped(
             B.stride(0),
             B.stride(2),
             B.stride(1),
             bs_u8.stride(0),
             bs_u8.stride(2),
             bs_u8.stride(1),
@@ -1569,7 +1580,8 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped(
     A:  (S, K) pre-quantized FP8 activations — rows addressed via ``gather_idx``
     B:  (num_experts, N, K) FP8 expert weights; under ``gate`` the (num_experts, 2N, K) stack
-    As: (S,) per-token activation scales
     Bs: (num_experts,) or (num_experts, 1, 1) per-expert weight scales
     expert_start: (num_experts_pow2 + 1,) int32 — cumulative sorted-row starts, S sentinel
     gather_idx: optional (S,) — sorted position -> source row of A; None = A is expert-sorted
@@ -1593,11 +1605,17 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped(
     # Normalize Bs to (num_experts, 1, 1) — one per-tensor scale (covers the gate|up stack)
     Bs = normalize_per_expert_scale(Bs, num_experts)
-    # A raw (As is None) -> quantize here (offline, per-token); else pre-quantized.
     output_dtype = resolve_output_dtype(output_dtype, A, As)
-    if As is None:
-        A, As = fp8_act_quant_tensor_wide(A, K)
-        # post-quant: trade the in-kernel gather for one packed-row copy where that wins
         A, _, gather_idx = expand_gather_below_parity(A, None, gather_idx, num_experts)
     C = A.new_empty(S, N, dtype=output_dtype)
     num_sms = sm_count(A.device.index)
@@ -1631,7 +1649,8 @@ def w8a8_tensor_dynamic_fp8_matmul_grouped(
             K,
             A.stride(0),
             A.stride(1),
-            As.stride(0),
             B.stride(0),
             B.stride(2),
             B.stride(1),
@@ -2234,12 +2253,15 @@ def matmul_grouped(
         "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_grouped(
             A,
             B,

 from .compat import add_op_namespace_prefix, FP8_DTYPE, NIBBLES_PER_BYTE, compile_time_only_triton_op, compile_time_only_triton_wrap, device_context, get_accelerator_autotuning_configs, sm_count, tl_dtype
 from .descriptors import build_grouped_operand_descriptors, rebind_grouped_descriptors, rebind_grouped_mx_descriptors
 from .formats import check_activation_format, global_scale_stride, is_per_expert_global, normalize_global_scale, e2m1_as_uint8, expert_weight_shape, is_mx, mx_scale_family, normalize_per_expert_scale, resolve_activation_format, resolve_output_dtype, routed_rows, tokens_per_expert_bucket, ue8m0_as_uint8, validate_dense_operands, weight_block_size, weight_format
+from .quant import fp8_act_quant_block_dynamic, MX_ACT_QUANT, mx_act_quant_grouped, quantize_routed_rows_per_expert, static_expert_act_operands, swizzle_grouped_mx_scales, tensor_wide_act_operands
 from .swizzle import swizzled_scale_descriptor
 from .mma import block_dynamic_dot, fp8_dot, mx_compute, mx_weight_only_compute, static_dot
+from .scheduling import build_tile_layout, expand_gather_below_parity, expand_regime, build_packed_schedule, load_packed_schedule, prefetch_packed_entry, resolve_grouped_tile, resolve_grouped_tile_packed
 from .loading.tiles import (
     load_act_block_dynamic,
     load_act_mx,
     weight_tile_ptrs,
 )
 from .epilogue import acc_init, bias_strides, gemm_epilogue
+from .pruners import PATH_ANCHOR_AXES, fp8_dot_warp_pruner, raw_activation_pointer_pruner, global_scale_warp_spec_pruner, warp_spec_memory_mode_pruner, packed_schedule_scope_pruner, affine_scale_warp_spec_pruner, block_dynamic_grouped_matmul_pruner, block_fits_dim_pruner, block_within_dim_pruner, compose_pruners, descriptor_box_pruner, gate_stacked_tmem_trap_pruner, gated_pointer_weight_warp_spec_pruner, mx_config_pruner, require_moe_dims_aligned, smem_pruner, swizzled_out_bm_pruner, swizzled_scale_config_pruner, swizzled_scales_bm_pruner, warp_spec_compile_guard_pruner
 @bayesian_autotune(
             fp8_dot_warp_pruner(),
             packed_schedule_scope_pruner(),
             block_dynamic_grouped_matmul_pruner(),
+            warp_spec_memory_mode_pruner(),
             descriptor_box_pruner(),
         )
     },
             packed_schedule_scope_pruner(),
             block_dynamic_grouped_matmul_pruner(),
             descriptor_box_pruner(),
+            raw_activation_pointer_pruner(),
         )
     },
 )
 def w8a8_block_static_fp8_matmul_grouped_kernel(
     A,  # (num_tokens, K) E4M3 activations (pre-quantized against the static scale by the wrapper)
     ADescriptor,  # host TMA descriptor over A (rows, K), box (BM, BK); read iff A_MEMORY_MODE != "pointer"
+    As,  # calibrated (static) activation scale: one value, or one per expert
     B,  # (num_experts, N, K) FP8 weights; under GATE the (num_experts, 2N, K) gate|up stack
     BDescriptor,  # host TMA descriptor over B viewed (E, 2N|N, K), box (1, (2|1)*BN, BK); read iff B_MEMORY_MODE != "pointer"
     Bs,  # (num_experts, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales (2N under GATE)
     stride_b_e,
     stride_b_k,
     stride_b_n,
+    stride_as_e,  # 0 = one calibrated scale shared by every expert
     stride_bs_e,
     stride_bs_k,
     stride_bs_n,
     weight scales apply per-K-tile (plain ``tl.dot`` + software rescale, ``accumulate("static")``),
     and the scalar activation scale multiplies the accumulator once after the loop. GATE=False is
     the plain grouped GEMM (down projection), bit-identical."""
     start_pid = tl.program_id(axis=0)
     if PACKED_SCHEDULE:
         total_m_tiles = load_packed_schedule(Schedule, BLOCK_SIZE_M)
         gate_s_ptr = Bs + expert_id64 * stride_bs_e + ((2 if GATE else 1) * pid_n) * stride_bs_n
         up_s_ptr = gate_s_ptr + stride_bs_n
+        # this tile's expert scale: stride 0 means one calibrated scale for every expert.
+        # It feeds the inline quant arm (raw A) and folds back onto the accumulator post-loop.
+        a_s_static = tl.load(As + expert_id64.to(tl.int32) * stride_as_e)
         acc = acc_init("dot", BLOCK_SIZE_M, (2 if GATE else 1) * BLOCK_SIZE_N, False)
         for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC):
             a, a_dead = load_act_static(
+                a_ptrs, ADescriptor, m_start, k * BLOCK_SIZE_K, row_mask, in_row, a_s_static,
                 A_MEMORY_MODE, GatherIdx is not None,
             )
             w, w_s = load_weight_static(
             warp_spec_compile_guard_pruner(),
             descriptor_box_pruner(),
             smem_pruner(),
+            raw_activation_pointer_pruner(),
         )
     },
 )
 def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel(
     A,  # (num_tokens, K) pre-quantized FP8 activations, any row order
     ADescriptor,  # host TMA descriptor over A (rows, K), box (BM, BK); read iff A_MEMORY_MODE != "pointer"
+    As,  # per-token activation scales (S,), or a calibrated (static) one: shared, or per expert
     B,  # (num_experts, N, K) FP8 weights; under GATE the (num_experts, 2N, K) gate|up stack
     BDescriptor,  # host TMA descriptor over B viewed (E, 2N|N, K), box (1, (2|1)*BN, BK); read iff B_MEMORY_MODE != "pointer"
     Bs,  # (num_experts, 1, 1) per-tensor weight scales (one scalar covers the gate|up stack)
     stride_a_m,
     stride_a_k,
     stride_as_m,
+    stride_as_e,  # 0 = the scale is per token; 1 = one calibrated scale per expert
     stride_b_e,
     stride_b_k,
     stride_b_n,
             stride_b_k,
             False,
         )
+        # per token, or the tile's own expert under a calibrated scale (stride_as_m 0); 1.0 on
+        # the masked rows so the static arm's in-register divide never sees a zero
+        a_s = tl.load(
+            As + in_row * stride_as_m + expert_id64.to(tl.int32) * stride_as_e, mask=row_mask, other=1.0
+        )
         b_s = tl.load(Bs + expert_id64 * stride_bs_e)
         acc = acc_init("dot", BLOCK_SIZE_M, (2 if GATE else 1) * BLOCK_SIZE_N, False)
         for k in tl.range(0, tl.cdiv(K, BLOCK_SIZE_K), warp_specialize=WARP_SPEC):
+            a, _as = load_act_static(
+                a_ptrs, ADescriptor, m_start, k * BLOCK_SIZE_K, row_mask, in_row, a_s[:, None],
                 A_MEMORY_MODE, GatherIdx is not None,
             )
             w, _ws = load_weight_plain(
             smem_pruner(min_sets=2),  # two full buffer sets: the grouped pipeline's floor (see smem_pruner)
             warp_spec_compile_guard_pruner(),
             affine_scale_warp_spec_pruner(),
+            warp_spec_memory_mode_pruner(),
         )
     },
 )
     path_anchor_axes=PATH_ANCHOR_AXES,
     prune_configs_by={
         "early_config_prune": compose_pruners(
+            warp_spec_memory_mode_pruner(weight_only=True),
             packed_schedule_scope_pruner(),
             # the tail tile slides back in bounds (see the K-loop) — BK only has to fit
             block_fits_dim_pruner("K"),
     A:  (S, 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,) for one calibrated scale, or (num_experts,) for a MoE that
+        calibrates each expert separately
     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)
+    A_q, As, as_stride = static_expert_act_operands(A, As, num_experts)
     bs_u8 = ue8m0_as_uint8(Bs)
     if requant:
         C = A.new_empty(S, N, dtype=FP8_DTYPE)
         # UE8M0 model (ue8m0 weights) -> UE8M0 intermediate scales; the epilogue infers the format
             B.stride(0),
             B.stride(2),
             B.stride(1),
+            as_stride,
             bs_u8.stride(0),
             bs_u8.stride(2),
             bs_u8.stride(1),
     A:  (S, K) pre-quantized FP8 activations — rows addressed via ``gather_idx``
     B:  (num_experts, N, K) FP8 expert weights; under ``gate`` the (num_experts, 2N, K) stack
+    As: (S,) per-token scales alongside a pre-quantized A, or — on a raw A — the calibrated
+        (static) activation scale: one value, or one per expert
     Bs: (num_experts,) or (num_experts, 1, 1) per-expert weight scales
     expert_start: (num_experts_pow2 + 1,) int32 — cumulative sorted-row starts, S sentinel
     gather_idx: optional (S,) — sorted position -> source row of A; None = A is expert-sorted
     # Normalize Bs to (num_experts, 1, 1) — one per-tensor scale (covers the gate|up stack)
     Bs = normalize_per_expert_scale(Bs, num_experts)
     output_dtype = resolve_output_dtype(output_dtype, A, As)
+    A, As, as_stride_m, as_stride_e = tensor_wide_act_operands(A, As, num_experts)
+    if as_stride_e and expand_regime(S, num_experts):
+        # a per-expert calibrated scale: lay the routed rows out quantized rather than read
+        # them raw through every N-tile — the same copy-amortizes-at-prefill law below
+        A = quantize_routed_rows_per_expert(A, As, gather_idx, expert_start, num_experts)
+        gather_idx = None  # spent: the rows are laid out now
+    if A.dtype == FP8_DTYPE:
+        # post-quant: trade the in-kernel gather for one packed-row copy where that wins. Keyed
+        # on A being quantized, which the dynamic quant, a shared calibrated scale and the
+        # per-expert layout above all reach; raw rows (decode) keep the gather.
         A, _, gather_idx = expand_gather_below_parity(A, None, gather_idx, num_experts)
     C = A.new_empty(S, N, dtype=output_dtype)
     num_sms = sm_count(A.device.index)
             K,
             A.stride(0),
             A.stride(1),
+            as_stride_m,
+            as_stride_e,
             B.stride(0),
             B.stride(2),
             B.stride(1),
         "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)"
     )
+    # a calibrated (static) scale on a raw A (see `tensor_wide_act_operands`): block-scale weights
+    # have a dedicated static kernel, per-tensor ones read the same As on the tensor-wide arm
+    static_act = As is not None and As.ndim <= 1 and As.numel() in (1, B.shape[0]) and A.dtype != FP8_DTYPE
+    if static_act:
+        assert Bs is not None and not is_mx(B, Bs), (
+            "a calibrated (static) activation scale is an FP8 form — MX activations carry a scale "
+            "per group, derived per call"
         )
+    if static_act and weight_block_size(B, Bs) is not None:
         out = w8a8_block_static_fp8_matmul_grouped(
             A,
             B,

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

@@ -268,15 +268,18 @@ 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

     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, through the same tile load in every memory mode: a pre-quantized
+    fp8 A is the rows-major MMA lhs as-is, raw bf16/fp16 is quantized in register against
+    ``a_s_static``. The raw arm is what a PER-EXPERT scale needs: one row can be routed to
+    several experts whose calibrated scales differ, so it has no single pre-quantized form. Its
+    own load would have to repeat the tile loader's masking and descriptor arms — a bare
+    pointer load reads past the tensor on a grouped tail tile. ``a_s`` = the values (the scale
+    is folded post-loop)."""
+    a = load_grouped_act_tile(
+        a_ptrs, a_descriptor, m_off, k_off, value_mask, gather_rows, A_MEMORY_MODE, A_GATHER
+    )
+    if a_ptrs.dtype.element_ty != tl.float8e4nv:  # raw rows: quantize against the tile's scale
+        a = (a.to(tl.float32) / a_s_static).to(tl.float8e4nv)
     a_s = a
     return a, a_s

build/torch-rocm/matmul.py CHANGED Viewed

@@ -25,7 +25,7 @@ from .compat import add_op_namespace_prefix, FP8_DTYPE, is_sm10x, NIBBLES_PER_BY
 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 (
@@ -299,7 +299,7 @@ def w8a8_block_dynamic_fp8_matmul_kernel(
 @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
@@ -1252,6 +1252,7 @@ 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",
@@ -1263,15 +1264,27 @@ def w8a8_tensor_dynamic_fp8_matmul(
     """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
@@ -1279,9 +1292,8 @@ def w8a8_tensor_dynamic_fp8_matmul(
     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)
@@ -1307,7 +1319,7 @@ def w8a8_tensor_dynamic_fp8_matmul(
             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),
@@ -1374,15 +1386,8 @@ def mx_dynamic_matmul(
     # 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),
@@ -1728,7 +1733,7 @@ def mx_weight_only_matmul_2d(
     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):
@@ -1738,6 +1743,14 @@ def mx_weight_only_matmul_2d(
 # 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,
@@ -1748,7 +1761,7 @@ def mx_weight_only_matmul_2d(
 _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,
@@ -1764,6 +1777,13 @@ def _scaled_mm_2d_op(
     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(),
@@ -1797,6 +1817,38 @@ def _smm_reject(reason):
     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
@@ -1950,10 +2002,9 @@ def matmul_2d(
         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:

 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, maybe_act_quant, mxfp8_act_quant, nvfp4_act_quant, quantize_rows_static, tensor_wide_act_operands
 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 (
 @triton.jit
 def w8a8_tensor_dynamic_fp8_matmul_kernel(
     A,  # (M, K) pre-quantized FP8 activations
+    As,  # (M,) per-token activation scales, or one calibrated (static) scale (stride_as_m 0)
     B,  # (N, K) FP8 weights
     Bs,  # scalar/(1,) per-tensor weight scale
     C,  # (M, N) output
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
+    As: torch.Tensor | None = None,
     output_dtype: torch.dtype | None = None,
     gate: bool = False,
     act_fn: str = "silu",
     """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)``, or
+        quantized against ``As`` when a calibrated (static) one is given.
     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.
+    As: the calibrated (static) activation scale, one value for the whole matmul; ``None``
+        derives one per row from the data.
     ``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)
+    # cuBLAS per-tensor fast path, INSIDE the op for the same reasons as the MX one below:
+    # every caller gets it, and a None falls through to the Triton launch unchanged.
+    if not gate and not simulate_unfused and bias is None:
+        _smm = _torch_scaled_mm_2d_tensor_wise(
+            A, B, As, Bs, None, resolve_output_dtype(output_dtype, A, None)
+        )
+        if _smm is not None:
+            return [_smm]
     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
     assert Bs.numel() == 1, f"Bs must be scalar or (1,), got {tuple(Bs.shape)}"
+    # one scale per row (derived here), or the calibrated one every row reads (stride 0)
+    qA, As, as_stride_m, _ = tensor_wide_act_operands(A, As)
     Bs = Bs.reshape(1)
     C = A.new_empty(A.shape[:-1] + (N,), dtype=output_dtype)
             int(M).bit_length(),  # m_bit_length key bucket
             qA.stride(-2),
             qA.stride(-1),
+            as_stride_m,
             B.stride(1),
             B.stride(0),
             C.stride(-2),
     # 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.
+    if not gate and not quantize_output and bias is None and not simulate_unfused:
         _smm = _torch_scaled_mm_2d(
             A, B, As, Bs, activation_format, a_global_scale, b_global_scale,
             resolve_output_dtype(output_dtype, A, None),
     return [C]
+# ── torch scaled_mm fast path (2D dense, 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):
 # 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.
 #
+# Per-tensor static FP8 wins at EVERY M, decode included, so that arm has no floor (B200,
+# N=12288 K=4096, `quantize_rows_static` feeding both arms, cudagraph):
+#     M=1 12.2us vs 14.0   M=8 9.8 vs 13.0   M=64 8.9 vs 12.0   M=512 23.2 vs 34.1
+#     M=8192 320.5us vs 413.0
+# GEMM alone that is 2807 TFLOP/s against our 1993, and cuBLAS lands within 1.4% of its own
+# big-square best case (2846) — the 2-CTA + TMA warp-specialized architecture we cannot reach
+# from Triton, not a tuning gap. Relerr 1.2e-5 against the Triton arm.
+#
 # 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,
 _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"), mutates_args=())
 def _scaled_mm_2d_op(
     Aq: torch.Tensor,
     As: torch.Tensor,
     cuBLAS fast path too, not just eager ones."""
     F = torch.nn.functional
     ST, SW = F.ScalingType, F.SwizzleType
+    if fmt == "fp8":
+        return F.scaled_mm(
+            Aq, B.t(),
+            As.reshape(()), ST.TensorWise,
+            Bs.reshape(()), ST.TensorWise,
+            output_dtype=torch.bfloat16,
+        )
     if fmt == "mxfp8":
         return F.scaled_mm(
             Aq, B.t(),
     return None
+def _torch_scaled_mm_2d_tensor_wise(A, B, As, Bs, activation_format, out_dtype):
+    """The scaled_mm route for per-tensor static FP8, or ``None`` to take the Triton ops.
+    One fp32 value per operand IS scaled_mm's ``TensorWise``, so this arm needs neither swizzled
+    scales nor a block-aligned K — only a STATIC activation scale, since a dynamic one would cost
+    the amax pass the Triton kernel fuses into its own loop.
+    """
+    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")
+    if activation_format not in (None, "fp8"):
+        return _smm_reject("activation_format")
+    if out_dtype is not torch.bfloat16:
+        return _smm_reject("out-dtype")
+    if As is None or As.numel() != 1:
+        return _smm_reject("not-static-per-tensor")
+    if As.dtype is not torch.float32 or Bs.dtype is not torch.float32:
+        # a UE8M0 container's scalar is an EXPONENT, which TensorWise would read as a multiplier
+        return _smm_reject("non-fp32 scale")
+    K = A.shape[-1]
+    if K % 16 or B.shape[0] % 16:
+        return _smm_reject("shape")
+    M = A.numel() // K
+    As = As.reshape(-1).float()
+    Aq = A.reshape(M, K) if A.dtype == FP8_DTYPE else quantize_rows_static(A.reshape(M, K), As)
+    out = _scaled_mm_2d_op(Aq, As, B, Bs.reshape(-1).float(), None, None, "fp8")
+    return out.reshape(*A.shape[:-1], B.shape[0])
 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
         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
         return _unwrap(
             w8a8_tensor_dynamic_fp8_matmul(
+                A, B, Bs, As, 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:

build/torch-rocm/metadata.json CHANGED Viewed

@@ -1,6 +1,6 @@
 {
   "name": "finegrained-kernels",
-  "id": "_finegrained_kernels_rocm_1348465",
   "version": 0,
   "license": "Apache-2.0",
   "python-depends": [],
@@ -11,25 +11,25 @@
     "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="
     }
@@ -41,7 +41,7 @@
       "dirty": false
     },
     "kernel": {
-      "sha": "1348465af00f2d4c596d37bdb262ae0875dc1d77",
       "dirty": false
     }
   }

 {
   "name": "finegrained-kernels",
+  "id": "_finegrained_kernels_rocm_b0c6080",
   "version": 0,
   "license": "Apache-2.0",
   "python-depends": [],
     "algorithm": "sha256",
     "files": {
       "__init__.py": "43nBPHxybam5ChWZmsz1Ji+pjJ6K6zF5lzwbqPcMwE4=",
+      "_ops.py": "EL67XyYKmJCWdvLz7dUF+v4b0yedbiW265F3nOEoZXw=",
       "backward.py": "ri9lKPhb9SmkYBvUS8KvuORiW5PcKViNamFNVWxb69g=",
+      "batched.py": "683kgUc1fTmvPZPwNha4+VRQMEdZWAULFtDH5/rs8pg=",
+      "bayesian_autotuner.py": "GoXX4MLMduf7L55t66joOgfMjRIt4NCQt3aGdon9qJU=",
+      "compat.py": "QUGxtD0TogkxhHX8nNNFtwOyjd8ECN72CUGefiwkL/M=",
       "descriptors.py": "4acRMazGPFlbydNt1Hbo8RpAqb3fOtnOQBv5yb4onqw=",
       "epilogue.py": "tGALTI/VMgO3Mz0sQ6H6wuRZpO4UKqM1/4h7d81SnVM=",
       "finegrained_kernels/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
       "formats.py": "V6F0PXdNhXQn/UBz0/frbvcp5Lgterh053XDrvQ/Cpw=",
+      "grouped.py": "14Zfza8MI033MXY2JFXVDfaNVbKHPqiYletmBazGwMo=",
       "loading/__init__.py": "so7kYmlVoiASfoOzemdU5fBTiP8yA85gVPNQExIq/KU=",
       "loading/scales.py": "8HJpflmImEXYJquA6hVLStUUEecvExAxBEYVb3qPPWE=",
+      "loading/tiles.py": "tcStORrCdZbFeoG4AXIE3YznJXGt8rIIBHnyNumZMCE=",
+      "matmul.py": "Wfmnyoez0xwplElVEd4zpzc0KQ75rr3m09eww6aOFBE=",
       "mma.py": "1fg2an7Dc8Vd6qpSmq+Q0781A6sNilOJHU0PqEbouT8=",
+      "moe.py": "Mj+n8CNqL0mOQ5xmqpopisUWGwKoayanfaAymXiGhRk=",
       "norm.py": "VvE6GUE0Oztf7PDcviLCcB4fhnXt7/zgdE+aApv6DBM=",
+      "pruners.py": "/uWNcx3KqQwJGG1FAeJOSY7Jha9PwBsczVfKYQb5GEI=",
+      "quant.py": "Wyw3yCpgybFHswP5UEsl13BQ1Hqgm61Z4/XL2TGhvfk=",
       "scheduling.py": "UGgQ5uVdm4cWPnH+4w6Ffou/PSX3iUarMeRo2E04BhY=",
       "swizzle.py": "ueIJ9FsO8QZhVt3nMFlA+8b+U9/ZHabHFlhufsONV04="
     }
       "dirty": false
     },
     "kernel": {
+      "sha": "b0c60801e54949c92345c0a86a32f36a1fd2885f",
       "dirty": false
     }
   }

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

@@ -1 +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="}}

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json", "verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUZ81snzjoHnJyxgJ4p7A9yWtxEb4wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwOTIxMTYyOTUxWhcNMjYwOTIxMTYzOTUxWjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE7W5TLpJgVdBM2F/NFhNhCgI3KCd12XeXtQ2rv9vspmSCzZ4GQL4edPNf8mvnZpN2AZvlzrluj3wmXKXB/bpyLKOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQU+ywNib8rPeO1AnFU7Bs2vkwfuoAwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoYjBjNjA4MDFlNTQ5NDljOTIzNDVjMGE4NmEzMmYzNmExZmQyODg1ZjATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoYjBjNjA4MDFlNTQ5NDljOTIzNDVjMGE4NmEzMmYzNmExZmQyODg1ZjAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoYjBjNjA4MDFlNTQ5NDljOTIzNDVjMGE4NmEzMmYzNmExZmQyODg1ZjAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKGIwYzYwODAxZTU0OTQ5YzkyMzQ1YzBhODZhMzJmMzZhMWZkMjg4NWYwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMzU2MjUzMjM5MTcvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABoMTNaZ4AAAQDAEcwRQIgP5gxSfLy4FLP9vZ4T4C2HBN01Vp/ZNNmXDiNJ3WZD6ACIQCMqPG8YUNpvBoo+si9YGfqD0VwIvi7PJSA1P2G8UfREDAKBggqhkjOPQQDAwNoADBlAjEA1D95XvsUGbnqjYBmuXYNdygeRfJhWSWI3AhvR46OC/CYXaqPHkFguedK8W5xUYhIAjAHWbzokgYBbsMSWbI621lxdeUfGGh4BKLr4z6Ka17axaCjQqK1dha5RQusUp4bzL4="}, "tlogEntries":[{"logIndex":"2906397162", "logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="}, "kindVersion":{"kind":"hashedrekord", "version":"0.0.1"}, "integratedTime":"1790008191", "inclusionPromise":{"signedEntryTimestamp":"MEQCICAibnJuErv4oTr2p9bLcj0cUuI4sYBjdKMeUNBmjEdMAiB956LiLQAgS3owSEq8A2dIW5RobiBhPv3UJjU94ulNWg=="}, "inclusionProof":{"logIndex":"2784492900", "rootHash":"ruuRuA2bOwvUCTUwjql5lLPMMxBRbVBujq7C52EY+xI=", "treeSize":"2784492912", "hashes":["W4aHnZgxOxaEX8nN0MvEbfXKv21ls3qgPIO9h5ANlQw=", "YPmXuUxpSz1W8ielAHHAKbM60+hu3794xNkvr1tqwd0=", "d2ml3/jgADqEi8yzSiS9BGj6oH8Ud81IHkRiyTHT9A0=", "dQijL74kFpEfFyGTF9IGOMoU04sZsc791MgeiLDR/ng=", "a2ewThrYCNlBctYvOMWwas2uP6zmyXk9g6mMYpGjqRA=", "vgq7Vwl4M42Xa6zosYvUzr3OX6U14I8aMd2kpNxaRAs=", "KGl94H7+f+B/yxu/2qgaWzTvKr9tDUarFyMOCuFM25c=", "n5rW0Uy02cMAl76x6jWiwI7hKyUcCC7Fr0zU6EJwOdA=", "a9uxtQgboWs7ecf47zfjx0t0FQE/PqTjpq+xX55YJSo=", "uMrI4K593Uc4YIUn568kQQU3P5WI4DcAn+rS3jG0Loo=", "NxuiinDjnbF2IyjYZfiDnZUwafSidtG8YO4Q8hE7VLk=", "0NTU/2u7wVmA0DVpAJ84n4sVvSCho8WHC+DkEdg92JU=", "zkeHshHFXmbLQPJRTS95jeMY35dwTUsGQe3ERLPQ7RE=", "XU+M8SHgYBPE6zkZMa/88OBoHCVM5MhtjfLPXnLvW7g=", "jsRQr2n8GzEgko08H6qJKNeX3a7bkPYvDuB55qbMm7g=", "4N5aDWlTGjrYD2P9BkugCVYpeek+qqUf//YGGQZkIsE=", "zA+3WHfvsHMHgoq3paLzdllP+3iYxHkVPro8CUBuZa4=", "GfSJjUNAa3XTMUgDFMrv4pTzWdeXIUdphT8VAloU9uM=", "/2a58lzrUQfoAuFNUx07Dmjo61/cNyxjKuxmZfV/uSU=", "tyO4Oj1UrXex15tzrEEO6sSnE8y8obZWghZ2XixBPtE=", "rjDiNJHGnboMYYUH7OSezbKWj3wMXRTHGCLa/SaZ6O4=", "mfWjSdNaEfNXxDww6AAOw/NoK9e6z1yKaZm5xvmMabs=", "qxzHanAzz57SDdmJe0B7bJK72NTIBbwEMGBBKvDOROw=", "xH/DCseLHr9eKoYT8qsORZK7zVdEGYWHuVtsVrD95wY="], "checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n2784492912\nruuRuA2bOwvUCTUwjql5lLPMMxBRbVBujq7C52EY+xI=\n\n— rekor.sigstore.dev wNI9ajBFAiAEwubxDVjUNcv+3VF0UEaDoWvSKfocgRHNTlIm7SbBIAIhAK5uuPxv1Jv8MIC9e5jdi/yn2wE+dcljjKFJA2yCOvEc\n"}}, "canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI2YmM4NDE1NDE2MTYyN2MyMTJlYjg5YmM5MTMzMTI5NDVjYmRkMDQxMDAyYWFhY2I5YWZkNGMxMmRkODFlNzRhIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJQ0VJOE0vdlVKZEJVTXNUbTVoU2p1TGEvTVNJNlRkdGpZdkp2c0pWYjVvSEFpRUFnaTFGdGNnRDFpQm5oQjVtWEI2Z1g1VkVvK2hqbnpCam8xQW5iY3k1MWtnPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVldqZ3hjMjU2YW05SWJrcDVlR2RLTkhBM1FUbDVWM1I0UldJMGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDlVU1hoTlZGbDVUMVJWZUZkb1kwNU5hbGwzVDFSSmVFMVVXWHBQVkZWNFYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVUzVnpWVVRIQktaMVprUWsweVJpOU9SbWhPYUVOblNUTkxRMlF4TWxobFdIUlJNbklLZGpsMmMzQnRVME42V2pSSFVVdzBaV1JRVG1ZNGJYWnVXbkJPTWtGYWRteDZjbXgxYWpOM2JWaExXRUl2WW5CNVRFdFBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlVyZVhkT0NtbGlPSEpRWlU4eFFXNUdWVGRDY3pKMmEzZG1kVzlCZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVpha0pxVG1wQk5FMUVSbXhPVkZFMVRrUnNhazlVU1hwT1JGWnFDazFIUlRST2JVVjZUVzFaZWs1dFJYaGFiVkY1VDBSbk1WcHFRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMWxxUW1wT2FrRTBUVVJHYkU1VVVUVk9SR3hxVDFSSmVrNUVWbXBOUjBVMFRtMUZlazF0V1hvS1RtMUZlRnB0VVhsUFJHY3hXbXBCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMWxxUW1wT2FrRTBUVVJHYkU1VVVUVk9SR3hxVDFSSmVrNUVWbW9LVFVkRk5FNXRSWHBOYlZsNlRtMUZlRnB0VVhsUFJHY3hXbXBCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwZEpkMWw2V1hjS1QwUkJlRnBVVlRCUFZGRTFXWHByZVUxNlVURlpla0pvVDBSYWFFMTZTbTFOZWxwb1RWZGFhMDFxWnpST1YxbDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOZWxVeVRXcFZlazFxVFRWTlZHTjJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtOU5WRTVoV2pSQlFVRlJSQXBCUldOM1VsRkpaMUExWjNoVFpreDVORVpNVURsMldqUlVORU15U0VKT01ERldjQzlhVGs1dFdFUnBUa296VjFwRU5rRkRTVkZEVFhGUVJ6aFpWVTV3Q25aQ2IyOHJjMms1V1VkbWNVUXdWbmRKZG1rM1VFcFRRVEZRTWtjNFZXWlNSVVJCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEZRVEZFT1RVS1dIWnpWVWRpYm5GcVdVSnRkVmhaVG1SNVoyVlNaa3BvVjFOWFNUTkJhSFpTTkRaUFF5OURXVmhoY1ZCSWEwWm5kV1ZrU3poWE5YaFZXV2hKUVdwQlNBcFhZbnB2YTJkWlFtSnpUVk5YWWtrMk1qRnNlR1JsVldaSFIyZzBRa3RNY2pSNk5rdGhNVGRoZUdGRGFsRnhTekZrYUdFMVVsRjFjMVZ3TkdKNlREUTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}], "timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyjADAgEAMIICwQYJKoZIhvcNAQcCoIICsjCCAq4CAQMxDTALBglghkgBZQMEAgEwgbgGCyqGSIb3DQEJEAEEoIGoBIGlMIGiAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgDM0ZH/BFkL7abcswmvCpzKpHqo/dDY4oQpTE2rSegcYCFQCEQUOszsswNLq5HwXLmDQPi0uCoBgPMjAyNjA5MjExNjI5NTFaMAMCAQGgMqQwMC4xFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEVMBMGA1UEAxMMc2lnc3RvcmUtdHNhoAAxggHbMIIB1wIBATBRMDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAsGCWCGSAFlAwQCAaCB/DAaBgkqhkiG9w0BCQMxDQYLKoZIhvcNAQkQAQQwHAYJKoZIhvcNAQkFMQ8XDTI2MDkyMTE2Mjk1MVowLwYJKoZIhvcNAQkEMSIEINZVjoB6Ftn+joLxXvUK/3YRG/5j5NvBF1Bn3jv91UpoMIGOBgsqhkiG9w0BCRACLzF/MH0wezB5BCCF+Se8B6tiysO0Q1bBDvyBssaIP9p6uebYcNnROs0FtzBVMD2kOzA5MRUwEwYDVQQKEwxzaWdzdG9yZS5kZXYxIDAeBgNVBAMTF3NpZ3N0b3JlLXRzYS1zZWxmc2lnbmVkAhQ6E1QvDJBh7rzBQy/Lio6LKiOLDDAKBggqhkjOPQQDAgRnMGUCMHuLT4V9Ga1MtwJ9bd3spYCretPRwz5W9+Qb9Vn2yXVXqo5BcKaoV9k9a6ad7UZwogIxAI8j8+evxPw2yj9+byUuefD5+AGCAppHEIPdzwM0Gn04tiTfiWKL4FeqQXga/7aB4w=="}]}}, "messageSignature":{"messageDigest":{"algorithm":"SHA2_256", "digest":"a8hBVBYWJ8IS64m8kTMSlFy90EEAKqrLmv1MEt2B50o="}, "signature":"MEUCICEI8M/vUJdBUMsTm5hSjuLa/MSI6TdtjYvJvsJVb5oHAiEAgi1FtcgD1iBnhB5mXB6gX5VEo+hjnzBjo1Anbcy51kg="}}

build/torch-rocm/moe.py CHANGED Viewed

@@ -47,7 +47,15 @@ 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
@@ -320,6 +328,8 @@ def moe_fused_grouped(
     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
@@ -363,9 +373,11 @@ def moe_fused_grouped(
     # 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,
@@ -376,9 +388,9 @@ def moe_fused_grouped(
         **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,
     )
@@ -392,7 +404,7 @@ def moe_fused_grouped(
     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,
@@ -429,6 +441,8 @@ def moe_fused_batched(
     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
@@ -472,9 +486,11 @@ def moe_fused_batched(
     # 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,
@@ -488,10 +504,14 @@ def moe_fused_batched(
         # (``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,
@@ -506,7 +526,7 @@ def moe_fused_batched(
     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,
@@ -545,6 +565,8 @@ def moe_unfused_grouped(
     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
@@ -562,8 +584,10 @@ def moe_unfused_grouped(
     ``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)
@@ -582,6 +606,7 @@ def moe_unfused_grouped(
     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,
@@ -597,6 +622,7 @@ def moe_unfused_grouped(
     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,
@@ -624,6 +650,8 @@ def moe_torch_grouped(
     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
@@ -653,6 +681,10 @@ def moe_torch_grouped(
     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)"
     )
@@ -664,9 +696,10 @@ def moe_torch_grouped(
         "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).
@@ -803,6 +836,8 @@ def moe_unfused_batched(
     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
@@ -819,8 +854,10 @@ def moe_unfused_batched(
     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)
@@ -836,6 +873,7 @@ def moe_unfused_batched(
     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,
@@ -850,6 +888,7 @@ def moe_unfused_batched(
     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,

 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,
+    ScalingType,
+    SwizzleType,
+    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
     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,
+    gate_up_proj_activation_scale: torch.Tensor | None = None,
+    down_proj_activation_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
     # 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.
+    static_act = gate_up_proj_activation_scale is not None or down_proj_activation_scale is not None
     gate_up_out = matmul_grouped(
         hidden_states,
         gate_up_proj,
+        As=gate_up_proj_activation_scale,
         Bs=gate_up_proj_scale_inv,
         a_global_scale=gate_up_proj_input_global_scale,
         b_global_scale=gate_up_proj_weight_global_scale,
         **glu,
         bias=gate_up_proj_bias,
         # fmt is the resolved format; "bf16" (weight-only) leaves the GLU intermediate bf16,
+        # no requant. So does static, whose epilogue has no calibrated scale to requant against.
         activation_format=fmt,
+        quantize_output=bool(glu) and fmt != "bf16" and not static_act,
         output_dtype=hidden_states.dtype,
         gather_idx=gather_idx,
     )
     down_out = matmul_grouped(
         inter,
         down_proj,
+        As=down_proj_activation_scale if static_act else inter_scale,
         Bs=down_proj_scale_inv,
         a_global_scale=down_proj_input_global_scale,
         b_global_scale=down_proj_weight_global_scale,
     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,
+    gate_up_proj_activation_scale: torch.Tensor | None = None,
+    down_proj_activation_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
     # 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
+    static_act = gate_up_proj_activation_scale is not None or down_proj_activation_scale is not None
     gate_up_out = matmul_batched(
         hidden_states,
         gate_up_proj,
+        As=gate_up_proj_activation_scale,
         Bs=gate_up_proj_scale_inv,
         a_global_scale=gate_up_proj_input_global_scale,
         b_global_scale=gate_up_proj_weight_global_scale,
         # (``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. Static keeps bf16 either way, its epilogue
+        # having no calibrated scale to requant against.
         activation_format=fmt,
         quantize_output=(
+            bool(glu)
+            and fmt != "bf16"
+            and (fmt != "fp8" or expert_ids.numel() <= GATE_UNSTACK_MAX_S)
+            and not static_act
         ),
         output_dtype=hidden_states.dtype,
         gather_idx=gather_idx,
     down_out = matmul_batched(
         inter,
         down_proj,
+        As=down_proj_activation_scale if static_act else inter_scale,
         Bs=down_proj_scale_inv,
         a_global_scale=down_proj_input_global_scale,
         b_global_scale=down_proj_weight_global_scale,
     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,
+    gate_up_proj_activation_scale: torch.Tensor | None = None,
+    down_proj_activation_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
     ``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``, and a ``*_activation_scale`` quantizes it
+    against that calibrated scale instead of a runtime one. ``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)
     gate_up_out = matmul_grouped(
         hidden_states,
         gate_up_proj,
+        As=gate_up_proj_activation_scale,
         Bs=gate_up_proj_scale_inv,
         a_global_scale=gate_up_proj_input_global_scale,
         b_global_scale=gate_up_proj_weight_global_scale,
     down_out = matmul_grouped(
         inter,
         down_proj,
+        As=down_proj_activation_scale,
         Bs=down_proj_scale_inv,
         a_global_scale=down_proj_input_global_scale,
         b_global_scale=down_proj_weight_global_scale,
     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,
+    gate_up_proj_activation_scale: torch.Tensor | None = None,
+    down_proj_activation_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
     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 ScalingType is not None, (
+        "this torch has no torch.nn.functional.ScalingType — the baseline's scaled_grouped_mm "
+        "scaling enums arrived with the op itself, so an older torch cannot run it"
+    )
     assert gate_up_proj.dtype in (torch.int8, torch.float8_e4m3fn), (
         "torch grouped baseline is MX-only (packed E2M1 or E4M3 weights)"
     )
         "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"
     )
+    assert gate_up_proj_activation_scale is None and down_proj_activation_scale is None, (
+        "the torch baseline quantizes each activation against a scale it derives per call — a "
+        "calibrated (static) scale is not representable here"
+    )
     # torchao >= 0.18 required with cutlass-dsl >= 4.6 (0.17 imports a helper path 4.6
     # removed; fixed upstream in pytorch/ao).
     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,
+    gate_up_proj_activation_scale: torch.Tensor | None = None,
+    down_proj_activation_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
     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``, and a ``*_activation_scale`` quantizes it against that calibrated
+    scale instead of a runtime one. ``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)
     gate_up_out = matmul_batched(
         hidden_states,
         gate_up_proj,
+        As=gate_up_proj_activation_scale,
         Bs=gate_up_proj_scale_inv,
         a_global_scale=gate_up_proj_input_global_scale,
         b_global_scale=gate_up_proj_weight_global_scale,
     down_out = matmul_batched(
         inter,
         down_proj,
+        As=down_proj_activation_scale,
         Bs=down_proj_scale_inv,
         a_global_scale=down_proj_input_global_scale,
         b_global_scale=down_proj_weight_global_scale,

build/torch-rocm/pruners.py CHANGED Viewed

@@ -15,7 +15,7 @@
 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 ────────────────────────────────────────────────────────────
@@ -101,6 +101,10 @@ from .mma import MMA_N_ATOM_WIDTH
 # 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
@@ -334,6 +338,10 @@ def mx_config_pruner(k_arg: str, n_arg: str | None = None, block_within_k: bool
     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
@@ -474,6 +482,18 @@ def mx_config_pruner(k_arg: str, n_arg: str | None = None, block_within_k: bool
             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
@@ -495,6 +515,7 @@ def mx_config_pruner(k_arg: str, n_arg: str | None = None, block_within_k: bool
     # 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
@@ -714,49 +735,65 @@ def packed_schedule_scope_pruner(min_bm: int = 128):
     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())
@@ -776,6 +813,37 @@ def weight_only_swap_scope_pruner():
     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
@@ -804,6 +872,14 @@ def mx_2d_swap_scope_pruner(max_m: int = 16):
             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)
@@ -911,9 +987,9 @@ def swizzled_scale_config_pruner(allow_gate_subblock=False):
     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(

 import torch
 import triton
+from .compat import FP8_DTYPE, get_active_device_type, is_sm10x, is_sm90, sm_count, sm_shared_memory_limit
 from .mma import MMA_N_ATOM_WIDTH
 # ── config pruners ────────────────────────────────────────────────────────────
 # it fit shared memory and failed only as benign launch-time smem overflows.
 SM10X_SCALED_MMA_MAX_N = 256
+# Widest N tile the raw-activation (per-expert calibrated) arm survives; see
+# raw_activation_pointer_pruner for the bisect.
+RAW_ACT_MAX_BN = 128
 # 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
     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.
+    A weight-only launch whose raw activation is fp32 drops ``dot_scaled`` outright: the op's
+    lhs must be bf16/fp16, so the arm cannot compile at all there (arch-independent, unlike the
+    shape gates below).
     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
             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 dot_scaled_act_dtype_ok(c, args):
+        # `tl.dot_scaled` takes a bf16/fp16 lhs. A weight-only launch hands it the RAW
+        # activation, so an fp32 one cannot reach the arm at all: every dot_scaled config dies
+        # in "Unexpected dtype for bf16. Got fp32" and the tuner books the whole arm as compile
+        # FAILURES rather than as a declared fence (87 of them on one grouped launch). A model
+        # running in fp32 reaches this through the transformers integration. The dot and scalar
+        # arms upcast the weight and take any float lhs, so they serve the launch — drop rather
+        # than raise. Not arch-gated: the lhs dtype is the op's contract, not an sm_10x quirk.
+        if c.kwargs.get("COMPUTE_MODE") != "dot_scaled" or acts_are_scaled(args):
+            return True
+        return getattr(args.get("A"), "dtype", None) != torch.float32
     def scales_are_e4m3(args):
         return getattr(args.get("Bs"), "dtype", None) == torch.float8_e4m3fn
     # re-admitting the single-trip/width traps it existed to remove).
     return compose_pruners(
         *stages,
+        config_filter(dot_scaled_act_dtype_ok),
         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
     return config_filter(ok)
+# Pointer/pointer warp counts that trap on the WEIGHT-ONLY kernel under Triton < 3.8. Charted
+# 2026-08-26 and REPRODUCED on 3.7.1 by the 2026-09-17 two-version sweep, so these rows are real
+# measurements of that compiler, not the mis-chart the descriptor rows turned out to be. Triton
+# 3.8 fixes every one of them, which is why `warp_spec_memory_mode_pruner` only consults the
+# table below that version.
+_WEIGHT_ONLY_POINTER_TRAPS = {
+    ("dot", "pointer"): {16},
+    ("dot_scaled", "pointer"): {4, 8, 16},
+}
+def _warp_spec_needs_descriptor_tile() -> bool:
+    """Triton >= 3.8, where a TMA-descriptor operand narrows WS to one tile (see the pruner)."""
+    return tuple(int(n) for n in triton.__version__.split(".")[:2]) >= (3, 8)
+def warp_spec_memory_mode_pruner(weight_only: bool = False):
+    """``early_config_prune`` dropping ``warp_specialize`` configs the WS passes cannot lower
+    (``TritonGPUOptimizePartitionWarps`` / ``RelayoutTritonGPU`` -> ``PassManager::run failed``).
+    The rule is COMPILER-DEPENDENT, and the two versions are near mirror images. Charted
+    2026-09-17 on B200 by filtering each kernel's OWN tuner configs (never hand-built — that skips
+    the descriptor pre_hooks and fails every descriptor cell for the wrong reason), BN held at 128,
+    BM and ``num_warps`` both axes, gpt-oss 2048x2048, on three kernels across two families:
+    grouped MX weight-only, grouped MX dynamic, and block-dynamic FP8.
+        operands          triton 3.8.0              triton 3.7.1
+        pointer/pointer   every warp count and BM   mostly FAILS (see the table above)
+        desc/desc         ONLY w4 AND BM == 128     most cells lower; w2 never does
+        mixed             identical to desc/desc    identical to desc/desc
+    Cell for cell across all three kernels, 3.8 rescues NO descriptor cell that 3.7.1 lowered
+    (0 of 96) and 3.7.1 rescues NO pointer cell that 3.8 lowered (0 of 80): descriptor+WS
+    REGRESSED in 3.8 and pointer+WS was FIXED. So each version gets the rule its own sweep
+    measured, rather than an intersection that would fence off a whole memory mode on both.
+    ``weight_only`` adds the rows that kernel charted for itself and the sweep did not re-probe:
+    ``BN == 32`` with WS, which failed 9/9 cells there, and (below 3.8) its trapping
+    pointer/pointer warp counts.
+    NOT covered here, deliberately: ``mx_dynamic``'s ``dot`` arm dies with a descriptor operand on
+    BOTH versions, and it is not a WS failure at all — it raises ``descriptor gather of uint8 must
+    have at least 32 columns, but got 16`` with ``warp_specialize`` on AND off. That is a gather
+    width law and belongs wherever gather width is decided."""
     def ok(c, args):
         if not c.kwargs.get("WARP_SPEC"):
             return True
+        if weight_only and config_dim(c, args, "BLOCK_SIZE_N") == 32:
             return False
         a, b = c.kwargs.get("A_MEMORY_MODE"), c.kwargs.get("B_MEMORY_MODE")
+        descriptor = a != "pointer" or b != "pointer"
+        if _warp_spec_needs_descriptor_tile():
+            # 3.8: one surviving descriptor cell, and pointer/pointer is wide open
+            return not descriptor or (c.num_warps == 4 and config_dim(c, args, "BLOCK_SIZE_M") == 128)
+        if descriptor:
             return True
+        traps = _WEIGHT_ONLY_POINTER_TRAPS if weight_only else {}
+        return c.num_warps not in traps.get((c.kwargs.get("COMPUTE_MODE"), a), set())
     return config_filter(ok, when=lambda args: is_sm10x())
     return config_filter(ok)
+def raw_activation_pointer_pruner():
+    """``early_config_prune`` scoping the RAW-activation arm to the pointer load and away from
+    warp specialization. A calibrated scale held per expert leaves ``A`` unquantized for the
+    kernel to quantize per tile (a gathered row serves several experts, so there is no single
+    pre-quantized form), which puts a quantize INSIDE the K-loop and changes its structure:
+    - the TMA gather cannot serve it at all (``async_tma_gather`` wants 4 contiguous elements
+      per thread and the lowering fails), which at a deployment shape left the tuner no config;
+    - ``BLOCK_SIZE_N`` above 128 traps the device at prefill scale — a sticky misaligned
+      address. Bisected on B200 / Triton 3.8 with a COLD tune cache (a warm one replays a safe
+      crown and hides it): holding BN <= 128 runs clean, while fencing warp specialization, the
+      packed schedule, or BK < 128 each still trap. So the tile width is the variable, not the
+      schedule, the memory mode or WS.
+    A pre-quantized ``A`` keeps the full grid, unchanged — which is what the grouped prefill
+    now hands it (``quantize_routed_rows_per_expert`` lays the routed rows out quantized once
+    the row count pays for the copy), so the raw arm is reached below that regime, or above it
+    when ``FINEGRAINED_FORCE_GATHER`` holds the gather. The BN fence is load-bearing there."""
+    def ok(c, args):
+        a = args.get("A")
+        if getattr(a, "dtype", None) == FP8_DTYPE or args.get("As") is None:
+            return True
+        return (
+            c.kwargs.get("A_MEMORY_MODE", "pointer") == "pointer"
+            and config_dim(c, args, "BLOCK_SIZE_N") <= RAW_ACT_MAX_BN
+        )
+    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
             config_dim(c, args, "BLOCK_SIZE_M") == 1
             and getattr(args.get("Bs"), "dtype", None) == torch.float8_e4m3fn
             and args["M"] <= max_m
+            # An A-side TMA descriptor under SWAP_AB traps the device on Triton 3.8 — a sticky
+            # misaligned address, after which every later launch reports it wherever it came
+            # from. Observed on B200 in a GLM-5.2-NVFP4 forward (16 tokens; clean at 4) at
+            # dot_scaled, BM=1, BN ∈ {128, 256}, BK ∈ {64, 128}: the tile WIDTH is not the
+            # variable (BN=128 traps too), the A descriptor is. Narrowed by probe — fencing the
+            # B side as well was NOT needed, so the weight descriptor stays available to the
+            # crown, and pointer-mode A keeps the decode arm this pruner exists for.
+            and c.kwargs.get("A_MEMORY_MODE", "pointer") == "pointer"
         )
     return config_filter(ok)
     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 "
+            f"tile; N={args.get('N')}, K={args.get('K')}, GATE={bool(args.get('GATE'))}) — say "
+            "which dim is short rather than guessing: pass affine (row-major) scales for it."
         )
     return config_filter(

build/torch-rocm/quant.py CHANGED Viewed

@@ -19,7 +19,7 @@ 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
@@ -927,9 +927,121 @@ def _fp8_act_quant_kernel(
     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]:

 from .bayesian_autotuner import bayesian_autotune
 from .formats import global_scale_stride, is_per_expert_global
+from .compat import FP8_DTYPE, MX_SCALE_GROUP_K, NVFP4_SCALE_GROUP_K, 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
     tl.store(s_ptr + pid, s)
+def static_expert_act_operands(A, As, num_experts):
+    """Activation operands of the per-expert STATIC arms: ``(A_q, As, as_stride)``.
+    One calibrated scale for the whole matmul, or one per expert — a MoE calibrates each
+    separately. ``as_stride`` 0 makes every expert read the one scale, 1 gives each its own.
+    Per expert the kernel quantizes in register against the tile's own scale, because one token
+    routes to several experts whose scales differ and so has no single pre-quantized form; one
+    scale for all of them pre-quantizes once here instead.
+    """
+    As = As.reshape(-1).to(torch.float32)
+    as_stride = int(As.numel() == num_experts)
+    A_q = A if as_stride else (A.to(torch.float32) / As).to(FP8_DTYPE)
+    return A_q, As, as_stride
+# Host-side operand marshalling, NOT a custom op: it returns its inputs unchanged on the
+# pre-quantized and per-expert branches, which `torch.library.custom_op` forbids (the output may
+# not alias an input), and it hands back Python strides. The quant kernels it calls carry their
+# own compile handling.
+def tensor_wide_act_operands(
+    A: torch.Tensor, As: torch.Tensor | None, num_experts: int | None = None
+) -> tuple[torch.Tensor, torch.Tensor, int, int]:
+    """Activation operands of the tensor-scale FP8 arms: ``(A, As, stride_m, stride_e)``.
+    No ``As`` derives one scale per row here; a pre-quantized (E4M3) ``A`` brings its own. An
+    ``As`` on a raw ``A`` is the CALIBRATED (static) scale that replaces them: one value, which
+    ``A`` quantizes against here so every row reads the one entry, or — where the op routes —
+    one per expert, which stays raw for the kernel to quantize each tile against its own expert
+    — a gathered row serves several experts, so it has no single pre-quantized form until the
+    routed rows are laid out (``quantize_routed_rows_per_expert``). ``num_experts`` ``None`` on
+    the 2D op, whose rows belong to no expert.
+    """
+    if As is None:
+        A, As = fp8_act_quant_tensor_wide(A, A.shape[-1])
+        As = As.reshape(-1)
+        return A, As, As.stride(0), 0
+    if A.dtype == FP8_DTYPE:
+        return A, As, As.stride(0), 0
+    As = As.reshape(-1).float()
+    assert As.numel() in (1, num_experts), (
+        f"a calibrated activation scale is one value, or one per expert where the op routes; "
+        f"got {As.numel()}"
+    )
+    if As.numel() == 1:
+        return quantize_rows_static(A, As), As, 0, 0
+    return A, As, 0, 1
+@triton.jit
+def _static_row_scale(S, ExpertStart, row, stride_s, NUM_EXPERTS: tl.constexpr, PER_EXPERT: tl.constexpr):
+    """This row's calibrated scale. ``PER_EXPERT`` resolves it from the routing itself — the
+    rows are expert-sorted, so the count of expert ends at or below the row IS its expert,
+    which costs one ``(NUM_EXPERTS,)`` compare and spares the host materializing a per-row
+    vector (a gather or a repeat_interleave, either of which outweighs this whole kernel)."""
+    if PER_EXPERT:
+        ends = tl.load(ExpertStart + 1 + tl.arange(0, NUM_EXPERTS))
+        s = tl.load(S + tl.sum((ends <= row).to(tl.int32)))
+    else:
+        s = tl.load(S + row * stride_s)
+    return s
+@triton.jit
+def _fp8_quant_static_kernel(
+    X, Y, S, GatherIdx, ExpertStart, stride_x_m, stride_y_m, stride_s, K,
+    NUM_EXPERTS: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
+    GATHER: tl.constexpr, PER_EXPERT: tl.constexpr,
+):
+    """One row tile against a PROVIDED scale: gather, divide in fp32, store E4M3. The fp32
+    divide is the in-register static arm's arithmetic, so both forms of ``A`` round alike.
+    ``stride_s`` 0 broadcasts one scale over every row; ``GATHER`` reads ``X`` through
+    ``GatherIdx`` so the routed rows land laid out."""
+    row = tl.program_id(0)
+    src = tl.load(GatherIdx + row) if GATHER else row
+    offs = tl.program_id(1) * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
+    mask = offs < K
+    x = tl.load(X + src.to(tl.int64) * stride_x_m + offs, mask=mask, other=0.0)
+    y = (x.to(tl.float32) / _static_row_scale(S, ExpertStart, row, stride_s, NUM_EXPERTS, PER_EXPERT)).to(tl.float8e4nv)
+    tl.store(Y + row * stride_y_m + offs, y, mask=mask)
+def quantize_rows_static(A, scales, gather_idx=None, expert_start=None, num_experts=None):
+    """``A`` as E4M3 against ``scales`` — one per row, one for all of them, or (with
+    ``expert_start``) one per expert resolved from the routing. ``gather_idx`` lays the routed
+    rows out in the same pass. One read and one write, where the torch spelling costs four."""
+    assert A.stride(-1) == 1, "the static row quant reads K contiguous"
+    rows = gather_idx.shape[0] if gather_idx is not None else A.shape[0]
+    K = A.shape[-1]
+    y = A.new_empty(rows, K, dtype=FP8_DTYPE)
+    block_k = min(triton.next_power_of_2(K), 1024)
+    with device_context(A.device):
+        compile_time_only_triton_wrap(_fp8_quant_static_kernel)[(rows, triton.cdiv(K, block_k))](
+            A, y, scales, gather_idx, expert_start,
+            A.stride(0), y.stride(0), 0 if scales.numel() == 1 else scales.stride(0), K,
+            NUM_EXPERTS=num_experts or 1, BLOCK_SIZE_K=block_k,
+            GATHER=gather_idx is not None, PER_EXPERT=expert_start is not None,
+        )
+    return y
+def quantize_routed_rows_per_expert(A, As, gather_idx, expert_start, num_experts):
+    """The routed rows as E4M3, each quantized against its own expert's calibrated scale.
+    A per-expert scale otherwise hands the kernel raw rows to quantize per tile, which reads
+    ``A`` at bf16 width through every N-tile and holds the tile at ``RAW_ACT_MAX_BN``. Laying
+    the rows out once buys back both, at the rounding the in-register arm produces, because a
+    row's expert is known once the rows are expert-sorted — ``gather_idx`` orders them, or the
+    intermediate already is. ``As`` is untouched: the kernel dequantizes per expert whichever
+    form ``A`` arrives in.
+    Returns the laid-out rows; the caller's ``gather_idx`` is spent by that layout.
+    """
+    return quantize_rows_static(A, As, gather_idx, expert_start, num_experts)
 def fp8_act_quant_tensor_wide(
     x: torch.Tensor, block_size: int = 128
 ) -> tuple[torch.Tensor, torch.Tensor]: