Kernels:

kernels-community
/

finegrained-fp8

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on Jun 29

Commit

1a264fc

verified ·

1 Parent(s): 14efb56

Uploaded using `kernel-builder`.

Browse files

Files changed (30) hide show

build/torch-cuda/__init__.py +14 -26
build/torch-cuda/_ops.py +1 -1
build/torch-cuda/act_quant.py +0 -73
build/torch-cuda/batched.py +374 -146
build/torch-cuda/bayesian_autotuner.py +358 -0
build/torch-cuda/fused_batched.py +853 -0
build/torch-cuda/fused_grouped.py +1110 -0
build/torch-cuda/grouped.py +480 -198
build/torch-cuda/matmul.py +574 -168
build/torch-cuda/metadata.json +17 -2
build/torch-cuda/metadata.json.sigstore +1 -0
build/torch-cuda/utils.py +755 -1
build/torch-rocm/__init__.py +0 -32
build/torch-rocm/_ops.py +0 -38
build/torch-rocm/act_quant.py +0 -73
build/torch-rocm/batched.py +0 -390
build/torch-rocm/finegrained_fp8/__init__.py +0 -26
build/torch-rocm/grouped.py +0 -477
build/torch-rocm/matmul.py +0 -411
build/torch-rocm/metadata.json +0 -10
build/torch-rocm/utils.py +0 -13
build/torch-xpu/__init__.py +0 -32
build/torch-xpu/_ops.py +0 -38
build/torch-xpu/act_quant.py +0 -73
build/torch-xpu/batched.py +0 -398
build/torch-xpu/finegrained_fp8/__init__.py +0 -26
build/torch-xpu/grouped.py +0 -477
build/torch-xpu/matmul.py +0 -411
build/torch-xpu/metadata.json +0 -10
build/torch-xpu/utils.py +0 -13

build/torch-cuda/__init__.py CHANGED Viewed

@@ -1,32 +1,20 @@
-from .act_quant import fp8_act_quant
-from .batched import (
-    w8a8_fp8_matmul_batched,
-    w8a8_block_fp8_matmul_batched,
-    w8a8_tensor_fp8_matmul_batched,
-)
-from .grouped import (
-    w8a8_fp8_matmul_grouped,
-    w8a8_block_fp8_matmul_grouped,
-    w8a8_tensor_fp8_matmul_grouped,
-)
-from .matmul import (
-    w8a8_fp8_matmul,
-    w8a8_block_fp8_matmul,
-    w8a8_tensor_fp8_matmul,
-)
 __all__ = [
     "fp8_act_quant",
-    # Single matmul
-    "w8a8_fp8_matmul",
-    "w8a8_block_fp8_matmul",
-    "w8a8_tensor_fp8_matmul",
     # Batched matmul
-    "w8a8_fp8_matmul_batched",
-    "w8a8_block_fp8_matmul_batched",
-    "w8a8_tensor_fp8_matmul_batched",
     # Grouped matmul
-    "w8a8_fp8_matmul_grouped",
-    "w8a8_block_fp8_matmul_grouped",
-    "w8a8_tensor_fp8_matmul_grouped",
 ]

+from .utils import fp8_act_quant
+from .matmul import matmul_2d
+from .batched import matmul_batched
+from .grouped import matmul_grouped
+from .fused_batched import moe_fused_batched
+from .fused_grouped import moe_fused_grouped
 __all__ = [
     "fp8_act_quant",
+    # 2D matmul
+    "matmul_2d",
     # Batched matmul
+    "matmul_batched",
     # Grouped matmul
+    "matmul_grouped",
+    # Fused batched matmul
+    "moe_fused_batched",
+    # Fused grouped matmul
+    "moe_fused_grouped",
 ]

build/torch-cuda/_ops.py CHANGED Viewed

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

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

build/torch-cuda/act_quant.py DELETED Viewed

@@ -1,73 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-import triton
-import triton.language as tl
-from torch.library import triton_op, wrap_triton
-from .utils import device_context
-_FP8_DTYPE = torch.float8_e4m3fn
-# Copied from https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
-@triton.jit
-def _fp8_act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr):
-    pid = tl.program_id(axis=0)
-    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    x = tl.load(x_ptr + offs).to(tl.float32)
-    s = tl.max(tl.abs(x)) / 448.0  # float8_e4m3fn max
-    y = (x / s).to(y_ptr.dtype.element_ty)
-    tl.store(y_ptr + offs, y)
-    tl.store(s_ptr + pid, s)
-@triton_op("finegrained_fp8::fp8_act_quant", mutates_args=())
-def _fp8_act_quant(
-    x: torch.Tensor, block_size: int = 128
-) -> tuple[torch.Tensor, torch.Tensor]:
-    assert x.is_contiguous()
-    assert x.shape[-1] % block_size == 0
-    y = torch.empty_like(x, dtype=_FP8_DTYPE)
-    grid = (triton.cdiv(x.numel(), block_size),)
-    s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
-    with device_context(x.device):
-        wrap_triton(_fp8_act_quant_kernel)[grid](x, y, s, BLOCK_SIZE=block_size)
-    return y, s
-def fp8_act_quant(
-    x: torch.Tensor, block_size: int = 128
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Quantize activations to FP8 with per-block dynamic scaling.
-    Splits the last dimension of ``x`` into blocks of ``block_size`` elements,
-    computes ``scale = max(|x_block|) / 448`` per block, and quantizes to
-    ``float8_e4m3fn``.
-    Args:
-        x: Input tensor in bf16/fp16/fp32. Last dimension must be divisible by
-            ``block_size`` and the tensor must be contiguous.
-        block_size: Number of elements per quantization block (default: 128).
-    Returns:
-        A tuple ``(quantized, scales)`` where ``quantized`` has dtype
-        ``float8_e4m3fn`` with the same shape as ``x``, and ``scales`` has
-        shape ``(*x.shape[:-1], x.shape[-1] // block_size)`` in float32.
-    """
-    return torch.ops.finegrained_fp8.fp8_act_quant(x, block_size)

build/torch-cuda/batched.py CHANGED Viewed

@@ -15,31 +15,101 @@
 import torch
 import triton
 import triton.language as tl
-from .act_quant import fp8_act_quant
 from torch.library import triton_op, wrap_triton
-from .utils import device_context
 @triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
-    key=["N", "K"],
 )
 @triton.jit
-def w8a8_block_fp8_matmul_batched_kernel(
-    A,  # (S, K)  raw BF16/FP16 activations
-    B,  # (E, N, K) FP8 weight matrices
-    C,  # (S, N)  output
-    Bs,  # (E, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales
     ExpertIds,  # (S,) — which expert each batch element routes to
     # Shape
     S,
     N,
     K,
     stride_am,
     stride_ak,
     stride_be,
@@ -50,28 +120,24 @@ def w8a8_block_fp8_matmul_batched_kernel(
     stride_bs_e,
     stride_bs_k,
     stride_bs_n,
     # Meta-parameters
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
 ):
     """Block-scale batched FP8 expert matmul kernel.
     Each program handles one routed token row and one N-tile, looks up the
     owning expert from ``ExpertIds``, and applies fused activation quantization.
     """
-    batch_id = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    # Cast expert_id to int64 to prevent int32 overflow when computing
-    # expert_id * stride_Eb (e.g. 255 * 9_437_184 > 2^31 for 256 experts of
-    # 3072×3072 FP8 weights).
-    expert_id = tl.load(ExpertIds + batch_id).to(tl.int64)
-    A = A + batch_id * stride_am
-    B = B + expert_id * stride_be
-    C = C + batch_id * stride_cm
-    Bs = Bs + expert_id * stride_bs_e
     offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
     offs_k = tl.arange(0, BLOCK_SIZE_K)
@@ -82,55 +148,35 @@ def w8a8_block_fp8_matmul_batched_kernel(
     accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        # ---- fused fp8_act_quant ----
         a_raw = tl.load(a_ptrs).to(tl.float32)
-        a_s = tl.max(tl.abs(a_raw)) / 448.0
-        a = (a_raw / tl.maximum(a_s, 1e-12)).to(tl.float8e4nv)
-        # ---- matmul ----
         b = tl.load(b_ptrs)
-        b_s = tl.load(bs_ptrs + k * stride_bs_k)
-        accumulator += tl.dot(a, b) * a_s * b_s[None, :]
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + offs_cm[:, None] * 0 + stride_cn * offs_cn[None, :]
-    # Fake-batch trick aliases all BLOCK_SIZE_M lanes to the same C row, so emitting
-    # `tl.store(c_ptrs, c)` issues BLOCK_SIZE_M duplicate-address stores. On NVIDIA
-    # WGMMA this is usually benign (last-write-wins of identical bytes), but on Intel
-    # XPU the duplicate-address store has hardware-undefined behavior and corrupts the
-    # output. Mask so only lane 0 stores — the (M, N) accumulator rows are
-    # mathematically identical (same A row × same B), so lane 0 holds the right value.
-    tl.store(c_ptrs, c, mask=(offs_cm == 0)[:, None])
 @triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
     key=["N", "K"],
 )
 @triton.jit
-def w8a8_tensor_fp8_matmul_batched_kernel(
     A,  # (S, K) pre-quantized FP8 activations
-    B,  # (E, N, K) FP8 weight matrices
     C,  # (S, N) output
-    As,  # (S, 1) per-tensor activation scales
-    Bs,  # (E, 1, 1) per-tensor weight scales
-    ExpertIds,
     S,
     N,
     K,
     stride_am,
     stride_ak,
     stride_be,
@@ -140,24 +186,24 @@ def w8a8_tensor_fp8_matmul_batched_kernel(
     stride_cn,
     stride_as_m,
     stride_bs_e,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
 ):
     """Tensor-scale batched FP8 expert matmul kernel.
     Activations are already quantized; the kernel applies per-token activation
     scales and per-expert tensor weight scales.
     """
-    batch_id = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    expert_id = tl.load(ExpertIds + batch_id).to(tl.int64)
-    A = A + batch_id * stride_am
-    B = B + expert_id * stride_be
-    C = C + batch_id * stride_cm
-    Bs = Bs + expert_id * stride_bs_e
     offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
     offs_k = tl.arange(0, BLOCK_SIZE_K)
@@ -177,45 +223,126 @@ def w8a8_tensor_fp8_matmul_batched_kernel(
     accumulator = accumulator * a_s * b_s
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + offs_cm[:, None] * 0 + stride_cn * offs_cn[None, :]
-    # See block-FP8 kernel above: BLOCK_SIZE_M lanes alias the same C row;
-    # mask so only lane 0 stores to avoid hardware-undefined duplicate writes on XPU.
-    tl.store(c_ptrs, c, mask=(offs_cm == 0)[:, None])
-@triton_op("finegrained_fp8::w8a8_block_fp8_matmul_batched", mutates_args=())
-def _w8a8_block_fp8_matmul_batched(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     expert_ids: torch.Tensor,
     block_size: list[int],
 ) -> torch.Tensor:
     """Block-scale batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T, with fused act quant.
     A:  (S, K) raw bf16/fp16 activations
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
     """
     assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
     assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
     assert B.is_contiguous(), "B must be contiguous"
     assert A.shape[1] == B.shape[2], (
         f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
     )
     S, K = A.shape
-    E, N, _ = B.shape
     assert len(block_size) == 2, (
         f"block_size must be [block_n, block_k], got {block_size}"
@@ -225,21 +352,18 @@ def _w8a8_block_fp8_matmul_batched(
     assert N % block_n == 0, f"N ({N}) must be divisible by block_n ({block_n})"
     assert K % block_k == 0, f"K ({K}) must be divisible by block_k ({block_k})"
     assert Bs.ndim == 3, (
-        f"Bs must be 3D (E, N//block_n, K//block_k), got ndim={Bs.ndim}"
     )
-    assert Bs.shape == (E, N // block_n, K // block_k), (
-        f"Bs shape {tuple(Bs.shape)} != expected ({E}, {N // block_n}, {K // block_k})"
     )
-    C = A.new_empty(S, N)
-    # Adaptive BLOCK_SIZE_M: smallest power-of-2 >= M, floored at 16, capped at 128.
-    # Matches the WGMMA tile to the actual row count — smaller tiles use less
-    # register pressure and a better-matched FP8 WGMMA instruction, improving
-    # both accuracy and performance for small M (decode).
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
     grid = (S, triton.cdiv(N, block_n))
     with device_context(A.device):
-        wrap_triton(w8a8_block_fp8_matmul_batched_kernel)[grid](
             A,
             B,
             C,
@@ -258,60 +382,63 @@ def _w8a8_block_fp8_matmul_batched(
             Bs.stride(0),
             Bs.stride(2),
             Bs.stride(1),
             BLOCK_SIZE_N=block_n,
             BLOCK_SIZE_K=block_k,
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
         )
     return C
-@triton_op("finegrained_fp8::w8a8_tensor_fp8_matmul_batched", mutates_args=())
-def _w8a8_tensor_fp8_matmul_batched(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     expert_ids: torch.Tensor,
 ) -> torch.Tensor:
     """Tensor-scale batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T, with fused act quant.
     A:  (S, K) raw bf16/fp16 activations
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
     """
     assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
     assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
     assert B.is_contiguous(), "B must be contiguous"
     assert A.shape[1] == B.shape[2], (
         f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
     )
     S, K = A.shape
-    E, N, _ = B.shape
-    # Normalize Bs to (E, 1, 1)
     if Bs.ndim == 1:
-        assert Bs.shape[0] == E, f"Bs shape {tuple(Bs.shape)} != expected ({E},)"
-        Bs = Bs.reshape(E, 1, 1)
     else:
-        assert Bs.shape == (E, 1, 1), (
-            f"Bs shape {tuple(Bs.shape)} != expected ({E}, 1, 1)"
         )
-    BLOCK_SIZE_N = 128
-    BLOCK_SIZE_K = 128
-    C = A.new_empty(S, N)
     qA, As = fp8_act_quant(A, K)
-    grid = (S, triton.cdiv(N, BLOCK_SIZE_N))
-    # Adaptive BLOCK_SIZE_M: smallest power-of-2 >= M, floored at 16, capped at 128.
-    # Matches the WGMMA tile to the actual row count — smaller tiles use less
-    # register pressure and a better-matched FP8 WGMMA instruction, improving
-    # both accuracy and performance for small M (decode).
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
-    grid = (S, triton.cdiv(N, BLOCK_SIZE_N))
     with device_context(A.device):
-        wrap_triton(w8a8_tensor_fp8_matmul_batched_kernel)[grid](
             qA,
             B,
             C,
@@ -330,69 +457,170 @@ def _w8a8_tensor_fp8_matmul_batched(
             C.stride(1),
             As.stride(0),
             Bs.stride(0),
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
         )
     return C
-def w8a8_block_fp8_matmul_batched(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     expert_ids: torch.Tensor,
     block_size: list[int],
 ) -> torch.Tensor:
     """Block-scale batched FP8 matmul with fused activation quantization.
     A:  (S, K) raw activations, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
     """
-    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_batched(
-        A, B, Bs, expert_ids, block_size
     )
-def w8a8_tensor_fp8_matmul_batched(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     expert_ids: torch.Tensor,
 ) -> torch.Tensor:
     """Tensor-scale batched FP8 matmul with fused activation quantization.
     A:  (S, K) raw activations, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
     """
-    return torch.ops.finegrained_fp8.w8a8_tensor_fp8_matmul_batched(
-        A, B, Bs, expert_ids
     )
-def w8a8_fp8_matmul_batched(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     expert_ids: torch.Tensor,
     block_size: list[int] | None,
 ) -> torch.Tensor:
-    """Unified batched W8A8 FP8 matmul dispatcher.
-    Dispatch rules:
-    - tensor mode when ``block_size is None``
-    - tensor mode when ``block_size == [N, K]``
-    - otherwise block mode
-    Returns:
-        Output tensor ``[S, N]`` in the same dtype as ``A``.
     """
-    if block_size is None or (
-        block_size[0] == B.size(1) and block_size[1] == B.size(2)
-    ):
-        return w8a8_tensor_fp8_matmul_batched(A, B, Bs, expert_ids)
-    return w8a8_block_fp8_matmul_batched(A, B, Bs, expert_ids, block_size)

 import torch
 import triton
 import triton.language as tl
 from torch.library import triton_op, wrap_triton
+from ._ops import add_op_namespace_prefix, ops
+from .bayesian_autotuner import bayesian_autotune
+from .utils import (
+    MX_SCALE_GROUP_K,
+    NIBBLES_PER_BYTE,
+    DECODE_BLOCK_SIZE_M,
+    decode_ue8m0_scale,
+    device_context,
+    mx_dot_rescale,
+    mx_dot_scaled,
+    mx_scalar_reduce,
+    mxfp_act_quant_inline,
+    mxfp4_e2m1_to_e4m3,
+    fp8_act_quant,
+    fp8_act_quant_inline,
+    get_accelerator_autotuning_configs,
+    get_mxfp_autotuning_configs,
+    is_mxfp,
+    is_tensor_wide,
+    e2m1_as_uint8,
+    ue8m0_as_uint8,
+)
+@triton.jit
+def expert_setup(
+    A,
+    B,
+    C,
+    Bs,
+    ExpertIds,
+    stride_am,
+    stride_be,
+    stride_cm,
+    stride_bs_e,
+    stride_eid,
+):
+    """Per-(row, expert) prologue shared by the batched kernels: read the program
+    ids, look up the routed expert, and advance the A/B/C/Bs base pointers to this
+    row's slice. Returns ``(batch_id, pid_n, expert_id, A, B, C, Bs)``.
+    The caller must early-return on the EP sentinel (``expert_id >= num_experts``)
+    before any load — the pointer arithmetic itself is harmless, only the loads on a
+    non-local expert would be out of bounds."""
+    batch_id = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    # Cast to int64 to prevent overflow on expert_id * stride_be.
+    expert_id = tl.load(ExpertIds + batch_id * stride_eid).to(tl.int64)
+    A = A + batch_id * stride_am
+    B = B + expert_id * stride_be
+    C = C + batch_id * stride_cm
+    Bs = Bs + expert_id * stride_bs_e
+    return batch_id, pid_n, expert_id, A, B, C, Bs
+@triton.jit
+def store_row(
+    C,
+    accumulator,
+    pid_n,
+    stride_cn,
+    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)
+    c_ptrs = C + offs_cm[:, None] * 0 + stride_cn * offs_cn[None, :]
+    tl.store(c_ptrs, c, mask=(offs_cm == 0)[:, None])
 @triton.autotune(
+    configs=get_accelerator_autotuning_configs(),
+    key=["N", "K", "S"],
 )
 @triton.jit
+def w8a8_block_dynamic_fp8_matmul_batched_kernel(
+    A,  # (S, K) raw BF16/FP16 activations
+    B,  # (num_experts, N, K) FP8 weight matrices
+    C,  # (S, N) output
+    Bs,  # (num_experts, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales
     ExpertIds,  # (S,) — which expert each batch element routes to
     # Shape
     S,
     N,
     K,
+    # Strides
     stride_am,
     stride_ak,
     stride_be,
     stride_bs_e,
     stride_bs_k,
     stride_bs_n,
+    stride_eid,
+    num_experts,
     # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
 ):
     """Block-scale batched FP8 expert matmul kernel.
     Each program handles one routed token row and one N-tile, looks up the
     owning expert from ``ExpertIds``, and applies fused activation quantization.
     """
+    _, pid_n, expert_id, A, B, C, Bs = expert_setup(
+        A, B, C, Bs, ExpertIds, stride_am, stride_be, stride_cm, stride_bs_e, stride_eid
+    )
+    # EP sentinel: row routed to a non-local expert; output is left uninit.
+    if expert_id >= num_experts:
+        return
     offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
     offs_k = tl.arange(0, BLOCK_SIZE_K)
     accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a_raw = tl.load(a_ptrs).to(tl.float32)
+        a, a_s = fp8_act_quant_inline(a_raw)
         b = tl.load(b_ptrs)
+        b_s = decode_ue8m0_scale(tl.load(bs_ptrs))
+        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
+        bs_ptrs += stride_bs_k
+    store_row(C, accumulator, pid_n, stride_cn, BLOCK_SIZE_M, BLOCK_SIZE_N)
 @triton.autotune(
+    configs=get_accelerator_autotuning_configs(tune_block_nk=True),
     key=["N", "K"],
 )
 @triton.jit
+def w8a8_tensor_dynamic_fp8_matmul_batched_kernel(
     A,  # (S, K) pre-quantized FP8 activations
+    B,  # (num_experts, N, K) FP8 weight matrices
     C,  # (S, N) output
+    As,  # (S,) per-token activation scales
+    Bs,  # (num_experts, 1, 1) per-tensor weight scales
+    ExpertIds,  # (S,) — which expert each batch element routes to
+    # Shape
     S,
     N,
     K,
+    # Strides
     stride_am,
     stride_ak,
     stride_be,
     stride_cn,
     stride_as_m,
     stride_bs_e,
+    stride_eid,
+    num_experts,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
 ):
     """Tensor-scale batched FP8 expert matmul kernel.
     Activations are already quantized; the kernel applies per-token activation
     scales and per-expert tensor weight scales.
     """
+    batch_id, pid_n, expert_id, A, B, C, Bs = expert_setup(
+        A, B, C, Bs, ExpertIds, stride_am, stride_be, stride_cm, stride_bs_e, stride_eid
+    )
+    # EP sentinel: row routed to a non-local expert; output is left uninit.
+    if expert_id >= num_experts:
+        return
     offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
     offs_k = tl.arange(0, BLOCK_SIZE_K)
     accumulator = accumulator * a_s * b_s
+    store_row(C, accumulator, pid_n, stride_cn, BLOCK_SIZE_M, BLOCK_SIZE_N)
+@bayesian_autotune(get_mxfp_autotuning_configs(), ["N", "K", "S"], n_trials=60)
+@triton.jit
+def mxfp_dynamic_matmul_batched_kernel(
+    A,  # (S, K) raw BF16/FP16 activations
+    B,  # (num_experts, N, K) E4M3 (MXFP8) or (num_experts, N, K // 2) packed E2M1 (MXFP4) expert weights
+    C,  # (S, N) output
+    Bs,  # (num_experts, N, K // SCALE_GROUP_K) UE8M0 weight scales
+    ExpertIds,  # (S,) — which expert each routed row uses
+    # Shape
+    S,
+    N,
+    K,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_be,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_bs_e,
+    stride_bs_k,
+    stride_bs_n,
+    stride_eid,
+    num_experts,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    VALUES_PER_BYTE: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    COMPUTE_MODE: tl.constexpr,
+):
+    """Unified batched MXFP4/MXFP8 (W4A8/W8A8) expert matmul with fused act quant.
+    One routed row + one N-tile per program; expert looked up from ``ExpertIds``. ``A`` is
+    quantized to E4M3 per K-group inline (UE8M0 scale). ``VALUES_PER_BYTE`` picks the
+    weight format (2 = packed E2M1 / MXFP4, 1 = unpacked E4M3 / MXFP8); ``COMPUTE_MODE``
+    picks ``tl.dot_scaled`` (native M=128) vs fp8 ``tl.dot`` + per-group rescale (wins at
+    decode where the scaled MMA's M→128 pad is waste; FP4 unpacks E2M1->E4M3, lossless).
+    """
+    _, pid_n, expert_id, A, B, C, Bs = expert_setup(
+        A, B, C, Bs, ExpertIds, stride_am, stride_be, stride_cm, stride_bs_e, stride_eid
+    )
+    # EP sentinel: row routed to a non-local expert; output is left uninit.
+    if expert_id >= num_experts:
+        return
+    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    offs_kb = tl.arange(0, BLOCK_SIZE_K // VALUES_PER_BYTE)
+    offs_sf = tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K)
+    a_ptrs = A + tl.arange(0, BLOCK_SIZE_M)[:, None] * 0 + offs_k[None, :] * stride_ak
+    b_ptrs = B + (offs_kb[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+    bs_ptrs = Bs + offs_bn[:, None] * stride_bs_n + offs_sf[None, :] * stride_bs_k
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a_raw = tl.load(a_ptrs).to(tl.float32)
+        a, a_scale = mxfp_act_quant_inline(
+            a_raw, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K
+        )
+        b = tl.load(b_ptrs)
+        b_s = tl.load(bs_ptrs).to(tl.uint8)
+        bq = mxfp4_e2m1_to_e4m3(b) if VALUES_PER_BYTE == 2 else b
+        if COMPUTE_MODE == "dot_scaled":
+            accumulator = mx_dot_scaled(
+                a, a_scale, b, b_s, accumulator, VALUES_PER_BYTE
+            )
+        elif COMPUTE_MODE == "dot":
+            accumulator = mx_dot_rescale(accumulator, a, bq, a_scale, b_s)
+        else:  # scalar
+            accumulator = mx_scalar_reduce(
+                accumulator,
+                a,
+                a_scale,
+                bq,
+                b_s,
+                BLOCK_SIZE_M,
+                BLOCK_SIZE_N,
+                BLOCK_SIZE_K,
+                SCALE_GROUP_K,
+            )
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += (BLOCK_SIZE_K // VALUES_PER_BYTE) * stride_bk
+        bs_ptrs += (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_bs_k
+    store_row(C, accumulator, pid_n, stride_cn, BLOCK_SIZE_M, BLOCK_SIZE_N)
+@triton_op(
+    add_op_namespace_prefix("w8a8_block_dynamic_fp8_matmul_batched"), mutates_args=()
+)
+def _w8a8_block_dynamic_fp8_matmul_batched(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     expert_ids: torch.Tensor,
     block_size: list[int],
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Block-scale batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T, with fused act quant.
     A:  (S, K) raw bf16/fp16 activations
+    B:  (num_experts, N, K) FP8 expert weights
+    Bs: (num_experts, N // block_n, K // block_k) per-block weight scales
     """
     assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
     assert A.is_contiguous(), "A must be contiguous"
+    assert B.ndim == 3, f"B must be 3D (num_experts, N, K), got ndim={B.ndim}"
     assert B.is_contiguous(), "B must be contiguous"
     assert A.shape[1] == B.shape[2], (
         f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
     )
     S, K = A.shape
+    num_experts, N, _ = B.shape
     assert len(block_size) == 2, (
         f"block_size must be [block_n, block_k], got {block_size}"
     assert N % block_n == 0, f"N ({N}) must be divisible by block_n ({block_n})"
     assert K % block_k == 0, f"K ({K}) must be divisible by block_k ({block_k})"
     assert Bs.ndim == 3, (
+        f"Bs must be 3D (num_experts, N//block_n, K//block_k), got ndim={Bs.ndim}"
     )
+    assert Bs.shape == (num_experts, N // block_n, K // block_k), (
+        f"Bs shape {tuple(Bs.shape)} != expected ({num_experts}, {N // block_n}, {K // block_k})"
     )
+    Bs = ue8m0_as_uint8(Bs)
     grid = (S, triton.cdiv(N, block_n))
+    C = A.new_empty(S, N, dtype=output_dtype)
     with device_context(A.device):
+        wrap_triton(w8a8_block_dynamic_fp8_matmul_batched_kernel)[grid](
             A,
             B,
             C,
             Bs.stride(0),
             Bs.stride(2),
             Bs.stride(1),
+            expert_ids.stride(0),
             BLOCK_SIZE_N=block_n,
             BLOCK_SIZE_K=block_k,
+            BLOCK_SIZE_M=DECODE_BLOCK_SIZE_M,
+            num_experts=num_experts,
         )
     return C
+@triton_op(
+    add_op_namespace_prefix("w8a8_tensor_dynamic_fp8_matmul_batched"), mutates_args=()
+)
+def _w8a8_tensor_dynamic_fp8_matmul_batched(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     expert_ids: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Tensor-scale batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T, with fused act quant.
     A:  (S, K) raw bf16/fp16 activations
+    B:  (num_experts, N, K) FP8 expert weights
+    Bs: (num_experts,) or (num_experts, 1, 1) per-expert weight scales
     """
     assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
     assert A.is_contiguous(), "A must be contiguous"
+    assert B.ndim == 3, f"B must be 3D (num_experts, N, K), got ndim={B.ndim}"
     assert B.is_contiguous(), "B must be contiguous"
     assert A.shape[1] == B.shape[2], (
         f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
     )
     S, K = A.shape
+    num_experts, N, _ = B.shape
+    # Normalize Bs to (num_experts, 1, 1)
     if Bs.ndim == 1:
+        assert Bs.shape[0] == num_experts, (
+            f"Bs shape {tuple(Bs.shape)} != expected ({num_experts},)"
+        )
+        Bs = Bs.reshape(num_experts, 1, 1)
     else:
+        assert Bs.shape == (num_experts, 1, 1), (
+            f"Bs shape {tuple(Bs.shape)} != expected ({num_experts}, 1, 1)"
         )
+    Bs = ue8m0_as_uint8(Bs)
     qA, As = fp8_act_quant(A, K)
+    C = A.new_empty(S, N, dtype=output_dtype)
+    def grid(META):
+        return (S, triton.cdiv(N, META["BLOCK_SIZE_N"]))
     with device_context(A.device):
+        wrap_triton(w8a8_tensor_dynamic_fp8_matmul_batched_kernel)[grid](
             qA,
             B,
             C,
             C.stride(1),
             As.stride(0),
             Bs.stride(0),
+            expert_ids.stride(0),
+            BLOCK_SIZE_M=DECODE_BLOCK_SIZE_M,
+            num_experts=num_experts,
         )
     return C
+def w8a8_block_dynamic_fp8_matmul_batched(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     expert_ids: torch.Tensor,
     block_size: list[int],
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Block-scale batched FP8 matmul with fused activation quantization.
     A:  (S, K) raw activations, bf16/fp16/fp32
+    B:  (num_experts, N, K) FP8 expert weights
+    Bs: (num_experts, N // block_n, K // block_k) per-block weight scales
+    expert_ids: (S,) — kernel loads stride-aware, any int dtype works
+    output_dtype: defaults to ``A.dtype``
     """
+    return ops.w8a8_block_dynamic_fp8_matmul_batched(
+        A, B, Bs, expert_ids, block_size, output_dtype
     )
+def w8a8_tensor_dynamic_fp8_matmul_batched(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     expert_ids: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Tensor-scale batched FP8 matmul with fused activation quantization.
     A:  (S, K) raw activations, bf16/fp16/fp32
+    B:  (num_experts, N, K) FP8 expert weights
+    Bs: (num_experts,) or (num_experts, 1, 1) per-expert weight scales
+    output_dtype: defaults to ``A.dtype``
+    """
+    return ops.w8a8_tensor_dynamic_fp8_matmul_batched(
+        A, B, Bs, expert_ids, output_dtype
+    )
+@triton_op(add_op_namespace_prefix("mxfp_dynamic_matmul_batched"), mutates_args=())
+def _mxfp_dynamic_matmul_batched(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    expert_ids: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Batched MX matmul ``C[s] = A[s] @ B[expert_ids[s]].T`` with fused act quant.
+    Weight format is detected from ``B.dtype``: ``int8`` → packed E2M1 (MXFP4, ``B`` is
+    ``(num_experts, N, K//2)``); ``float8_e4m3fn`` → unpacked E4M3 (MXFP8, ``(num_experts, N, K)``). Both use
+    UE8M0 group-32 scales ``(num_experts, N, K//32)``; tile + dot path are autotuned.
+    A:  (S, K) raw activations, bf16/fp16/fp32 (quantized inline to E4M3)
+    expert_ids: (S,) which expert each routed row uses
     """
+    assert A.ndim == 2 and B.ndim == 3 and Bs.ndim == 3
+    assert expert_ids.ndim == 1
+    assert B.dtype in (torch.int8, torch.float8_e4m3fn), (
+        f"B must be int8 (packed E2M1) or float8_e4m3fn (E4M3), got {B.dtype}"
+    )
+    assert Bs.dtype == torch.float8_e8m0fnu, (
+        f"Bs must be float8_e8m0fnu, got {Bs.dtype}"
+    )
+    VALUES_PER_BYTE = NIBBLES_PER_BYTE if B.dtype == torch.int8 else 1
+    S, K = A.shape
+    num_experts, N, K_b = B.shape
+    assert expert_ids.shape[0] == S
+    assert K == VALUES_PER_BYTE * K_b, (
+        f"K (={K}) must equal {VALUES_PER_BYTE} * B.shape[2] (={K_b})"
+    )
+    assert K % MX_SCALE_GROUP_K == 0, (
+        f"K (={K}) must be a multiple of {MX_SCALE_GROUP_K}"
     )
+    assert Bs.shape == (num_experts, N, K // MX_SCALE_GROUP_K), (
+        f"Bs shape {tuple(Bs.shape)} != ({num_experts}, {N}, {K // MX_SCALE_GROUP_K})"
+    )
+    B = e2m1_as_uint8(B)
+    bs_u8 = ue8m0_as_uint8(Bs)
+    C = A.new_empty((S, N), dtype=output_dtype)
+    def grid(META):
+        return (S, triton.cdiv(N, META["BLOCK_SIZE_N"]))
+    with device_context(A.device):
+        wrap_triton(mxfp_dynamic_matmul_batched_kernel)[grid](
+            A,
+            B,
+            C,
+            bs_u8,
+            expert_ids,
+            S,
+            N,
+            K,
+            A.stride(0),
+            A.stride(1),
+            B.stride(0),
+            B.stride(2),
+            B.stride(1),
+            C.stride(0),
+            C.stride(1),
+            bs_u8.stride(0),
+            bs_u8.stride(2),
+            bs_u8.stride(1),
+            expert_ids.stride(0),
+            # Meta-parameters (BLOCK_SIZE_N, BLOCK_SIZE_K come from autotune Config)
+            BLOCK_SIZE_M=DECODE_BLOCK_SIZE_M,
+            VALUES_PER_BYTE=VALUES_PER_BYTE,
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+            num_experts=num_experts,
+        )
+    return C
+def mxfp_dynamic_matmul_batched(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    expert_ids: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Batched MX (MXFP4/MXFP8) matmul with fused act quant; the weight format
+    (packed E2M1 vs E4M3) is detected from ``B.dtype``. Tile + dot path autotuned."""
+    return ops.mxfp_dynamic_matmul_batched(A, B, Bs, expert_ids, output_dtype)
+def matmul_batched(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     expert_ids: torch.Tensor,
     block_size: list[int] | None,
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
+    """Batched quantized matmul dispatcher (W8A8 FP8 or W4A8 FP4). Routes one
+    routed row per program to ``B[expert_ids[s]]``.
+    ``output_dtype`` defaults to ``A.dtype``.
+    Routes by weight dtype and ``block_size``:
+    - MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0
+      group-32 ``Bs`` (shape ``[num_experts, N, K//32]``) → ``mxfp_dynamic_matmul_batched``
+      (``block_size`` ignored; tile + dot path autotuned).
+    - ``block_size`` None or full ``[N, K]`` → ``w8a8_tensor_dynamic_fp8_matmul_batched``.
+    - otherwise → ``w8a8_block_dynamic_fp8_matmul_batched``.
     """
+    if is_mxfp(B, Bs):
+        return mxfp_dynamic_matmul_batched(A, B, Bs, expert_ids, output_dtype)
+    if is_tensor_wide(block_size, B):
+        return w8a8_tensor_dynamic_fp8_matmul_batched(
+            A, B, Bs, expert_ids, output_dtype
+        )
+    return w8a8_block_dynamic_fp8_matmul_batched(
+        A, B, Bs, expert_ids, block_size, output_dtype
+    )

build/torch-cuda/bayesian_autotuner.py ADDED Viewed

	@@ -0,0 +1,358 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Bayesian-optimization autotuner: benches a budgeted sample of the config grid instead of
+the full grid, via a Tree-structured Parzen Estimator (TPE) over the config *dimensions*
+(``num_warps`` / ``num_stages`` / tile sizes / flags). After a short random seed phase it
+models the good (top-``gamma`` by measured time) vs bad per-dimension value densities and
+benches the unmeasured config that maximizes ``l(x)/g(x)`` — the Expected-Improvement proxy
+— then a coordinate-descent pass polishes the best. Subclasses Triton's ``Autotuner`` and
+uses only the Python stdlib (no external optimizer dependency).
+Each new key's search is warm-started from the most recently tuned key's best config (nearby
+workloads share tile-shape preferences). Grids smaller than ``n_trials`` defer to the stock
+exhaustive bench-all.
+"""
+from __future__ import annotations
+import hashlib
+import json
+import math
+import os
+import random
+import time
+from collections import defaultdict
+from typing import Dict, List
+from triton.runtime.autotuner import (
+    Autotuner,
+    Config,
+    JITFunction,
+    driver,
+    get_cache_invalidating_env_vars,
+    get_cache_manager,
+    knobs,
+    triton_key,
+)
+class BayesianAutotuner(Autotuner):
+    """Drop-in replacement for ``triton.runtime.autotuner.Autotuner`` that
+    benches ~``n_trials`` configs per key via TPE Bayesian optimization +
+    coordinate-descent refinement, instead of the full grid."""
+    def __init__(
+        self,
+        *args,
+        n_trials: int = 80,
+        n_startup_trials: int = 12,
+        gamma: float = 0.25,
+        refine: bool = True,
+        max_refine_iters: int = 5,
+        log_path: str | None = None,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        # Bayesian trial budget — the per-decorator default, overridable via the
+        # FINEGRAINED_AUTOTUNE_TRIALS env var (quick sweeps / exhaustive runs without touching
+        # the decorators; set it >= grid size to fall back to stock exhaustive bench-all).
+        self.n_trials = int(os.environ.get("FINEGRAINED_AUTOTUNE_TRIALS") or n_trials)
+        self.n_startup_trials = n_startup_trials
+        # top fraction of measured configs the TPE treats as "good"
+        self.gamma = gamma
+        self.refine = refine
+        self.max_refine_iters = max_refine_iters
+        # JSONL log of every benched (key, config, ms) — set here or via the
+        # FINEGRAINED_AUTOTUNE_LOG env var. Analyse offline to prune bad configs.
+        self.log_path = log_path or os.environ.get("FINEGRAINED_AUTOTUNE_LOG")
+    def run(self, *args, **kwargs):
+        # Small grid → defer to parent (stock exhaustive bench-all).
+        if len(self.configs) <= 1 or self.n_trials >= len(self.configs):
+            return super().run(*args, **kwargs)
+        self.nargs = dict(zip(self.arg_names, args))
+        all_args = {**self.nargs, **kwargs}
+        _args = {k: v for k, v in all_args.items() if k in self.arg_names}
+        key = [_args[k] for k in self.keys if k in _args]
+        key.extend(str(v.dtype) for v in _args.values() if hasattr(v, "dtype"))
+        key = tuple(key)
+        if key not in self.cache:
+            pruned = self.prune_configs(kwargs)
+            def benchmark():
+                t0 = time.time()
+                self.cache[key] = self._bayesian_search(pruned, args, kwargs, key)
+                self.bench_time = time.time() - t0
+                if knobs.autotuning.print:
+                    fn_name = getattr(self.fn, "__name__", str(self.fn))
+                    print(
+                        f"[bayesian-autotune] {fn_name} tuned "
+                        f"{len(self.configs_timings)} configs in {self.bench_time:.1f}s — "
+                        f"key={key}, best={self.cache[key].all_kwargs()}"
+                    )
+                full_nargs = {**self.nargs, **kwargs, **self.cache[key].all_kwargs()}
+                self.pre_hook(full_nargs, reset_only=True)
+            if self.cache_results:
+                self.check_disk_cache(key, pruned, benchmark)
+            else:
+                benchmark()
+        config = self.cache[key]
+        self.best_config = config
+        if config.pre_hook is not None:
+            config.pre_hook({**self.nargs, **kwargs, **config.all_kwargs()})
+        ret = self.fn.run(*args, **kwargs, **config.all_kwargs())
+        self.nargs = None
+        return ret
+    def _bayesian_search(self, configs: List[Config], args, kwargs, key) -> Config:
+        timings: Dict[int, float] = {}
+        sigs = [tuple(sorted(c.all_kwargs().items())) for c in configs]
+        def bench_idx(idx: int) -> float:
+            if idx in timings:
+                return timings[idx]
+            try:
+                ms = self._bench(*args, config=configs[idx], **kwargs)
+                if isinstance(ms, (tuple, list)):
+                    ms = ms[0]
+                timings[idx] = float(ms)
+            except Exception:
+                timings[idx] = float("inf")
+            self._log_result(key, configs[idx], timings[idx])
+            return timings[idx]
+        # Distinct values per config dimension, for the TPE's Laplace smoothing.
+        dim_vals: Dict = defaultdict(set)
+        for sig in sigs:
+            for d, v in sig:
+                dim_vals[d].add(v)
+        # Seed phase: a few random configs (seeded → deterministic), warm-started from the
+        # most recent cached key's best, to give the TPE an initial good/bad split.
+        n_startup = max(2, min(self.n_startup_trials, self.n_trials))
+        order = list(range(len(configs)))
+        random.Random(0).shuffle(order)
+        warm_idx = self._warm_start_index(configs)
+        if warm_idx is not None:
+            order = [warm_idx] + [i for i in order if i != warm_idx]
+        for idx in order[:n_startup]:
+            bench_idx(idx)
+        # TPE: split measured configs into good (top-gamma) / bad, build per-dimension value
+        # densities for each, and bench the unmeasured config maximizing log l(x) - log g(x)
+        # (Expected-Improvement proxy), updating the model after each measurement.
+        while len(timings) < self.n_trials:
+            ranked = sorted(timings, key=timings.get)
+            n_good = max(1, round(self.gamma * len(ranked)))
+            good_c: Dict = defaultdict(lambda: defaultdict(int))
+            bad_c: Dict = defaultdict(lambda: defaultdict(int))
+            for j, i in enumerate(ranked):
+                tgt = good_c if j < n_good else bad_c
+                for d, v in sigs[i]:
+                    tgt[d][v] += 1
+            n_bad = len(ranked) - n_good
+            best_i, best_score = None, -math.inf
+            for i in range(len(configs)):
+                if i in timings:
+                    continue
+                score = 0.0
+                for d, v in sigs[i]:
+                    V = len(dim_vals[d])
+                    lp = (good_c[d][v] + 1.0) / (
+                        n_good + V
+                    )  # P(v | good), Laplace-smoothed
+                    gp = (bad_c[d][v] + 1.0) / (n_bad + V)  # P(v | bad)
+                    score += math.log(lp) - math.log(gp)
+                if score > best_score:
+                    best_score, best_i = score, i
+            if best_i is None:
+                break
+            bench_idx(best_i)
+        # Coordinate-descent refinement: try single-dim perturbations around
+        # the current best until no neighbor improves.
+        if self.refine:
+            for _ in range(self.max_refine_iters):
+                best_idx = min(timings, key=timings.get)
+                best_sig = sigs[best_idx]
+                best_ms = timings[best_idx]
+                improved = False
+                for i, s in enumerate(sigs):
+                    if i in timings:
+                        continue
+                    diff = sum(1 for (_, a), (_, b) in zip(best_sig, s) if a != b)
+                    if diff != 1:
+                        continue
+                    if bench_idx(i) < best_ms:
+                        improved = True
+                if not improved:
+                    break
+        self.configs_timings = {configs[i]: t for i, t in timings.items()}
+        return configs[min(timings, key=timings.get)]
+    def _log_result(self, key, config: Config, ms: float):
+        """Append one ``(key, config, ms)`` record as JSONL for offline analysis —
+        e.g. pruning configs that are consistently far off the per-key best. ``inf`` ms
+        marks a config that failed to compile/run (out of resources, etc.)."""
+        if not self.log_path:
+            return
+        try:
+            rec = {
+                "timestamp": time.time(),
+                "fn_name": getattr(self.fn, "__name__", str(self.fn)),
+                "key": list(key),
+                "kwargs": config.kwargs,
+                "num_warps": config.num_warps,
+                "num_stages": config.num_stages,
+                "ms": ms,
+            }
+            with open(self.log_path, "a") as f:
+                f.write(json.dumps(rec, default=str) + "\n")
+        except Exception:
+            pass
+    def _warm_start_index(self, configs: List[Config]):
+        """Return the index in ``configs`` matching the most recently cached
+        key's best config (or ``None`` if no prior tune or no match in the
+        current pruned list)."""
+        if not self.cache:
+            return None
+        # Python 3.7+ dicts preserve insertion order; last entry = most recent tune.
+        prev_best = next(reversed(self.cache.values()))
+        prev_kwargs = prev_best.all_kwargs()
+        for i, c in enumerate(configs):
+            if c.all_kwargs() == prev_kwargs:
+                return i
+        return None
+    def check_disk_cache(self, tuning_key, configs, bench_fn):
+        """Persist the tuned best config to Triton's on-disk cache so a later run (or process)
+        skips the whole minutes-long search+compile — most of which is Triton codegen, not
+        benching. Keyed exactly like Triton's own autotune cache (triton version + backend +
+        kernel source hash + invalidating env + key + the full config grid), so any of those
+        changing re-tunes.
+        Unlike Triton's stock version we do NOT bail when configs carry a ``pre_hook`` (ours
+        always do): we persist each config's ``all_kwargs()`` — serializable and pre_hook-free —
+        and on a hit re-match it to the live ``Config`` object, so the pre_hook is preserved from
+        code rather than lost in (de)serialization. A hit launches only the winning config."""
+        if not tuning_key:
+            bench_fn()
+            return False
+        from triton.compiler.compiler import make_backend
+        fn = self.fn
+        while not isinstance(fn, JITFunction):
+            fn = fn.fn
+        env_vars = get_cache_invalidating_env_vars()
+        cache_key = [
+            triton_key(),
+            make_backend(driver.active.get_current_target()).hash(),
+            fn.cache_key,
+            str(sorted(env_vars.items())),
+            str(tuning_key),
+        ] + [str(c) for c in configs]  # str(Config) is pre_hook-free and process-stable
+        cache_key = hashlib.sha256("-".join(cache_key).encode("utf-8")).hexdigest()
+        cache = get_cache_manager(cache_key)
+        file_name = f"{fn.__name__[:150]}.bayes_autotune.json"
+        # signature -> live Config (carries the pre_hook); used to re-match the cached winner
+        by_sig = {tuple(sorted(c.all_kwargs().items())): c for c in configs}
+        path = cache.get_file(file_name)
+        if path:
+            try:
+                with open(path) as f:
+                    data = json.load(f)
+                best = by_sig.get(tuple(sorted(data["best"].items())))
+                if best is not None:
+                    self.cache[tuning_key] = best
+                    self.configs_timings = {
+                        by_sig[s]: t
+                        for kw, t in data["timings"]
+                        for s in (tuple(sorted(kw.items())),)
+                        if s in by_sig
+                    }
+                    return True
+            except Exception:
+                pass  # corrupt/stale cache file → fall through and re-tune
+        bench_fn()
+        try:
+            cache.put(
+                json.dumps(
+                    {
+                        "best": self.cache[tuning_key].all_kwargs(),
+                        "timings": [
+                            (c.all_kwargs(), t) for c, t in self.configs_timings.items()
+                        ],
+                    }
+                ),
+                file_name,
+                binary=False,
+            )
+        except Exception:
+            pass
+        return False
+def bayesian_autotune(
+    configs,
+    key,
+    *,
+    n_trials: int = 80,
+    n_startup_trials: int = 12,
+    gamma: float = 0.25,
+    refine: bool = True,
+    max_refine_iters: int = 5,
+    log_path: str | None = None,
+    cache_results: bool = True,
+    reset_to_zero=None,
+    restore_value=None,
+    **kwargs,
+):
+    """Decorator mirroring ``@triton.autotune``. Extra kwargs:
+    n_trials:                 total configs benched per key (TPE budget)
+    n_startup_trials:         random seed configs before the TPE model kicks in
+    gamma:                    top fraction of measured configs treated as "good"
+    refine, max_refine_iters: coordinate-descent refinement after the TPE
+    log_path:                 JSONL of benched configs (or FINEGRAINED_AUTOTUNE_LOG)
+    cache_results:            persist the tuned best config to disk (on by default) so later
+                              runs skip the search+compile — see BayesianAutotuner.check_disk_cache"""
+    def decorator(fn):
+        return BayesianAutotuner(
+            fn,
+            fn.arg_names,
+            configs,
+            key,
+            reset_to_zero,
+            restore_value,
+            n_trials=n_trials,
+            n_startup_trials=n_startup_trials,
+            gamma=gamma,
+            refine=refine,
+            max_refine_iters=max_refine_iters,
+            log_path=log_path,
+            cache_results=cache_results,
+            **kwargs,
+        )
+    return decorator

build/torch-cuda/fused_batched.py ADDED Viewed

	@@ -0,0 +1,853 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Fused batched MoE: two deterministic kernels (no atomics, no sort), per routed token.
+  Kernel 1 (S x N-tiles): gather -> gate_up GEMM -> SiLU -> FP8 requant -> fp8 intermediate
+  Kernel 2 (S x H-tiles): fp8 intermediate -> down GEMM -> routing weight -> output
+Recipe-named to mirror ``batched.py``; ``moe_fused_batched`` is the neutral dispatcher.
+"""
+import torch
+import triton
+import triton.language as tl
+from torch.library import triton_op, wrap_triton
+from ._ops import add_op_namespace_prefix, ops
+from .utils import (
+    FP8_DTYPE,
+    MX_SCALE_GROUP_K,
+    NIBBLES_PER_BYTE,
+    DECODE_BLOCK_SIZE_M,
+    device_context,
+    fp8_act_quant_inline,
+    get_mxfp_autotuning_configs,
+    is_mxfp,
+    is_mxfp4,
+    mxfp_act_quant_inline,
+    mx_compute,
+    mx_compute_gate_up,
+    glu,
+    topk_reduce_kernel,
+    TOPK_REDUCE_BLOCK_H,
+    e2m1_as_uint8,
+    ue8m0_as_uint8,
+)
+from .batched import store_row
+from .bayesian_autotuner import bayesian_autotune
+# ── Batched fused: two-kernel approach (no sorting, no atomics) ──────────────
+#
+# Same two-kernel architecture as grouped fused but with per-token dispatch:
+# Kernel 1: (S, N-tiles) — gate_up + SiLU + FP8 quant → intermediate buffer
+# Kernel 2: (S, H-tiles) — fp8 intermediate → down proj → output
+# No sorting needed — expert lookup is per-token via ExpertIds.
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=w, num_stages=s)
+        for w in [2, 4, 8, 16]
+        for s in [2, 3, 4, 5, 6]
+    ],
+    # num_routed_tokens (= grid axis-0, runtime) is in the key so decode re-tunes per
+    # batch size — a config tuned at S=8 is wrong at S=256 (GPU unsaturated at S=8). The
+    # constexpr dims auto-partition the cache, so they don't need to be in the key.
+    key=["num_routed_tokens"],
+)
+@triton.jit
+def w8a8_block_dynamic_fp8_moe_batched_gate_up_kernel(
+    HiddenStates,
+    GateUp,
+    GateUpScale,
+    Intermediate,
+    IntermediateScale,
+    ExpertIds,
+    stride_a_m,
+    stride_a_k,
+    stride_gu_e,
+    stride_gu_n,
+    stride_gu_k,
+    stride_gus_e,
+    stride_gus_n,
+    stride_gus_k,
+    num_routed_tokens,  # tuning knob: in the autotune key so decode re-tunes per batch
+    NUM_EXPERTS: tl.constexpr,
+    NUM_TOP_K: tl.constexpr,
+    HIDDEN_DIM: tl.constexpr,
+    NUM_N_TILES: tl.constexpr,
+    INTERMEDIATE_DIM: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    SIMULATE_UNFUSED: tl.constexpr,
+    ACT_FN: tl.constexpr,
+):
+    """Batched kernel 1: per-token gate_up + SiLU + FP8 quant. Grid: (S, N-tiles)."""
+    batch_id = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    # Each token is replicated NUM_TOP_K times in the routed (S,) layout, so the
+    # source row is batch_id // NUM_TOP_K — gathered from the unexpanded activations
+    # (no token-index tensor, no top_k-replicated copy). int64 against offset overflow.
+    token = (batch_id // NUM_TOP_K).to(tl.int64)
+    expert_id = tl.load(ExpertIds + batch_id).to(tl.int64)
+    # EP sentinel: row routed to a non-local expert — skip. The intermediate row is left
+    # uninit; the down kernel writes zeros for this row, so the output stays well-defined.
+    if expert_id >= NUM_EXPERTS:
+        return
+    a_ptr = (
+        HiddenStates
+        + (token + tl.arange(0, BLOCK_SIZE_M))[:, None] * stride_a_m
+        + tl.arange(0, BLOCK_SIZE_K)[None, :] * stride_a_k
+    )
+    b_gate_ptr = (
+        GateUp
+        + expert_id * stride_gu_e
+        + tl.arange(0, BLOCK_SIZE_K)[:, None] * stride_gu_k
+        + (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))[None, :] * stride_gu_n
+    )
+    b_up_ptr = (
+        GateUp
+        + expert_id * stride_gu_e
+        + tl.arange(0, BLOCK_SIZE_K)[:, None] * stride_gu_k
+        + (INTERMEDIATE_DIM + pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))[
+            None, :
+        ]
+        * stride_gu_n
+    )
+    bs_gate_ptr = GateUpScale + expert_id * stride_gus_e + pid_n * stride_gus_n
+    bs_up_ptr = (
+        GateUpScale + expert_id * stride_gus_e + (NUM_N_TILES + pid_n) * stride_gus_n
+    )
+    acc_gate = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    acc_up = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for _ in range(0, tl.cdiv(HIDDEN_DIM, BLOCK_SIZE_K)):
+        a_raw = tl.load(a_ptr).to(tl.float32)
+        a, a_s = fp8_act_quant_inline(a_raw)
+        w_gate = tl.load(b_gate_ptr)
+        w_up = tl.load(b_up_ptr)
+        w_s_gate = tl.load(bs_gate_ptr)
+        w_s_up = tl.load(bs_up_ptr)
+        acc_gate += tl.dot(a, w_gate) * a_s[:, None] * w_s_gate
+        acc_up += tl.dot(a, w_up) * a_s[:, None] * w_s_up
+        a_ptr += BLOCK_SIZE_K * stride_a_k
+        b_gate_ptr += BLOCK_SIZE_K * stride_gu_k
+        b_up_ptr += BLOCK_SIZE_K * stride_gu_k
+        bs_gate_ptr += stride_gus_k
+        bs_up_ptr += stride_gus_k
+    intermediate = glu(
+        acc_gate, acc_up, ACT_FN, SIMULATE_UNFUSED, HiddenStates.dtype.element_ty
+    )
+    # Requant the intermediate to FP8 — the same inline per-row act quant as the inputs;
+    # with BLOCK_SIZE_M=1 the per-row scale is the single per-tile scalar we store.
+    inter, inter_s = fp8_act_quant_inline(intermediate)
+    store_row(
+        Intermediate + batch_id * INTERMEDIATE_DIM,
+        inter,
+        pid_n,
+        1,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+    )
+    tl.store(IntermediateScale + batch_id * NUM_N_TILES + pid_n, tl.max(inter_s))
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=w, num_stages=s)
+        for w in [2, 4, 8, 16]
+        for s in [2, 3, 4, 5, 6]
+    ],
+    key=["num_routed_tokens"],
+)
+@triton.jit
+def w8a8_block_dynamic_fp8_moe_batched_down_kernel(
+    Intermediate,
+    IntermediateScale,
+    Down,
+    DownScale,
+    ExpertIds,
+    SampleWeights,
+    Out,
+    stride_down_e,
+    stride_down_n,
+    stride_down_k,
+    stride_downs_e,
+    stride_downs_n,
+    stride_downs_k,
+    num_routed_tokens,  # tuning knob: in the autotune key so decode re-tunes per batch
+    NUM_EXPERTS: tl.constexpr,
+    HIDDEN_DIM: tl.constexpr,
+    INTERMEDIATE_DIM: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_H: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    NUM_N_TILES: tl.constexpr,
+    SIMULATE_UNFUSED: tl.constexpr,
+):
+    """Batched kernel 2: fp8 intermediate → down proj → output. Grid: (S, H-tiles)."""
+    batch_id = tl.program_id(axis=0)
+    pid_h = tl.program_id(axis=1)
+    expert_id = tl.load(ExpertIds + batch_id).to(tl.int64)
+    # EP sentinel: row routed to a non-local expert. The program is already launched, so
+    # write its zero tile here (skipping the weight load) — cheaper than a host-side mask
+    # pass, and it leaves the output fully defined for a plain top-k sum.
+    if expert_id >= NUM_EXPERTS:
+        z = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_H), dtype=tl.float32)
+        store_row(Out + batch_id * HIDDEN_DIM, z, pid_h, 1, BLOCK_SIZE_M, BLOCK_SIZE_H)
+        return
+    w_down_ptr = (
+        Down
+        + expert_id * stride_down_e
+        + tl.arange(0, BLOCK_SIZE_N)[:, None] * stride_down_k
+        + (pid_h * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H))[None, :] * stride_down_n
+    )
+    inter_ptr = (
+        Intermediate
+        + (batch_id + tl.arange(0, BLOCK_SIZE_M))[:, None] * INTERMEDIATE_DIM
+        + tl.arange(0, BLOCK_SIZE_N)[None, :]
+    )
+    inter_s_ptr = (
+        IntermediateScale
+        + (batch_id + tl.arange(0, BLOCK_SIZE_M))[:, None] * NUM_N_TILES
+    )
+    ws_down_ptr = DownScale + expert_id * stride_downs_e + pid_h * stride_downs_n
+    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_H), dtype=tl.float32)
+    for _ in range(0, NUM_N_TILES):
+        inter = tl.load(inter_ptr)
+        inter_s = tl.load(inter_s_ptr)
+        w_s_down = tl.load(ws_down_ptr)
+        w_down = tl.load(w_down_ptr)
+        acc += tl.dot(inter, w_down) * inter_s * w_s_down
+        inter_ptr += BLOCK_SIZE_N
+        inter_s_ptr += 1
+        ws_down_ptr += stride_downs_k
+        w_down_ptr += BLOCK_SIZE_N * stride_down_k
+    if SIMULATE_UNFUSED:
+        acc = acc.to(Out.dtype.element_ty).to(tl.float32)
+    acc = acc * tl.load(SampleWeights + batch_id)
+    store_row(Out + batch_id * HIDDEN_DIM, acc, pid_h, 1, BLOCK_SIZE_M, BLOCK_SIZE_H)
+@triton_op(
+    add_op_namespace_prefix("w8a8_block_dynamic_fp8_moe_batched"), mutates_args=()
+)
+def _w8a8_block_dynamic_fp8_moe_batched(
+    hidden_states: torch.Tensor,
+    gate_up_proj: torch.Tensor,
+    gate_up_proj_scale: torch.Tensor,
+    down_proj: torch.Tensor,
+    down_proj_scale: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    block_size: list[int],
+    output_dtype: torch.dtype,
+    simulate_unfused: bool = False,
+    act_fn: str = "silu",
+) -> torch.Tensor:
+    """Block-dynamic FP8 batched fused MoE in ONE op: gate_up + SiLU + FP8 requant →
+    grouped down → routing-weighted per-(token, expert) output. gate_up gathers each routed
+    row directly from the unexpanded ``hidden_states`` (source row ``s // NUM_TOP_K``, no
+    replicated copy). ``inter``/``inter_scales`` are internal; the caller reduces over top-k."""
+    device = hidden_states.device
+    HIDDEN_DIM = hidden_states.size(1)
+    NUM_EXPERTS = gate_up_proj.size(0)
+    num_tokens = hidden_states.size(0)
+    num_routed_tokens = top_k_index.numel()
+    INTERMEDIATE_DIM = gate_up_proj.size(1) // 2
+    NUM_TOP_K = num_routed_tokens // hidden_states.size(0)
+    BLOCK_SIZE_N, BLOCK_SIZE_K = block_size
+    NUM_N_TILES = triton.cdiv(INTERMEDIATE_DIM, BLOCK_SIZE_N)
+    NUM_H_TILES = triton.cdiv(HIDDEN_DIM, BLOCK_SIZE_N)
+    inter = torch.empty(
+        num_routed_tokens, INTERMEDIATE_DIM, device=device, dtype=FP8_DTYPE
+    )
+    inter_scales = torch.empty(
+        num_routed_tokens, NUM_N_TILES, device=device, dtype=torch.float32
+    )
+    out = torch.empty(num_routed_tokens, HIDDEN_DIM, device=device, dtype=output_dtype)
+    reduced = torch.empty(num_tokens, HIDDEN_DIM, device=device, dtype=output_dtype)
+    with device_context(device):
+        wrap_triton(w8a8_block_dynamic_fp8_moe_batched_gate_up_kernel)[
+            (num_routed_tokens, NUM_N_TILES)
+        ](
+            hidden_states,
+            gate_up_proj,
+            gate_up_proj_scale,
+            inter,
+            inter_scales,
+            top_k_index,
+            hidden_states.stride(0),
+            hidden_states.stride(1),
+            gate_up_proj.stride(0),
+            gate_up_proj.stride(1),
+            gate_up_proj.stride(2),
+            gate_up_proj_scale.stride(0),
+            gate_up_proj_scale.stride(1),
+            gate_up_proj_scale.stride(2),
+            num_routed_tokens=num_routed_tokens,
+            NUM_TOP_K=NUM_TOP_K,
+            NUM_EXPERTS=NUM_EXPERTS,
+            HIDDEN_DIM=HIDDEN_DIM,
+            INTERMEDIATE_DIM=INTERMEDIATE_DIM,
+            BLOCK_SIZE_N=BLOCK_SIZE_N,
+            BLOCK_SIZE_K=BLOCK_SIZE_K,
+            BLOCK_SIZE_M=DECODE_BLOCK_SIZE_M,
+            NUM_N_TILES=NUM_N_TILES,
+            SIMULATE_UNFUSED=simulate_unfused,
+            ACT_FN=act_fn,
+        )
+        wrap_triton(w8a8_block_dynamic_fp8_moe_batched_down_kernel)[
+            (num_routed_tokens, NUM_H_TILES)
+        ](
+            inter,
+            inter_scales,
+            down_proj,
+            down_proj_scale,
+            top_k_index,
+            top_k_weights,
+            out,
+            down_proj.stride(0),
+            down_proj.stride(1),
+            down_proj.stride(2),
+            down_proj_scale.stride(0),
+            down_proj_scale.stride(1),
+            down_proj_scale.stride(2),
+            num_routed_tokens=num_routed_tokens,
+            NUM_EXPERTS=NUM_EXPERTS,
+            HIDDEN_DIM=HIDDEN_DIM,
+            INTERMEDIATE_DIM=INTERMEDIATE_DIM,
+            BLOCK_SIZE_N=BLOCK_SIZE_N,
+            BLOCK_SIZE_H=BLOCK_SIZE_N,
+            BLOCK_SIZE_M=DECODE_BLOCK_SIZE_M,
+            NUM_N_TILES=NUM_N_TILES,
+            SIMULATE_UNFUSED=simulate_unfused,
+        )
+        wrap_triton(topk_reduce_kernel)[
+            (num_tokens, triton.cdiv(HIDDEN_DIM, TOPK_REDUCE_BLOCK_H))
+        ](
+            out,
+            reduced,
+            HIDDEN_DIM,
+            out.stride(0),
+            out.stride(1),
+            reduced.stride(0),
+            reduced.stride(1),
+            NUM_TOP_K=NUM_TOP_K,
+            BLOCK_H=TOPK_REDUCE_BLOCK_H,
+        )
+    return reduced
+def w8a8_block_dynamic_fp8_moe_batched(
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    gate_up_proj: torch.Tensor,
+    down_proj: torch.Tensor,
+    gate_up_proj_scale: torch.Tensor,
+    down_proj_scale: torch.Tensor,
+    block_size: list[int],
+    simulate_unfused: bool = False,
+    act_fn: str = "silu",
+) -> torch.Tensor:
+    """Batched fused MoE (deterministic, no sorting, no atomics): a single ``triton_op``
+    runs gate_up + down (gathering each routed row from the unexpanded ``hidden_states``,
+    source row ``s // num_top_k``, no replicated copy). The top-k reduce stays plain torch
+    so ``torch.compile`` can fuse it with the surrounding model graph."""
+    out = ops.w8a8_block_dynamic_fp8_moe_batched(
+        hidden_states,
+        gate_up_proj,
+        gate_up_proj_scale,
+        down_proj,
+        down_proj_scale,
+        top_k_index,
+        top_k_weights,
+        block_size,
+        hidden_states.dtype,
+        simulate_unfused,
+        act_fn,
+    )
+    return out
+# ── MXFP8 fused (tl.dot_scaled, UE8M0 group-32 scales, tunable tiles) ─────────
+#
+# Mirrors the block-dynamic kernels but: activations/intermediate quantize via the MX
+# group-32 inline quant, weights/scales feed tl.dot_scaled, and the tile (BLOCK_SIZE_N/K)
+# is free to autotune — the group-32 scale runs along K independently of the compute
+# tile (the whole reason MXFP8 can tune the tile that the block recipe must lock to 128).
+@bayesian_autotune(
+    get_mxfp_autotuning_configs(),
+    ["num_routed_tokens"],
+    n_trials=60,
+)
+@triton.jit
+def mxfp_dynamic_moe_batched_gate_up_kernel(
+    HiddenStates,
+    GateUp,
+    GateUpScale,
+    Intermediate,
+    IntermediateScale,
+    ExpertIds,
+    stride_a_m,
+    stride_a_k,
+    stride_gu_e,
+    stride_gu_n,
+    stride_gu_k,
+    stride_gus_e,
+    stride_gus_n,
+    stride_gus_k,
+    num_routed_tokens,  # tuning knob: in the autotune key so decode re-tunes per batch
+    NUM_EXPERTS: tl.constexpr,
+    NUM_TOP_K: tl.constexpr,
+    HIDDEN_DIM: tl.constexpr,
+    INTERMEDIATE_DIM: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    VALUES_PER_BYTE: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    SIMULATE_UNFUSED: tl.constexpr,
+    COMPUTE_MODE: tl.constexpr,
+    ACT_FN: tl.constexpr,
+):
+    """MXFP4/MXFP8 kernel 1: gate_up + SiLU + MXFP8 requant. N = intermediate
+    (output) tile, K = hidden (contraction) tile — both tunable. Grid: (S, N-tiles)."""
+    batch_id = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    token = (batch_id // NUM_TOP_K).to(tl.int64)
+    expert_id = tl.load(ExpertIds + batch_id).to(tl.int64)
+    if expert_id >= NUM_EXPERTS:  # EP sentinel; down kernel zeros the output row
+        return
+    a_ptr = (
+        HiddenStates
+        + (token + tl.arange(0, BLOCK_SIZE_M))[:, None] * stride_a_m
+        + tl.arange(0, BLOCK_SIZE_K)[None, :] * stride_a_k
+    )
+    b_gate_ptr = (
+        GateUp
+        + expert_id * stride_gu_e
+        + tl.arange(0, BLOCK_SIZE_K // VALUES_PER_BYTE)[:, None] * stride_gu_k
+        + (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))[None, :] * stride_gu_n
+    )
+    b_up_ptr = (
+        GateUp
+        + expert_id * stride_gu_e
+        + tl.arange(0, BLOCK_SIZE_K // VALUES_PER_BYTE)[:, None] * stride_gu_k
+        + (INTERMEDIATE_DIM + pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))[
+            None, :
+        ]
+        * stride_gu_n
+    )
+    bs_gate_ptr = (
+        GateUpScale
+        + expert_id * stride_gus_e
+        + (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))[:, None] * stride_gus_n
+        + tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K)[None, :] * stride_gus_k
+    )
+    bs_up_ptr = (
+        GateUpScale
+        + expert_id * stride_gus_e
+        + (INTERMEDIATE_DIM + pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))[
+            :, None
+        ]
+        * stride_gus_n
+        + tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K)[None, :] * stride_gus_k
+    )
+    acc_gate = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    acc_up = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for _ in range(0, tl.cdiv(HIDDEN_DIM, BLOCK_SIZE_K)):
+        a_raw = tl.load(a_ptr).to(tl.float32)
+        a, a_scale = mxfp_act_quant_inline(
+            a_raw, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K
+        )
+        b_gate = tl.load(b_gate_ptr)
+        b_up = tl.load(b_up_ptr)
+        bs_gate = tl.load(bs_gate_ptr)
+        bs_up = tl.load(bs_up_ptr)
+        acc_gate, acc_up = mx_compute_gate_up(
+            acc_gate,
+            acc_up,
+            a,
+            a_scale,
+            b_gate,
+            b_up,
+            bs_gate,
+            bs_up,
+            COMPUTE_MODE,
+            VALUES_PER_BYTE,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            SCALE_GROUP_K,
+        )
+        a_ptr += BLOCK_SIZE_K * stride_a_k
+        b_gate_ptr += (BLOCK_SIZE_K // VALUES_PER_BYTE) * stride_gu_k
+        b_up_ptr += (BLOCK_SIZE_K // VALUES_PER_BYTE) * stride_gu_k
+        bs_gate_ptr += (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_gus_k
+        bs_up_ptr += (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_gus_k
+    intermediate = glu(
+        acc_gate, acc_up, ACT_FN, SIMULATE_UNFUSED, HiddenStates.dtype.element_ty
+    )
+    # MXFP8 requant of the intermediate (E4M3 + UE8M0 group-32 along this N-tile).
+    inter, inter_scale = mxfp_act_quant_inline(
+        intermediate, BLOCK_SIZE_M, BLOCK_SIZE_N, SCALE_GROUP_K
+    )
+    store_row(
+        Intermediate + batch_id * INTERMEDIATE_DIM,
+        inter,
+        pid_n,
+        1,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+    )
+    offs_sc = pid_n * (BLOCK_SIZE_N // SCALE_GROUP_K) + tl.arange(
+        0, BLOCK_SIZE_N // SCALE_GROUP_K
+    )
+    sc_ptrs = (
+        IntermediateScale
+        + batch_id * (INTERMEDIATE_DIM // SCALE_GROUP_K)
+        + offs_sc[None, :]
+    )
+    tl.store(sc_ptrs, inter_scale)
+@bayesian_autotune(
+    get_mxfp_autotuning_configs(),
+    ["num_routed_tokens"],
+    n_trials=60,
+)
+@triton.jit
+def mxfp_dynamic_moe_batched_down_kernel(
+    Intermediate,
+    IntermediateScale,
+    Down,
+    DownScale,
+    ExpertIds,
+    SampleWeights,
+    Out,
+    stride_down_e,
+    stride_down_n,
+    stride_down_k,
+    stride_downs_e,
+    stride_downs_n,
+    stride_downs_k,
+    num_routed_tokens,  # tuning knob: in the autotune key so decode re-tunes per batch
+    NUM_EXPERTS: tl.constexpr,
+    HIDDEN_DIM: tl.constexpr,
+    INTERMEDIATE_DIM: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    VALUES_PER_BYTE: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    SIMULATE_UNFUSED: tl.constexpr,
+    COMPUTE_MODE: tl.constexpr,
+):
+    """MXFP4/MXFP8 kernel 2: MXFP8 intermediate ��� down proj. N = hidden
+    (output) tile, K = intermediate (contraction) tile. Grid: (S, H-tiles)."""
+    batch_id = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    expert_id = tl.load(ExpertIds + batch_id).to(tl.int64)
+    if expert_id >= NUM_EXPERTS:  # EP sentinel: zero the output row for the top-k sum
+        z = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+        store_row(Out + batch_id * HIDDEN_DIM, z, pid_n, 1, BLOCK_SIZE_M, BLOCK_SIZE_N)
+        return
+    a_ptr = (
+        Intermediate
+        + (batch_id + tl.arange(0, BLOCK_SIZE_M))[:, None] * INTERMEDIATE_DIM
+        + tl.arange(0, BLOCK_SIZE_K)[None, :]
+    )
+    as_ptr = (
+        IntermediateScale
+        + (batch_id + tl.arange(0, BLOCK_SIZE_M))[:, None]
+        * (INTERMEDIATE_DIM // SCALE_GROUP_K)
+        + tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K)[None, :]
+    )
+    w_down_ptr = (
+        Down
+        + expert_id * stride_down_e
+        + tl.arange(0, BLOCK_SIZE_K // VALUES_PER_BYTE)[:, None] * stride_down_k
+        + (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))[None, :] * stride_down_n
+    )
+    ws_down_ptr = (
+        DownScale
+        + expert_id * stride_downs_e
+        + (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))[:, None] * stride_downs_n
+        + tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K)[None, :] * stride_downs_k
+    )
+    acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for _ in range(0, tl.cdiv(INTERMEDIATE_DIM, BLOCK_SIZE_K)):
+        a = tl.load(a_ptr)
+        a_scale = tl.load(as_ptr)
+        w = tl.load(w_down_ptr)
+        w_scale = tl.load(ws_down_ptr)
+        acc = mx_compute(
+            acc,
+            a,
+            a_scale,
+            w,
+            w_scale,
+            COMPUTE_MODE,
+            VALUES_PER_BYTE,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            SCALE_GROUP_K,
+        )
+        a_ptr += BLOCK_SIZE_K
+        as_ptr += BLOCK_SIZE_K // SCALE_GROUP_K
+        w_down_ptr += (BLOCK_SIZE_K // VALUES_PER_BYTE) * stride_down_k
+        ws_down_ptr += (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_downs_k
+    if SIMULATE_UNFUSED:
+        acc = acc.to(Out.dtype.element_ty).to(tl.float32)
+    acc = acc * tl.load(SampleWeights + batch_id)
+    store_row(Out + batch_id * HIDDEN_DIM, acc, pid_n, 1, BLOCK_SIZE_M, BLOCK_SIZE_N)
+@triton_op(add_op_namespace_prefix("mxfp_dynamic_moe_batched"), mutates_args=())
+def _mxfp_dynamic_moe_batched(
+    hidden_states: torch.Tensor,
+    gate_up_proj: torch.Tensor,
+    gate_up_proj_scale: torch.Tensor,
+    down_proj: torch.Tensor,
+    down_proj_scale: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    output_dtype: torch.dtype,
+    simulate_unfused: bool = False,
+    act_fn: str = "silu",
+) -> torch.Tensor:
+    """MXFP4/MXFP8 batched fused MoE in ONE op: gate_up + SiLU + MXFP8 requant → grouped
+    down → routing-weighted per-(token, expert) output. gate_up and down must share the same MX
+    format. ``inter``/``inter_scales`` are internal; caller reduces over top-k."""
+    gate_up_is_fp4 = is_mxfp4(gate_up_proj, gate_up_proj_scale)
+    down_is_fp4 = is_mxfp4(down_proj, down_proj_scale)
+    if gate_up_is_fp4 != down_is_fp4:
+        raise ValueError(
+            "gate_up_proj and down_proj must use the same MX format (both MXFP4 or both MXFP8)."
+        )
+    device = hidden_states.device
+    HIDDEN_DIM = hidden_states.size(1)
+    NUM_EXPERTS = gate_up_proj.size(0)
+    num_tokens = hidden_states.size(0)
+    num_routed_tokens = top_k_index.numel()
+    INTERMEDIATE_DIM = gate_up_proj.size(1) // 2
+    NUM_TOP_K = num_routed_tokens // hidden_states.size(0)
+    VALUES_PER_BYTE = NIBBLES_PER_BYTE if gate_up_is_fp4 else 1
+    gate_up_proj_u8 = e2m1_as_uint8(gate_up_proj)
+    gate_up_proj_scale_u8 = ue8m0_as_uint8(gate_up_proj_scale)
+    down_proj_u8 = e2m1_as_uint8(down_proj)
+    down_proj_scale_u8 = ue8m0_as_uint8(down_proj_scale)
+    inter = torch.empty(
+        num_routed_tokens, INTERMEDIATE_DIM, device=device, dtype=FP8_DTYPE
+    )
+    inter_scales = torch.empty(
+        num_routed_tokens,
+        INTERMEDIATE_DIM // MX_SCALE_GROUP_K,
+        device=device,
+        dtype=torch.uint8,
+    )
+    out = torch.empty(num_routed_tokens, HIDDEN_DIM, device=device, dtype=output_dtype)
+    reduced = torch.empty(num_tokens, HIDDEN_DIM, device=device, dtype=output_dtype)
+    def gate_up_grid(META):
+        return (num_routed_tokens, triton.cdiv(INTERMEDIATE_DIM, META["BLOCK_SIZE_N"]))
+    def down_grid(META):
+        return (num_routed_tokens, triton.cdiv(HIDDEN_DIM, META["BLOCK_SIZE_N"]))
+    with device_context(device):
+        wrap_triton(mxfp_dynamic_moe_batched_gate_up_kernel)[gate_up_grid](
+            hidden_states,
+            gate_up_proj_u8,
+            gate_up_proj_scale_u8,
+            inter,
+            inter_scales,
+            top_k_index,
+            hidden_states.stride(0),
+            hidden_states.stride(1),
+            gate_up_proj_u8.stride(0),
+            gate_up_proj_u8.stride(1),
+            gate_up_proj_u8.stride(2),
+            gate_up_proj_scale_u8.stride(0),
+            gate_up_proj_scale_u8.stride(1),
+            gate_up_proj_scale_u8.stride(2),
+            num_routed_tokens=num_routed_tokens,
+            NUM_TOP_K=NUM_TOP_K,
+            NUM_EXPERTS=NUM_EXPERTS,
+            HIDDEN_DIM=HIDDEN_DIM,
+            INTERMEDIATE_DIM=INTERMEDIATE_DIM,
+            BLOCK_SIZE_M=DECODE_BLOCK_SIZE_M,
+            VALUES_PER_BYTE=VALUES_PER_BYTE,
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+            SIMULATE_UNFUSED=simulate_unfused,
+            ACT_FN=act_fn,
+        )
+        wrap_triton(mxfp_dynamic_moe_batched_down_kernel)[down_grid](
+            inter,
+            inter_scales,
+            down_proj_u8,
+            down_proj_scale_u8,
+            top_k_index,
+            top_k_weights,
+            out,
+            down_proj_u8.stride(0),
+            down_proj_u8.stride(1),
+            down_proj_u8.stride(2),
+            down_proj_scale_u8.stride(0),
+            down_proj_scale_u8.stride(1),
+            down_proj_scale_u8.stride(2),
+            num_routed_tokens=num_routed_tokens,
+            NUM_EXPERTS=NUM_EXPERTS,
+            HIDDEN_DIM=HIDDEN_DIM,
+            INTERMEDIATE_DIM=INTERMEDIATE_DIM,
+            BLOCK_SIZE_M=DECODE_BLOCK_SIZE_M,
+            VALUES_PER_BYTE=VALUES_PER_BYTE,
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+            SIMULATE_UNFUSED=simulate_unfused,
+        )
+        wrap_triton(topk_reduce_kernel)[
+            (num_tokens, triton.cdiv(HIDDEN_DIM, TOPK_REDUCE_BLOCK_H))
+        ](
+            out,
+            reduced,
+            HIDDEN_DIM,
+            out.stride(0),
+            out.stride(1),
+            reduced.stride(0),
+            reduced.stride(1),
+            NUM_TOP_K=NUM_TOP_K,
+            BLOCK_H=TOPK_REDUCE_BLOCK_H,
+        )
+    return reduced
+def mxfp_dynamic_moe_batched(
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    gate_up_proj: torch.Tensor,
+    down_proj: torch.Tensor,
+    gate_up_proj_scale: torch.Tensor,
+    down_proj_scale: torch.Tensor,
+    simulate_unfused: bool = False,
+    act_fn: str = "silu",
+) -> torch.Tensor:
+    """Two-kernel batched fused MX MoE — MXFP4 or MXFP8 weights (UE8M0 group-32), the
+    format picked per-weight by the ops. Same structure as the block-dynamic path but
+    with a tunable tile and an MXFP8 group-32 intermediate; ``block_size`` is unused."""
+    out = ops.mxfp_dynamic_moe_batched(
+        hidden_states,
+        gate_up_proj,
+        gate_up_proj_scale,
+        down_proj,
+        down_proj_scale,
+        top_k_index,
+        top_k_weights,
+        hidden_states.dtype,
+        simulate_unfused,
+        act_fn,
+    )
+    return out
+# ── Dispatcher ────────────────────────────────────────────────────────────────
+def moe_fused_batched(
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    gate_up_proj: torch.Tensor,
+    down_proj: torch.Tensor,
+    gate_up_proj_scale_inv: torch.Tensor,
+    down_proj_scale_inv: torch.Tensor,
+    block_size: list[int] | None,
+    simulate_unfused: bool = False,
+    act_fn: str = "silu",
+) -> torch.Tensor:
+    """Fused batched-MoE dispatcher — routes to the recipe matching the weight dtype /
+    scale layout, mirroring ``matmul_batched``. Implemented: block-dynamic FP8 and MXFP8
+    (tensor-dynamic) and MXFP4. ``simulate_unfused`` (testing) rounds each step through the
+    activation dtype so the output matches the unfused reference to reduce order."""
+    gate_up_is_mx = is_mxfp(gate_up_proj, gate_up_proj_scale_inv)
+    down_is_mx = is_mxfp(down_proj, down_proj_scale_inv)
+    if gate_up_is_mx != down_is_mx:
+        raise ValueError(
+            "gate_up_proj and down_proj must use the same recipe (both MX or both block-dynamic FP8)."
+        )
+    if gate_up_is_mx:
+        return mxfp_dynamic_moe_batched(
+            hidden_states,
+            top_k_index,
+            top_k_weights,
+            gate_up_proj,
+            down_proj,
+            gate_up_proj_scale_inv,
+            down_proj_scale_inv,
+            simulate_unfused,
+            act_fn,
+        )
+    if block_size is None:
+        raise ValueError("block_size is required for block-dynamic FP8 weights.")
+    return w8a8_block_dynamic_fp8_moe_batched(
+        hidden_states,
+        top_k_index,
+        top_k_weights,
+        gate_up_proj,
+        down_proj,
+        gate_up_proj_scale_inv,
+        down_proj_scale_inv,
+        block_size,
+        simulate_unfused,
+        act_fn,
+    )

build/torch-cuda/fused_grouped.py ADDED Viewed

	@@ -0,0 +1,1110 @@

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+"""Fused grouped MoE (prefill counterpart to the fused batched/decode path).
+Two **persistent** kernels over expert-grouped M-tiles, no ``A_sorted`` materialization:
+  1. Routing (on-device): an atomic counting-sort scatter (``_scatter_kernel``, O(S)) →
+     ``perm`` / ``perm_token`` + an ``expert_start`` exclusive offset vector — the sort is
+     just an index; activations stay unsorted.
+  2. Phase 1 — gate_up, **gathering** each tile's hidden rows via ``perm_token``, + SiLU,
+     then **FP8 requant** of the intermediate — same trick as ``fused_batched``: the down
+     kernel reads FP8 directly instead of a bf16 round-trip.
+  3. Phase 2 — grouped down over the FP8 intermediate, then **fuses the routing-weight ×
+     reorder** (no atomics): each row is scaled by its weight and **scattered** to its flat
+     ``(token, slot)`` row via ``perm``, leaving the host only a top-k sum back to (T, H).
+Scheduling is persistent: a fixed ``NUM_SMS`` programs grid-stride over all tiles. The
+expert-tile layout is built once into registers from ``expert_start`` (``_build_tile_layout``)
+and each tile's owner is resolved inline (``_resolve_tile_inline``) — no padded grid, no
+per-tile binary search, no host-side tile-offset precompute.
+Recipe-named to mirror ``fused_batched``: ``w8a8_block_dynamic_fp8_moe_grouped`` (128x128
+block scales, FP8 intermediate) and ``mxfp_dynamic_moe_grouped`` (MXFP4/MXFP8, UE8M0
+group-32 scales, MXFP8 intermediate); ``moe_fused_grouped`` is the neutral dispatcher.
+"""
+import torch
+import triton
+import triton.language as tl
+from triton.tools.tensor_descriptor import TensorDescriptor
+from .bayesian_autotuner import bayesian_autotune
+from .grouped import store_tile
+from .utils import (
+    FP8_DTYPE,
+    MX_SCALE_GROUP_K,
+    NIBBLES_PER_BYTE,
+    decode_ue8m0_scale,
+    device_context,
+    sm_count,
+    fp8_act_quant_inline,
+    topk_reduce_kernel,
+    TOPK_REDUCE_BLOCK_H,
+    get_accelerator_autotuning_configs,
+    get_mxfp_autotuning_configs,
+    is_mxfp,
+    is_mxfp4,
+    mxfp_act_quant_inline,
+    mx_compute,
+    smem_config_pruner,
+    glu,
+    e2m1_as_uint8,
+    ue8m0_as_uint8,
+)
+# ── Persistent tile scheduling (register-resident expert-tile layout) ─────────
+# Flat-slot tile per program for the O(S) routing kernels (count + scatter). These are small
+# latency-bound atomic kernels that want many programs: a sweep over {256..4096} x prefill shapes
+# put 256 best (or within ~1%) for both, with 1024 up to ~1.5x slower. The grid derives from it
+# so the two can't drift. Power of 2.
+_ROUTING_BLOCK_SIZE = 256
+@triton.jit
+def _build_tile_layout(
+    ExpertStart, NUM_EXPERTS: tl.constexpr, BLOCK_SIZE_M: tl.constexpr
+):
+    """Load ``expert_start`` once and derive the per-BM tile layout vectors (kept in
+    registers for the whole persistent loop): per-expert first sorted row, token count,
+    exclusive tile-start cumsum, and the total M-tile count. ``ExpertStart`` is
+    ``(NUM_EXPERTS + 1,)`` with a trailing ``S`` sentinel (``expert_start[E] == S``)."""
+    e_offs = tl.arange(0, NUM_EXPERTS)
+    exp_start = tl.load(ExpertStart + e_offs)
+    exp_end = tl.load(ExpertStart + e_offs + 1)
+    freqs = exp_end - exp_start
+    tiles_per_e = (freqs + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M
+    tile_start_excl = (
+        tl.cumsum(tiles_per_e, 0) - tiles_per_e
+    )  # first tile index of expert e
+    total_m_tiles = tl.sum(tiles_per_e, 0)
+    return exp_start, freqs, tile_start_excl, total_m_tiles, e_offs
+@triton.jit
+def _resolve_tile_inline(
+    pid_m, exp_start, freqs, tile_start_excl, e_offs, BLOCK_SIZE_M: tl.constexpr
+):
+    """Map an M-tile id to its owning expert + the tile's sorted row range, from the
+    register-resident layout (no global loads). Returns ``(expert_id, sorted_indices,
+    row_mask)``."""
+    # Bucketize via the exclusive tile cumsum: #experts whose tile-start <= pid_m, minus 1.
+    expert_id = tl.sum((tile_start_excl <= pid_m).to(tl.int32), 0) - 1
+    sel = (
+        e_offs == expert_id
+    )  # scalar-index the E-vectors via mask-sum (no dynamic index)
+    e_start = tl.sum(tl.where(sel, exp_start, 0), 0)
+    e_tile_start = tl.sum(tl.where(sel, tile_start_excl, 0), 0)
+    freq = tl.sum(tl.where(sel, freqs, 0), 0)
+    within = pid_m - e_tile_start
+    m_start = e_start + within * BLOCK_SIZE_M
+    offs = tl.arange(0, BLOCK_SIZE_M)
+    row_mask = offs < freq - within * BLOCK_SIZE_M
+    sorted_indices = tl.max_contiguous(m_start + offs, BLOCK_SIZE_M)
+    return expert_id, sorted_indices, row_mask
+@triton.jit
+def _exclusive_offsets_kernel(
+    ExpertFreq, ExpertStart, Counters, NUM_EXPERTS: tl.constexpr
+):
+    """Exclusive cumsum of per-expert token counts → ``expert_start`` (leading 0, trailing
+    S), and zero the scatter counters — one launch."""
+    offs = tl.arange(0, NUM_EXPERTS)
+    incl = tl.cumsum(tl.load(ExpertFreq + offs), 0)
+    tl.store(ExpertStart, 0)
+    tl.store(ExpertStart + 1 + offs, incl)
+    tl.store(Counters + offs, tl.zeros([NUM_EXPERTS], tl.int32))
+@triton.jit
+def _scatter_kernel(
+    ExpertIds,
+    Perm,
+    PermToken,
+    ExpertStart,
+    Counters,
+    S,
+    NUM_TOP_K: tl.constexpr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    """Counting-sort scatter: each flat slot atomically claims the next slot of its expert
+    (``expert_start[e] + counter[e]++``). O(S), replaces an O(S·logS) argsort. Within-expert
+    order is arbitrary (atomic race) — fine, the per-token reduce is order-invariant."""
+    offs = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offs < S
+    expert_id = tl.load(ExpertIds + offs, mask=mask, other=0)
+    dest = tl.load(ExpertStart + expert_id, mask=mask, other=0) + tl.atomic_add(
+        Counters + expert_id, 1, mask=mask
+    )
+    tl.store(Perm + dest, offs, mask=mask)
+    tl.store(PermToken + dest, offs // NUM_TOP_K, mask=mask)
+@triton.jit
+def _count_kernel(
+    ExpertIds, ExpertFreq, S, NUM_EXPERTS: tl.constexpr, BLOCK_SIZE: tl.constexpr
+):
+    """Per-expert token count via atomics — replaces ``torch.histc`` (no float cast), fixed
+    ``(NUM_EXPERTS,)`` output stays CUDA-graph friendly. ``ExpertFreq`` is pre-zeroed."""
+    offs = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offs < S
+    expert_id = tl.load(ExpertIds + offs, mask=mask, other=NUM_EXPERTS)
+    tl.atomic_add(ExpertFreq + expert_id, 1, mask=mask & (expert_id < NUM_EXPERTS))
+def _grouped_routing(expert_ids: torch.Tensor, num_experts: int, num_top_k: int):
+    """On-device routing: expert-sorted index (no copy of the activations) via two Triton
+    launches — exclusive offsets + an atomic counting-sort scatter (replaces host
+    ``argsort``). Returns ``(perm_token, perm, expert_start, num_experts, num_routed_tokens)``.
+    ``perm`` is the sorted-position → flat
+    ``(t*K + j)`` map (gate_up gathers via ``perm_token = perm // K``, down scatters via
+    ``perm``); ``expert_start`` is ``(E+1,)`` padded with S so the kernels build the tile
+    layout in-register (E is a power of 2)."""
+    device = expert_ids.device
+    expert_ids = expert_ids.int()
+    num_routed_tokens = expert_ids.numel()  # S = num_tokens * num_top_k
+    expert_freq = torch.zeros(num_experts, dtype=torch.int32, device=device)
+    expert_start = torch.empty(num_experts + 1, dtype=torch.int32, device=device)
+    counters = torch.empty(num_experts, dtype=torch.int32, device=device)
+    perm = torch.empty(num_routed_tokens, dtype=torch.int32, device=device)
+    perm_token = torch.empty(num_routed_tokens, dtype=torch.int32, device=device)
+    with device_context(device):
+        _count_kernel[(triton.cdiv(num_routed_tokens, _ROUTING_BLOCK_SIZE),)](
+            expert_ids,
+            expert_freq,
+            num_routed_tokens,
+            NUM_EXPERTS=num_experts,
+            BLOCK_SIZE=_ROUTING_BLOCK_SIZE,
+        )
+        _exclusive_offsets_kernel[(1,)](
+            expert_freq,
+            expert_start,
+            counters,
+            NUM_EXPERTS=num_experts,
+        )
+        _scatter_kernel[(triton.cdiv(num_routed_tokens, _ROUTING_BLOCK_SIZE),)](
+            expert_ids,
+            perm,
+            perm_token,
+            expert_start,
+            counters,
+            num_routed_tokens,
+            NUM_TOP_K=num_top_k,
+            BLOCK_SIZE=_ROUTING_BLOCK_SIZE,
+        )
+    return perm_token, perm, expert_start, num_experts, num_routed_tokens
+# MX MoE autotune axes: the MMA flavor and the weight-load mechanism. The kernels implement all
+# three memory modes; get_mxfp_autotuning_configs drops the device-inapplicable descriptor
+# flavor (host on XPU, device on CUDA).
+_MX_COMPUTE_MODES = ("dot_scaled", "dot")
+_MX_MEMORY_MODES = ("host_descriptor", "device_descriptor", "pointer")
+# ── Block-dynamic FP8 ────────────────────────────────────────────────────────
+@bayesian_autotune(
+    get_accelerator_autotuning_configs(tune_block_m=True),
+    ["INTERMEDIATE_DIM", "HIDDEN_DIM", "tokens_per_sm_bit_length"],
+    n_trials=60,
+)
+@triton.jit
+def w8a8_block_dynamic_fp8_moe_grouped_gate_up_kernel(
+    Hidden,  # (T, H) raw activations, UNSORTED
+    PermToken,  # (S,) int32 — sorted position -> source token id
+    GateUp,  # (E, 2I, H) FP8
+    GateUpScale,  # (E, 2I//bn, H//bk) UE8M0 block scales
+    ExpertStart,  # (NUM_EXPERTS+1,) int32 — exclusive sorted-row start per expert (pad S)
+    Inter,  # (S, I) FP8 — output (expert-sorted)
+    InterScale,  # (S, NUM_I_TILES) fp32 — per-row, per-I-tile activation scale
+    stride_h_t,
+    stride_h_k,
+    stride_gu_e,
+    stride_gu_n,
+    stride_gu_k,
+    stride_gus_e,
+    stride_gus_n,
+    stride_gus_k,
+    stride_int_m,
+    stride_int_n,
+    stride_is_m,
+    stride_is_n,
+    stride_pt,
+    tokens_per_sm_bit_length,  # autotune key only (log2 tokens-per-SM bucket); unused in body
+    HIDDEN_DIM: tl.constexpr,
+    INTERMEDIATE_DIM: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    NUM_EXPERTS: tl.constexpr,
+    NUM_SMS: tl.constexpr,
+    SIMULATE_UNFUSED: tl.constexpr,
+    ACT_FN: tl.constexpr,
+):
+    """Phase 1: persistent grid-stride over (M-tile, I-tile). Gather hidden rows per expert
+    M-tile, gate + up block-FP8 matmuls, SiLU-combine, FP8-requant the intermediate."""
+    start_pid = tl.program_id(axis=0)
+    exp_start, freqs, tile_start_excl, total_m_tiles, e_offs = _build_tile_layout(
+        ExpertStart, NUM_EXPERTS, BLOCK_SIZE_M
+    )
+    num_n_tiles = tl.cdiv(INTERMEDIATE_DIM, BLOCK_SIZE_N)
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    for tile_id in tl.range(start_pid, total_m_tiles * num_n_tiles, NUM_SMS):
+        pid_m = tile_id // num_n_tiles
+        pid_n = tile_id % num_n_tiles
+        expert_id, offs_global_m, row_mask = _resolve_tile_inline(
+            pid_m, exp_start, freqs, tile_start_excl, e_offs, BLOCK_SIZE_M
+        )
+        offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+        token = tl.load(PermToken + offs_global_m * stride_pt, mask=row_mask, other=0)
+        a_ptrs = Hidden + token[:, None] * stride_h_t + offs_k[None, :] * stride_h_k
+        gate_ptr = (
+            GateUp
+            + expert_id * stride_gu_e
+            + tl.arange(0, BLOCK_SIZE_K)[:, None] * stride_gu_k
+            + (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))[None, :] * stride_gu_n
+        )
+        up_ptr = (
+            GateUp
+            + expert_id * stride_gu_e
+            + tl.arange(0, BLOCK_SIZE_K)[:, None] * stride_gu_k
+            + (INTERMEDIATE_DIM + pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))[
+                None, :
+            ]
+            * stride_gu_n
+        )
+        gate_s_ptr = GateUpScale + expert_id * stride_gus_e + pid_n * stride_gus_n
+        up_s_ptr = (
+            GateUpScale
+            + expert_id * stride_gus_e
+            + (num_n_tiles + pid_n) * stride_gus_n
+        )
+        acc_gate = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+        acc_up = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+        for _ in range(0, tl.cdiv(HIDDEN_DIM, BLOCK_SIZE_K)):
+            a_raw = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0).to(tl.float32)
+            a, a_s = fp8_act_quant_inline(a_raw)
+            w_gate = tl.load(gate_ptr)
+            w_up = tl.load(up_ptr)
+            w_s_gate = decode_ue8m0_scale(tl.load(gate_s_ptr))
+            w_s_up = decode_ue8m0_scale(tl.load(up_s_ptr))
+            acc_gate += tl.dot(a, w_gate) * a_s[:, None] * w_s_gate
+            acc_up += tl.dot(a, w_up) * a_s[:, None] * w_s_up
+            a_ptrs += BLOCK_SIZE_K * stride_h_k
+            gate_ptr += BLOCK_SIZE_K * stride_gu_k
+            up_ptr += BLOCK_SIZE_K * stride_gu_k
+            gate_s_ptr += stride_gus_k
+            up_s_ptr += stride_gus_k
+        intermediate = glu(
+            acc_gate, acc_up, ACT_FN, SIMULATE_UNFUSED, Hidden.dtype.element_ty
+        )
+        inter, inter_s = fp8_act_quant_inline(
+            intermediate
+        )  # FP8 requant, per-row scale
+        int_ptrs = (
+            Inter
+            + offs_global_m[:, None] * stride_int_m
+            + offs_bn[None, :] * stride_int_n
+        )
+        tl.store(int_ptrs, inter, mask=row_mask[:, None])
+        tl.store(
+            InterScale + offs_global_m * stride_is_m + pid_n * stride_is_n,
+            inter_s,
+            mask=row_mask,
+        )
+@bayesian_autotune(
+    get_accelerator_autotuning_configs(tune_block_m=True),
+    ["INTERMEDIATE_DIM", "HIDDEN_DIM", "tokens_per_sm_bit_length"],
+    n_trials=60,
+)
+@triton.jit
+def w8a8_block_dynamic_fp8_moe_grouped_down_kernel(
+    Inter,  # (S, I) FP8 — expert-sorted intermediate
+    InterScale,  # (S, NUM_I_TILES) fp32 — per-row, per-I-tile scale
+    Down,  # (E, H, I) FP8
+    DownScale,  # (E, H//bn, I//bk) UE8M0 block scales
+    ExpertStart,  # (NUM_EXPERTS+1,) int32 — exclusive sorted-row start per expert (pad S)
+    Perm,  # (S,) int32 — sorted position -> flat (token, slot) row
+    SampleWeights,  # (S,) routing weight per flat row
+    ProjOut,  # (S, H) bf16, flat (token, slot) order
+    stride_int_m,
+    stride_int_n,
+    stride_is_m,
+    stride_is_n,
+    stride_down_e,
+    stride_down_h,
+    stride_down_i,
+    stride_downs_e,
+    stride_downs_h,
+    stride_downs_i,
+    stride_po_m,
+    stride_po_n,
+    stride_perm,
+    tokens_per_sm_bit_length,  # autotune key only (log2 tokens-per-SM bucket); unused in body
+    HIDDEN_DIM: tl.constexpr,
+    INTERMEDIATE_DIM: tl.constexpr,
+    NUM_I_TILES: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_H: tl.constexpr,
+    BLOCK_SIZE_I: tl.constexpr,
+    NUM_EXPERTS: tl.constexpr,
+    NUM_SMS: tl.constexpr,
+    SIMULATE_UNFUSED: tl.constexpr,
+):
+    """Phase 2: persistent grid-stride over (M-tile, H-tile). Grouped down over the FP8
+    intermediate, then routing-weight × scatter to the flat (token, slot) row."""
+    start_pid = tl.program_id(axis=0)
+    exp_start, freqs, tile_start_excl, total_m_tiles, e_offs = _build_tile_layout(
+        ExpertStart, NUM_EXPERTS, BLOCK_SIZE_M
+    )
+    num_h_tiles = tl.cdiv(HIDDEN_DIM, BLOCK_SIZE_H)
+    offs_i = tl.arange(0, BLOCK_SIZE_I)
+    for tile_id in tl.range(start_pid, total_m_tiles * num_h_tiles, NUM_SMS):
+        pid_m = tile_id // num_h_tiles
+        pid_h = tile_id % num_h_tiles
+        expert_id, offs_global_m, row_mask = _resolve_tile_inline(
+            pid_m, exp_start, freqs, tile_start_excl, e_offs, BLOCK_SIZE_M
+        )
+        offs_h = pid_h * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H)
+        w_down_ptr = (
+            Down
+            + expert_id * stride_down_e
+            + tl.arange(0, BLOCK_SIZE_I)[:, None] * stride_down_i
+            + (pid_h * BLOCK_SIZE_H + tl.arange(0, BLOCK_SIZE_H))[None, :]
+            * stride_down_h
+        )
+        ws_down_ptr = DownScale + expert_id * stride_downs_e + pid_h * stride_downs_h
+        acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_H), dtype=tl.float32)
+        for i_tile in range(0, NUM_I_TILES):
+            i_off = i_tile * BLOCK_SIZE_I
+            inter = tl.load(
+                Inter
+                + offs_global_m[:, None] * stride_int_m
+                + (i_off + offs_i)[None, :] * stride_int_n,
+                mask=row_mask[:, None],
+                other=0.0,
+            )
+            inter_s = tl.load(
+                InterScale + offs_global_m * stride_is_m + i_tile * stride_is_n,
+                mask=row_mask,
+                other=0.0,
+            )
+            w_s_down = decode_ue8m0_scale(tl.load(ws_down_ptr))
+            w_down = tl.load(w_down_ptr)
+            acc += tl.dot(inter, w_down) * inter_s[:, None] * w_s_down
+            w_down_ptr += BLOCK_SIZE_I * stride_down_i
+            ws_down_ptr += stride_downs_i
+        if SIMULATE_UNFUSED:
+            acc = acc.to(ProjOut.dtype.element_ty).to(tl.float32)
+        # Fused routing-weight × top-k reorder (no atomics): scale each row by its weight
+        # and scatter to its flat (token, slot) row; the host then just sums over slots.
+        flat = tl.load(Perm + offs_global_m * stride_perm, mask=row_mask, other=0)
+        weight = tl.load(SampleWeights + flat, mask=row_mask, other=0.0)
+        acc = acc * weight[:, None]
+        store_tile(ProjOut, acc, flat, offs_h, row_mask, stride_po_m, stride_po_n)
+def w8a8_block_dynamic_fp8_moe_grouped(
+    hidden_states: torch.Tensor,  # (T, H)
+    top_k_index: torch.Tensor,  # (T, K) int
+    top_k_weights: torch.Tensor,  # (T, K)
+    gate_up_proj: torch.Tensor,  # (E, 2I, H) FP8
+    down_proj: torch.Tensor,  # (E, H, I) FP8
+    gate_up_proj_scale: torch.Tensor,
+    down_proj_scale: torch.Tensor,
+    block_size: list[int],
+    simulate_unfused: bool = False,
+    act_fn: str = "silu",
+) -> torch.Tensor:
+    """Block-dynamic FP8 fused grouped MoE: gather gate_up+SiLU → FP8 intermediate →
+    grouped down → routing-weighted top-k reduce. Returns ``(num_tokens, hidden_dim)``."""
+    device = hidden_states.device
+    num_tokens = hidden_states.size(0)
+    num_top_k = top_k_index.size(-1)
+    HIDDEN_DIM = hidden_states.size(1)
+    INTERMEDIATE_DIM = down_proj.size(2)
+    BLOCK_SIZE_N, BLOCK_SIZE_K = block_size
+    NUM_I_TILES = INTERMEDIATE_DIM // BLOCK_SIZE_N
+    perm_token, perm, expert_start, NUM_EXPERTS, num_routed_tokens = _grouped_routing(
+        top_k_index, gate_up_proj.size(0), num_top_k
+    )
+    num_sms = sm_count(device.index)
+    tokens_per_sm_bit_length = (num_routed_tokens // num_sms).bit_length()
+    gate_up_scale_u8 = ue8m0_as_uint8(gate_up_proj_scale)
+    down_scale_u8 = ue8m0_as_uint8(down_proj_scale)
+    inter = torch.empty(
+        num_routed_tokens, INTERMEDIATE_DIM, device=device, dtype=FP8_DTYPE
+    )
+    inter_scale = torch.empty(
+        num_routed_tokens, NUM_I_TILES, device=device, dtype=torch.float32
+    )
+    out = torch.empty(
+        num_routed_tokens, HIDDEN_DIM, device=device, dtype=hidden_states.dtype
+    )
+    reduced = torch.empty(
+        num_tokens, HIDDEN_DIM, device=device, dtype=hidden_states.dtype
+    )
+    with device_context(device):
+        w8a8_block_dynamic_fp8_moe_grouped_gate_up_kernel[(num_sms,)](
+            hidden_states,
+            perm_token,
+            gate_up_proj,
+            gate_up_scale_u8,
+            expert_start,
+            inter,
+            inter_scale,
+            hidden_states.stride(0),
+            hidden_states.stride(1),
+            gate_up_proj.stride(0),
+            gate_up_proj.stride(1),
+            gate_up_proj.stride(2),
+            gate_up_scale_u8.stride(0),
+            gate_up_scale_u8.stride(1),
+            gate_up_scale_u8.stride(2),
+            inter.stride(0),
+            inter.stride(1),
+            inter_scale.stride(0),
+            inter_scale.stride(1),
+            perm_token.stride(0),
+            tokens_per_sm_bit_length=tokens_per_sm_bit_length,
+            HIDDEN_DIM=HIDDEN_DIM,
+            INTERMEDIATE_DIM=INTERMEDIATE_DIM,
+            BLOCK_SIZE_N=BLOCK_SIZE_N,
+            BLOCK_SIZE_K=BLOCK_SIZE_K,
+            NUM_EXPERTS=NUM_EXPERTS,
+            NUM_SMS=num_sms,
+            SIMULATE_UNFUSED=simulate_unfused,
+            ACT_FN=act_fn,
+        )
+        w8a8_block_dynamic_fp8_moe_grouped_down_kernel[(num_sms,)](
+            inter,
+            inter_scale,
+            down_proj,
+            down_scale_u8,
+            expert_start,
+            perm,
+            top_k_weights,
+            out,
+            inter.stride(0),
+            inter.stride(1),
+            inter_scale.stride(0),
+            inter_scale.stride(1),
+            down_proj.stride(0),
+            down_proj.stride(1),
+            down_proj.stride(2),
+            down_scale_u8.stride(0),
+            down_scale_u8.stride(1),
+            down_scale_u8.stride(2),
+            out.stride(0),
+            out.stride(1),
+            perm.stride(0),
+            tokens_per_sm_bit_length=tokens_per_sm_bit_length,
+            HIDDEN_DIM=HIDDEN_DIM,
+            INTERMEDIATE_DIM=INTERMEDIATE_DIM,
+            NUM_I_TILES=NUM_I_TILES,
+            BLOCK_SIZE_H=BLOCK_SIZE_N,
+            BLOCK_SIZE_I=BLOCK_SIZE_K,
+            NUM_EXPERTS=NUM_EXPERTS,
+            NUM_SMS=num_sms,
+            SIMULATE_UNFUSED=simulate_unfused,
+        )
+        topk_reduce_kernel[(num_tokens, triton.cdiv(HIDDEN_DIM, TOPK_REDUCE_BLOCK_H))](
+            out,
+            reduced,
+            HIDDEN_DIM,
+            out.stride(0),
+            out.stride(1),
+            reduced.stride(0),
+            reduced.stride(1),
+            NUM_TOP_K=num_top_k,
+            BLOCK_H=TOPK_REDUCE_BLOCK_H,
+        )
+    return reduced
+# ── MXFP4/MXFP8 (UE8M0 group-32, tunable tile, MXFP8 intermediate) ────────────
+def _set_gate_up_descriptor(nargs):
+    """Per-config pre_hook: set the gate_up TMA descriptor box to the autotuned
+    [2 (gate|up), BLOCK_SIZE_N, BK//VALUES_PER_BYTE] over the (2E, I, H/vpb) weight view.
+    MUST mutate ``block_shape`` in place, not rebind ``nargs[...]`` to a fresh descriptor:
+    the autotuner launches with the original ``*args`` object, so a rebind never reaches the
+    kernel (the placeholder box survives → all but BN=32,BK=32 configs fail the reshape). An
+    in-place edit of the passed object propagates to both compile-time specialization and the
+    launch-time tensor-map fill. ``shape``/``strides`` are already correct from the wrapper's
+    ``from_tensor`` placeholder; only the per-config box changes. No-op unless MEMORY_MODE is host_descriptor (pointer /
+    device_descriptor don't use the host-built descriptor)."""
+    if nargs["MEMORY_MODE"] != "host_descriptor":
+        return
+    nargs["GateUpDescriptor"].block_shape = [
+        2,
+        nargs["BLOCK_SIZE_N"],
+        nargs["BLOCK_SIZE_K"] // nargs["VALUES_PER_BYTE"],
+    ]
+@bayesian_autotune(
+    get_mxfp_autotuning_configs(
+        pre_hook=_set_gate_up_descriptor,
+        compute_modes=_MX_COMPUTE_MODES,
+        memory_modes=_MX_MEMORY_MODES,
+        tune_block_m=True,
+    ),  # prefill: no scalar; combined gate∪up dot ([2*BN,K] reshape) — TMA vs block-ptr load
+    ["INTERMEDIATE_DIM", "HIDDEN_DIM", "tokens_per_sm_bit_length"],
+    n_trials=60,
+    # bf16 activation tile + fused gate|up weight tiles
+    prune_configs_by={
+        "early_config_prune": smem_config_pruner(act_bytes=2, n_weight_tiles=2)
+    },
+)
+@triton.jit
+def mxfp_dynamic_moe_grouped_gate_up_kernel(
+    Hidden,  # (T, H) raw activations, UNSORTED
+    PermToken,  # (S,) int32 — sorted position -> source token id
+    GateUp,  # (E, 2I, H//VALUES_PER_BYTE) MXFP4/MXFP8
+    GateUpScale,  # (E, 2I, H//32) UE8M0 group-32 scales
+    GateUpDescriptor,  # TMA descriptor over the (2E, I, H//VALUES_PER_BYTE) weight view
+    ExpertStart,  # (NUM_EXPERTS+1,) int32 — exclusive sorted-row start per expert (pad S)
+    Inter,  # (S, I) E4M3 — MXFP8 intermediate (expert-sorted)
+    InterScale,  # (S, I//32) UE8M0 group-32 scales
+    stride_h_t,
+    stride_h_k,
+    stride_gu_e,
+    stride_gu_n,
+    stride_gu_k,
+    stride_gus_e,
+    stride_gus_n,
+    stride_gus_k,
+    stride_int_m,
+    stride_int_n,
+    stride_is_m,
+    stride_is_n,
+    stride_pt,
+    tokens_per_sm_bit_length,  # autotune key only (log2 tokens-per-SM bucket); unused in body
+    HIDDEN_DIM: tl.constexpr,
+    INTERMEDIATE_DIM: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    VALUES_PER_BYTE: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    NUM_EXPERTS: tl.constexpr,
+    NUM_SMS: tl.constexpr,
+    SIMULATE_UNFUSED: tl.constexpr,
+    COMPUTE_MODE: tl.constexpr,
+    MEMORY_MODE: tl.constexpr,
+    ACT_FN: tl.constexpr,
+):
+    """MXFP4/MXFP8 phase 1 (persistent): gather hidden rows per expert M-tile, gate + up MX
+    matmuls (``tl.dot_scaled`` or fp8 ``tl.dot`` + per-group rescale), SiLU, MXFP8-requant."""
+    start_pid = tl.program_id(axis=0)
+    exp_start, freqs, tile_start_excl, total_m_tiles, e_offs = _build_tile_layout(
+        ExpertStart, NUM_EXPERTS, BLOCK_SIZE_M
+    )
+    num_n_tiles = tl.cdiv(INTERMEDIATE_DIM, BLOCK_SIZE_N)
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    for tile_id in tl.range(start_pid, total_m_tiles * num_n_tiles, NUM_SMS):
+        pid_m = tile_id // num_n_tiles
+        pid_n = tile_id % num_n_tiles
+        expert_id, offs_global_m, row_mask = _resolve_tile_inline(
+            pid_m, exp_start, freqs, tile_start_excl, e_offs, BLOCK_SIZE_M
+        )
+        offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+        token = tl.load(PermToken + offs_global_m * stride_pt, mask=row_mask, other=0)
+        a_ptrs = Hidden + token[:, None] * stride_h_t + offs_k[None, :] * stride_h_k
+        n_off = pid_n * BLOCK_SIZE_N
+        # Load gate (row 2e) + up (row 2e+1) of the (2E, I, H) view as ONE combined tile, then
+        # split — this is the shared loop. MEMORY_MODE picks the LOAD only (decoupled from COMPUTE_MODE):
+        # a host/device descriptor (TMA on NVIDIA) vs explicit rank-3 pointers. COMPUTE_MODE picks the compute on the
+        # loaded tile: scaled-MMA (dot_scaled) or fp8 dot + per-group software rescale (dot).
+        gu_row = expert_id * 2
+        # (E, 2I, H//vpb) reinterpreted as (2E, I, H//vpb); host_descriptor uses the passed
+        # GateUpDescriptor, device_descriptor builds one in-kernel, pointer indexes it directly.
+        if MEMORY_MODE == "pointer":
+            gu_ptr = (
+                GateUp
+                + (gu_row + tl.arange(0, 2))[:, None, None]
+                * (INTERMEDIATE_DIM * stride_gu_n)
+                + (n_off + tl.arange(0, BLOCK_SIZE_N))[None, :, None] * stride_gu_n
+                + tl.arange(0, BLOCK_SIZE_K // VALUES_PER_BYTE)[None, None, :]
+                * stride_gu_k
+            )
+        elif MEMORY_MODE == "device_descriptor":
+            gu_desc = tl.make_tensor_descriptor(
+                GateUp,
+                shape=(
+                    2 * NUM_EXPERTS,
+                    INTERMEDIATE_DIM,
+                    HIDDEN_DIM // VALUES_PER_BYTE,
+                ),
+                strides=(INTERMEDIATE_DIM * stride_gu_n, stride_gu_n, stride_gu_k),
+                block_shape=(2, BLOCK_SIZE_N, BLOCK_SIZE_K // VALUES_PER_BYTE),
+            )
+        gu_scale_ptr = (
+            GateUpScale
+            + expert_id * stride_gus_e
+            + tl.arange(0, 2)[:, None, None] * (INTERMEDIATE_DIM * stride_gus_n)
+            + (n_off + tl.arange(0, BLOCK_SIZE_N))[None, :, None] * stride_gus_n
+            + tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K)[None, None, :] * stride_gus_k
+        )
+        acc = tl.zeros((BLOCK_SIZE_M, 2 * BLOCK_SIZE_N), dtype=tl.float32)
+        for k_off in tl.range(0, HIDDEN_DIM, BLOCK_SIZE_K):
+            a_raw = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0).to(tl.float32)
+            a, a_scale = mxfp_act_quant_inline(
+                a_raw, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K
+            )
+            if MEMORY_MODE == "host_descriptor":
+                gu = tl.reshape(
+                    GateUpDescriptor.load([gu_row, n_off, k_off // VALUES_PER_BYTE]),
+                    [2 * BLOCK_SIZE_N, BLOCK_SIZE_K // VALUES_PER_BYTE],
+                )
+            elif MEMORY_MODE == "device_descriptor":
+                gu = tl.reshape(
+                    gu_desc.load([gu_row, n_off, k_off // VALUES_PER_BYTE]),
+                    [2 * BLOCK_SIZE_N, BLOCK_SIZE_K // VALUES_PER_BYTE],
+                )
+            else:
+                gu = tl.reshape(
+                    tl.load(gu_ptr),
+                    [2 * BLOCK_SIZE_N, BLOCK_SIZE_K // VALUES_PER_BYTE],
+                )
+                gu_ptr += (BLOCK_SIZE_K // VALUES_PER_BYTE) * stride_gu_k
+            gu_scale = tl.reshape(
+                tl.load(gu_scale_ptr + (k_off // SCALE_GROUP_K) * stride_gus_k),
+                [2 * BLOCK_SIZE_N, BLOCK_SIZE_K // SCALE_GROUP_K],
+            )
+            gu_t = tl.trans(
+                gu
+            )  # [BK, 2*BN] gate∪up weight; mx_compute decodes it for the fp8 path
+            acc = mx_compute(
+                acc,
+                a,
+                a_scale,
+                gu_t,
+                gu_scale,
+                COMPUTE_MODE,
+                VALUES_PER_BYTE,
+                BLOCK_SIZE_M,
+                BLOCK_SIZE_N,
+                BLOCK_SIZE_K,
+                SCALE_GROUP_K,
+            )
+            a_ptrs += BLOCK_SIZE_K * stride_h_k
+        acc_3d = tl.permute(tl.reshape(acc, [BLOCK_SIZE_M, 2, BLOCK_SIZE_N]), (0, 2, 1))
+        acc_gate, acc_up = tl.split(acc_3d)
+        intermediate = glu(
+            acc_gate, acc_up, ACT_FN, SIMULATE_UNFUSED, Hidden.dtype.element_ty
+        )
+        # MXFP8 requant of the intermediate (E4M3 + UE8M0 group-32 along this N-tile).
+        inter, inter_scale = mxfp_act_quant_inline(
+            intermediate, BLOCK_SIZE_M, BLOCK_SIZE_N, SCALE_GROUP_K
+        )
+        int_ptrs = (
+            Inter
+            + offs_global_m[:, None] * stride_int_m
+            + offs_bn[None, :] * stride_int_n
+        )
+        tl.store(int_ptrs, inter, mask=row_mask[:, None])
+        offs_sc = pid_n * (BLOCK_SIZE_N // SCALE_GROUP_K) + tl.arange(
+            0, BLOCK_SIZE_N // SCALE_GROUP_K
+        )
+        sc_ptrs = (
+            InterScale
+            + offs_global_m[:, None] * stride_is_m
+            + offs_sc[None, :] * stride_is_n
+        )
+        tl.store(sc_ptrs, inter_scale, mask=row_mask[:, None])
+def _set_down_descriptor(nargs):
+    """Per-config pre_hook: in-place set the down TMA descriptor box to
+    [BLOCK_SIZE_N, BK//VALUES_PER_BYTE] over the (E*H, I/vpb) weight view. In-place, not
+    rebind — see _set_gate_up_descriptor. No-op unless MEMORY_MODE is host_descriptor."""
+    if nargs["MEMORY_MODE"] != "host_descriptor":
+        return
+    nargs["DownDescriptor"].block_shape = [
+        nargs["BLOCK_SIZE_N"],
+        nargs["BLOCK_SIZE_K"] // nargs["VALUES_PER_BYTE"],
+    ]
+@bayesian_autotune(
+    get_mxfp_autotuning_configs(
+        pre_hook=_set_down_descriptor,
+        compute_modes=_MX_COMPUTE_MODES,
+        memory_modes=_MX_MEMORY_MODES,
+        tune_block_m=True,
+    ),  # prefill: no scalar; TMA vs block-ptr load
+    ["INTERMEDIATE_DIM", "HIDDEN_DIM", "tokens_per_sm_bit_length"],
+    n_trials=60,
+    # fp8 intermediate activation tile + single down weight tile
+    prune_configs_by={
+        "early_config_prune": smem_config_pruner(act_bytes=1, n_weight_tiles=1)
+    },
+)
+@triton.jit
+def mxfp_dynamic_moe_grouped_down_kernel(
+    Inter,  # (S, I) E4M3 — MXFP8 intermediate (expert-sorted)
+    InterScale,  # (S, I//32) UE8M0 group-32 scales
+    Down,  # (E, H, I//VALUES_PER_BYTE) MXFP4/MXFP8
+    DownScale,  # (E, H, I//32) UE8M0 group-32 scales
+    DownDescriptor,  # TMA descriptor over the (E*H, I//VALUES_PER_BYTE) weight view
+    ExpertStart,  # (NUM_EXPERTS+1,) int32 — exclusive sorted-row start per expert (pad S)
+    Perm,  # (S,) int32 — sorted position -> flat (token, slot) row
+    SampleWeights,  # (S,) routing weight per flat row
+    ProjOut,  # (S, H) bf16, flat (token, slot) order
+    stride_int_m,
+    stride_int_n,
+    stride_is_m,
+    stride_is_n,
+    stride_down_e,
+    stride_down_n,
+    stride_down_k,
+    stride_downs_e,
+    stride_downs_n,
+    stride_downs_k,
+    stride_po_m,
+    stride_po_n,
+    stride_perm,
+    tokens_per_sm_bit_length,  # autotune key only (log2 tokens-per-SM bucket); unused in body
+    HIDDEN_DIM: tl.constexpr,
+    INTERMEDIATE_DIM: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    VALUES_PER_BYTE: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    NUM_EXPERTS: tl.constexpr,
+    NUM_SMS: tl.constexpr,
+    SIMULATE_UNFUSED: tl.constexpr,
+    COMPUTE_MODE: tl.constexpr,
+    MEMORY_MODE: tl.constexpr,
+):
+    """MXFP4/MXFP8 phase 2 (persistent): MXFP8 intermediate → down proj, then routing-weight
+    × scatter to the flat (token, slot) row. N = hidden tile, K = intermediate tile."""
+    start_pid = tl.program_id(axis=0)
+    exp_start, freqs, tile_start_excl, total_m_tiles, e_offs = _build_tile_layout(
+        ExpertStart, NUM_EXPERTS, BLOCK_SIZE_M
+    )
+    num_n_tiles = tl.cdiv(HIDDEN_DIM, BLOCK_SIZE_N)
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    offs_sf = tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K)
+    for tile_id in tl.range(start_pid, total_m_tiles * num_n_tiles, NUM_SMS):
+        pid_m = tile_id // num_n_tiles
+        pid_n = tile_id % num_n_tiles
+        expert_id, offs_global_m, row_mask = _resolve_tile_inline(
+            pid_m, exp_start, freqs, tile_start_excl, e_offs, BLOCK_SIZE_M
+        )
+        offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+        a_ptrs = (
+            Inter
+            + offs_global_m[:, None] * stride_int_m
+            + offs_k[None, :] * stride_int_n
+        )
+        as_ptrs = (
+            InterScale
+            + offs_global_m[:, None] * stride_is_m
+            + offs_sf[None, :] * stride_is_n
+        )
+        n_off = pid_n * BLOCK_SIZE_N
+        ws_down_ptr = (
+            DownScale
+            + expert_id * stride_downs_e
+            + (n_off + tl.arange(0, BLOCK_SIZE_N))[:, None] * stride_downs_n
+            + tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K)[None, :] * stride_downs_k
+        )
+        # Down weight tile [BK//vpb, BN] loaded once per K-chunk: MEMORY_MODE picks the LOAD (a
+        # host/device descriptor over the (E*H, I//vpb) view vs explicit pointers), COMPUTE_MODE
+        # picks the compute. Scales stay a pointer load (a descriptor needs >=16B inner; the
+        # BK//32 row is too narrow).
+        if MEMORY_MODE == "pointer":
+            w_down_ptr = (
+                Down
+                + expert_id * stride_down_e
+                + tl.arange(0, BLOCK_SIZE_K // VALUES_PER_BYTE)[:, None] * stride_down_k
+                + (n_off + tl.arange(0, BLOCK_SIZE_N))[None, :] * stride_down_n
+            )
+        elif MEMORY_MODE == "device_descriptor":
+            down_desc = tl.make_tensor_descriptor(
+                Down,  # (E, H, I//vpb) flattened to the (E*H, I//vpb) view
+                shape=(NUM_EXPERTS * HIDDEN_DIM, INTERMEDIATE_DIM // VALUES_PER_BYTE),
+                strides=(stride_down_n, stride_down_k),
+                block_shape=(BLOCK_SIZE_N, BLOCK_SIZE_K // VALUES_PER_BYTE),
+            )
+        acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+        for k_off in tl.range(0, INTERMEDIATE_DIM, BLOCK_SIZE_K):
+            a = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0)
+            a_scale = tl.load(as_ptrs, mask=row_mask[:, None], other=0)
+            if MEMORY_MODE == "host_descriptor":
+                w = tl.trans(
+                    tl.reshape(
+                        DownDescriptor.load(
+                            [expert_id * HIDDEN_DIM + n_off, k_off // VALUES_PER_BYTE]
+                        ),
+                        [BLOCK_SIZE_N, BLOCK_SIZE_K // VALUES_PER_BYTE],
+                    )
+                )
+            elif MEMORY_MODE == "device_descriptor":
+                w = tl.trans(
+                    tl.reshape(
+                        down_desc.load(
+                            [expert_id * HIDDEN_DIM + n_off, k_off // VALUES_PER_BYTE]
+                        ),
+                        [BLOCK_SIZE_N, BLOCK_SIZE_K // VALUES_PER_BYTE],
+                    )
+                )
+            else:
+                w = tl.load(w_down_ptr)
+                w_down_ptr += (BLOCK_SIZE_K // VALUES_PER_BYTE) * stride_down_k
+            w_scale = tl.load(ws_down_ptr + (k_off // SCALE_GROUP_K) * stride_downs_k)
+            acc = mx_compute(
+                acc,
+                a,
+                a_scale,
+                w,
+                w_scale,
+                COMPUTE_MODE,
+                VALUES_PER_BYTE,
+                BLOCK_SIZE_M,
+                BLOCK_SIZE_N,
+                BLOCK_SIZE_K,
+                SCALE_GROUP_K,
+            )
+            a_ptrs += BLOCK_SIZE_K * stride_int_n
+            as_ptrs += (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_is_n
+        if SIMULATE_UNFUSED:
+            acc = acc.to(ProjOut.dtype.element_ty).to(tl.float32)
+        # Fused routing-weight × top-k reorder (no atomics): scale each row by its weight
+        # and scatter to its flat (token, slot) row; the host then just sums over slots.
+        flat = tl.load(Perm + offs_global_m * stride_perm, mask=row_mask, other=0)
+        weight = tl.load(SampleWeights + flat, mask=row_mask, other=0.0)
+        acc = acc * weight[:, None]
+        store_tile(ProjOut, acc, flat, offs_bn, row_mask, stride_po_m, stride_po_n)
+def mxfp_dynamic_moe_grouped(
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    gate_up_proj: torch.Tensor,
+    down_proj: torch.Tensor,
+    gate_up_proj_scale: torch.Tensor,
+    down_proj_scale: torch.Tensor,
+    simulate_unfused: bool = False,
+    act_fn: str = "silu",
+) -> torch.Tensor:
+    """MXFP4/MXFP8 fused grouped MoE (UE8M0 group-32); gate_up and down must share the same MX
+    format. Same structure as the block-dynamic path but with a tunable tile and MXFP8 intermediate."""
+    gate_up_is_fp4 = is_mxfp4(gate_up_proj, gate_up_proj_scale)
+    down_is_fp4 = is_mxfp4(down_proj, down_proj_scale)
+    if gate_up_is_fp4 != down_is_fp4:
+        raise ValueError(
+            "gate_up_proj and down_proj must use the same MX format (both MXFP4 or both MXFP8)."
+        )
+    device = hidden_states.device
+    num_sms = sm_count(device.index)
+    num_top_k = top_k_index.size(-1)
+    num_tokens = hidden_states.size(0)
+    HIDDEN_DIM = hidden_states.size(1)
+    INTERMEDIATE_DIM = gate_up_proj.size(1) // 2
+    perm_token, perm, expert_start, NUM_EXPERTS, num_routed_tokens = _grouped_routing(
+        top_k_index, gate_up_proj.size(0), num_top_k
+    )
+    tokens_per_sm_bit_length = (num_routed_tokens // num_sms).bit_length()
+    VALUES_PER_BYTE = NIBBLES_PER_BYTE if gate_up_is_fp4 else 1
+    gate_up_proj_u8 = e2m1_as_uint8(gate_up_proj)
+    down_proj_u8 = e2m1_as_uint8(down_proj)
+    gate_up_scale_u8 = ue8m0_as_uint8(gate_up_proj_scale)
+    down_scale_u8 = ue8m0_as_uint8(down_proj_scale)
+    inter = torch.empty(
+        num_routed_tokens, INTERMEDIATE_DIM, device=device, dtype=FP8_DTYPE
+    )
+    inter_scale = torch.empty(
+        num_routed_tokens,
+        INTERMEDIATE_DIM // MX_SCALE_GROUP_K,
+        device=device,
+        dtype=torch.uint8,
+    )
+    out = torch.empty(
+        num_routed_tokens, HIDDEN_DIM, device=device, dtype=hidden_states.dtype
+    )
+    reduced = torch.empty(
+        num_tokens, HIDDEN_DIM, device=device, dtype=hidden_states.dtype
+    )
+    # Host-built descriptors for host_descriptor mode: views (not reshapes) of the weight
+    # buffers — gate_up (E, 2I, H/vpb) → (2E, I, H/vpb) so one [2, BN, BK/vpb] box loads gate
+    # (row 2e) + up (2e+1); down (E, H, I/vpb) → (E*H, I/vpb), one [BN, BK/vpb] box per (expert,
+    # hidden-tile). Cheap to build, so always created (device_descriptor/pointer ignore them).
+    gate_up_2e = gate_up_proj_u8.view(
+        2 * gate_up_proj_u8.size(0), INTERMEDIATE_DIM, gate_up_proj_u8.size(2)
+    )
+    gate_up_descriptor = TensorDescriptor.from_tensor(
+        gate_up_2e, [2, 32, 32 // VALUES_PER_BYTE]
+    )
+    down_eh = down_proj_u8.view(down_proj_u8.size(0) * HIDDEN_DIM, down_proj_u8.size(2))
+    down_descriptor = TensorDescriptor.from_tensor(down_eh, [32, 32 // VALUES_PER_BYTE])
+    with device_context(device):
+        mxfp_dynamic_moe_grouped_gate_up_kernel[(num_sms,)](
+            hidden_states,
+            perm_token,
+            gate_up_proj_u8,
+            gate_up_scale_u8,
+            gate_up_descriptor,
+            expert_start,
+            inter,
+            inter_scale,
+            hidden_states.stride(0),
+            hidden_states.stride(1),
+            gate_up_proj_u8.stride(0),
+            gate_up_proj_u8.stride(1),
+            gate_up_proj_u8.stride(2),
+            gate_up_scale_u8.stride(0),
+            gate_up_scale_u8.stride(1),
+            gate_up_scale_u8.stride(2),
+            inter.stride(0),
+            inter.stride(1),
+            inter_scale.stride(0),
+            inter_scale.stride(1),
+            perm_token.stride(0),
+            tokens_per_sm_bit_length=tokens_per_sm_bit_length,
+            HIDDEN_DIM=HIDDEN_DIM,
+            INTERMEDIATE_DIM=INTERMEDIATE_DIM,
+            VALUES_PER_BYTE=VALUES_PER_BYTE,
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+            NUM_EXPERTS=NUM_EXPERTS,
+            NUM_SMS=num_sms,
+            SIMULATE_UNFUSED=simulate_unfused,
+            ACT_FN=act_fn,
+        )
+        mxfp_dynamic_moe_grouped_down_kernel[(num_sms,)](
+            inter,
+            inter_scale,
+            down_proj_u8,
+            down_scale_u8,
+            down_descriptor,
+            expert_start,
+            perm,
+            top_k_weights,
+            out,
+            inter.stride(0),
+            inter.stride(1),
+            inter_scale.stride(0),
+            inter_scale.stride(1),
+            down_proj_u8.stride(0),
+            down_proj_u8.stride(1),
+            down_proj_u8.stride(2),
+            down_scale_u8.stride(0),
+            down_scale_u8.stride(1),
+            down_scale_u8.stride(2),
+            out.stride(0),
+            out.stride(1),
+            perm.stride(0),
+            tokens_per_sm_bit_length=tokens_per_sm_bit_length,
+            HIDDEN_DIM=HIDDEN_DIM,
+            INTERMEDIATE_DIM=INTERMEDIATE_DIM,
+            VALUES_PER_BYTE=VALUES_PER_BYTE,
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+            NUM_EXPERTS=NUM_EXPERTS,
+            NUM_SMS=num_sms,
+            SIMULATE_UNFUSED=simulate_unfused,
+        )
+        topk_reduce_kernel[(num_tokens, triton.cdiv(HIDDEN_DIM, TOPK_REDUCE_BLOCK_H))](
+            out,
+            reduced,
+            HIDDEN_DIM,
+            out.stride(0),
+            out.stride(1),
+            reduced.stride(0),
+            reduced.stride(1),
+            NUM_TOP_K=num_top_k,
+            BLOCK_H=TOPK_REDUCE_BLOCK_H,
+        )
+    return reduced
+# ── Dispatcher ────────────────────────────────────────────────────────────────
+def moe_fused_grouped(
+    hidden_states: torch.Tensor,
+    top_k_index: torch.Tensor,
+    top_k_weights: torch.Tensor,
+    gate_up_proj: torch.Tensor,
+    down_proj: torch.Tensor,
+    gate_up_proj_scale_inv: torch.Tensor,
+    down_proj_scale_inv: torch.Tensor,
+    block_size: list[int] | None,
+    simulate_unfused: bool = False,
+    act_fn: str = "silu",
+) -> torch.Tensor:
+    """Fused grouped-MoE dispatcher — routes to the recipe matching the weight dtype /
+    scale layout, mirroring ``moe_fused_batched``. Implemented: block-dynamic FP8 and MXFP8 /
+    MXFP4 (UE8M0 group-32). ``simulate_unfused`` (testing) rounds each step through the
+    activation dtype so the output matches the unfused reference to reduce order."""
+    gate_up_is_mx = is_mxfp(gate_up_proj, gate_up_proj_scale_inv)
+    down_is_mx = is_mxfp(down_proj, down_proj_scale_inv)
+    if gate_up_is_mx != down_is_mx:
+        raise ValueError(
+            "gate_up_proj and down_proj must use the same recipe (both MX or both block-dynamic FP8)."
+        )
+    if gate_up_is_mx:
+        return mxfp_dynamic_moe_grouped(
+            hidden_states,
+            top_k_index,
+            top_k_weights,
+            gate_up_proj,
+            down_proj,
+            gate_up_proj_scale_inv,
+            down_proj_scale_inv,
+            simulate_unfused,
+            act_fn,
+        )
+    if block_size is None:
+        raise ValueError("block_size is required for block-dynamic FP8 weights.")
+    return w8a8_block_dynamic_fp8_moe_grouped(
+        hidden_states,
+        top_k_index,
+        top_k_weights,
+        gate_up_proj,
+        down_proj,
+        gate_up_proj_scale_inv,
+        down_proj_scale_inv,
+        block_size,
+        simulate_unfused,
+        act_fn,
+    )

build/torch-cuda/grouped.py CHANGED Viewed

@@ -15,28 +15,138 @@
 import torch
 import triton
 import triton.language as tl
-from .act_quant import fp8_act_quant
 from torch.library import triton_op, wrap_triton
-from .utils import device_context
 @triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
     key=["N", "K", "BLOCK_SIZE_M"],
 )
 @triton.jit
-def w8a8_block_fp8_matmul_grouped_kernel(
-    A,  # (S, K)  raw BF16/FP16 activations, sorted/grouped by expert id
-    B,  # (E, N, K) FP8 weight matrices
-    C,  # (S, N)  output
-    Bs,  # (E, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales
-    Offsets,  # (E,) int32 — cumulative row-end per expert
-    TileOffsets,  # (E,) int32 — cumulative tile-end per expert
     # Shape
     S,
     N,
@@ -52,11 +162,13 @@ def w8a8_block_fp8_matmul_grouped_kernel(
     stride_bs_e,
     stride_bs_k,
     stride_bs_n,
     # Meta-parameters
-    NUM_EXPERTS: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
     NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
 ):
     """Block-scale grouped FP8 expert matmul kernel.
@@ -67,47 +179,23 @@ def w8a8_block_fp8_matmul_grouped_kernel(
     pid_m = tl.program_id(axis=0)
     pid_n = tl.program_id(axis=1)
-    # Exit early for programs beyond the actual tile count.
-    total_tiles = tl.load(TileOffsets + NUM_EXPERTS - 1)
     if pid_m >= total_tiles:
         return
-    # Binary search in TileOffsets to find the owning expert.
-    # Finds the smallest e such that TileOffsets[e] > pid_m (upper_bound semantics),
-    # which is the expert whose tile range contains pid_m.
-    # O(log2(NUM_EXPERTS)) loads instead of the O(NUM_EXPERTS) linear scan.
-    # NUM_EXPERTS_BIT_LENGTH is ceil(log2(E))+1 for powers-of-two, giving one
-    # harmless extra iteration when lo==hi; it's a compile-time constant so the
-    # loop is fully unrolled by the compiler.
-    lo = 0
-    hi = NUM_EXPERTS
-    for _ in tl.static_range(NUM_EXPERTS_BIT_LENGTH):
-        mid = (lo + hi) >> 1
-        mid_val = tl.load(TileOffsets + mid)
-        is_left = mid_val <= pid_m
-        lo = tl.where(is_left, mid + 1, lo)
-        hi = tl.where(is_left, hi, mid)
-    # Cast expert_id to int64 to prevent int32 overflow when computing
-    # expert_id * stride_be (e.g. 255 * 9_437_184 > 2^31 for 256 experts of
-    # 3072×3072 FP8 weights).
-    expert_id = lo.to(tl.int64)
-    prev_eid = tl.maximum(expert_id - 1, 0)
-    expert_start = tl.where(expert_id == 0, 0, tl.load(Offsets + prev_eid))
-    expert_end = tl.load(Offsets + expert_id)
-    M_expert = expert_end - expert_start
-    expert_tile_start = tl.where(expert_id == 0, 0, tl.load(TileOffsets + prev_eid))
-    local_tile = pid_m - expert_tile_start
-    m_off = local_tile * BLOCK_SIZE_M
-    offs_am = m_off + tl.arange(0, BLOCK_SIZE_M)
-    row_mask = offs_am < M_expert
-    offs_global_m = expert_start + offs_am
-    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
     a_ptrs = A + offs_global_m[:, None] * stride_am + offs_k[None, :] * stride_ak
     b_ptrs = (
@@ -120,49 +208,37 @@ def w8a8_block_fp8_matmul_grouped_kernel(
     accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        # ---- fused fp8_act_quant ----
         a_raw = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0).to(tl.float32)
-        a_s = tl.max(tl.abs(a_raw), axis=1) / 448.0
-        a = (a_raw / tl.maximum(a_s[:, None], 1e-12)).to(tl.float8e4nv)
-        # ---- matmul ----
         b = tl.load(b_ptrs)
-        b_s = tl.load(bs_ptrs + k * stride_bs_k)
         accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    c_ptrs = C + stride_cm * offs_global_m[:, None] + stride_cn * offs_bn[None, :]
-    c_mask = row_mask[:, None]
-    tl.store(c_ptrs, c, mask=c_mask)
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
-    key=["N", "K", "BLOCK_SIZE_M"],
 )
 @triton.jit
-def w8a8_tensor_fp8_matmul_grouped_kernel(
-    A,  # (S, K) raw BF16/FP16 activations, sorted/grouped by expert idc
-    B,  # (E, N, K) FP8 weight matrices
     C,  # (S, N) output
-    As,  # (S, 1) activation scales
-    Bs,  # (E, 1, 1) per-tensor weight scales
-    Offsets,
-    TileOffsets,
     S,
     N,
     K,
     stride_am,
     stride_ak,
     stride_be,
@@ -172,10 +248,13 @@ def w8a8_tensor_fp8_matmul_grouped_kernel(
     stride_cn,
     stride_as_m,
     stride_bs_e,
-    NUM_EXPERTS: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
     NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
 ):
     """Tensor-scale grouped FP8 expert matmul kernel.
@@ -186,35 +265,23 @@ def w8a8_tensor_fp8_matmul_grouped_kernel(
     pid_m = tl.program_id(axis=0)
     pid_n = tl.program_id(axis=1)
-    total_tiles = tl.load(TileOffsets + NUM_EXPERTS - 1)
     if pid_m >= total_tiles:
         return
-    lo = 0
-    hi = NUM_EXPERTS
-    for _ in tl.static_range(NUM_EXPERTS_BIT_LENGTH):
-        mid = (lo + hi) >> 1
-        mid_val = tl.load(TileOffsets + mid)
-        is_left = mid_val <= pid_m
-        lo = tl.where(is_left, mid + 1, lo)
-        hi = tl.where(is_left, hi, mid)
-    expert_id = lo.to(tl.int64)
-    prev_eid = tl.maximum(expert_id - 1, 0)
-    expert_start = tl.where(expert_id == 0, 0, tl.load(Offsets + prev_eid))
-    expert_end = tl.load(Offsets + expert_id)
-    M_expert = expert_end - expert_start
-    expert_tile_start = tl.where(expert_id == 0, 0, tl.load(TileOffsets + prev_eid))
-    local_tile = pid_m - expert_tile_start
-    m_off = local_tile * BLOCK_SIZE_M
-    offs_am = m_off + tl.arange(0, BLOCK_SIZE_M)
-    row_mask = offs_am < M_expert
-    offs_global_m = expert_start + offs_am
-    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
     a_ptrs = A + offs_global_m[:, None] * stride_am + offs_k[None, :] * stride_ak
     b_ptrs = (
@@ -231,50 +298,151 @@ def w8a8_tensor_fp8_matmul_grouped_kernel(
     for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0)
         b = tl.load(b_ptrs)
         accumulator += tl.dot(a, b)
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
     accumulator = accumulator * a_s[:, None] * b_s
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    c_ptrs = C + stride_cm * offs_global_m[:, None] + stride_cn * offs_bn[None, :]
-    c_mask = row_mask[:, None]
-    tl.store(c_ptrs, c, mask=c_mask)
-@triton_op("finegrained_fp8::w8a8_block_fp8_matmul_grouped", mutates_args=())
-def _w8a8_block_fp8_matmul_grouped(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     offsets: torch.Tensor,
     tokens_per_expert: torch.Tensor,
     block_size: list[int],
 ) -> torch.Tensor:
     """Block-scale grouped FP8 matmul: C = A @ B.T per expert, with fused act quant.
     A:  (S, K) raw bf16/fp16 activations, sorted by expert
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
     """
     assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
     assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
     assert B.is_contiguous(), "B must be contiguous"
     assert A.shape[1] == B.shape[2], (
         f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
     )
     S, K = A.shape
-    E, N, _ = B.shape
     assert len(block_size) == 2, (
         f"block_size must be [block_n, block_k], got {block_size}"
@@ -284,25 +452,22 @@ def _w8a8_block_fp8_matmul_grouped(
     assert N % block_n == 0, f"N ({N}) must be divisible by block_n ({block_n})"
     assert K % block_k == 0, f"K ({K}) must be divisible by block_k ({block_k})"
     assert Bs.ndim == 3, (
-        f"Bs must be 3D (E, N//block_n, K//block_k), got ndim={Bs.ndim}"
     )
-    assert Bs.shape == (E, N // block_n, K // block_k), (
-        f"Bs shape {tuple(Bs.shape)} != expected ({E}, {N // block_n}, {K // block_k})"
     )
-    C = A.new_empty(S, N)
-    # Adaptive BLOCK_SIZE_M: match tile to average tokens per expert.
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
-    tiles_per_expert = (tokens_per_expert + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M
-    tile_offsets = torch.cumsum(tiles_per_expert, dim=0).to(torch.int32)
-    # Upper bound on M-tiles: sum_e ceil(M_e / BLOCK_M) <= ceil(S / BLOCK_M) + E.
-    # Programs beyond the real tile count exit immediately via the early-return
-    # guard inside the kernel. This is faster than syncing for the exact count
-    # and keeps the grid size data-independent (cuda-graph / torch.compile safe).
-    max_M_tiles = triton.cdiv(S, BLOCK_SIZE_M) + E
-    grid = (max_M_tiles, triton.cdiv(N, block_n))
     with device_context(A.device):
-        wrap_triton(w8a8_block_fp8_matmul_grouped_kernel)[grid](
             A,
             B,
             C,
@@ -322,66 +487,70 @@ def _w8a8_block_fp8_matmul_grouped(
             Bs.stride(0),
             Bs.stride(2),
             Bs.stride(1),
             # Meta-parameters
-            NUM_EXPERTS=E,
             BLOCK_SIZE_N=block_n,
             BLOCK_SIZE_K=block_k,
             BLOCK_SIZE_M=BLOCK_SIZE_M,
-            NUM_EXPERTS_BIT_LENGTH=E.bit_length(),
         )
     return C
-@triton_op("finegrained_fp8::w8a8_tensor_fp8_matmul_grouped", mutates_args=())
-def _w8a8_tensor_fp8_matmul_grouped(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     offsets: torch.Tensor,
     tokens_per_expert: torch.Tensor,
 ) -> torch.Tensor:
     """Tensor-scale grouped FP8 matmul: C = A @ B.T per expert, with fused act quant.
     A:  (S, K) raw bf16/fp16 activations, sorted by expert
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
     """
     assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
     assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
     assert B.is_contiguous(), "B must be contiguous"
     assert A.shape[1] == B.shape[2], (
         f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
     )
     S, K = A.shape
-    E, N, _ = B.shape
-    # Normalize Bs to (E, 1, 1)
     if Bs.ndim == 1:
-        assert Bs.shape[0] == E, f"Bs shape {tuple(Bs.shape)} != expected ({E},)"
-        Bs = Bs.reshape(E, 1, 1)
     else:
-        assert Bs.shape == (E, 1, 1), (
-            f"Bs shape {tuple(Bs.shape)} != expected ({E}, 1, 1)"
         )
-    BLOCK_SIZE_N = 128
-    BLOCK_SIZE_K = 128
-    C = A.new_empty(S, N)
     qA, As = fp8_act_quant(A, K)
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
-    tiles_per_expert = (tokens_per_expert + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M
-    tile_offsets = torch.cumsum(tiles_per_expert, dim=0).to(torch.int32)
-    # Upper bound on M-tiles: sum_e ceil(M_e / BLOCK_M) <= ceil(S / BLOCK_M) + E.
-    # Programs beyond the real tile count exit immediately via the early-return
-    # guard inside the kernel. This is faster than syncing for the exact count
-    # and keeps the grid size data-independent (cuda-graph / torch.compile safe).
-    max_M_tiles = triton.cdiv(S, BLOCK_SIZE_M) + E
-    grid = (max_M_tiles, triton.cdiv(N, BLOCK_SIZE_N))
     with device_context(A.device):
-        wrap_triton(w8a8_tensor_fp8_matmul_grouped_kernel)[grid](
             qA,
             B,
             C,
@@ -401,77 +570,190 @@ def _w8a8_tensor_fp8_matmul_grouped(
             C.stride(1),
             As.stride(0),
             Bs.stride(0),
-            # Meta-parameters
-            NUM_EXPERTS=E,
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
             BLOCK_SIZE_M=BLOCK_SIZE_M,
-            NUM_EXPERTS_BIT_LENGTH=E.bit_length(),
         )
     return C
-def w8a8_block_fp8_matmul_grouped(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     offsets: torch.Tensor,
     tokens_per_expert: torch.Tensor,
     block_size: list[int],
 ) -> torch.Tensor:
     """Block-scale grouped FP8 matmul with fused activation quantization.
     A:  (S, K) raw activations sorted by expert, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
     """
-    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_grouped(
-        A, B, Bs, offsets, tokens_per_expert, block_size
     )
-def w8a8_tensor_fp8_matmul_grouped(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     offsets: torch.Tensor,
     tokens_per_expert: torch.Tensor,
 ) -> torch.Tensor:
     """Tensor-scale grouped FP8 matmul with fused activation quantization.
     A:  (S, K) raw activations sorted by expert, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
     """
-    return torch.ops.finegrained_fp8.w8a8_tensor_fp8_matmul_grouped(
-        A, B, Bs, offsets, tokens_per_expert
     )
-def w8a8_fp8_matmul_grouped(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     offsets: torch.Tensor,
     tokens_per_expert: torch.Tensor,
     block_size: list[int] | None,
 ) -> torch.Tensor:
-    """Unified grouped W8A8 FP8 matmul dispatcher.
-    Dispatch rules:
-    - tensor mode when ``block_size is None``
-    - tensor mode when ``block_size == [N, K]``
-    - otherwise block mode
-    Returns:
-        Output tensor ``[S, N]`` in the same dtype as ``A``, in expert-sorted order.
     """
-    if block_size is None or (
-        block_size[0] == B.size(1) and block_size[1] == B.size(2)
-    ):
-        return w8a8_tensor_fp8_matmul_grouped(A, B, Bs, offsets, tokens_per_expert)
-    return w8a8_block_fp8_matmul_grouped(
-        A, B, Bs, offsets, tokens_per_expert, block_size
     )

 import torch
 import triton
 import triton.language as tl
 from torch.library import triton_op, wrap_triton
+from ._ops import add_op_namespace_prefix, ops
+from .bayesian_autotuner import bayesian_autotune
+from .utils import (
+    MX_SCALE_GROUP_K,
+    NIBBLES_PER_BYTE,
+    adaptive_block_size_m,
+    device_context,
+    mxfp_act_quant_inline,
+    mxfp4_e2m1_to_e4m3,
+    fp8_act_quant,
+    fp8_act_quant_inline,
+    get_accelerator_autotuning_configs,
+    get_mxfp_autotuning_configs,
+    is_mxfp,
+    is_tensor_wide,
+    e2m1_as_uint8,
+    ue8m0_as_uint8,
+    decode_ue8m0_scale,
+    mx_dot_rescale,
+    mx_dot_scaled,
+)
+def _grouped_tile_layout(
+    tokens_per_expert: torch.Tensor,
+    block_size_m: int,
+    S: int,
+    num_experts: int,
+) -> tuple[torch.Tensor, int]:
+    """Compute the M-tile layout for the grouped kernels.
+    Returns ``(tile_offsets, max_m_tiles)``:
+    - ``tile_offsets``: int32 (num_experts,) cumulative tile-end per expert, used by
+      ``_grouped_tile_setup`` to locate an M-tile's owning expert.
+    - ``max_m_tiles``: upper bound on total M-tiles, used as the grid axis-0
+      size. Real tile count <= this; surplus programs early-return inside the
+      kernel. Keeps the grid data-independent (cuda-graph / torch.compile safe).
+    """
+    tiles_per_expert = (tokens_per_expert + block_size_m - 1) // block_size_m
+    tile_offsets = torch.cumsum(tiles_per_expert, dim=0).to(torch.int32)
+    max_m_tiles = triton.cdiv(S, block_size_m) + num_experts
+    return tile_offsets, max_m_tiles
+@triton.jit
+def _grouped_tile_setup(
+    pid_m,
+    pid_n,
+    Offsets,
+    TileOffsets,
+    stride_offs,
+    stride_tile,
+    num_experts,
+    NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    """Map a grouped M-tile to its expert and build the per-tile offset vectors —
+    the prologue shared by every grouped kernel.
+    Returns ``(expert_id, offs_global_m, row_mask, offs_bn, offs_k)``:
+    - ``expert_id``: int64 owning expert
+    - ``offs_global_m``: ``(BLOCK_SIZE_M,)`` global row indices into A
+    - ``row_mask``: ``(BLOCK_SIZE_M,)`` validity mask within the expert's M
+    - ``offs_bn``: ``(BLOCK_SIZE_N,)`` output column offsets
+    - ``offs_k``: ``(BLOCK_SIZE_K,)`` K range
+    Caller must have early-returned if ``pid_m`` exceeds total_tiles
+    (``TileOffsets[(num_experts - 1) * stride_tile]``) — the ``Offsets`` load below
+    is out of bounds for an out-of-range tile otherwise.
+    """
+    # Binary search: upper_bound(TileOffsets, pid_m). NUM_EXPERTS_BIT_LENGTH is
+    # ceil(log2(num_experts))+1, giving one harmless extra iteration; constexpr so the
+    # loop unrolls.
+    lo = 0
+    hi = num_experts
+    for _ in tl.static_range(NUM_EXPERTS_BIT_LENGTH):
+        mid = (lo + hi) >> 1
+        mid_val = tl.load(TileOffsets + mid * stride_tile)
+        is_left = mid_val <= pid_m
+        lo = tl.where(is_left, mid + 1, lo)
+        hi = tl.where(is_left, hi, mid)
+    # Cast to int64 so ``expert_id * stride_be`` doesn't overflow for large num_experts
+    # × large weight matrices (e.g. 255 * 9_437_184 > 2^31).
+    expert_id = lo.to(tl.int64)
+    prev_eid = tl.maximum(expert_id - 1, 0)
+    expert_start = tl.where(
+        expert_id == 0, 0, tl.load(Offsets + prev_eid * stride_offs)
+    )
+    expert_end = tl.load(Offsets + expert_id * stride_offs)
+    M_expert = expert_end - expert_start
+    expert_tile_start = tl.where(
+        expert_id == 0, 0, tl.load(TileOffsets + prev_eid * stride_tile)
+    )
+    local_tile = pid_m - expert_tile_start
+    m_off = local_tile * BLOCK_SIZE_M
+    offs_am = m_off + tl.arange(0, BLOCK_SIZE_M)
+    row_mask = offs_am < M_expert
+    offs_global_m = expert_start + offs_am
+    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    return expert_id, offs_global_m, row_mask, offs_bn, offs_k
+@triton.jit
+def store_tile(C, accumulator, offs_global_m, offs_bn, row_mask, stride_cm, stride_cn):
+    """Output epilogue shared by the grouped kernels: cast the fp32 accumulator to
+    ``C``'s dtype and store the tile at expert-sorted global rows ``offs_global_m`` ×
+    columns ``offs_bn``, masked to the expert's valid rows (``row_mask``)."""
+    c = accumulator.to(C.dtype.element_ty)
+    c_ptrs = C + stride_cm * offs_global_m[:, None] + stride_cn * offs_bn[None, :]
+    tl.store(c_ptrs, c, mask=row_mask[:, None])
 @triton.autotune(
+    configs=get_accelerator_autotuning_configs(),
     key=["N", "K", "BLOCK_SIZE_M"],
 )
 @triton.jit
+def w8a8_block_dynamic_fp8_matmul_grouped_kernel(
+    A,  # (S, K) raw BF16/FP16 activations, sorted/grouped by expert id
+    B,  # (num_experts, N, K) FP8 weight matrices
+    C,  # (S, N) output
+    Bs,  # (num_experts, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales
+    Offsets,  # (num_experts,) int32 — cumulative row-end per expert
+    TileOffsets,  # (num_experts,) int32 — cumulative tile-end per expert
     # Shape
     S,
     N,
     stride_bs_e,
     stride_bs_k,
     stride_bs_n,
+    stride_offs,
+    stride_tile,
+    num_experts,
     # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
     NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
 ):
     """Block-scale grouped FP8 expert matmul kernel.
     pid_m = tl.program_id(axis=0)
     pid_n = tl.program_id(axis=1)
+    total_tiles = tl.load(TileOffsets + (num_experts - 1) * stride_tile)
     if pid_m >= total_tiles:
         return
+    expert_id, offs_global_m, row_mask, offs_bn, offs_k = _grouped_tile_setup(
+        pid_m,
+        pid_n,
+        Offsets,
+        TileOffsets,
+        stride_offs,
+        stride_tile,
+        num_experts,
+        NUM_EXPERTS_BIT_LENGTH,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+        BLOCK_SIZE_K,
+    )
     a_ptrs = A + offs_global_m[:, None] * stride_am + offs_k[None, :] * stride_ak
     b_ptrs = (
     accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a_raw = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0).to(tl.float32)
+        a, a_s = fp8_act_quant_inline(a_raw)
         b = tl.load(b_ptrs)
+        b_s = decode_ue8m0_scale(tl.load(bs_ptrs))
         accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
+        bs_ptrs += stride_bs_k
+    store_tile(C, accumulator, offs_global_m, offs_bn, row_mask, stride_cm, stride_cn)
+@bayesian_autotune(
+    get_accelerator_autotuning_configs(tune_block_nk=True),
+    ["N", "K", "BLOCK_SIZE_M"],
+    n_trials=60,
 )
 @triton.jit
+def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel(
+    A,  # (S, K) pre-quantized FP8 activations, sorted/grouped by expert id
+    B,  # (num_experts, N, K) FP8 weight matrices
     C,  # (S, N) output
+    As,  # (S,) per-token activation scales
+    Bs,  # (num_experts, 1, 1) per-tensor weight scales
+    Offsets,  # (num_experts,) int32 — cumulative row-end per expert
+    TileOffsets,  # (num_experts,) int32 — cumulative tile-end per expert
+    # Shape
     S,
     N,
     K,
+    # Strides
     stride_am,
     stride_ak,
     stride_be,
     stride_cn,
     stride_as_m,
     stride_bs_e,
+    stride_offs,
+    stride_tile,
+    num_experts,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
     NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
 ):
     """Tensor-scale grouped FP8 expert matmul kernel.
     pid_m = tl.program_id(axis=0)
     pid_n = tl.program_id(axis=1)
+    total_tiles = tl.load(TileOffsets + (num_experts - 1) * stride_tile)
     if pid_m >= total_tiles:
         return
+    expert_id, offs_global_m, row_mask, offs_bn, offs_k = _grouped_tile_setup(
+        pid_m,
+        pid_n,
+        Offsets,
+        TileOffsets,
+        stride_offs,
+        stride_tile,
+        num_experts,
+        NUM_EXPERTS_BIT_LENGTH,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+        BLOCK_SIZE_K,
+    )
     a_ptrs = A + offs_global_m[:, None] * stride_am + offs_k[None, :] * stride_ak
     b_ptrs = (
     for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0)
         b = tl.load(b_ptrs)
         accumulator += tl.dot(a, b)
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
     accumulator = accumulator * a_s[:, None] * b_s
+    store_tile(C, accumulator, offs_global_m, offs_bn, row_mask, stride_cm, stride_cn)
+@bayesian_autotune(
+    get_mxfp_autotuning_configs(
+        compute_modes=("dot_scaled", "dot")
+    ),  # prefill: no scalar branch
+    ["N", "K", "BLOCK_SIZE_M"],
+    n_trials=60,
+)
+@triton.jit
+def mxfp_dynamic_matmul_grouped_kernel(
+    A,  # (S, K) raw BF16/FP16 activations, sorted by expert id
+    B,  # (num_experts, N, K) E4M3 (MXFP8) or (num_experts, N, K // 2) packed E2M1 (MXFP4) expert weights
+    C,  # (S, N) output
+    Bs,  # (num_experts, N, K // SCALE_GROUP_K) UE8M0 weight scales
+    Offsets,  # (num_experts,) int32 — cumulative row-end per expert
+    TileOffsets,  # (num_experts,) int32 — cumulative tile-end per expert
+    # Shape
+    S,
+    N,
+    K,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_be,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_bs_e,
+    stride_bs_k,
+    stride_bs_n,
+    stride_offs,
+    stride_tile,
+    num_experts,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
+    VALUES_PER_BYTE: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    COMPUTE_MODE: tl.constexpr,
+):
+    """Unified grouped MXFP4/MXFP8 (W4A8/W8A8) expert matmul with fused act quant.
+    Tokens sorted by expert; each M-tile maps to its expert via ``TileOffsets``. ``A`` is
+    quantized to E4M3 per K-group inline (UE8M0 scale). ``VALUES_PER_BYTE`` picks the
+    weight format (2 = packed E2M1 / MXFP4, 1 = unpacked E4M3 / MXFP8); ``COMPUTE_MODE``
+    picks ``tl.dot_scaled`` vs fp8 ``tl.dot`` + per-group rescale (decode; FP4 unpacks
+    E2M1->E4M3 first, lossless).
+    """
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    total_tiles = tl.load(TileOffsets + (num_experts - 1) * stride_tile)
+    if pid_m >= total_tiles:
+        return
+    expert_id, offs_global_m, row_mask, offs_bn, offs_k = _grouped_tile_setup(
+        pid_m,
+        pid_n,
+        Offsets,
+        TileOffsets,
+        stride_offs,
+        stride_tile,
+        num_experts,
+        NUM_EXPERTS_BIT_LENGTH,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+        BLOCK_SIZE_K,
+    )
+    offs_kb = tl.arange(0, BLOCK_SIZE_K // VALUES_PER_BYTE)
+    offs_sf = tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K)
+    a_ptrs = A + offs_global_m[:, None] * stride_am + offs_k[None, :] * stride_ak
+    b_ptrs = (
+        B
+        + expert_id * stride_be
+        + offs_kb[:, None] * stride_bk
+        + offs_bn[None, :] * stride_bn
+    )
+    bs_ptrs = (
+        Bs
+        + expert_id * stride_bs_e
+        + offs_bn[:, None] * stride_bs_n
+        + offs_sf[None, :] * stride_bs_k
+    )
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        a_raw = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0).to(tl.float32)
+        a, a_scale = mxfp_act_quant_inline(
+            a_raw, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K
+        )
+        b = tl.load(b_ptrs)
+        b_s = tl.load(bs_ptrs).to(tl.uint8)
+        bq = mxfp4_e2m1_to_e4m3(b) if VALUES_PER_BYTE == 2 else b
+        if COMPUTE_MODE == "dot_scaled":
+            accumulator = mx_dot_scaled(
+                a, a_scale, b, b_s, accumulator, VALUES_PER_BYTE
+            )
+        else:  # dot
+            accumulator = mx_dot_rescale(accumulator, a, bq, a_scale, b_s)
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += (BLOCK_SIZE_K // VALUES_PER_BYTE) * stride_bk
+        bs_ptrs += (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_bs_k
+    store_tile(C, accumulator, offs_global_m, offs_bn, row_mask, stride_cm, stride_cn)
+@triton_op(
+    add_op_namespace_prefix("w8a8_block_dynamic_fp8_matmul_grouped"), mutates_args=()
+)
+def _w8a8_block_dynamic_fp8_matmul_grouped(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     offsets: torch.Tensor,
     tokens_per_expert: torch.Tensor,
     block_size: list[int],
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Block-scale grouped FP8 matmul: C = A @ B.T per expert, with fused act quant.
     A:  (S, K) raw bf16/fp16 activations, sorted by expert
+    B:  (num_experts, N, K) FP8 expert weights
+    Bs: (num_experts, N // block_n, K // block_k) per-block weight scales
     """
     assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
     assert A.is_contiguous(), "A must be contiguous"
+    assert B.ndim == 3, f"B must be 3D (num_experts, N, K), got ndim={B.ndim}"
     assert B.is_contiguous(), "B must be contiguous"
     assert A.shape[1] == B.shape[2], (
         f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
     )
     S, K = A.shape
+    num_experts, N, _ = B.shape
     assert len(block_size) == 2, (
         f"block_size must be [block_n, block_k], got {block_size}"
     assert N % block_n == 0, f"N ({N}) must be divisible by block_n ({block_n})"
     assert K % block_k == 0, f"K ({K}) must be divisible by block_k ({block_k})"
     assert Bs.ndim == 3, (
+        f"Bs must be 3D (num_experts, N//block_n, K//block_k), got ndim={Bs.ndim}"
     )
+    assert Bs.shape == (num_experts, N // block_n, K // block_k), (
+        f"Bs shape {tuple(Bs.shape)} != expected ({num_experts}, {N // block_n}, {K // block_k})"
     )
+    Bs = ue8m0_as_uint8(Bs)
+    C = A.new_empty(S, N, dtype=output_dtype)
+    BLOCK_SIZE_M = adaptive_block_size_m((S + num_experts - 1) // num_experts)
+    tile_offsets, max_m_tiles = _grouped_tile_layout(
+        tokens_per_expert, BLOCK_SIZE_M, S, num_experts
+    )
+    grid = (max_m_tiles, triton.cdiv(N, block_n))
     with device_context(A.device):
+        wrap_triton(w8a8_block_dynamic_fp8_matmul_grouped_kernel)[grid](
             A,
             B,
             C,
             Bs.stride(0),
             Bs.stride(2),
             Bs.stride(1),
+            offsets.stride(0),
+            tile_offsets.stride(0),
             # Meta-parameters
+            num_experts=num_experts,
             BLOCK_SIZE_N=block_n,
             BLOCK_SIZE_K=block_k,
             BLOCK_SIZE_M=BLOCK_SIZE_M,
+            NUM_EXPERTS_BIT_LENGTH=num_experts.bit_length(),
         )
     return C
+@triton_op(
+    add_op_namespace_prefix("w8a8_tensor_dynamic_fp8_matmul_grouped"), mutates_args=()
+)
+def _w8a8_tensor_dynamic_fp8_matmul_grouped(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     offsets: torch.Tensor,
     tokens_per_expert: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Tensor-scale grouped FP8 matmul: C = A @ B.T per expert, with fused act quant.
     A:  (S, K) raw bf16/fp16 activations, sorted by expert
+    B:  (num_experts, N, K) FP8 expert weights
+    Bs: (num_experts,) or (num_experts, 1, 1) per-expert weight scales
     """
     assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
     assert A.is_contiguous(), "A must be contiguous"
+    assert B.ndim == 3, f"B must be 3D (num_experts, N, K), got ndim={B.ndim}"
     assert B.is_contiguous(), "B must be contiguous"
     assert A.shape[1] == B.shape[2], (
         f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
     )
     S, K = A.shape
+    num_experts, N, _ = B.shape
+    # Normalize Bs to (num_experts, 1, 1)
     if Bs.ndim == 1:
+        assert Bs.shape[0] == num_experts, (
+            f"Bs shape {tuple(Bs.shape)} != expected ({num_experts},)"
+        )
+        Bs = Bs.reshape(num_experts, 1, 1)
     else:
+        assert Bs.shape == (num_experts, 1, 1), (
+            f"Bs shape {tuple(Bs.shape)} != expected ({num_experts}, 1, 1)"
         )
     qA, As = fp8_act_quant(A, K)
+    C = A.new_empty(S, N, dtype=output_dtype)
+    BLOCK_SIZE_M = adaptive_block_size_m((S + num_experts - 1) // num_experts)
+    tile_offsets, max_m_tiles = _grouped_tile_layout(
+        tokens_per_expert, BLOCK_SIZE_M, S, num_experts
+    )
+    def grid(META):
+        return (max_m_tiles, triton.cdiv(N, META["BLOCK_SIZE_N"]))
     with device_context(A.device):
+        wrap_triton(w8a8_tensor_dynamic_fp8_matmul_grouped_kernel)[grid](
             qA,
             B,
             C,
             C.stride(1),
             As.stride(0),
             Bs.stride(0),
+            offsets.stride(0),
+            tile_offsets.stride(0),
+            # Meta-parameters (BLOCK_SIZE_N, BLOCK_SIZE_K come from autotune Config)
+            num_experts=num_experts,
             BLOCK_SIZE_M=BLOCK_SIZE_M,
+            NUM_EXPERTS_BIT_LENGTH=num_experts.bit_length(),
         )
     return C
+def w8a8_block_dynamic_fp8_matmul_grouped(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     offsets: torch.Tensor,
     tokens_per_expert: torch.Tensor,
     block_size: list[int],
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Block-scale grouped FP8 matmul with fused activation quantization.
     A:  (S, K) raw activations sorted by expert, bf16/fp16/fp32
+    B:  (num_experts, N, K) FP8 expert weights
+    Bs: (num_experts, N // block_n, K // block_k) per-block weight scales
+    output_dtype: defaults to ``A.dtype``
     """
+    return ops.w8a8_block_dynamic_fp8_matmul_grouped(
+        A, B, Bs, offsets, tokens_per_expert, block_size, output_dtype
     )
+def w8a8_tensor_dynamic_fp8_matmul_grouped(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     offsets: torch.Tensor,
     tokens_per_expert: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
     """Tensor-scale grouped FP8 matmul with fused activation quantization.
     A:  (S, K) raw activations sorted by expert, bf16/fp16/fp32
+    B:  (num_experts, N, K) FP8 expert weights
+    Bs: (num_experts,) or (num_experts, 1, 1) per-expert weight scales
+    output_dtype: defaults to ``A.dtype``
+    """
+    return ops.w8a8_tensor_dynamic_fp8_matmul_grouped(
+        A, B, Bs, offsets, tokens_per_expert, output_dtype
+    )
+@triton_op(add_op_namespace_prefix("mxfp_dynamic_matmul_grouped"), mutates_args=())
+def _mxfp_dynamic_matmul_grouped(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    offsets: torch.Tensor,
+    tokens_per_expert: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Grouped MX matmul with fused act quant — per-expert ``C[s] = A[s] @ B[e].T`` over
+    contiguous, expert-sorted rows. Weight format detected from ``B.dtype``: ``int8`` →
+    packed E2M1 (MXFP4, ``B`` is ``(num_experts, N, K//2)``); ``float8_e4m3fn`` → unpacked E4M3
+    (MXFP8, ``(num_experts, N, K)``). UE8M0 group-32 scales ``(num_experts, N, K//32)``; tile + dot autotuned.
+    A:  (S, K) raw activations, bf16/fp16/fp32, expert-sorted (quantized inline to E4M3)
+    offsets: (num_experts,) — exclusive prefix of expert token counts (cumsum)
+    tokens_per_expert: (num_experts,) — per-expert row count
     """
+    assert A.ndim == 2 and B.ndim == 3 and Bs.ndim == 3
+    assert offsets.ndim == 1 and tokens_per_expert.ndim == 1
+    assert B.dtype in (torch.int8, torch.float8_e4m3fn), (
+        f"B must be int8 (packed E2M1) or float8_e4m3fn (E4M3), got {B.dtype}"
+    )
+    assert Bs.dtype == torch.float8_e8m0fnu, (
+        f"Bs must be float8_e8m0fnu, got {Bs.dtype}"
+    )
+    VALUES_PER_BYTE = NIBBLES_PER_BYTE if B.dtype == torch.int8 else 1
+    S, K = A.shape
+    num_experts, N, K_b = B.shape
+    assert offsets.shape[0] == num_experts and tokens_per_expert.shape[0] == num_experts
+    assert K == VALUES_PER_BYTE * K_b, (
+        f"K (={K}) must equal {VALUES_PER_BYTE} * B.shape[2] (={K_b})"
+    )
+    assert K % MX_SCALE_GROUP_K == 0, (
+        f"K (={K}) must be a multiple of {MX_SCALE_GROUP_K}"
+    )
+    assert Bs.shape == (num_experts, N, K // MX_SCALE_GROUP_K), (
+        f"Bs shape {tuple(Bs.shape)} != ({num_experts}, {N}, {K // MX_SCALE_GROUP_K})"
+    )
+    B = e2m1_as_uint8(B)
+    bs_u8 = ue8m0_as_uint8(Bs)
+    C = A.new_empty((S, N), dtype=output_dtype)
+    BLOCK_SIZE_M = adaptive_block_size_m((S + num_experts - 1) // num_experts)
+    tile_offsets, max_m_tiles = _grouped_tile_layout(
+        tokens_per_expert, BLOCK_SIZE_M, S, num_experts
+    )
+    def grid(META):
+        return (max_m_tiles, triton.cdiv(N, META["BLOCK_SIZE_N"]))
+    with device_context(A.device):
+        wrap_triton(mxfp_dynamic_matmul_grouped_kernel)[grid](
+            A,
+            B,
+            C,
+            bs_u8,
+            offsets,
+            tile_offsets,
+            S,
+            N,
+            K,
+            A.stride(0),
+            A.stride(1),
+            B.stride(0),
+            B.stride(2),
+            B.stride(1),
+            C.stride(0),
+            C.stride(1),
+            bs_u8.stride(0),
+            bs_u8.stride(2),
+            bs_u8.stride(1),
+            offsets.stride(0),
+            tile_offsets.stride(0),
+            # Meta-parameters (BLOCK_SIZE_N, BLOCK_SIZE_K come from autotune Config)
+            num_experts=num_experts,
+            BLOCK_SIZE_M=BLOCK_SIZE_M,
+            NUM_EXPERTS_BIT_LENGTH=num_experts.bit_length(),
+            VALUES_PER_BYTE=VALUES_PER_BYTE,
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+        )
+    return C
+def mxfp_dynamic_matmul_grouped(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    offsets: torch.Tensor,
+    tokens_per_expert: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Grouped MX (MXFP4/MXFP8) matmul with fused act quant — per-expert
+    ``C[s] = A[s] @ B[e].T`` over expert-sorted rows; weight format detected from
+    ``B.dtype``. Tile + dot path autotuned."""
+    return ops.mxfp_dynamic_matmul_grouped(
+        A, B, Bs, offsets, tokens_per_expert, output_dtype
     )
+def matmul_grouped(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     offsets: torch.Tensor,
     tokens_per_expert: torch.Tensor,
     block_size: list[int] | None,
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
+    """Grouped quantized matmul dispatcher (W8A8 FP8 or W4A8 FP4). Tokens must
+    be sorted by expert; M-tiles are mapped to experts via ``offsets``.
+    ``output_dtype`` defaults to ``A.dtype``.
+    Routes by weight dtype and ``block_size``:
+    - MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0
+      group-32 ``Bs`` (shape ``[num_experts, N, K//32]``) → ``mxfp_dynamic_matmul_grouped``
+      (``block_size`` ignored; tile + dot path autotuned).
+    - ``block_size`` None or full ``[N, K]`` → ``w8a8_tensor_dynamic_fp8_matmul_grouped``.
+    - otherwise → ``w8a8_block_dynamic_fp8_matmul_grouped``.
     """
+    if is_mxfp(B, Bs):
+        return mxfp_dynamic_matmul_grouped(
+            A, B, Bs, offsets, tokens_per_expert, output_dtype
+        )
+    if is_tensor_wide(block_size, B):
+        return w8a8_tensor_dynamic_fp8_matmul_grouped(
+            A, B, Bs, offsets, tokens_per_expert, output_dtype
+        )
+    return w8a8_block_dynamic_fp8_matmul_grouped(
+        A, B, Bs, offsets, tokens_per_expert, block_size, output_dtype
     )

build/torch-cuda/matmul.py CHANGED Viewed

@@ -17,38 +17,103 @@ import triton
 import triton.language as tl
 from torch.library import triton_op, wrap_triton
-from .utils import device_context
-# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py
 @triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4]
-    ],
     key=["N", "K", "BLOCK_SIZE_M"],
 )
 @triton.jit
-def w8a8_block_fp8_matmul_kernel(
-    # Pointers to inputs and output
-    A,
-    B,
-    C,
-    As,
-    Bs,
-    # Shape for matmul
     M,
     N,
     K,
     stride_am,
     stride_ak,
     stride_bk,
     stride_bn,
     stride_cm,
     stride_cn,
-    stride_as_m,
-    stride_as_k,
     stride_bs_k,
     stride_bs_n,
     # Meta-parameters
@@ -57,72 +122,65 @@ def w8a8_block_fp8_matmul_kernel(
     BLOCK_SIZE_K: tl.constexpr,
     GROUP_SIZE_M: tl.constexpr,
 ):
-    """Block-scale FP8 GEMM kernel.
-    Computes ``C = A @ B.T`` with block-wise activation/weight scales.
-    Uses a 2D grid with swizzle for L2 cache locality on B tiles.
     """
-    pid_m = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-    pid_m, pid_n = tl.swizzle2d(pid_m, pid_n, num_pid_m, num_pid_n, GROUP_SIZE_M)
-    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
-    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
     a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
     b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
-    as_ptrs = As + offs_am * stride_as_m
     offs_bsn = offs_bn // BLOCK_SIZE_N
     bs_ptrs = Bs + offs_bsn * stride_bs_n
     accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         k_remaining = K - k * BLOCK_SIZE_K
-        a = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
         b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
-        a_s = tl.load(as_ptrs + k * stride_as_k)
-        b_s = tl.load(bs_ptrs + k * stride_bs_k)
         accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-    tl.store(c_ptrs, c, mask=c_mask)
 @triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4]
-    ],
     key=["N", "K", "BLOCK_SIZE_M"],
 )
 @triton.jit
-def w8a8_tensor_fp8_matmul_kernel(
-    A,
-    B,
-    C,
-    As,
-    Bs,
     M,
     N,
     K,
     stride_am,
     stride_ak,
     stride_bk,
@@ -130,6 +188,7 @@ def w8a8_tensor_fp8_matmul_kernel(
     stride_cm,
     stride_cn,
     stride_as_m,
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
@@ -141,16 +200,9 @@ def w8a8_tensor_fp8_matmul_kernel(
     weight scale for the full matrix.
     Uses a 2D grid with swizzle for L2 cache locality on B tiles.
     """
-    pid_m = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-    pid_m, pid_n = tl.swizzle2d(pid_m, pid_n, num_pid_m, num_pid_n, GROUP_SIZE_M)
-    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
-    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
     a_ptrs = A + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
     b_ptrs = B + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn
@@ -169,33 +221,201 @@ def w8a8_tensor_fp8_matmul_kernel(
     accumulator = accumulator * a_s[:, None] * b_s
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-    tl.store(c_ptrs, c, mask=c_mask)
-@triton_op("finegrained_fp8::w8a8_block_fp8_matmul", mutates_args=())
-def _w8a8_block_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
-    As: torch.Tensor,
     Bs: torch.Tensor,
     block_size: list[int],
-    output_dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
-    """Block-scale FP8 matmul: C = A @ B.T with per-block scales.
-    As: (M, K // block_k) — per-token-group activation scales
-    Bs: (N // block_n, K // block_k) — per-block weight scales
     """
     assert len(block_size) == 2, (
         f"block_size must be [block_n, block_k], got {block_size}"
@@ -212,10 +432,6 @@ def _w8a8_block_fp8_matmul(
     N, K = B.shape
     M = A.numel() // A.shape[-1]
-    assert As.ndim >= 2, f"As must be at least 2D, got ndim={As.ndim}"
-    assert As.shape[-1] == triton.cdiv(K, block_k), (
-        f"As last dim {As.shape[-1]} != expected {triton.cdiv(K, block_k)} (cdiv(K={K}, block_k={block_k}))"
-    )
     assert Bs.ndim == 2, f"Bs must be 2D (N//block_n, K//block_k), got ndim={Bs.ndim}"
     assert Bs.shape == (triton.cdiv(N, block_n), triton.cdiv(K, block_k)), (
         f"Bs shape {tuple(Bs.shape)} != expected ({triton.cdiv(N, block_n)}, {triton.cdiv(K, block_k)})"
@@ -223,20 +439,17 @@ def _w8a8_block_fp8_matmul(
     BLOCK_SIZE_K = block_k
     BLOCK_SIZE_N = block_n
     C_shape = A.shape[:-1] + (N,)
     C = A.new_empty(C_shape, dtype=output_dtype)
-    # Adaptive BLOCK_SIZE_M: smallest power-of-2 >= M, floored at 16, capped at 128.
-    # Matches the WGMMA tile to the actual row count — smaller tiles use less
-    # register pressure and a better-matched FP8 WGMMA instruction, improving
-    # both accuracy and performance for small M (decode).
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2(M), 16), 128)
     grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
     with device_context(A.device):
-        wrap_triton(w8a8_block_fp8_matmul_kernel)[grid](
             A,
             B,
             C,
-            As,
             Bs,
             M,
             N,
@@ -247,64 +460,73 @@ def _w8a8_block_fp8_matmul(
             B.stride(0),
             C.stride(-2),
             C.stride(-1),
-            As.stride(-2),
-            As.stride(-1),
             Bs.stride(1),
             Bs.stride(0),
             # Meta-parameters
             BLOCK_SIZE_M=BLOCK_SIZE_M,
             BLOCK_SIZE_N=BLOCK_SIZE_N,
             BLOCK_SIZE_K=BLOCK_SIZE_K,
-            GROUP_SIZE_M=8,
         )
     return C
-@triton_op("finegrained_fp8::w8a8_tensor_fp8_matmul", mutates_args=())
-def _w8a8_tensor_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
-    As: torch.Tensor,
     Bs: torch.Tensor,
-    output_dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
-    """Tensor-scale FP8 matmul: C = A @ B.T with per-row / per-tensor scales.
-    As: scalar, (M,), or (M, 1) — per-row activation scales
-    Bs: scalar, (1,), or (1, 1) — single weight scale
     """
     assert A.shape[-1] == B.shape[-1], (
         f"K mismatch: A has K={A.shape[-1]}, B has K={B.shape[-1]}"
     )
     assert A.is_contiguous(), "A must be contiguous"
     assert B.ndim == 2, f"B must be 2D (N, K), got ndim={B.ndim}"
     assert B.is_contiguous(), "B must be contiguous"
     N, K = B.shape
     M = A.numel() // A.shape[-1]
-    # Normalize As to (M,)
-    if As.numel() == 1:
-        As = As.reshape(1).expand(M).contiguous()
-    elif As.ndim == 2:
-        As = As.reshape(M)
-    assert As.ndim == 1 and As.shape[0] == M, (
-        f"As must be scalar, (M,), or (M,1) with M={M}, got {tuple(As.shape)}"
     )
-    # Normalize Bs to (1,)
-    assert Bs.numel() == 1, f"Bs must be scalar or (1,), got {tuple(Bs.shape)}"
-    Bs = Bs.reshape(1)
-    BLOCK_SIZE_N = 128
-    BLOCK_SIZE_K = 128
     C_shape = A.shape[:-1] + (N,)
     C = A.new_empty(C_shape, dtype=output_dtype)
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2(M), 16), 128)
-    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
     with device_context(A.device):
-        wrap_triton(w8a8_tensor_fp8_matmul_kernel)[grid](
             A,
             B,
             C,
@@ -319,93 +541,277 @@ def _w8a8_tensor_fp8_matmul(
             B.stride(0),
             C.stride(-2),
             C.stride(-1),
-            As.stride(0),
             # Meta-parameters
             BLOCK_SIZE_M=BLOCK_SIZE_M,
             BLOCK_SIZE_N=BLOCK_SIZE_N,
             BLOCK_SIZE_K=BLOCK_SIZE_K,
-            GROUP_SIZE_M=8,
         )
     return C
-def w8a8_block_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
-    As: torch.Tensor,
     Bs: torch.Tensor,
-    block_size: list[int],
-    output_dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
-    """Block-wise W8A8 FP8 matrix multiplication.
-    Computes ``C = A @ B.T`` where both operands are pre-quantized to
-    ``float8_e4m3fn`` with per-block scales, and accumulates in float32
-    before casting to ``output_dtype``.
     Args:
-        A: Quantized activation tensor ``[M, K]`` in ``float8_e4m3fn``.
-        B: Quantized weight tensor ``[N, K]`` in ``float8_e4m3fn``.
-        As: Per-token-group activation scales ``[M, K // block_size[1]]``.
         Bs: Per-block weight scales ``[N // block_size[0], K // block_size[1]]``.
-        block_size: ``[block_n, block_k]`` quantization block dimensions, e.g. ``[128, 128]``.
         output_dtype: dtype of the returned tensor (default: ``torch.float32``).
     Returns:
-        Output tensor ``[M, N]`` in ``output_dtype``.
     """
-    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul(
-        A, B, As, Bs, block_size, output_dtype
-    )
-def w8a8_tensor_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
     As: torch.Tensor,
     Bs: torch.Tensor,
-    output_dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
-    """Tensor-scale W8A8 FP8 matrix multiplication.
-    Computes ``C = A @ B.T`` in tensor-scale mode using pre-quantized FP8
-    activations/weights and tensor scales.
     Args:
-        A: Quantized activation tensor ``[M, K]`` in ``float8_e4m3fn``.
-        B: Quantized weight tensor ``[N, K]`` in ``float8_e4m3fn``.
-        As: Per-row activation scales ``[M]``.
         Bs: Single weight scale, scalar or ``[1]``.
         output_dtype: dtype of the returned tensor.
-    Returns:
-        Output tensor ``[M, N]`` in ``output_dtype``.
     """
-    return torch.ops.finegrained_fp8.w8a8_tensor_fp8_matmul(A, B, As, Bs, output_dtype)
-def w8a8_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
-    As: torch.Tensor,
     Bs: torch.Tensor,
     block_size: list[int] | None,
-    output_dtype: torch.dtype = torch.float32,
 ) -> torch.Tensor:
-    """Unified W8A8 FP8 matmul dispatcher.
-    Dispatch rules:
-    - tensor mode when ``block_size is None``
-    - tensor mode when ``block_size == [N, K]``
-    - otherwise block mode
-    Returns:
-        Output tensor ``[M, N]`` in ``output_dtype``.
-    """
-    if block_size is None or (
-        block_size[0] == B.size(0) and block_size[1] == B.size(1)
-    ):
-        return w8a8_tensor_fp8_matmul(A, B, As, Bs, output_dtype)
-    return w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype)

 import triton.language as tl
 from torch.library import triton_op, wrap_triton
+from ._ops import add_op_namespace_prefix, ops
+from .bayesian_autotuner import bayesian_autotune
+from .utils import (
+    NIBBLES_PER_BYTE,
+    MX_SCALE_GROUP_K,
+    adaptive_block_size_m,
+    decode_ue8m0_scale,
+    mx_dot_rescale,
+    mx_dot_scaled,
+    device_context,
+    fp8_act_quant,
+    fp8_act_quant_inline,
+    get_accelerator_autotuning_configs,
+    get_mxfp_autotuning_configs,
+    is_mxfp,
+    is_tensor_wide,
+    mxfp_act_quant_inline,
+    mxfp4_e2m1_to_e4m3,
+    e2m1_as_uint8,
+    ue8m0_as_uint8,
+)
+# Swizzle group size for the 2D-grid kernels' L2-locality tiling (``_swizzle_offsets``) —
+# a perf knob passed as the ``GROUP_SIZE_M`` constexpr, not a correctness parameter.
+GROUP_SIZE_M = 8
+@triton.jit
+def _swizzle_offsets(
+    M,
+    N,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+):
+    """2D-grid tile scheduling shared by the kernels below: swizzle the
+    ``(pid_m, pid_n)`` program ids for L2 locality on B, then build the operand
+    offset vectors. Returns ``(pid_m, pid_n, offs_am, offs_bn, offs_k)`` — the
+    swizzled ids (reused by the output store) and the ``%``-wrapped row/col offsets
+    plus the K range."""
+    pid_m = tl.program_id(axis=0)
+    pid_n = tl.program_id(axis=1)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    pid_m, pid_n = tl.swizzle2d(pid_m, pid_n, num_pid_m, num_pid_n, GROUP_SIZE_M)
+    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    return pid_m, pid_n, offs_am, offs_bn, offs_k
+@triton.jit
+def _store_masked(
+    C,
+    accumulator,
+    pid_m,
+    pid_n,
+    M,
+    N,
+    stride_cm,
+    stride_cn,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+):
+    """Shared output epilogue of the kernels below: cast the fp32 accumulator to
+    ``C``'s dtype and store the ``(BLOCK_SIZE_M, BLOCK_SIZE_N)`` tile at the swizzled
+    ``(pid_m, pid_n)``, masked to the ``(M, N)`` bounds."""
+    c = accumulator.to(C.dtype.element_ty)
+    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
+    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
+    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, c, mask=c_mask)
 @triton.autotune(
+    configs=get_accelerator_autotuning_configs(),
     key=["N", "K", "BLOCK_SIZE_M"],
 )
 @triton.jit
+def w8a8_block_dynamic_fp8_matmul_kernel(
+    A,  # (M, K) raw BF16/FP16 activations
+    B,  # (N, K) FP8 weights
+    C,  # (M, N) output
+    Bs,  # (N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales (fp32 or uint8/UE8M0)
+    # Shape
     M,
     N,
     K,
+    # Strides
     stride_am,
     stride_ak,
     stride_bk,
     stride_bn,
     stride_cm,
     stride_cn,
     stride_bs_k,
     stride_bs_n,
     # Meta-parameters
     BLOCK_SIZE_K: tl.constexpr,
     GROUP_SIZE_M: tl.constexpr,
 ):
+    """Block-scale FP8 GEMM kernel with fused activation quantization.
+    Computes ``C = A @ B.T`` with bf16/fp16 ``A`` quantized to FP8 per-K-tile
+    inline (one scale per M-row per BLOCK_SIZE_K) and pre-quantized FP8 weights
+    with per-block scales. 2D grid with swizzle for L2 cache locality on B.
     """
+    pid_m, pid_n, offs_am, offs_bn, offs_k = _swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M
+    )
     a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
     b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
     offs_bsn = offs_bn // BLOCK_SIZE_N
     bs_ptrs = Bs + offs_bsn * stride_bs_n
     accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         k_remaining = K - k * BLOCK_SIZE_K
+        a_raw = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0).to(
+            tl.float32
+        )
+        a, a_s = fp8_act_quant_inline(a_raw)
         b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
+        b_s = decode_ue8m0_scale(tl.load(bs_ptrs))
         accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
+        bs_ptrs += stride_bs_k
+    _store_masked(
+        C,
+        accumulator,
+        pid_m,
+        pid_n,
+        M,
+        N,
+        stride_cm,
+        stride_cn,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+    )
 @triton.autotune(
+    configs=get_accelerator_autotuning_configs(tune_block_nk=True),
     key=["N", "K", "BLOCK_SIZE_M"],
 )
 @triton.jit
+def w8a8_tensor_dynamic_fp8_matmul_kernel(
+    A,  # (M, K) pre-quantized FP8 activations
+    B,  # (N, K) FP8 weights
+    C,  # (M, N) output
+    As,  # (M,) per-token activation scales
+    Bs,  # scalar/(1,) per-tensor weight scale
+    # Shape
     M,
     N,
     K,
+    # Strides
     stride_am,
     stride_ak,
     stride_bk,
     stride_cm,
     stride_cn,
     stride_as_m,
+    # Meta-parameters
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
     weight scale for the full matrix.
     Uses a 2D grid with swizzle for L2 cache locality on B tiles.
     """
+    pid_m, pid_n, offs_am, offs_bn, offs_k = _swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M
+    )
     a_ptrs = A + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
     b_ptrs = B + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn
     accumulator = accumulator * a_s[:, None] * b_s
+    _store_masked(
+        C,
+        accumulator,
+        pid_m,
+        pid_n,
+        M,
+        N,
+        stride_cm,
+        stride_cn,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+    )
+@triton.autotune(
+    configs=get_accelerator_autotuning_configs(),
+    key=["N", "K", "BLOCK_SIZE_M"],
+)
+@triton.jit
+def w8a8_block_static_fp8_matmul_kernel(
+    A,  # (M, K) raw BF16/FP16 activations
+    B,  # (N, K) FP8 weights
+    C,  # (M, N) output
+    As,  # scalar — static per-tensor activation scale (calibration-time)
+    Bs,  # (N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales (fp32 or uint8/UE8M0)
+    # Shape
+    M,
+    N,
+    K,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_bs_k,
+    stride_bs_n,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+):
+    """Block-scale FP8 GEMM with static (per-tensor) activation scale.
+    ``A`` is raw bf16/fp16; the kernel divides by the scalar ``As`` and casts
+    to FP8 inline. Per-block weight scales apply per-K-tile during
+    accumulation; the scalar activation scale factors out of the loop and
+    is applied once at the end.
+    """
+    pid_m, pid_n, offs_am, offs_bn, offs_k = _swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M
+    )
+    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
+    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+    offs_bsn = offs_bn // BLOCK_SIZE_N
+    bs_ptrs = Bs + offs_bsn * stride_bs_n
+    a_s_static = tl.load(As)
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        k_remaining = K - k * BLOCK_SIZE_K
+        a_raw = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0).to(
+            tl.float32
+        )
+        a = (a_raw / a_s_static).to(tl.float8e4nv)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
+        b_s = decode_ue8m0_scale(tl.load(bs_ptrs))
+        accumulator += tl.dot(a, b) * b_s[None, :]
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * stride_bk
+        bs_ptrs += stride_bs_k
+    accumulator = accumulator * a_s_static
+    _store_masked(
+        C,
+        accumulator,
+        pid_m,
+        pid_n,
+        M,
+        N,
+        stride_cm,
+        stride_cn,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+    )
+@bayesian_autotune(
+    get_mxfp_autotuning_configs(
+        compute_modes=("dot_scaled", "dot")
+    ),  # no scalar branch here
+    ["N", "K", "BLOCK_SIZE_M"],
+    n_trials=60,
+)
+@triton.jit
+def mxfp_dynamic_matmul_kernel(
+    A,  # (M, K) raw BF16/FP16 activations
+    B,  # (N, K) E4M3 (MXFP8) or (N, K // 2) packed E2M1 (MXFP4) weights
+    C,  # (M, N) output
+    Bs,  # (N, K // SCALE_GROUP_K) UE8M0 weight scales
+    # Shape
+    M,
+    N,
+    K,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_bs_k,
+    stride_bs_n,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+    VALUES_PER_BYTE: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    COMPUTE_MODE: tl.constexpr,
+):
+    """Unified MXFP4/MXFP8 (W4A8/W8A8) GEMM with fused activation quantization.
+    ``C = A @ B.T`` with bf16/fp16 ``A`` quantized to E4M3 per K-group inline (UE8M0
+    scale). ``VALUES_PER_BYTE`` picks the weight format: 2 = packed E2M1 (MXFP4),
+    1 = unpacked E4M3 (MXFP8). ``COMPUTE_MODE`` picks the MMA: ``tl.dot_scaled``
+    (native M=128 scaled MMA) vs fp8 ``tl.dot`` + per-group software rescale (wins at
+    decode; FP4 unpacks E2M1->E4M3 first — lossless). 2D grid with swizzle for L2 reuse.
+    """
+    pid_m, pid_n, offs_am, offs_bn, offs_k = _swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M
+    )
+    offs_kb = tl.arange(0, BLOCK_SIZE_K // VALUES_PER_BYTE)
+    offs_sf = tl.arange(0, BLOCK_SIZE_K // SCALE_GROUP_K)
+    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
+    b_ptrs = B + (offs_kb[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+    bs_ptrs = Bs + (offs_bn[:, None] * stride_bs_n + offs_sf[None, :] * stride_bs_k)
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        k_remaining = K - k * BLOCK_SIZE_K
+        a_raw = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0).to(
+            tl.float32
+        )
+        a, a_scale = mxfp_act_quant_inline(
+            a_raw, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K
+        )
+        b = tl.load(
+            b_ptrs, mask=offs_kb[:, None] < k_remaining // VALUES_PER_BYTE, other=0.0
+        )
+        b_s = tl.load(
+            bs_ptrs, mask=offs_sf[None, :] < k_remaining // SCALE_GROUP_K, other=0
+        ).to(tl.uint8)
+        bq = mxfp4_e2m1_to_e4m3(b) if VALUES_PER_BYTE == 2 else b
+        if COMPUTE_MODE == "dot_scaled":
+            accumulator = mx_dot_scaled(
+                a, a_scale, b, b_s, accumulator, VALUES_PER_BYTE
+            )
+        else:  # dot
+            accumulator = mx_dot_rescale(accumulator, a, bq, a_scale, b_s)
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += (BLOCK_SIZE_K // VALUES_PER_BYTE) * stride_bk
+        bs_ptrs += (BLOCK_SIZE_K // SCALE_GROUP_K) * stride_bs_k
+    _store_masked(
+        C,
+        accumulator,
+        pid_m,
+        pid_n,
+        M,
+        N,
+        stride_cm,
+        stride_cn,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+    )
+@triton_op(add_op_namespace_prefix("w8a8_block_dynamic_fp8_matmul"), mutates_args=())
+def _w8a8_block_dynamic_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     block_size: list[int],
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
+    """Block-scale FP8 matmul: ``C = A @ B.T`` with fused activation quantization.
+    A:  (..., K) raw activations, bf16/fp16/fp32 (quantized inline to FP8)
+    B:  (N, K) FP8 weights
+    Bs: (N // block_n, K // block_k) per-block weight scales
     """
     assert len(block_size) == 2, (
         f"block_size must be [block_n, block_k], got {block_size}"
     N, K = B.shape
     M = A.numel() // A.shape[-1]
     assert Bs.ndim == 2, f"Bs must be 2D (N//block_n, K//block_k), got ndim={Bs.ndim}"
     assert Bs.shape == (triton.cdiv(N, block_n), triton.cdiv(K, block_k)), (
         f"Bs shape {tuple(Bs.shape)} != expected ({triton.cdiv(N, block_n)}, {triton.cdiv(K, block_k)})"
     BLOCK_SIZE_K = block_k
     BLOCK_SIZE_N = block_n
+    Bs = ue8m0_as_uint8(Bs)
     C_shape = A.shape[:-1] + (N,)
+    BLOCK_SIZE_M = adaptive_block_size_m(M)
     C = A.new_empty(C_shape, dtype=output_dtype)
     grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
     with device_context(A.device):
+        wrap_triton(w8a8_block_dynamic_fp8_matmul_kernel)[grid](
             A,
             B,
             C,
             Bs,
             M,
             N,
             B.stride(0),
             C.stride(-2),
             C.stride(-1),
             Bs.stride(1),
             Bs.stride(0),
             # Meta-parameters
             BLOCK_SIZE_M=BLOCK_SIZE_M,
             BLOCK_SIZE_N=BLOCK_SIZE_N,
             BLOCK_SIZE_K=BLOCK_SIZE_K,
+            GROUP_SIZE_M=GROUP_SIZE_M,
         )
     return C
+@triton_op(add_op_namespace_prefix("w8a8_block_static_fp8_matmul"), mutates_args=())
+def _w8a8_block_static_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
+    As: torch.Tensor,
+    block_size: list[int],
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
+    """Block-scale FP8 matmul with static (per-tensor) activation quantization.
+    A:  (..., K) raw bf16/fp16 activations — quantized to FP8 inline against ``As``
+    B:  (N, K) FP8 weights
+    Bs: (N // block_n, K // block_k) per-block weight scales
+    As: scalar / (1,) — per-tensor static activation scale
     """
+    assert len(block_size) == 2, (
+        f"block_size must be [block_n, block_k], got {block_size}"
+    )
+    block_n, block_k = block_size[0], block_size[1]
+    assert B.dtype != torch.int8, (
+        "static activation quant is not supported on the FP4 path"
+    )
+    assert not (block_n == B.size(0) and block_k == B.size(1)), (
+        "static activation quant requires block-wise weights, not tensor-mode"
+    )
     assert A.shape[-1] == B.shape[-1], (
         f"K mismatch: A has K={A.shape[-1]}, B has K={B.shape[-1]}"
     )
     assert A.is_contiguous(), "A must be contiguous"
     assert B.ndim == 2, f"B must be 2D (N, K), got ndim={B.ndim}"
     assert B.is_contiguous(), "B must be contiguous"
+    assert As.numel() == 1, f"As must be scalar or (1,), got {tuple(As.shape)}"
     N, K = B.shape
     M = A.numel() // A.shape[-1]
+    assert Bs.ndim == 2, f"Bs must be 2D (N//block_n, K//block_k), got ndim={Bs.ndim}"
+    assert Bs.shape == (triton.cdiv(N, block_n), triton.cdiv(K, block_k)), (
+        f"Bs shape {tuple(Bs.shape)} != expected ({triton.cdiv(N, block_n)}, {triton.cdiv(K, block_k)})"
     )
+    BLOCK_SIZE_K = block_k
+    BLOCK_SIZE_N = block_n
+    BLOCK_SIZE_M = adaptive_block_size_m(M)
+    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
+    Bs = ue8m0_as_uint8(Bs)
     C_shape = A.shape[:-1] + (N,)
     C = A.new_empty(C_shape, dtype=output_dtype)
+    As = As.reshape(1).to(torch.float32)
     with device_context(A.device):
+        wrap_triton(w8a8_block_static_fp8_matmul_kernel)[grid](
             A,
             B,
             C,
             B.stride(0),
             C.stride(-2),
             C.stride(-1),
+            Bs.stride(1),
+            Bs.stride(0),
             # Meta-parameters
             BLOCK_SIZE_M=BLOCK_SIZE_M,
             BLOCK_SIZE_N=BLOCK_SIZE_N,
             BLOCK_SIZE_K=BLOCK_SIZE_K,
+            GROUP_SIZE_M=GROUP_SIZE_M,
         )
     return C
+@triton_op(add_op_namespace_prefix("w8a8_tensor_dynamic_fp8_matmul"), mutates_args=())
+def _w8a8_tensor_dynamic_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Tensor-scale FP8 matmul: ``C = A @ B.T`` with fused activation quantization.
+    A:  (..., K) raw activations, bf16/fp16/fp32 (flattened to (M, K)
+        internally) — per-row scales computed via ``fp8_act_quant(A, K)``.
+    B:  (N, K) FP8 weights.
+    Bs: scalar, (1,), or (1, 1) — single tensor-scale weight scale.
+    """
+    assert A.shape[-1] == B.shape[-1], (
+        f"K mismatch: A has K={A.shape[-1]}, B has K={B.shape[-1]}"
+    )
+    assert A.is_contiguous(), "A must be contiguous"
+    assert B.ndim == 2, f"B must be 2D (N, K), got ndim={B.ndim}"
+    assert B.is_contiguous(), "B must be contiguous"
+    N, K = B.shape
+    M = A.numel() // A.shape[-1]
+    assert Bs.numel() == 1, f"Bs must be scalar or (1,), got {tuple(Bs.shape)}"
+    # Per-row scalar activation scale (one per token).
+    qA, As = fp8_act_quant(A, K)
+    As = As.reshape(M)
+    Bs = Bs.reshape(1)
+    C_shape = A.shape[:-1] + (N,)
+    C = A.new_empty(C_shape, dtype=output_dtype)
+    BLOCK_SIZE_M = adaptive_block_size_m(M)
+    def grid(META):
+        return (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, META["BLOCK_SIZE_N"]))
+    with device_context(A.device):
+        wrap_triton(w8a8_tensor_dynamic_fp8_matmul_kernel)[grid](
+            qA,
+            B,
+            C,
+            As,
+            Bs,
+            M,
+            N,
+            K,
+            qA.stride(-2),
+            qA.stride(-1),
+            B.stride(1),
+            B.stride(0),
+            C.stride(-2),
+            C.stride(-1),
+            As.stride(0),
+            # Meta-parameters (BLOCK_SIZE_N, BLOCK_SIZE_K come from autotune Config)
+            BLOCK_SIZE_M=BLOCK_SIZE_M,
+            GROUP_SIZE_M=GROUP_SIZE_M,
+        )
+    return C
+@triton_op(add_op_namespace_prefix("mxfp_dynamic_matmul"), mutates_args=())
+def _mxfp_dynamic_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
+    """MX matmul ``C = A @ B.T`` with fused activation quant. Weight format detected
+    from ``B.dtype``: ``int8`` → packed E2M1 (MXFP4, ``B`` is ``(N, K//2)``);
+    ``float8_e4m3fn`` → unpacked E4M3 (MXFP8, ``(N, K)``). Both use UE8M0 group-32 scales
+    ``(N, K//32)``; tile + dot path are autotuned (scale granularity fixed at 32).
+    A:  (M, K) raw activations, bf16/fp16/fp32 (quantized inline to E4M3)
+    """
+    assert A.ndim == 2 and B.ndim == 2 and Bs.ndim == 2
+    assert B.dtype in (torch.int8, torch.float8_e4m3fn), (
+        f"B must be int8 (packed E2M1) or float8_e4m3fn (E4M3), got {B.dtype}"
+    )
+    assert Bs.dtype == torch.float8_e8m0fnu, (
+        f"Bs must be float8_e8m0fnu, got {Bs.dtype}"
+    )
+    assert A.is_contiguous(), "A must be contiguous"
+    assert B.is_contiguous(), "B must be contiguous"
+    VALUES_PER_BYTE = NIBBLES_PER_BYTE if B.dtype == torch.int8 else 1
+    M, K = A.shape
+    N, K_b = B.shape
+    assert K == VALUES_PER_BYTE * K_b, (
+        f"K (={K}) must equal {VALUES_PER_BYTE} * B.shape[1] (={K_b})"
+    )
+    assert K % MX_SCALE_GROUP_K == 0, (
+        f"K (={K}) must be a multiple of {MX_SCALE_GROUP_K}"
+    )
+    assert Bs.shape == (N, K // MX_SCALE_GROUP_K), (
+        f"Bs shape {tuple(Bs.shape)} != ({N}, {K // MX_SCALE_GROUP_K})"
+    )
+    B = e2m1_as_uint8(B)
+    bs_u8 = ue8m0_as_uint8(Bs)
+    C = A.new_empty((M, N), dtype=output_dtype)
+    BLOCK_SIZE_M = adaptive_block_size_m(M)
+    def grid(META):
+        return (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, META["BLOCK_SIZE_N"]))
+    with device_context(A.device):
+        wrap_triton(mxfp_dynamic_matmul_kernel)[grid](
+            A,
+            B,
+            C,
+            bs_u8,
+            M,
+            N,
+            K,
+            A.stride(0),
+            A.stride(1),
+            B.stride(1),
+            B.stride(0),
+            C.stride(0),
+            C.stride(1),
+            bs_u8.stride(1),
+            bs_u8.stride(0),
+            # Meta-parameters (BLOCK_SIZE_N, BLOCK_SIZE_K come from autotune Config)
+            BLOCK_SIZE_M=BLOCK_SIZE_M,
+            GROUP_SIZE_M=GROUP_SIZE_M,
+            VALUES_PER_BYTE=VALUES_PER_BYTE,
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+        )
+    return C
+def w8a8_block_dynamic_fp8_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    block_size: list[int],
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Block-wise W8A8 FP8 matrix multiplication with fused activation quantization.
     Args:
+        A: Raw activation tensor ``[..., M, K]`` in bf16/fp16/fp32 — quantized
+            inline to ``float8_e4m3fn`` against per-K-tile per-row UE8 scales.
+        B: FP8 weight tensor ``[N, K]`` in ``float8_e4m3fn``.
         Bs: Per-block weight scales ``[N // block_size[0], K // block_size[1]]``.
+        block_size: ``[block_n, block_k]`` weight quantization block dimensions.
         output_dtype: dtype of the returned tensor (default: ``torch.float32``).
     Returns:
+        Output tensor ``[..., M, N]`` in ``output_dtype``.
     """
+    return ops.w8a8_block_dynamic_fp8_matmul(A, B, Bs, block_size, output_dtype)
+def w8a8_block_static_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
+    Bs: torch.Tensor,
     As: torch.Tensor,
+    block_size: list[int],
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Block-wise W8A8 FP8 matmul with static (per-tensor) activation quantization.
+    Args:
+        A: Raw activation tensor ``[..., M, K]`` in bf16/fp16/fp32 — quantized
+            inline against the per-tensor scalar ``As``.
+        B: FP8 weight tensor ``[N, K]`` in ``float8_e4m3fn``.
+        Bs: Per-block weight scales ``[N // block_size[0], K // block_size[1]]``.
+        As: Static per-tensor activation scale (scalar or ``[1]``).
+        block_size: ``[block_n, block_k]`` weight quantization block dimensions.
+        output_dtype: dtype of the returned tensor (default: ``torch.float32``).
+    Returns:
+        Output tensor ``[..., M, N]`` in ``output_dtype``.
+    """
+    return ops.w8a8_block_static_fp8_matmul(A, B, Bs, As, block_size, output_dtype)
+def w8a8_tensor_dynamic_fp8_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
     Bs: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
+    """Tensor-scale W8A8 FP8 matmul with fused activation quantization.
+    Computes ``C = A @ B.T`` with raw bf16/fp16/fp32 ``A`` quantized to FP8
+    per-row (one scale per token) before the dot.
     Args:
+        A: Raw activation tensor ``[M, K]`` in bf16/fp16/fp32.
+        B: FP8 weight tensor ``[N, K]`` in ``float8_e4m3fn``.
         Bs: Single weight scale, scalar or ``[1]``.
         output_dtype: dtype of the returned tensor.
+    """
+    return ops.w8a8_tensor_dynamic_fp8_matmul(A, B, Bs, output_dtype)
+def mxfp_dynamic_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """MX (MXFP4/MXFP8) matmul ``C = A @ B.T`` with fused activation quant; the weight
+    format is detected from ``B.dtype``. Tile + dot path are autotuned; MX scale
+    granularity is fixed at 32 (the MX-format spec).
+    Args:
+        A: Raw activations ``[M, K]`` in bf16/fp16/fp32.
+        B: Packed E2M1 ``[N, K // 2]`` (``int8``) or E4M3 ``[N, K]`` (``float8_e4m3fn``).
+        Bs: UE8M0 weight scales ``[N, K // 32]`` (``float8_e8m0fnu``).
+        output_dtype: dtype of the returned tensor (default ``bfloat16``).
     """
+    return ops.mxfp_dynamic_matmul(A, B, Bs, output_dtype)
+def matmul_2d(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
     block_size: list[int] | None,
+    output_dtype: torch.dtype | None = None,
+    activation_scale: torch.Tensor | None = None,
 ) -> torch.Tensor:
+    """Quantized matmul dispatcher (W8A8 FP8 or W4A8 FP4).
+    ``A`` is always raw bf16/fp16/fp32; quantization is fused into every path.
+    With ``activation_scale`` set, the kernel uses that per-tensor scalar
+    (static quant); otherwise it computes its own scale from ``A`` (dynamic).
+    ``output_dtype`` defaults to ``A.dtype``.
+    Routes by weight dtype and ``block_size``:
+    - MX weights — ``int8`` (packed E2M1) or ``float8_e4m3fn`` (E4M3) with UE8M0
+      group-32 ``Bs`` (shape ``[N, K//32]``) → ``mxfp_dynamic_matmul`` (``block_size``
+      ignored, ``activation_scale`` unsupported; scale granularity fixed at 32, autotuned).
+    - ``block_size`` None or full ``[N, K]`` → ``w8a8_tensor_dynamic_fp8_matmul``.
+    - otherwise → ``w8a8_block_dynamic_fp8_matmul`` (or its static variant when
+      ``activation_scale`` is given).
+    """
+    if is_mxfp(B, Bs):
+        if activation_scale is not None:
+            raise NotImplementedError(
+                "activation_scale (static activation quant) is not supported for MX weights — "
+                "the MX path quantizes activations dynamically per group. Omit activation_scale."
+            )
+        return mxfp_dynamic_matmul(A, B, Bs, output_dtype)
+    if is_tensor_wide(block_size, B):
+        return w8a8_tensor_dynamic_fp8_matmul(A, B, Bs, output_dtype)
+    # Block-wise FP8: static when a per-tensor activation scale is supplied, else dynamic.
+    if activation_scale is not None:
+        return w8a8_block_static_fp8_matmul(
+            A, B, Bs, activation_scale, block_size, output_dtype
+        )
+    return w8a8_block_dynamic_fp8_matmul(A, B, Bs, block_size, output_dtype)

build/torch-cuda/metadata.json CHANGED Viewed

@@ -1,10 +1,25 @@
 {
   "name": "finegrained-fp8",
-  "id": "_finegrained_fp8_cuda_7c5619e",
-  "version": 1,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda"
   }
 }

 {
   "name": "finegrained-fp8",
+  "id": "_finegrained_fp8_cuda_93d7641",
+  "version": 4,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "zhfXqFHpZ1lPv0hUwaOhJHxmcDq2QcE9R+9cgVfLTI4=",
+      "_ops.py": "v4IWQ4TB7LvLjTio3GctOEdzLSpc45Z3OBcWILsDeSM=",
+      "batched.py": "siX6kWh3I9GzioEw58Y66p1aZX1m6kSAf4E/asiLnVA=",
+      "bayesian_autotuner.py": "jBUYvyfpEHsTBgFlpt2Ts5emz+CSMyvjc1TxUfZKVrg=",
+      "finegrained_fp8/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "fused_batched.py": "09a9bN8fzGBG5CapCwvIUL5589F7Ve+CbT+S1+4+BAI=",
+      "fused_grouped.py": "foyMOTru53VKJ1d7ASJpEGWqiF3OqbSQT91tTPXcWKk=",
+      "grouped.py": "Gq1YS9LXnFuKWQ+M0XW6WuP2Fg7zwKwwT50o6YhQXnw=",
+      "matmul.py": "PEdwUJGaJetm4lWF6de3eYHBHtnHsZhevoDQd6fJYKw=",
+      "utils.py": "aUUYq5/POJbaZKmo8QgM89XQrJN3HNHTuJOsPhgYxQw="
+    }
   }
 }

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

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtCgAwIBAgIUHLiFbbNYXDYPtqIG48HGoOqA7RcwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjI5MTcxNDQ5WhcNMjYwNjI5MTcyNDQ5WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAEdD3jvq2C4tEZZAE37ydzk+UsPgoh++aHLcQELDY+a5KZgEqJqxpiW8OHSGEe07m85EHlRXB5sB9BmHJTBi7AuKOCBe8wggXrMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUUaXQbM7F2hw5Au49d8kdUNd+gTkwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoOTNkNzY0MTZlOTY1YTE0Y2IzY2Y1ZGFlOWNjYTdhNTNmNzBjNWU1YjATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoOTNkNzY0MTZlOTY1YTE0Y2IzY2Y1ZGFlOWNjYTdhNTNmNzBjNWU1YjAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoOTNkNzY0MTZlOTY1YTE0Y2IzY2Y1ZGFlOWNjYTdhNTNmNzBjNWU1YjAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDkzZDc2NDE2ZTk2NWExNGNiM2NmNWRhZTljY2E3YTUzZjcwYzVlNWIwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjgzODk2ODc4NDkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiQYKKwYBBAHWeQIEAgR7BHkAdwB1AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABnxRgZgUAAAQDAEYwRAIga3GJTEHsVlnuoI1kG25kPC3sc4On1G1CH7jgrcznaq4CIAEDXVUDkXPlbdYZztMVhEPydBetfSt9x44bOu1mhEqmMAoGCCqGSM49BAMDA2kAMGYCMQD2YFLb+M5svHbxwKLjbe7VlXaa7CuLvdujtpVpwR4WQYwAjZh/NsYfHRTK9YF5d6YCMQDPbSsy3rBtAI4QppFB48c8/PC8BS+YwTaWq1guriywQSt9JXkgOD7TlWMPSrvG5LE="},"tlogEntries":[{"logIndex":"2011412491","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1782753290","inclusionPromise":{"signedEntryTimestamp":"MEQCIBu7joAyJ1Scv9JziiQLgyQ7crhnWnP1MdOfZVZU/belAiATlq9G4be8PKp5mBqEjIZDvKaDOmj4wdqGHFq4ov+8eg=="},"inclusionProof":{"logIndex":"1889508229","rootHash":"Qc/hbi1qBcXqZet3VqITiwbaJ5KYuZZp4AZzYzDFmPY=","treeSize":"1889508650","hashes":["tWE1dKnr+xyVR9DqZl179d4a2pzILWEWoG3QQ016pQA=","eJDrwh6X93nkRtJXPluc2cbg01QfYAH0gkvcZW5fZdg=","2Bvz7mTmsTUXVkKOwxb/onqiaThvKrmBkjB83YlfoYk=","u97qIngM8q1QfW5Lo/9csT+WAgu8+rpNQjqarIgDrnE=","7dbKejwhyjxkywj03FXGECoQ7QTz1yxa3/gxCieD8WU=","RGsebap7/2pdY7rrT3oY8CUkcbwkB+6gHwnx1/CsGAs=","NYulN1dObgLker5XjMR4jAFzRdx6Sy0SXp6X3KRf0m4=","zsh1PjbsqQJic+3EcBbfMKAlRareth4KaBthTv/3/SA=","0Lm4vgxuYggxvPJg+Rhp0EJTjbPoCTRGIaGqkTkZ4Lo=","X9YQHakf+wqAGVCUCm2cEW/JGbf4dyp/Db+Ku1v1Du8=","PORSs03e4HDJYPwDuaM6rQrziDBN3e5Y4SuajVQFHg0=","DGp8UqFvSdPpdWfi6+/wVsDiGX858w1Jtjle0C4QGVc=","uxLbTK4c/VFXh7/w8pb8eAk6P1ppLRLvvQ5tU6i9MFg=","tvcQx4ohRLYBU5QOWeu4jhR3lEKztnlyYTspOSiOMcA=","pu3boNdI/2WREJSDmOgBX37bljG4IyCsSsPMDAzAB8Q=","f2GuSxsX8CM80Bzrdr+iaFcvdbcdx8PsOGfP1A5iZmg=","wznNTIJf/erZkw1aAaotGts1BhEB9zNmnJSKo0zdCeM=","+hhERR0matid5NvHk7ddsPOWJ5JAyyn11f+ruZpMW38=","Mp0blBzkxwj6OGAogVNNNMwQM4XxRTFCrHD5ZZy2wYg=","dIBftsHOWhWti/c7pucPNfamsTNyNuam3ViDx67krLA=","+/VZ56MsIPxMiyLAodzKXo5TEWdQp36z89qLhpzloAo=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1889508650\nQc/hbi1qBcXqZet3VqITiwbaJ5KYuZZp4AZzYzDFmPY=\n\n— rekor.sigstore.dev wNI9ajBFAiEAmeN2VwOoIwnC9dJ2+KWMPB3BmzlDx3rjop2jipH6bJECIEFFujbiD917hnqHY5pLSHj0TL9j6+R9/irZ4NLnl49k\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJkOTQxNDY1ZjQ1Y2FkOGE3M2Q5MjhiNWY1YWUxYmIyZDhjOGRlN2Y3ZjI5MDIzMTdkMGZiODUyOWI1ZjMyMmE0In19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUUNGam5tbFJQbG4rUG52T2VNUTN2b1BCR0ZkaDRmMDcvcjJtY05JN2V5eDNBSWdXajhTck42N0FITWFDMHcxYmpEbG9IRXhGSmxCTkhNUUZtUWcyZ01OcHNzPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SRFowRjNTVUpCWjBsVlNFeHBSbUppVGxsWVJGbFFkSEZKUnpRNFNFZHZUM0ZCTjFKamQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxU1RWTlZHTjRUa1JSTlZkb1kwNU5hbGwzVG1wSk5VMVVZM2xPUkZFMVYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZrUkROcWRuRXlRelIwUlZwYVFVVXpOM2xrZW1zclZYTlFaMjlvS3l0aFNFeGpVVVVLVEVSWksyRTFTMXBuUlhGS2NYaHdhVmM0VDBoVFIwVmxNRGR0T0RWRlNHeFNXRUkxYzBJNVFtMUlTbFJDYVRkQmRVdFBRMEpsT0hkbloxaHlUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZWWVZoUkNtSk5OMFl5YUhjMVFYVTBPV1E0YTJSVlRtUXJaMVJyZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVBWRTVyVG5wWk1FMVVXbXhQVkZreFdWUkZNRmt5U1hwWk1sa3hDbHBIUm14UFYwNXFXVlJrYUU1VVRtMU9la0pxVGxkVk1WbHFRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDlVVG10T2Vsa3dUVlJhYkU5VVdURlpWRVV3V1RKSmVsa3lXVEZhUjBac1QxZE9hbGxVWkdnS1RsUk9iVTU2UW1wT1YxVXhXV3BCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDlVVG10T2Vsa3dUVlJhYkU5VVdURlpWRVV3V1RKSmVsa3lXVEVLV2tkR2JFOVhUbXBaVkdSb1RsUk9iVTU2UW1wT1YxVXhXV3BCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUnJlbHBFWXpJS1RrUkZNbHBVYXpKT1YwVjRUa2RPYVUweVRtMU9WMUpvV2xSc2Fsa3lSVE5aVkZWNldtcGpkMWw2Vm14T1YwbDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1kNlQwUnJNazlFWXpST1JHdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwVVZsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTTjBKSWEwRUtaSGRDTVVGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTRVbWRhWjFWQlFVRlJSQXBCUlZsM1VrRkpaMkV6UjBwVVJVaHpWbXh1ZFc5Sk1XdEhNalZyVUVNemMyTTBUMjR4UnpGRFNEZHFaM0pqZW01aGNUUkRTVUZGUkZoV1ZVUnJXRkJzQ21Ka1dWcDZkRTFXYUVWUWVXUkNaWFJtVTNRNWVEUTBZazkxTVcxb1JYRnRUVUZ2UjBORGNVZFRUVFE1UWtGTlJFRXlhMEZOUjFsRFRWRkVNbGxHVEdJS0swMDFjM1pJWW5oM1MweHFZbVUzVm14WVlXRTNRM1ZNZG1SMWFuUndWbkIzVWpSWFVWbDNRV3BhYUM5T2MxbG1TRkpVU3psWlJqVmtObGxEVFZGRVVBcGlVM041TTNKQ2RFRkpORkZ3Y0VaQ05EaGpPQzlRUXpoQ1V5dFpkMVJoVjNFeFozVnlhWGwzVVZOME9VcFlhMmRQUkRkVWJGZE5VRk55ZGtjMVRFVTlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQguFU3H+mook948dfQKNgaoxHAoycH3GX2xe+B/+CeMJACFFbYYWssAU46XsBRISWFOQ8Li12eGA8yMDI2MDYyOTE3MTQ0OVowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNjI5MTcxNDQ5WjAvBgkqhkiG9w0BCQQxIgQguvirpEmIhGn3xr9LWY2ECPSMWhGKtDD4yC7y/bDh1KswgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIwDmJq7ErbKctLqPhm2iwOCTJDHIm8Je4xVwCY/fXufe2r+Z9Dgh8AW+PjeW6Rj9hZAjEAhiq+ySoRLE6b8eGDhPKHxDczbFbxVhtusLTA9bdwkTUDzb2bPMZKHP+wYEQ5La8d"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"2UFGX0XK2Kc9kotfWuG7LYyN5/fykCMX0PuFKbXzIqQ="},"signature":"MEUCIQCFjnmlRPln+PnvOeMQ3voPBGFdh4f07/r2mcNI7eyx3AIgWj8SrN67AHMaC0w1bjDloHExFJlBNHMQFmQg2gMNpss="}}

build/torch-cuda/utils.py CHANGED Viewed

@@ -1,6 +1,40 @@
-import torch
 from contextlib import contextmanager
 @contextmanager
 def device_context(device: torch.device):
@@ -11,3 +45,723 @@ def device_context(device: torch.device):
             yield
     else:
         yield

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import functools
 from contextlib import contextmanager
+import torch
+import triton
+import triton.language as tl
+from torch.library import triton_op, wrap_triton
+from ._ops import add_op_namespace_prefix, ops
+# ── Format constants ──────────────────────────────────────────────────────────
+# FP8 (E4M3) is the main format for weights and activations;
+FP8_DTYPE = torch.float8_e4m3fn
+# FP4 (E2M1) packs two 4-bit nibbles per byte. MX formats (MXFP4 weights, MXFP8
+# E4M3 weights/activations) share one UE8M0 scale per 32-element K-group — the OCP
+# MX block size, consumed by ``tl.dot_scaled``. Format constants, not tunables.
+NIBBLES_PER_BYTE = 2
+MX_SCALE_GROUP_K = 32
+# ── Host-side helpers ─────────────────────────────────────────────────────────
 @contextmanager
 def device_context(device: torch.device):
             yield
     else:
         yield
+@functools.lru_cache(maxsize=8)
+def sm_count(device_index: int) -> int:
+    """Processor (CUDA SM / XPU Xe-core) count for a device (a constant) — cached so callers
+    don't re-query it each launch. Read from Triton's active driver so it works on any backend,
+    with a CUDA fallback."""
+    try:
+        return triton.runtime.driver.active.utils.get_device_properties(device_index)[
+            "multiprocessor_count"
+        ]
+    except Exception:
+        active_device = get_active_device_type()
+        if active_device == "cuda":
+            return torch.cuda.get_device_properties(device_index).multi_processor_count
+        elif active_device == "xpu":
+            return torch.xpu.get_device_properties(device_index).multi_processor_count
+        else:
+            raise RuntimeError(
+                f"Unsupported device type {active_device} for sm_count; only cuda/xpu are supported."
+            )
+# H-tile each topk_reduce program reduces (one (token, tile) program per BLOCK_H span). The
+# reduce is bandwidth-bound and saturates at BLOCK_H>=512: a sweep over {128..2048} x prefill
+# shapes put 512/1024/2048 within ~7% (mostly noise) while 128/256 lagged — not worth a tuning
+# axis for this small slice of runtime. 512 = smallest that saturates (best at low token counts,
+# where bigger tiles waste masked lanes). Power of 2 (tl.arange).
+TOPK_REDUCE_BLOCK_H = 512
+@triton.jit
+def topk_reduce_kernel(
+    ProjOut,
+    Out,
+    H,
+    stride_pm,
+    stride_ph,
+    stride_om,
+    stride_oh,
+    NUM_TOP_K: tl.constexpr,
+    BLOCK_H: tl.constexpr,
+):
+    """Sum the ``NUM_TOP_K`` flat ``(token, slot)`` rows of ``ProjOut`` into ``Out[token]`` —
+    a tight replacement for ``proj_out.view(T, K, H).sum(1)`` (~2.8x faster than torch's
+    generic reduce; fp32 accumulate, so bit-identical). Caller allocs Out and launches it."""
+    t = tl.program_id(0)
+    offs_h = tl.program_id(1) * BLOCK_H + tl.arange(0, BLOCK_H)
+    mask = offs_h < H
+    acc = tl.zeros((BLOCK_H,), tl.float32)
+    for k in tl.static_range(NUM_TOP_K):
+        acc += tl.load(
+            ProjOut + (t * NUM_TOP_K + k) * stride_pm + offs_h * stride_ph,
+            mask=mask,
+            other=0.0,
+        ).to(tl.float32)
+    tl.store(
+        Out + t * stride_om + offs_h * stride_oh,
+        acc.to(Out.dtype.element_ty),
+        mask=mask,
+    )
+def ue8m0_as_uint8(scale: torch.Tensor) -> torch.Tensor:
+    """View UE8M0 (``float8_e8m0fnu``) weight scales as ``uint8`` for the Triton
+    binder, which doesn't recognize the dtype; kernels decode ``2^(exp-127)``
+    inline. fp32 (non-UE8M0) scales pass through unchanged."""
+    return scale.view(torch.uint8) if scale.dtype == torch.float8_e8m0fnu else scale
+def e2m1_as_uint8(weight: torch.Tensor) -> torch.Tensor:
+    """View an ``int8``-stored MXFP4 (packed E2M1) weight as ``uint8`` — a zero-cost
+    reinterpret. ``tl.dot_scaled`` requires the packed rhs as ``uint8``, so do it once here
+    instead of casting in-kernel at every load. E4M3 (MXFP8) weights pass through unchanged."""
+    return weight.view(torch.uint8) if weight.dtype == torch.int8 else weight
+def is_mxfp8(weight: torch.Tensor, scale: torch.Tensor) -> bool:
+    """MXFP8 weight/scale pair: E4M3 weights with UE8M0 group-32 scales — last dim
+    ``scale.shape[-1] == weight.shape[-1] // MX_SCALE_GROUP_K``, matching leading dims.
+    Works for 2D ``(N, K)`` and 3D ``(E, N, K)`` weights. The group-32 layout is what
+    separates MXFP8 from 128-block FP8 (which may also carry UE8M0 scales)."""
+    return (
+        weight.dtype == torch.float8_e4m3fn
+        and scale.dtype == torch.float8_e8m0fnu
+        and scale.ndim == weight.ndim
+        and scale.shape[:-1] == weight.shape[:-1]
+        and scale.shape[-1] == weight.shape[-1] // MX_SCALE_GROUP_K
+    )
+def is_mxfp4(weight: torch.Tensor, scale: torch.Tensor) -> bool:
+    """MXFP4 weight/scale pair: packed E2M1 weights (``int8``, two codes/byte) with
+    UE8M0 group-32 scales — ``scale.shape[-1] == weight.shape[-1] * NIBBLES_PER_BYTE //
+    MX_SCALE_GROUP_K`` (unpacked K = ``2 * K_half``), matching leading dims. 2D or 3D."""
+    return (
+        weight.dtype == torch.int8
+        and scale.dtype == torch.float8_e8m0fnu
+        and scale.ndim == weight.ndim
+        and scale.shape[:-1] == weight.shape[:-1]
+        and scale.shape[-1] == weight.shape[-1] * NIBBLES_PER_BYTE // MX_SCALE_GROUP_K
+    )
+def is_mxfp(weight: torch.Tensor, scale: torch.Tensor) -> bool:
+    """Any MX weight/scale pair — MXFP8 (``float8_e4m3fn``, one value/byte) or MXFP4
+    (``int8``, two E2M1 codes/byte), both with UE8M0 group-32 scales. ``values_per_byte``
+    folds the two cases: the scale's last dim covers the unpacked K
+    (``weight.shape[-1] * values_per_byte``) in groups of ``MX_SCALE_GROUP_K``. The
+    dispatchers route on this; the op picks the format from ``weight.dtype``."""
+    if weight.dtype == torch.float8_e4m3fn:
+        values_per_byte = 1
+    elif weight.dtype == torch.int8:
+        values_per_byte = NIBBLES_PER_BYTE
+    else:
+        return False
+    return (
+        scale.dtype == torch.float8_e8m0fnu
+        and scale.ndim == weight.ndim
+        and scale.shape[:-1] == weight.shape[:-1]
+        and scale.shape[-1] == weight.shape[-1] * values_per_byte // MX_SCALE_GROUP_K
+    )
+def is_tensor_wide(block_size, weight: torch.Tensor) -> bool:
+    """True when ``block_size`` selects per-tensor (tensor-dynamic) scaling: ``None`` or
+    equal to the weight's full ``(N, K)`` — one scale block spanning the whole matrix.
+    Handles 2D ``(N, K)`` and 3D ``(E, N, K)`` weights via the last two dims."""
+    return block_size is None or (
+        block_size[0] == weight.shape[-2] and block_size[1] == weight.shape[-1]
+    )
+# The per-token batched/fused kernels are decode-shaped — one routed row per program, so
+# BLOCK_SIZE_M is 1 (a larger tile would just recompute the same row). The 2D matmul sizes
+# its M tile to the workload instead, via ``adaptive_block_size_m``.
+DECODE_BLOCK_SIZE_M = 1
+def adaptive_block_size_m(target_m: int) -> int:
+    """Smallest power-of-2 >= ``target_m``, floored at 16 and capped at 128.
+    Used by all matmul wrappers (single / batched / grouped) to size the M tile
+    to the workload — small per-expert M wants smaller tiles, large M caps out
+    at 128 to keep register pressure bounded. Pass ``M`` for single matmul, or
+    ``(S + E - 1) // E`` (avg tokens per expert) for batched/grouped.
+    """
+    return min(max(triton.next_power_of_2(target_m), 16), 128)
+@functools.cache
+def get_active_device_type() -> str:
+    """Active torch device type for the current Triton backend (``"cuda"``, ``"xpu"``, ...).
+    Falls back to ``"cuda"`` when no driver is loaded — Triton raises
+    ``RuntimeError: 0 active drivers ([])`` on driverless build boxes, and the
+    autotune-config builder is evaluated at module-import time under the
+    ``@triton.autotune`` decorator (no kernel launches there, so the default is
+    only used to shape the config list).
+    """
+    try:
+        return triton.runtime.driver.active.get_active_torch_device().type
+    except RuntimeError:
+        return "cuda"
+def get_accelerator_autotuning_configs(
+    *,
+    tune_block_m: bool = False,
+    tune_block_nk: bool = False,
+):
+    """Autotune search grid for the current accelerator.
+    ``num_warps``, ``num_stages`` and ``blocks`` (the ``(BLOCK_SIZE_N, BLOCK_SIZE_K)``
+    tile shapes) are fixed up front from ``(is_xpu, tune_block_nk)``, then crossed
+    into the config list.
+    ``tune_block_nk=True`` sweeps the tile: used by kernels that have no caller
+    ``block_size`` to fix it — the MX ``tl.dot_scaled`` paths AND the tensor-dynamic
+    FP8 paths. ``tune_block_nk=False`` emits a single empty meta-dict (block-scaled
+    kernels take the tile from the caller's ``block_size``). ``tune_block_m=True`` also crosses
+    in BLOCK_SIZE_M ∈ ``(16, 32, 64, 128)`` (fused grouped, which lets the tuner pick M instead
+    of computing it via ``adaptive_block_size_m``).
+    The CUDA tile set is a data-driven prune of a B200 sweep across single (BM=128),
+    grouped MoE (BM=16/64) and decode (BM=1): winners only ever used these 4 tiles,
+    num_warps in {4,8,16}, num_stages in {2,3} (warps=2, stages=4 and 256x256 never
+    won) — 108 → 24. (Tuned on dot_scaled; tensor-dynamic tl.dot reuses it.)
+    """
+    is_xpu = get_active_device_type() == "xpu"
+    num_warps = [8, 16] if is_xpu else [2, 4, 8, 16]
+    num_stages = [2, 3, 4]
+    if tune_block_nk:
+        tiles = (
+            [(128, 128)] if is_xpu else [(128, 128), (256, 128), (128, 64), (64, 256)]
+        )
+        blocks = [{"BLOCK_SIZE_N": bn, "BLOCK_SIZE_K": bk} for bn, bk in tiles]
+    else:
+        blocks = [{}]
+    if tune_block_m:
+        blocks = [{**b, "BLOCK_SIZE_M": bm} for b in blocks for bm in (16, 32, 64, 128)]
+    return [
+        triton.Config(b, num_warps=w, num_stages=s)
+        for b in blocks
+        for w in num_warps
+        for s in num_stages
+    ]
+def get_mxfp_autotuning_configs(
+    pre_hook=None,
+    compute_modes=("dot_scaled", "dot", "scalar"),
+    memory_modes=(None,),
+    tune_block_m=False,
+):
+    """Autotune grid for the MXFP8 MoE/matmul kernels. ``COMPUTE_MODE`` (a constexpr string) picks
+    the multiply-accumulate, and the Bayesian tuner selects it per workload (token count is
+    in the key):
+    - ``"dot_scaled"``: native group-32 scaled MMA (wide K, ``BLOCK_SIZE_K`` ∈ {128, 256}) —
+      wins once the grid saturates (prefill, ~S≥32).
+    - ``"dot"``: fp8 ``tl.dot`` + per-group software rescale, one scale group per K-step
+      (``BLOCK_SIZE_K == 32``).
+    - ``"scalar"``: scalar CUDA-core FMA reduction, no tensor core — same per-group structure
+      as ``"dot"`` (hence bit-exact with it) but avoids the scaled-MMA's M→16 pad, so it wins
+      for memory-bound **decode** (M=1; ≈2× over ``dot_scaled`` at S≥32 in a sweep). Only emit
+      for kernels that implement the branch.
+    ``num_warps=16`` and ``BLOCK_SIZE_K=64`` are omitted (dead at M=1 per a sweep).
+    ``memory_modes`` (fused_grouped only) is the ``MEMORY_MODE`` axis for the weight load — e.g.
+    ``("host_descriptor", "pointer")`` on CUDA, ``("device_descriptor", "pointer")`` on XPU
+    (a host/device tensor descriptor — TMA on NVIDIA — vs explicit pointers); the
+    ``(None,)`` (default) emits no axis. Pass all the modes a kernel implements; this drops
+    the device-inapplicable descriptor flavor (host-side is NVIDIA-TMA-only, device-side is the
+    XPU win). a given pre_hook is attached to every config and self-guards on ``MEMORY_MODE``."""
+    drop = (
+        "host_descriptor" if get_active_device_type() == "xpu" else "device_descriptor"
+    )
+    memory_modes = tuple(m for m in memory_modes if m != drop)
+    if "device_descriptor" in memory_modes:
+        # device-side descriptors build in-kernel tensormaps that need a scratch allocator (the
+        # XPU path; CUDA keeps host-built descriptors and never reaches this).
+        triton.set_allocator(
+            lambda size, alignment, stream: torch.empty(
+                size, dtype=torch.int8, device="xpu"
+            )
+        )
+    bms = (16, 32, 64, 128) if tune_block_m else (None,)
+    return [
+        triton.Config(
+            {
+                "COMPUTE_MODE": mode,
+                "BLOCK_SIZE_N": bn,
+                "BLOCK_SIZE_K": bk,
+                **({"MEMORY_MODE": mm} if mm is not None else {}),
+                **({"BLOCK_SIZE_M": bm} if bm is not None else {}),
+            },
+            num_warps=w,
+            num_stages=s,
+            # Always attached when given; the hook self-guards on MEMORY_MODE (only
+            # host_descriptor configs use the host-built descriptor).
+            pre_hook=pre_hook,
+        )
+        for mode in compute_modes
+        for mm in memory_modes
+        for bn in [32, 64, 128, 256]
+        # only "dot" pins BK to the group (one scale group per fp8 tl.dot step); "dot_scaled"
+        # and "scalar" handle multiple groups per K-chunk, so they use a wide BK.
+        for bk in ([32] if mode == "dot" else [128, 256])
+        for s in [2, 3, 4, 5, 6]
+        for w in [2, 4, 8]
+        for bm in bms
+    ]
+@functools.lru_cache(maxsize=None)
+def sm_shared_memory_limit(device_index: int) -> int:
+    """Max dynamic shared memory per block (bytes) for a CUDA device — the cap Triton
+    reports as the 'Hardware limit' on an ``out of resource: shared memory`` failure
+    (~232 KB on B200, ~227 KB on H100, much less on older/consumer parts). Queried from
+    the driver so the prune adapts to the hardware instead of hardcoding one GPU."""
+    try:
+        return triton.runtime.driver.active.utils.get_device_properties(device_index)[
+            "max_shared_mem"
+        ]
+    except Exception:
+        if get_active_device_type() == "cuda":
+            return torch.cuda.get_device_properties(
+                device_index
+            ).shared_memory_per_block_optin
+        elif get_active_device_type() == "xpu":
+            return torch.xpu.get_device_properties(
+                device_index
+            ).shared_memory_per_block_optin
+        else:
+            raise RuntimeError(
+                f"Unsupported device type {get_active_device_type()} for sm_shared_memory_limit; only cuda/xpu are supported."
+            )
+def smem_config_pruner(act_bytes: int, n_weight_tiles: int, weight_bytes: int = 1):
+    """Build an ``early_config_prune`` that drops configs whose pipelined operand shared
+    memory would overflow the SM — the source of ``out of resource: shared memory`` autotune
+    failures (and the wasted compiles they cause).
+    Per-stage estimate (bytes) = ``BK · (act_bytes·BM + weight_bytes·n_weight_tiles·BN)``:
+    one activation tile ``[BM, BK]`` plus ``n_weight_tiles`` weight tiles ``[BN, BK]``
+    (gate_up fuses 2; down has 1), times ``num_stages``. ``BM`` is the routing-derived tile,
+    read from the launch args. MXFP4 packs 2 weights/byte, so ``weight_bytes=1`` (MXFP8) is a
+    safe upper bound. The limit is read from the active device. Never returns empty — keeps
+    the smallest-footprint config as a fallback."""
+    def prune(configs, named_args, **kwargs):
+        # The estimate needs all three tile dims. Each is either an autotuned config meta
+        # (MX) or a launch arg (block-dynamic), so look in both; raise clearly if missing.
+        all_args = {**named_args, **kwargs}
+        dev = triton.runtime.driver.active.get_active_torch_device()
+        limit = sm_shared_memory_limit(dev.index if dev.index is not None else 0)
+        def dim(c, name):
+            v = c.kwargs.get(name, all_args.get(name))
+            if v is None:
+                raise ValueError(
+                    f"smem_config_pruner needs {name} (autotune config meta or launch arg) "
+                    "to estimate shared memory; none found."
+                )
+            return v
+        def smem(c):
+            BM, BN, BK = (
+                dim(c, n) for n in ("BLOCK_SIZE_M", "BLOCK_SIZE_N", "BLOCK_SIZE_K")
+            )
+            return (
+                c.num_stages
+                * BK
+                * (act_bytes * BM + weight_bytes * n_weight_tiles * BN)
+            )
+        kept = [c for c in configs if smem(c) <= limit]
+        return kept or [min(configs, key=smem)]
+    return prune
+# ── Triton-side helpers (inlined by ``@triton.jit`` callers) ──────────────────
+@triton.jit
+def fp8_act_quant_inline(a_raw):
+    """Inline FP8 (E4M3) activation quant for the W8A8 block-scale path.
+    Per-row amax → fp32 scale ``amax/448`` (floored at 1e-12 against zero rows)
+    → cast values to FP8. Returns ``(a_fp8, a_s)`` with shapes ``(M, K)`` and
+    ``(M,)``.
+    """
+    a_s = tl.max(tl.abs(a_raw), axis=1) / 448.0
+    a_fp8 = (a_raw / tl.maximum(a_s[:, None], 1e-12)).to(tl.float8e4nv)
+    return a_fp8, a_s
+@triton.jit
+def mxfp_act_quant_inline(
+    a_raw,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+):
+    """Inline E4M3 activation quant for the MX paths (W4A8 MXFP4 / W8A8 MXFP8).
+    Per-row, per-K-group amax → UE8M0 scale (ceil to next power-of-2 via the
+    mantissa-nonzero bump trick) → cast values to FP8. Returns ``(a_fp8,
+    a_scale_u8)`` with shapes ``(M, K)`` and ``(M, K // SCALE_GROUP_K)``.
+    """
+    a_groups = tl.reshape(
+        a_raw, (BLOCK_SIZE_M, BLOCK_SIZE_K // SCALE_GROUP_K, SCALE_GROUP_K)
+    )
+    a_s_fp32 = tl.max(tl.abs(a_groups), axis=2) / 448.0
+    bits = a_s_fp32.to(tl.int32, bitcast=True)
+    # ceil_to_ue8m0: bump exponent by 1 when mantissa is non-zero.
+    exp_ceil = ((bits >> 23) & 0xFF) + ((bits & 0x7FFFFF) != 0).to(tl.int32)
+    exp_ceil = tl.minimum(tl.maximum(exp_ceil, 1), 254)
+    a_scale_u8 = exp_ceil.to(tl.uint8)
+    a_s_pow2 = (exp_ceil << 23).to(tl.float32, bitcast=True)
+    a_fp8 = tl.reshape(
+        a_groups / tl.maximum(a_s_pow2[:, :, None], 1e-12),
+        (BLOCK_SIZE_M, BLOCK_SIZE_K),
+    ).to(tl.float8e4nv)
+    return a_fp8, a_scale_u8
+@triton.jit
+def decode_ue8m0_scale(scale):
+    """Decode a UE8M0 weight scale to fp32: when it was loaded as ``uint8`` exponent
+    bits, ``value = 2^(exp - 127)``, built directly as the fp32 bit pattern. fp32
+    scales (block-dynamic with float scales) pass through. The dtype branch is a
+    compile-time constant, so only the taken path is emitted (single return — Triton
+    requires all ``return`` statements to share a type)."""
+    if scale.dtype == tl.uint8:
+        scale = (scale.to(tl.int32) << 23).to(tl.float32, bitcast=True)
+    return scale
+@triton.jit
+def expand_ue8m0_scale(
+    scale_u8, ROWS: tl.constexpr, BK: tl.constexpr, GK: tl.constexpr
+):
+    """Decode a UE8M0 group scale tile ``[ROWS, BK//GK]`` and broadcast each group scale
+    across its ``GK`` elements → ``[ROWS, BK]`` fp32 — one scale per K element. Used by the
+    scalar (CUDA-core FMA) MX paths to dequant the whole BK chunk before the reduction."""
+    s = decode_ue8m0_scale(scale_u8)
+    return tl.reshape(tl.broadcast_to(s[:, :, None], (ROWS, BK // GK, GK)), (ROWS, BK))
+@triton.jit
+def mx_dot_scaled(a, a_scale, w, w_scale, acc, VALUES_PER_BYTE: tl.constexpr):
+    """MX 'dot_scaled' path: scaled MMA folding the UE8M0 group scales into the tensor core —
+    ``a`` (E4M3) @ ``w`` (packed E2M1 if MXFP4, else E4M3). The rhs format is picked from
+    ``VALUES_PER_BYTE``. Caller pre-shapes ``w``/``w_scale`` (e.g. ``tl.trans(gu)``)."""
+    rhs_format: tl.constexpr = "e2m1" if VALUES_PER_BYTE == 2 else "e4m3"
+    return tl.dot_scaled(a, a_scale, "e4m3", w, w_scale, rhs_format, acc)
+@triton.jit
+def mx_dot_rescale(acc, a, wq, a_scale, w_scale):
+    """MX 'dot' path (BK == group): fp8 ``tl.dot`` + per-group software rescale, decoding both
+    UE8M0 scales internally and accumulating into ``acc`` (returned updated). Single weight — for
+    the gate∪up pair (shared activation scale) use ``mx_dot_rescale_gate_up``."""
+    return acc + tl.dot(a, wq) * decode_ue8m0_scale(a_scale) * tl.trans(
+        decode_ue8m0_scale(w_scale)
+    )
+@triton.jit
+def mx_dot_rescale_gate_up(
+    acc_gate, acc_up, a, w_gate, w_up, a_scale, gate_scale, up_scale
+):
+    """Gate∪up 'dot' path: decode the SHARED activation scale once, rescale both projections
+    and accumulate into ``acc_gate``/``acc_up`` (returned updated)."""
+    a_s = decode_ue8m0_scale(a_scale)
+    acc_gate += tl.dot(a, w_gate) * a_s * tl.trans(decode_ue8m0_scale(gate_scale))
+    acc_up += tl.dot(a, w_up) * a_s * tl.trans(decode_ue8m0_scale(up_scale))
+    return acc_gate, acc_up
+@triton.jit
+def mx_scalar_reduce(
+    acc,
+    a,
+    a_scale,
+    wq,
+    w_scale,
+    BLOCK_SIZE_M: tl.constexpr,
+    ROWS_W: tl.constexpr,
+    BK: tl.constexpr,
+    GK: tl.constexpr,
+):
+    """MX 'scalar' path: CUDA-core FMA GEMV, dequantizing activation + weight per-element by
+    their expanded group scales internally, then reducing and accumulating into ``acc`` (returned
+    updated). No tensor core (so no M→16 MMA pad) — wins for the memory-bound decode GEMV (M=1).
+    Single weight — for the gate∪up pair (shared activation dequant) use ``mx_scalar_reduce_gate_up``."""
+    a_deq = tl.trans(
+        a.to(tl.float32) * expand_ue8m0_scale(a_scale, BLOCK_SIZE_M, BK, GK)
+    )
+    w_exp = tl.trans(expand_ue8m0_scale(w_scale, ROWS_W, BK, GK))
+    return acc + tl.sum(a_deq * (wq.to(tl.float32) * w_exp), axis=0)[None, :]
+@triton.jit
+def mx_scalar_reduce_gate_up(
+    acc_gate,
+    acc_up,
+    a,
+    w_gate,
+    w_up,
+    a_scale,
+    gate_scale,
+    up_scale,
+    BLOCK_SIZE_M: tl.constexpr,
+    ROWS_W: tl.constexpr,
+    BK: tl.constexpr,
+    GK: tl.constexpr,
+):
+    """Gate∪up 'scalar' path: dequant the SHARED activation once, reduce both projections and
+    accumulate into ``acc_gate``/``acc_up`` (returned updated)."""
+    a_deq = tl.trans(
+        a.to(tl.float32) * expand_ue8m0_scale(a_scale, BLOCK_SIZE_M, BK, GK)
+    )
+    acc_gate += tl.sum(
+        a_deq
+        * (
+            w_gate.to(tl.float32)
+            * tl.trans(expand_ue8m0_scale(gate_scale, ROWS_W, BK, GK))
+        ),
+        axis=0,
+    )[None, :]
+    acc_up += tl.sum(
+        a_deq
+        * (
+            w_up.to(tl.float32) * tl.trans(expand_ue8m0_scale(up_scale, ROWS_W, BK, GK))
+        ),
+        axis=0,
+    )[None, :]
+    return acc_gate, acc_up
+@triton.jit
+def mx_compute_gate_up(
+    acc_gate,
+    acc_up,
+    a,
+    a_scale,
+    b_gate,
+    b_up,
+    gate_scale,
+    up_scale,
+    COMPUTE_MODE: tl.constexpr,
+    VALUES_PER_BYTE: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+):
+    """Gate∪up MMA step dispatched on ``COMPUTE_MODE``: scaled-MMA on the raw weights
+    (``b_gate``/``b_up``), or fp8 ``tl.dot`` + per-group rescale / scalar reduce on the
+    E4M3-decoded weights. Returns the updated ``(acc_gate, acc_up)``."""
+    w_gate = mxfp4_e2m1_to_e4m3(b_gate) if VALUES_PER_BYTE == 2 else b_gate
+    w_up = mxfp4_e2m1_to_e4m3(b_up) if VALUES_PER_BYTE == 2 else b_up
+    if COMPUTE_MODE == "dot_scaled":
+        acc_gate = mx_dot_scaled(
+            a, a_scale, b_gate, gate_scale, acc_gate, VALUES_PER_BYTE
+        )
+        acc_up = mx_dot_scaled(a, a_scale, b_up, up_scale, acc_up, VALUES_PER_BYTE)
+    elif COMPUTE_MODE == "dot":
+        acc_gate, acc_up = mx_dot_rescale_gate_up(
+            acc_gate, acc_up, a, w_gate, w_up, a_scale, gate_scale, up_scale
+        )
+    else:  # scalar
+        acc_gate, acc_up = mx_scalar_reduce_gate_up(
+            acc_gate,
+            acc_up,
+            a,
+            w_gate,
+            w_up,
+            a_scale,
+            gate_scale,
+            up_scale,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            SCALE_GROUP_K,
+        )
+    return acc_gate, acc_up
+@triton.jit
+def mx_compute(
+    acc,
+    a,
+    a_scale,
+    w,
+    w_scale,
+    COMPUTE_MODE: tl.constexpr,
+    VALUES_PER_BYTE: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+):
+    """Single-projection MMA step dispatched on ``COMPUTE_MODE``: scaled-MMA on the raw weight
+    (``w``), or fp8 ``tl.dot`` + per-group rescale / scalar reduce on the E4M3-decoded weight.
+    Returns the updated ``acc``."""
+    wq = mxfp4_e2m1_to_e4m3(w) if VALUES_PER_BYTE == 2 else w
+    if COMPUTE_MODE == "dot_scaled":
+        acc = mx_dot_scaled(a, a_scale, w, w_scale, acc, VALUES_PER_BYTE)
+    elif COMPUTE_MODE == "dot":
+        acc = mx_dot_rescale(acc, a, wq, a_scale, w_scale)
+    else:  # scalar
+        acc = mx_scalar_reduce(
+            acc,
+            a,
+            a_scale,
+            wq,
+            w_scale,
+            BLOCK_SIZE_M,
+            BLOCK_SIZE_N,
+            BLOCK_SIZE_K,
+            SCALE_GROUP_K,
+        )
+    return acc
+@triton.jit
+def _activation(x, ACT_FN: tl.constexpr):
+    """Pointwise GLU gate activation. ``ACT_FN`` in ``{"silu", "gelu", "relu"}``; gelu is exact
+    (erf), matching torch's default ``F.gelu``."""
+    if ACT_FN == "silu":
+        return x * tl.sigmoid(x)
+    elif ACT_FN == "gelu":
+        return 0.5 * x * (1.0 + tl.erf(x * 0.7071067811865476))
+    elif ACT_FN == "relu":
+        return tl.maximum(x, 0.0)
+    else:
+        tl.static_assert(
+            False, "unsupported ACT_FN; expected 'silu', 'gelu', or 'relu'"
+        )
+@triton.jit
+def glu(
+    acc_gate,
+    acc_up,
+    ACT_FN: tl.constexpr,
+    SIMULATE_UNFUSED: tl.constexpr,
+    OUT_DTYPE: tl.constexpr,
+):
+    """Gated linear unit: ``ACT_FN(gate) * up`` (SwiGLU / GeGLU / ReGLU). When ``SIMULATE_UNFUSED``,
+    round each op through ``OUT_DTYPE`` to match the numerics of the unfused (separate-kernel)
+    path, where every intermediate is materialized in that dtype."""
+    if SIMULATE_UNFUSED:
+        g = acc_gate.to(OUT_DTYPE).to(tl.float32)
+        u = acc_up.to(OUT_DTYPE).to(tl.float32)
+        act = _activation(g, ACT_FN).to(OUT_DTYPE).to(tl.float32)
+        return (act * u).to(OUT_DTYPE).to(tl.float32)
+    return _activation(acc_gate, ACT_FN) * acc_up
+@triton.jit
+def _e2m1_code_to_f32(code):
+    """One E2M1 4-bit code -> fp32. Layout ``[sign | exp(2) | mant(1)]``; the 8
+    magnitudes are ``{0, .5, 1, 1.5, 2, 3, 4, 6}`` (exp==0 is the 0/0.5 subnormal)."""
+    code = code.to(tl.int32)
+    s = (code >> 3) & 1
+    e = (code >> 1) & 3
+    m = (code & 1).to(tl.float32)
+    pow2 = tl.exp2((e - 1).to(tl.float32))  # e in 0..3 -> 0.5, 1, 2, 4
+    mag = tl.where(e == 0, 0.5 * m, (1.0 + 0.5 * m) * pow2)
+    return (1.0 - 2.0 * s.to(tl.float32)) * mag
+@triton.jit
+def mxfp4_e2m1_to_e4m3(b_packed):
+    """Unpack packed MXFP4 (E2M1, two nibbles/byte along K) to E4M3, doubling the K
+    (row) dim: ``(R, C) uint8 -> (2R, C) E4M3``. E2M1's 8 magnitudes are all exact in
+    E4M3, so this is lossless — it lets the FP8 ``tl.dot`` path stand in for
+    ``tl.dot_scaled`` at decode (avoiding its M->128 pad). K order is the low nibble
+    first: ``[byte0_lo, byte0_hi, byte1_lo, ...]``."""
+    lo = _e2m1_code_to_f32(b_packed & 0xF)
+    hi = _e2m1_code_to_f32(b_packed >> 4)
+    # interleave along the K (row) dim via trans -> interleave-last-dim -> trans back
+    unpacked = tl.trans(tl.interleave(tl.trans(lo), tl.trans(hi)))
+    return unpacked.to(tl.float8e4nv)
+# ── fp8_act_quant kernel (used by tensor-mode FP8 wrappers) ───────────────────
+@triton.jit
+def _fp8_act_quant_kernel(
+    x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr, PADDED_BLOCK: tl.constexpr
+):
+    # ``tl.arange`` needs a power-of-2 length, so iterate over PADDED_BLOCK (the next
+    # power of 2) and mask the tail — lets block_size be non-power-of-2 (e.g. a full
+    # row K=14336 in tensor-mode). Masked lanes load 0, which can't affect ``amax``.
+    pid = tl.program_id(axis=0)
+    cols = tl.arange(0, PADDED_BLOCK)
+    mask = cols < BLOCK_SIZE
+    offs = pid * BLOCK_SIZE + cols
+    x = tl.load(x_ptr + offs, mask=mask, other=0.0).to(tl.float32)
+    s = tl.max(tl.abs(x)) / 448.0  # float8_e4m3fn max
+    y = (x / tl.maximum(s, 1e-12)).to(y_ptr.dtype.element_ty)
+    tl.store(y_ptr + offs, y, mask=mask)
+    tl.store(s_ptr + pid, s)
+@triton_op(add_op_namespace_prefix("fp8_act_quant"), mutates_args=())
+def _fp8_act_quant(
+    x: torch.Tensor, block_size: int = 128
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert x.is_contiguous()
+    assert x.shape[-1] % block_size == 0
+    y = torch.empty_like(x, dtype=FP8_DTYPE)
+    grid = (triton.cdiv(x.numel(), block_size),)
+    s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
+    with device_context(x.device):
+        wrap_triton(_fp8_act_quant_kernel)[grid](
+            x,
+            y,
+            s,
+            BLOCK_SIZE=block_size,
+            PADDED_BLOCK=triton.next_power_of_2(block_size),
+        )
+    return y, s
+def fp8_act_quant(
+    x: torch.Tensor, block_size: int = 128
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Quantize activations to FP8 with per-block dynamic scaling.
+    Splits the last dimension of ``x`` into blocks of ``block_size`` elements,
+    computes ``scale = max(|x_block|) / 448`` per block, and quantizes to
+    ``float8_e4m3fn``.
+    Args:
+        x: Input tensor in bf16/fp16/fp32. Last dimension must be divisible by
+            ``block_size`` and the tensor must be contiguous.
+        block_size: Number of elements per quantization block (default: 128).
+    Returns:
+        A tuple ``(quantized, scales)`` where ``quantized`` has dtype
+        ``float8_e4m3fn`` with the same shape as ``x``, and ``scales`` has
+        shape ``(*x.shape[:-1], x.shape[-1] // block_size)`` in float32.
+    """
+    return ops.fp8_act_quant(x, block_size)

build/torch-rocm/__init__.py DELETED Viewed

@@ -1,32 +0,0 @@
-from .act_quant import fp8_act_quant
-from .batched import (
-    w8a8_fp8_matmul_batched,
-    w8a8_block_fp8_matmul_batched,
-    w8a8_tensor_fp8_matmul_batched,
-)
-from .grouped import (
-    w8a8_fp8_matmul_grouped,
-    w8a8_block_fp8_matmul_grouped,
-    w8a8_tensor_fp8_matmul_grouped,
-)
-from .matmul import (
-    w8a8_fp8_matmul,
-    w8a8_block_fp8_matmul,
-    w8a8_tensor_fp8_matmul,
-)
-__all__ = [
-    "fp8_act_quant",
-    # Single matmul
-    "w8a8_fp8_matmul",
-    "w8a8_block_fp8_matmul",
-    "w8a8_tensor_fp8_matmul",
-    # Batched matmul
-    "w8a8_fp8_matmul_batched",
-    "w8a8_block_fp8_matmul_batched",
-    "w8a8_tensor_fp8_matmul_batched",
-    # Grouped matmul
-    "w8a8_fp8_matmul_grouped",
-    "w8a8_block_fp8_matmul_grouped",
-    "w8a8_tensor_fp8_matmul_grouped",
-]

build/torch-rocm/_ops.py DELETED Viewed

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

build/torch-rocm/act_quant.py DELETED Viewed

@@ -1,73 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-import triton
-import triton.language as tl
-from torch.library import triton_op, wrap_triton
-from .utils import device_context
-_FP8_DTYPE = torch.float8_e4m3fn
-# Copied from https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
-@triton.jit
-def _fp8_act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr):
-    pid = tl.program_id(axis=0)
-    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    x = tl.load(x_ptr + offs).to(tl.float32)
-    s = tl.max(tl.abs(x)) / 448.0  # float8_e4m3fn max
-    y = (x / s).to(y_ptr.dtype.element_ty)
-    tl.store(y_ptr + offs, y)
-    tl.store(s_ptr + pid, s)
-@triton_op("finegrained_fp8::fp8_act_quant", mutates_args=())
-def _fp8_act_quant(
-    x: torch.Tensor, block_size: int = 128
-) -> tuple[torch.Tensor, torch.Tensor]:
-    assert x.is_contiguous()
-    assert x.shape[-1] % block_size == 0
-    y = torch.empty_like(x, dtype=_FP8_DTYPE)
-    grid = (triton.cdiv(x.numel(), block_size),)
-    s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
-    with device_context(x.device):
-        wrap_triton(_fp8_act_quant_kernel)[grid](x, y, s, BLOCK_SIZE=block_size)
-    return y, s
-def fp8_act_quant(
-    x: torch.Tensor, block_size: int = 128
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Quantize activations to FP8 with per-block dynamic scaling.
-    Splits the last dimension of ``x`` into blocks of ``block_size`` elements,
-    computes ``scale = max(|x_block|) / 448`` per block, and quantizes to
-    ``float8_e4m3fn``.
-    Args:
-        x: Input tensor in bf16/fp16/fp32. Last dimension must be divisible by
-            ``block_size`` and the tensor must be contiguous.
-        block_size: Number of elements per quantization block (default: 128).
-    Returns:
-        A tuple ``(quantized, scales)`` where ``quantized`` has dtype
-        ``float8_e4m3fn`` with the same shape as ``x``, and ``scales`` has
-        shape ``(*x.shape[:-1], x.shape[-1] // block_size)`` in float32.
-    """
-    return torch.ops.finegrained_fp8.fp8_act_quant(x, block_size)

build/torch-rocm/batched.py DELETED Viewed

@@ -1,390 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-import triton
-import triton.language as tl
-from .act_quant import fp8_act_quant
-from torch.library import triton_op, wrap_triton
-from .utils import device_context
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
-    key=["N", "K"],
-)
-@triton.jit
-def w8a8_block_fp8_matmul_batched_kernel(
-    A,  # (S, K)  raw BF16/FP16 activations
-    B,  # (E, N, K) FP8 weight matrices
-    C,  # (S, N)  output
-    Bs,  # (E, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales
-    ExpertIds,  # (S,) — which expert each batch element routes to
-    # Shape
-    S,
-    N,
-    K,
-    stride_am,
-    stride_ak,
-    stride_be,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_bs_e,
-    stride_bs_k,
-    stride_bs_n,
-    # Meta-parameters
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
-):
-    """Block-scale batched FP8 expert matmul kernel.
-    Each program handles one routed token row and one N-tile, looks up the
-    owning expert from ``ExpertIds``, and applies fused activation quantization.
-    """
-    batch_id = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    # Cast expert_id to int64 to prevent int32 overflow when computing
-    # expert_id * stride_Eb (e.g. 255 * 9_437_184 > 2^31 for 256 experts of
-    # 3072×3072 FP8 weights).
-    expert_id = tl.load(ExpertIds + batch_id).to(tl.int64)
-    A = A + batch_id * stride_am
-    B = B + expert_id * stride_be
-    C = C + batch_id * stride_cm
-    Bs = Bs + expert_id * stride_bs_e
-    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + tl.arange(0, BLOCK_SIZE_M)[:, None] * 0 + offs_k[None, :] * stride_ak
-    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
-    bs_ptrs = Bs + pid_n * stride_bs_n
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        # ---- fused fp8_act_quant ----
-        a_raw = tl.load(a_ptrs).to(tl.float32)
-        a_s = tl.max(tl.abs(a_raw)) / 448.0
-        a = (a_raw / tl.maximum(a_s, 1e-12)).to(tl.float8e4nv)
-        # ---- matmul ----
-        b = tl.load(b_ptrs)
-        b_s = tl.load(bs_ptrs + k * stride_bs_k)
-        accumulator += tl.dot(a, b) * a_s * b_s[None, :]
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + offs_cm[:, None] * 0 + stride_cn * offs_cn[None, :]
-    tl.store(c_ptrs, c)
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
-    key=["N", "K"],
-)
-@triton.jit
-def w8a8_tensor_fp8_matmul_batched_kernel(
-    A,  # (S, K) pre-quantized FP8 activations
-    B,  # (E, N, K) FP8 weight matrices
-    C,  # (S, N) output
-    As,  # (S, 1) per-tensor activation scales
-    Bs,  # (E, 1, 1) per-tensor weight scales
-    ExpertIds,
-    S,
-    N,
-    K,
-    stride_am,
-    stride_ak,
-    stride_be,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_as_m,
-    stride_bs_e,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
-):
-    """Tensor-scale batched FP8 expert matmul kernel.
-    Activations are already quantized; the kernel applies per-token activation
-    scales and per-expert tensor weight scales.
-    """
-    batch_id = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    expert_id = tl.load(ExpertIds + batch_id).to(tl.int64)
-    A = A + batch_id * stride_am
-    B = B + expert_id * stride_be
-    C = C + batch_id * stride_cm
-    Bs = Bs + expert_id * stride_bs_e
-    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + tl.arange(0, BLOCK_SIZE_M)[:, None] * 0 + offs_k[None, :] * stride_ak
-    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
-    b_s = tl.load(Bs)
-    a_s = tl.load(As + batch_id * stride_as_m)
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        a = tl.load(a_ptrs)
-        b = tl.load(b_ptrs)
-        accumulator += tl.dot(a, b)
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    accumulator = accumulator * a_s * b_s
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + offs_cm[:, None] * 0 + stride_cn * offs_cn[None, :]
-    tl.store(c_ptrs, c)
-@triton_op("finegrained_fp8::w8a8_block_fp8_matmul_batched", mutates_args=())
-def _w8a8_block_fp8_matmul_batched(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    expert_ids: torch.Tensor,
-    block_size: list[int],
-) -> torch.Tensor:
-    """Block-scale batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T, with fused act quant.
-    A:  (S, K) raw bf16/fp16 activations
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
-    """
-    assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    assert A.shape[1] == B.shape[2], (
-        f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
-    )
-    S, K = A.shape
-    E, N, _ = B.shape
-    assert len(block_size) == 2, (
-        f"block_size must be [block_n, block_k], got {block_size}"
-    )
-    block_n, block_k = block_size[0], block_size[1]
-    # MoE expert dimensions must be block-aligned; non-aligned N/K is not supported.
-    assert N % block_n == 0, f"N ({N}) must be divisible by block_n ({block_n})"
-    assert K % block_k == 0, f"K ({K}) must be divisible by block_k ({block_k})"
-    assert Bs.ndim == 3, (
-        f"Bs must be 3D (E, N//block_n, K//block_k), got ndim={Bs.ndim}"
-    )
-    assert Bs.shape == (E, N // block_n, K // block_k), (
-        f"Bs shape {tuple(Bs.shape)} != expected ({E}, {N // block_n}, {K // block_k})"
-    )
-    C = A.new_empty(S, N)
-    # Adaptive BLOCK_SIZE_M: smallest power-of-2 >= M, floored at 16, capped at 128.
-    # Matches the WGMMA tile to the actual row count — smaller tiles use less
-    # register pressure and a better-matched FP8 WGMMA instruction, improving
-    # both accuracy and performance for small M (decode).
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
-    grid = (S, triton.cdiv(N, block_n))
-    with device_context(A.device):
-        wrap_triton(w8a8_block_fp8_matmul_batched_kernel)[grid](
-            A,
-            B,
-            C,
-            Bs,
-            expert_ids,
-            S,
-            N,
-            K,
-            A.stride(0),
-            A.stride(1),
-            B.stride(0),
-            B.stride(2),
-            B.stride(1),
-            C.stride(0),
-            C.stride(1),
-            Bs.stride(0),
-            Bs.stride(2),
-            Bs.stride(1),
-            BLOCK_SIZE_N=block_n,
-            BLOCK_SIZE_K=block_k,
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-        )
-    return C
-@triton_op("finegrained_fp8::w8a8_tensor_fp8_matmul_batched", mutates_args=())
-def _w8a8_tensor_fp8_matmul_batched(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    expert_ids: torch.Tensor,
-) -> torch.Tensor:
-    """Tensor-scale batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T, with fused act quant.
-    A:  (S, K) raw bf16/fp16 activations
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
-    """
-    assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    assert A.shape[1] == B.shape[2], (
-        f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
-    )
-    S, K = A.shape
-    E, N, _ = B.shape
-    # Normalize Bs to (E, 1, 1)
-    if Bs.ndim == 1:
-        assert Bs.shape[0] == E, f"Bs shape {tuple(Bs.shape)} != expected ({E},)"
-        Bs = Bs.reshape(E, 1, 1)
-    else:
-        assert Bs.shape == (E, 1, 1), (
-            f"Bs shape {tuple(Bs.shape)} != expected ({E}, 1, 1)"
-        )
-    BLOCK_SIZE_N = 128
-    BLOCK_SIZE_K = 128
-    C = A.new_empty(S, N)
-    qA, As = fp8_act_quant(A, K)
-    grid = (S, triton.cdiv(N, BLOCK_SIZE_N))
-    # Adaptive BLOCK_SIZE_M: smallest power-of-2 >= M, floored at 16, capped at 128.
-    # Matches the WGMMA tile to the actual row count — smaller tiles use less
-    # register pressure and a better-matched FP8 WGMMA instruction, improving
-    # both accuracy and performance for small M (decode).
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
-    grid = (S, triton.cdiv(N, BLOCK_SIZE_N))
-    with device_context(A.device):
-        wrap_triton(w8a8_tensor_fp8_matmul_batched_kernel)[grid](
-            qA,
-            B,
-            C,
-            As,
-            Bs,
-            expert_ids,
-            S,
-            N,
-            K,
-            qA.stride(0),
-            qA.stride(1),
-            B.stride(0),
-            B.stride(2),
-            B.stride(1),
-            C.stride(0),
-            C.stride(1),
-            As.stride(0),
-            Bs.stride(0),
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-        )
-    return C
-def w8a8_block_fp8_matmul_batched(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    expert_ids: torch.Tensor,
-    block_size: list[int],
-) -> torch.Tensor:
-    """Block-scale batched FP8 matmul with fused activation quantization.
-    A:  (S, K) raw activations, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
-    """
-    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_batched(
-        A, B, Bs, expert_ids, block_size
-    )
-def w8a8_tensor_fp8_matmul_batched(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    expert_ids: torch.Tensor,
-) -> torch.Tensor:
-    """Tensor-scale batched FP8 matmul with fused activation quantization.
-    A:  (S, K) raw activations, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
-    """
-    return torch.ops.finegrained_fp8.w8a8_tensor_fp8_matmul_batched(
-        A, B, Bs, expert_ids
-    )
-def w8a8_fp8_matmul_batched(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    expert_ids: torch.Tensor,
-    block_size: list[int] | None,
-) -> torch.Tensor:
-    """Unified batched W8A8 FP8 matmul dispatcher.
-    Dispatch rules:
-    - tensor mode when ``block_size is None``
-    - tensor mode when ``block_size == [N, K]``
-    - otherwise block mode
-    Returns:
-        Output tensor ``[S, N]`` in the same dtype as ``A``.
-    """
-    if block_size is None or (
-        block_size[0] == B.size(1) and block_size[1] == B.size(2)
-    ):
-        return w8a8_tensor_fp8_matmul_batched(A, B, Bs, expert_ids)
-    return w8a8_block_fp8_matmul_batched(A, B, Bs, expert_ids, block_size)

build/torch-rocm/finegrained_fp8/__init__.py DELETED Viewed

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

build/torch-rocm/grouped.py DELETED Viewed

@@ -1,477 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-import triton
-import triton.language as tl
-from .act_quant import fp8_act_quant
-from torch.library import triton_op, wrap_triton
-from .utils import device_context
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
-    key=["N", "K", "BLOCK_SIZE_M"],
-)
-@triton.jit
-def w8a8_block_fp8_matmul_grouped_kernel(
-    A,  # (S, K)  raw BF16/FP16 activations, sorted/grouped by expert id
-    B,  # (E, N, K) FP8 weight matrices
-    C,  # (S, N)  output
-    Bs,  # (E, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales
-    Offsets,  # (E,) int32 — cumulative row-end per expert
-    TileOffsets,  # (E,) int32 — cumulative tile-end per expert
-    # Shape
-    S,
-    N,
-    K,
-    # Strides
-    stride_am,
-    stride_ak,
-    stride_be,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_bs_e,
-    stride_bs_k,
-    stride_bs_n,
-    # Meta-parameters
-    NUM_EXPERTS: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
-    NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
-):
-    """Block-scale grouped FP8 expert matmul kernel.
-    Tokens are assumed sorted by expert. The kernel maps each M-tile to its
-    owning expert via ``TileOffsets`` and applies fused activation quantization.
-    """
-    pid_m = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    # Exit early for programs beyond the actual tile count.
-    total_tiles = tl.load(TileOffsets + NUM_EXPERTS - 1)
-    if pid_m >= total_tiles:
-        return
-    # Binary search in TileOffsets to find the owning expert.
-    # Finds the smallest e such that TileOffsets[e] > pid_m (upper_bound semantics),
-    # which is the expert whose tile range contains pid_m.
-    # O(log2(NUM_EXPERTS)) loads instead of the O(NUM_EXPERTS) linear scan.
-    # NUM_EXPERTS_BIT_LENGTH is ceil(log2(E))+1 for powers-of-two, giving one
-    # harmless extra iteration when lo==hi; it's a compile-time constant so the
-    # loop is fully unrolled by the compiler.
-    lo = 0
-    hi = NUM_EXPERTS
-    for _ in tl.static_range(NUM_EXPERTS_BIT_LENGTH):
-        mid = (lo + hi) >> 1
-        mid_val = tl.load(TileOffsets + mid)
-        is_left = mid_val <= pid_m
-        lo = tl.where(is_left, mid + 1, lo)
-        hi = tl.where(is_left, hi, mid)
-    # Cast expert_id to int64 to prevent int32 overflow when computing
-    # expert_id * stride_be (e.g. 255 * 9_437_184 > 2^31 for 256 experts of
-    # 3072×3072 FP8 weights).
-    expert_id = lo.to(tl.int64)
-    prev_eid = tl.maximum(expert_id - 1, 0)
-    expert_start = tl.where(expert_id == 0, 0, tl.load(Offsets + prev_eid))
-    expert_end = tl.load(Offsets + expert_id)
-    M_expert = expert_end - expert_start
-    expert_tile_start = tl.where(expert_id == 0, 0, tl.load(TileOffsets + prev_eid))
-    local_tile = pid_m - expert_tile_start
-    m_off = local_tile * BLOCK_SIZE_M
-    offs_am = m_off + tl.arange(0, BLOCK_SIZE_M)
-    row_mask = offs_am < M_expert
-    offs_global_m = expert_start + offs_am
-    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + offs_global_m[:, None] * stride_am + offs_k[None, :] * stride_ak
-    b_ptrs = (
-        B
-        + expert_id * stride_be
-        + offs_k[:, None] * stride_bk
-        + offs_bn[None, :] * stride_bn
-    )
-    bs_ptrs = Bs + expert_id * stride_bs_e + pid_n * stride_bs_n
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        # ---- fused fp8_act_quant ----
-        a_raw = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0).to(tl.float32)
-        a_s = tl.max(tl.abs(a_raw), axis=1) / 448.0
-        a = (a_raw / tl.maximum(a_s[:, None], 1e-12)).to(tl.float8e4nv)
-        # ---- matmul ----
-        b = tl.load(b_ptrs)
-        b_s = tl.load(bs_ptrs + k * stride_bs_k)
-        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    c_ptrs = C + stride_cm * offs_global_m[:, None] + stride_cn * offs_bn[None, :]
-    c_mask = row_mask[:, None]
-    tl.store(c_ptrs, c, mask=c_mask)
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
-    key=["N", "K", "BLOCK_SIZE_M"],
-)
-@triton.jit
-def w8a8_tensor_fp8_matmul_grouped_kernel(
-    A,  # (S, K) raw BF16/FP16 activations, sorted/grouped by expert idc
-    B,  # (E, N, K) FP8 weight matrices
-    C,  # (S, N) output
-    As,  # (S, 1) activation scales
-    Bs,  # (E, 1, 1) per-tensor weight scales
-    Offsets,
-    TileOffsets,
-    S,
-    N,
-    K,
-    stride_am,
-    stride_ak,
-    stride_be,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_as_m,
-    stride_bs_e,
-    NUM_EXPERTS: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
-    NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
-):
-    """Tensor-scale grouped FP8 expert matmul kernel.
-    Uses grouped expert scheduling with pre-quantized activations plus
-    per-token activation scales and per-expert tensor weight scales.
-    """
-    pid_m = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    total_tiles = tl.load(TileOffsets + NUM_EXPERTS - 1)
-    if pid_m >= total_tiles:
-        return
-    lo = 0
-    hi = NUM_EXPERTS
-    for _ in tl.static_range(NUM_EXPERTS_BIT_LENGTH):
-        mid = (lo + hi) >> 1
-        mid_val = tl.load(TileOffsets + mid)
-        is_left = mid_val <= pid_m
-        lo = tl.where(is_left, mid + 1, lo)
-        hi = tl.where(is_left, hi, mid)
-    expert_id = lo.to(tl.int64)
-    prev_eid = tl.maximum(expert_id - 1, 0)
-    expert_start = tl.where(expert_id == 0, 0, tl.load(Offsets + prev_eid))
-    expert_end = tl.load(Offsets + expert_id)
-    M_expert = expert_end - expert_start
-    expert_tile_start = tl.where(expert_id == 0, 0, tl.load(TileOffsets + prev_eid))
-    local_tile = pid_m - expert_tile_start
-    m_off = local_tile * BLOCK_SIZE_M
-    offs_am = m_off + tl.arange(0, BLOCK_SIZE_M)
-    row_mask = offs_am < M_expert
-    offs_global_m = expert_start + offs_am
-    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + offs_global_m[:, None] * stride_am + offs_k[None, :] * stride_ak
-    b_ptrs = (
-        B
-        + expert_id * stride_be
-        + offs_k[:, None] * stride_bk
-        + offs_bn[None, :] * stride_bn
-    )
-    a_s = tl.load(As + offs_global_m * stride_as_m, mask=row_mask, other=0.0)
-    b_s = tl.load(Bs + expert_id * stride_bs_e)
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        a = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0)
-        b = tl.load(b_ptrs)
-        accumulator += tl.dot(a, b)
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    accumulator = accumulator * a_s[:, None] * b_s
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    c_ptrs = C + stride_cm * offs_global_m[:, None] + stride_cn * offs_bn[None, :]
-    c_mask = row_mask[:, None]
-    tl.store(c_ptrs, c, mask=c_mask)
-@triton_op("finegrained_fp8::w8a8_block_fp8_matmul_grouped", mutates_args=())
-def _w8a8_block_fp8_matmul_grouped(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    offsets: torch.Tensor,
-    tokens_per_expert: torch.Tensor,
-    block_size: list[int],
-) -> torch.Tensor:
-    """Block-scale grouped FP8 matmul: C = A @ B.T per expert, with fused act quant.
-    A:  (S, K) raw bf16/fp16 activations, sorted by expert
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
-    """
-    assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    assert A.shape[1] == B.shape[2], (
-        f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
-    )
-    S, K = A.shape
-    E, N, _ = B.shape
-    assert len(block_size) == 2, (
-        f"block_size must be [block_n, block_k], got {block_size}"
-    )
-    block_n, block_k = block_size[0], block_size[1]
-    # MoE expert dimensions must be block-aligned; non-aligned N/K is not supported.
-    assert N % block_n == 0, f"N ({N}) must be divisible by block_n ({block_n})"
-    assert K % block_k == 0, f"K ({K}) must be divisible by block_k ({block_k})"
-    assert Bs.ndim == 3, (
-        f"Bs must be 3D (E, N//block_n, K//block_k), got ndim={Bs.ndim}"
-    )
-    assert Bs.shape == (E, N // block_n, K // block_k), (
-        f"Bs shape {tuple(Bs.shape)} != expected ({E}, {N // block_n}, {K // block_k})"
-    )
-    C = A.new_empty(S, N)
-    # Adaptive BLOCK_SIZE_M: match tile to average tokens per expert.
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
-    tiles_per_expert = (tokens_per_expert + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M
-    tile_offsets = torch.cumsum(tiles_per_expert, dim=0).to(torch.int32)
-    # Upper bound on M-tiles: sum_e ceil(M_e / BLOCK_M) <= ceil(S / BLOCK_M) + E.
-    # Programs beyond the real tile count exit immediately via the early-return
-    # guard inside the kernel. This is faster than syncing for the exact count
-    # and keeps the grid size data-independent (cuda-graph / torch.compile safe).
-    max_M_tiles = triton.cdiv(S, BLOCK_SIZE_M) + E
-    grid = (max_M_tiles, triton.cdiv(N, block_n))
-    with device_context(A.device):
-        wrap_triton(w8a8_block_fp8_matmul_grouped_kernel)[grid](
-            A,
-            B,
-            C,
-            Bs,
-            offsets,
-            tile_offsets,
-            S,
-            N,
-            K,
-            A.stride(0),
-            A.stride(1),
-            B.stride(0),
-            B.stride(2),
-            B.stride(1),
-            C.stride(0),
-            C.stride(1),
-            Bs.stride(0),
-            Bs.stride(2),
-            Bs.stride(1),
-            # Meta-parameters
-            NUM_EXPERTS=E,
-            BLOCK_SIZE_N=block_n,
-            BLOCK_SIZE_K=block_k,
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-            NUM_EXPERTS_BIT_LENGTH=E.bit_length(),
-        )
-    return C
-@triton_op("finegrained_fp8::w8a8_tensor_fp8_matmul_grouped", mutates_args=())
-def _w8a8_tensor_fp8_matmul_grouped(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    offsets: torch.Tensor,
-    tokens_per_expert: torch.Tensor,
-) -> torch.Tensor:
-    """Tensor-scale grouped FP8 matmul: C = A @ B.T per expert, with fused act quant.
-    A:  (S, K) raw bf16/fp16 activations, sorted by expert
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
-    """
-    assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    assert A.shape[1] == B.shape[2], (
-        f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
-    )
-    S, K = A.shape
-    E, N, _ = B.shape
-    # Normalize Bs to (E, 1, 1)
-    if Bs.ndim == 1:
-        assert Bs.shape[0] == E, f"Bs shape {tuple(Bs.shape)} != expected ({E},)"
-        Bs = Bs.reshape(E, 1, 1)
-    else:
-        assert Bs.shape == (E, 1, 1), (
-            f"Bs shape {tuple(Bs.shape)} != expected ({E}, 1, 1)"
-        )
-    BLOCK_SIZE_N = 128
-    BLOCK_SIZE_K = 128
-    C = A.new_empty(S, N)
-    qA, As = fp8_act_quant(A, K)
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
-    tiles_per_expert = (tokens_per_expert + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M
-    tile_offsets = torch.cumsum(tiles_per_expert, dim=0).to(torch.int32)
-    # Upper bound on M-tiles: sum_e ceil(M_e / BLOCK_M) <= ceil(S / BLOCK_M) + E.
-    # Programs beyond the real tile count exit immediately via the early-return
-    # guard inside the kernel. This is faster than syncing for the exact count
-    # and keeps the grid size data-independent (cuda-graph / torch.compile safe).
-    max_M_tiles = triton.cdiv(S, BLOCK_SIZE_M) + E
-    grid = (max_M_tiles, triton.cdiv(N, BLOCK_SIZE_N))
-    with device_context(A.device):
-        wrap_triton(w8a8_tensor_fp8_matmul_grouped_kernel)[grid](
-            qA,
-            B,
-            C,
-            As,
-            Bs,
-            offsets,
-            tile_offsets,
-            S,
-            N,
-            K,
-            qA.stride(0),
-            qA.stride(1),
-            B.stride(0),
-            B.stride(2),
-            B.stride(1),
-            C.stride(0),
-            C.stride(1),
-            As.stride(0),
-            Bs.stride(0),
-            # Meta-parameters
-            NUM_EXPERTS=E,
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-            NUM_EXPERTS_BIT_LENGTH=E.bit_length(),
-        )
-    return C
-def w8a8_block_fp8_matmul_grouped(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    offsets: torch.Tensor,
-    tokens_per_expert: torch.Tensor,
-    block_size: list[int],
-) -> torch.Tensor:
-    """Block-scale grouped FP8 matmul with fused activation quantization.
-    A:  (S, K) raw activations sorted by expert, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
-    """
-    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_grouped(
-        A, B, Bs, offsets, tokens_per_expert, block_size
-    )
-def w8a8_tensor_fp8_matmul_grouped(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    offsets: torch.Tensor,
-    tokens_per_expert: torch.Tensor,
-) -> torch.Tensor:
-    """Tensor-scale grouped FP8 matmul with fused activation quantization.
-    A:  (S, K) raw activations sorted by expert, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
-    """
-    return torch.ops.finegrained_fp8.w8a8_tensor_fp8_matmul_grouped(
-        A, B, Bs, offsets, tokens_per_expert
-    )
-def w8a8_fp8_matmul_grouped(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    offsets: torch.Tensor,
-    tokens_per_expert: torch.Tensor,
-    block_size: list[int] | None,
-) -> torch.Tensor:
-    """Unified grouped W8A8 FP8 matmul dispatcher.
-    Dispatch rules:
-    - tensor mode when ``block_size is None``
-    - tensor mode when ``block_size == [N, K]``
-    - otherwise block mode
-    Returns:
-        Output tensor ``[S, N]`` in the same dtype as ``A``, in expert-sorted order.
-    """
-    if block_size is None or (
-        block_size[0] == B.size(1) and block_size[1] == B.size(2)
-    ):
-        return w8a8_tensor_fp8_matmul_grouped(A, B, Bs, offsets, tokens_per_expert)
-    return w8a8_block_fp8_matmul_grouped(
-        A, B, Bs, offsets, tokens_per_expert, block_size
-    )

build/torch-rocm/matmul.py DELETED Viewed

@@ -1,411 +0,0 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-import triton
-import triton.language as tl
-from torch.library import triton_op, wrap_triton
-from .utils import device_context
-# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4]
-    ],
-    key=["N", "K", "BLOCK_SIZE_M"],
-)
-@triton.jit
-def w8a8_block_fp8_matmul_kernel(
-    # Pointers to inputs and output
-    A,
-    B,
-    C,
-    As,
-    Bs,
-    # Shape for matmul
-    M,
-    N,
-    K,
-    stride_am,
-    stride_ak,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_as_m,
-    stride_as_k,
-    stride_bs_k,
-    stride_bs_n,
-    # Meta-parameters
-    BLOCK_SIZE_M: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    GROUP_SIZE_M: tl.constexpr,
-):
-    """Block-scale FP8 GEMM kernel.
-    Computes ``C = A @ B.T`` with block-wise activation/weight scales.
-    Uses a 2D grid with swizzle for L2 cache locality on B tiles.
-    """
-    pid_m = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-    pid_m, pid_n = tl.swizzle2d(pid_m, pid_n, num_pid_m, num_pid_n, GROUP_SIZE_M)
-    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
-    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
-    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
-    as_ptrs = As + offs_am * stride_as_m
-    offs_bsn = offs_bn // BLOCK_SIZE_N
-    bs_ptrs = Bs + offs_bsn * stride_bs_n
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        k_remaining = K - k * BLOCK_SIZE_K
-        a = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
-        b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
-        a_s = tl.load(as_ptrs + k * stride_as_k)
-        b_s = tl.load(bs_ptrs + k * stride_bs_k)
-        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-    tl.store(c_ptrs, c, mask=c_mask)
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4]
-    ],
-    key=["N", "K", "BLOCK_SIZE_M"],
-)
-@triton.jit
-def w8a8_tensor_fp8_matmul_kernel(
-    A,
-    B,
-    C,
-    As,
-    Bs,
-    M,
-    N,
-    K,
-    stride_am,
-    stride_ak,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_as_m,
-    BLOCK_SIZE_M: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    GROUP_SIZE_M: tl.constexpr,
-):
-    """Tensor-scale FP8 GEMM kernel.
-    Computes ``C = A @ B.T`` with one activation scale per row and one
-    weight scale for the full matrix.
-    Uses a 2D grid with swizzle for L2 cache locality on B tiles.
-    """
-    pid_m = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-    pid_m, pid_n = tl.swizzle2d(pid_m, pid_n, num_pid_m, num_pid_n, GROUP_SIZE_M)
-    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
-    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
-    b_ptrs = B + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn
-    a_s = tl.load(As + offs_am * stride_as_m)
-    b_s = tl.load(Bs)
-    # Accumulate raw dot products, apply scales once after the loop.
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        k_remaining = K - k * BLOCK_SIZE_K
-        a = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
-        b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
-        accumulator += tl.dot(a, b)
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    accumulator = accumulator * a_s[:, None] * b_s
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-    tl.store(c_ptrs, c, mask=c_mask)
-@triton_op("finegrained_fp8::w8a8_block_fp8_matmul", mutates_args=())
-def _w8a8_block_fp8_matmul(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    As: torch.Tensor,
-    Bs: torch.Tensor,
-    block_size: list[int],
-    output_dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """Block-scale FP8 matmul: C = A @ B.T with per-block scales.
-    As: (M, K // block_k) — per-token-group activation scales
-    Bs: (N // block_n, K // block_k) — per-block weight scales
-    """
-    assert len(block_size) == 2, (
-        f"block_size must be [block_n, block_k], got {block_size}"
-    )
-    block_n, block_k = block_size[0], block_size[1]
-    assert A.shape[-1] == B.shape[-1], (
-        f"K mismatch: A has K={A.shape[-1]}, B has K={B.shape[-1]}"
-    )
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 2, f"B must be 2D (N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    N, K = B.shape
-    M = A.numel() // A.shape[-1]
-    assert As.ndim >= 2, f"As must be at least 2D, got ndim={As.ndim}"
-    assert As.shape[-1] == triton.cdiv(K, block_k), (
-        f"As last dim {As.shape[-1]} != expected {triton.cdiv(K, block_k)} (cdiv(K={K}, block_k={block_k}))"
-    )
-    assert Bs.ndim == 2, f"Bs must be 2D (N//block_n, K//block_k), got ndim={Bs.ndim}"
-    assert Bs.shape == (triton.cdiv(N, block_n), triton.cdiv(K, block_k)), (
-        f"Bs shape {tuple(Bs.shape)} != expected ({triton.cdiv(N, block_n)}, {triton.cdiv(K, block_k)})"
-    )
-    BLOCK_SIZE_K = block_k
-    BLOCK_SIZE_N = block_n
-    C_shape = A.shape[:-1] + (N,)
-    C = A.new_empty(C_shape, dtype=output_dtype)
-    # Adaptive BLOCK_SIZE_M: smallest power-of-2 >= M, floored at 16, capped at 128.
-    # Matches the WGMMA tile to the actual row count — smaller tiles use less
-    # register pressure and a better-matched FP8 WGMMA instruction, improving
-    # both accuracy and performance for small M (decode).
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2(M), 16), 128)
-    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
-    with device_context(A.device):
-        wrap_triton(w8a8_block_fp8_matmul_kernel)[grid](
-            A,
-            B,
-            C,
-            As,
-            Bs,
-            M,
-            N,
-            K,
-            A.stride(-2),
-            A.stride(-1),
-            B.stride(1),
-            B.stride(0),
-            C.stride(-2),
-            C.stride(-1),
-            As.stride(-2),
-            As.stride(-1),
-            Bs.stride(1),
-            Bs.stride(0),
-            # Meta-parameters
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
-            GROUP_SIZE_M=8,
-        )
-    return C
-@triton_op("finegrained_fp8::w8a8_tensor_fp8_matmul", mutates_args=())
-def _w8a8_tensor_fp8_matmul(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    As: torch.Tensor,
-    Bs: torch.Tensor,
-    output_dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """Tensor-scale FP8 matmul: C = A @ B.T with per-row / per-tensor scales.
-    As: scalar, (M,), or (M, 1) — per-row activation scales
-    Bs: scalar, (1,), or (1, 1) — single weight scale
-    """
-    assert A.shape[-1] == B.shape[-1], (
-        f"K mismatch: A has K={A.shape[-1]}, B has K={B.shape[-1]}"
-    )
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 2, f"B must be 2D (N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    N, K = B.shape
-    M = A.numel() // A.shape[-1]
-    # Normalize As to (M,)
-    if As.numel() == 1:
-        As = As.reshape(1).expand(M).contiguous()
-    elif As.ndim == 2:
-        As = As.reshape(M)
-    assert As.ndim == 1 and As.shape[0] == M, (
-        f"As must be scalar, (M,), or (M,1) with M={M}, got {tuple(As.shape)}"
-    )
-    # Normalize Bs to (1,)
-    assert Bs.numel() == 1, f"Bs must be scalar or (1,), got {tuple(Bs.shape)}"
-    Bs = Bs.reshape(1)
-    BLOCK_SIZE_N = 128
-    BLOCK_SIZE_K = 128
-    C_shape = A.shape[:-1] + (N,)
-    C = A.new_empty(C_shape, dtype=output_dtype)
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2(M), 16), 128)
-    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
-    with device_context(A.device):
-        wrap_triton(w8a8_tensor_fp8_matmul_kernel)[grid](
-            A,
-            B,
-            C,
-            As,
-            Bs,
-            M,
-            N,
-            K,
-            A.stride(-2),
-            A.stride(-1),
-            B.stride(1),
-            B.stride(0),
-            C.stride(-2),
-            C.stride(-1),
-            As.stride(0),
-            # Meta-parameters
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
-            GROUP_SIZE_M=8,
-        )
-    return C
-def w8a8_block_fp8_matmul(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    As: torch.Tensor,
-    Bs: torch.Tensor,
-    block_size: list[int],
-    output_dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """Block-wise W8A8 FP8 matrix multiplication.
-    Computes ``C = A @ B.T`` where both operands are pre-quantized to
-    ``float8_e4m3fn`` with per-block scales, and accumulates in float32
-    before casting to ``output_dtype``.
-    Args:
-        A: Quantized activation tensor ``[M, K]`` in ``float8_e4m3fn``.
-        B: Quantized weight tensor ``[N, K]`` in ``float8_e4m3fn``.
-        As: Per-token-group activation scales ``[M, K // block_size[1]]``.
-        Bs: Per-block weight scales ``[N // block_size[0], K // block_size[1]]``.
-        block_size: ``[block_n, block_k]`` quantization block dimensions, e.g. ``[128, 128]``.
-        output_dtype: dtype of the returned tensor (default: ``torch.float32``).
-    Returns:
-        Output tensor ``[M, N]`` in ``output_dtype``.
-    """
-    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul(
-        A, B, As, Bs, block_size, output_dtype
-    )
-def w8a8_tensor_fp8_matmul(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    As: torch.Tensor,
-    Bs: torch.Tensor,
-    output_dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """Tensor-scale W8A8 FP8 matrix multiplication.
-    Computes ``C = A @ B.T`` in tensor-scale mode using pre-quantized FP8
-    activations/weights and tensor scales.
-    Args:
-        A: Quantized activation tensor ``[M, K]`` in ``float8_e4m3fn``.
-        B: Quantized weight tensor ``[N, K]`` in ``float8_e4m3fn``.
-        As: Per-row activation scales ``[M]``.
-        Bs: Single weight scale, scalar or ``[1]``.
-        output_dtype: dtype of the returned tensor.
-    Returns:
-        Output tensor ``[M, N]`` in ``output_dtype``.
-    """
-    return torch.ops.finegrained_fp8.w8a8_tensor_fp8_matmul(A, B, As, Bs, output_dtype)
-def w8a8_fp8_matmul(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    As: torch.Tensor,
-    Bs: torch.Tensor,
-    block_size: list[int] | None,
-    output_dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """Unified W8A8 FP8 matmul dispatcher.
-    Dispatch rules:
-    - tensor mode when ``block_size is None``
-    - tensor mode when ``block_size == [N, K]``
-    - otherwise block mode
-    Returns:
-        Output tensor ``[M, N]`` in ``output_dtype``.
-    """
-    if block_size is None or (
-        block_size[0] == B.size(0) and block_size[1] == B.size(1)
-    ):
-        return w8a8_tensor_fp8_matmul(A, B, As, Bs, output_dtype)
-    return w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype)

build/torch-rocm/metadata.json DELETED Viewed

@@ -1,10 +0,0 @@
-{
-  "name": "finegrained-fp8",
-  "id": "_finegrained_fp8_rocm_bbc7e0f",
-  "version": 1,
-  "license": "Apache-2.0",
-  "python-depends": [],
-  "backend": {
-    "type": "rocm"
-  }
-}

build/torch-rocm/utils.py DELETED Viewed

@@ -1,13 +0,0 @@
-import torch
-from contextlib import contextmanager
-@contextmanager
-def device_context(device: torch.device):
-    """Context manager that sets the active device for any backend (cuda, xpu, etc.)."""
-    backend = getattr(torch, device.type, None)
-    if backend is not None and hasattr(backend, "device"):
-        with backend.device(device):
-            yield
-    else:
-        yield

build/torch-xpu/__init__.py DELETED Viewed

@@ -1,32 +0,0 @@
-from .act_quant import fp8_act_quant
-from .batched import (
-    w8a8_fp8_matmul_batched,
-    w8a8_block_fp8_matmul_batched,
-    w8a8_tensor_fp8_matmul_batched,
-)
-from .grouped import (
-    w8a8_fp8_matmul_grouped,
-    w8a8_block_fp8_matmul_grouped,
-    w8a8_tensor_fp8_matmul_grouped,
-)
-from .matmul import (
-    w8a8_fp8_matmul,
-    w8a8_block_fp8_matmul,
-    w8a8_tensor_fp8_matmul,
-)
-__all__ = [
-    "fp8_act_quant",
-    # Single matmul
-    "w8a8_fp8_matmul",
-    "w8a8_block_fp8_matmul",
-    "w8a8_tensor_fp8_matmul",
-    # Batched matmul
-    "w8a8_fp8_matmul_batched",
-    "w8a8_block_fp8_matmul_batched",
-    "w8a8_tensor_fp8_matmul_batched",
-    # Grouped matmul
-    "w8a8_fp8_matmul_grouped",
-    "w8a8_block_fp8_matmul_grouped",
-    "w8a8_tensor_fp8_matmul_grouped",
-]

build/torch-xpu/_ops.py DELETED Viewed

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

build/torch-xpu/act_quant.py DELETED Viewed

@@ -1,73 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-import triton
-import triton.language as tl
-from torch.library import triton_op, wrap_triton
-from .utils import device_context
-_FP8_DTYPE = torch.float8_e4m3fn
-# Copied from https://huggingface.co/deepseek-ai/DeepSeek-V3/blob/main/inference/kernel.py
-@triton.jit
-def _fp8_act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr):
-    pid = tl.program_id(axis=0)
-    offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
-    x = tl.load(x_ptr + offs).to(tl.float32)
-    s = tl.max(tl.abs(x)) / 448.0  # float8_e4m3fn max
-    y = (x / s).to(y_ptr.dtype.element_ty)
-    tl.store(y_ptr + offs, y)
-    tl.store(s_ptr + pid, s)
-@triton_op("finegrained_fp8::fp8_act_quant", mutates_args=())
-def _fp8_act_quant(
-    x: torch.Tensor, block_size: int = 128
-) -> tuple[torch.Tensor, torch.Tensor]:
-    assert x.is_contiguous()
-    assert x.shape[-1] % block_size == 0
-    y = torch.empty_like(x, dtype=_FP8_DTYPE)
-    grid = (triton.cdiv(x.numel(), block_size),)
-    s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
-    with device_context(x.device):
-        wrap_triton(_fp8_act_quant_kernel)[grid](x, y, s, BLOCK_SIZE=block_size)
-    return y, s
-def fp8_act_quant(
-    x: torch.Tensor, block_size: int = 128
-) -> tuple[torch.Tensor, torch.Tensor]:
-    """Quantize activations to FP8 with per-block dynamic scaling.
-    Splits the last dimension of ``x`` into blocks of ``block_size`` elements,
-    computes ``scale = max(|x_block|) / 448`` per block, and quantizes to
-    ``float8_e4m3fn``.
-    Args:
-        x: Input tensor in bf16/fp16/fp32. Last dimension must be divisible by
-            ``block_size`` and the tensor must be contiguous.
-        block_size: Number of elements per quantization block (default: 128).
-    Returns:
-        A tuple ``(quantized, scales)`` where ``quantized`` has dtype
-        ``float8_e4m3fn`` with the same shape as ``x``, and ``scales`` has
-        shape ``(*x.shape[:-1], x.shape[-1] // block_size)`` in float32.
-    """
-    return torch.ops.finegrained_fp8.fp8_act_quant(x, block_size)

build/torch-xpu/batched.py DELETED Viewed

@@ -1,398 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-import triton
-import triton.language as tl
-from .act_quant import fp8_act_quant
-from torch.library import triton_op, wrap_triton
-from .utils import device_context
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
-    key=["N", "K"],
-)
-@triton.jit
-def w8a8_block_fp8_matmul_batched_kernel(
-    A,  # (S, K)  raw BF16/FP16 activations
-    B,  # (E, N, K) FP8 weight matrices
-    C,  # (S, N)  output
-    Bs,  # (E, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales
-    ExpertIds,  # (S,) — which expert each batch element routes to
-    # Shape
-    S,
-    N,
-    K,
-    stride_am,
-    stride_ak,
-    stride_be,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_bs_e,
-    stride_bs_k,
-    stride_bs_n,
-    # Meta-parameters
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
-):
-    """Block-scale batched FP8 expert matmul kernel.
-    Each program handles one routed token row and one N-tile, looks up the
-    owning expert from ``ExpertIds``, and applies fused activation quantization.
-    """
-    batch_id = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    # Cast expert_id to int64 to prevent int32 overflow when computing
-    # expert_id * stride_Eb (e.g. 255 * 9_437_184 > 2^31 for 256 experts of
-    # 3072×3072 FP8 weights).
-    expert_id = tl.load(ExpertIds + batch_id).to(tl.int64)
-    A = A + batch_id * stride_am
-    B = B + expert_id * stride_be
-    C = C + batch_id * stride_cm
-    Bs = Bs + expert_id * stride_bs_e
-    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + tl.arange(0, BLOCK_SIZE_M)[:, None] * 0 + offs_k[None, :] * stride_ak
-    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
-    bs_ptrs = Bs + pid_n * stride_bs_n
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        # ---- fused fp8_act_quant ----
-        a_raw = tl.load(a_ptrs).to(tl.float32)
-        a_s = tl.max(tl.abs(a_raw)) / 448.0
-        a = (a_raw / tl.maximum(a_s, 1e-12)).to(tl.float8e4nv)
-        # ---- matmul ----
-        b = tl.load(b_ptrs)
-        b_s = tl.load(bs_ptrs + k * stride_bs_k)
-        accumulator += tl.dot(a, b) * a_s * b_s[None, :]
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + offs_cm[:, None] * 0 + stride_cn * offs_cn[None, :]
-    # Fake-batch trick aliases all BLOCK_SIZE_M lanes to the same C row, so emitting
-    # `tl.store(c_ptrs, c)` issues BLOCK_SIZE_M duplicate-address stores. On NVIDIA
-    # WGMMA this is usually benign (last-write-wins of identical bytes), but on Intel
-    # XPU the duplicate-address store has hardware-undefined behavior and corrupts the
-    # output. Mask so only lane 0 stores — the (M, N) accumulator rows are
-    # mathematically identical (same A row × same B), so lane 0 holds the right value.
-    tl.store(c_ptrs, c, mask=(offs_cm == 0)[:, None])
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
-    key=["N", "K"],
-)
-@triton.jit
-def w8a8_tensor_fp8_matmul_batched_kernel(
-    A,  # (S, K) pre-quantized FP8 activations
-    B,  # (E, N, K) FP8 weight matrices
-    C,  # (S, N) output
-    As,  # (S, 1) per-tensor activation scales
-    Bs,  # (E, 1, 1) per-tensor weight scales
-    ExpertIds,
-    S,
-    N,
-    K,
-    stride_am,
-    stride_ak,
-    stride_be,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_as_m,
-    stride_bs_e,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
-):
-    """Tensor-scale batched FP8 expert matmul kernel.
-    Activations are already quantized; the kernel applies per-token activation
-    scales and per-expert tensor weight scales.
-    """
-    batch_id = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    expert_id = tl.load(ExpertIds + batch_id).to(tl.int64)
-    A = A + batch_id * stride_am
-    B = B + expert_id * stride_be
-    C = C + batch_id * stride_cm
-    Bs = Bs + expert_id * stride_bs_e
-    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + tl.arange(0, BLOCK_SIZE_M)[:, None] * 0 + offs_k[None, :] * stride_ak
-    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
-    b_s = tl.load(Bs)
-    a_s = tl.load(As + batch_id * stride_as_m)
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        a = tl.load(a_ptrs)
-        b = tl.load(b_ptrs)
-        accumulator += tl.dot(a, b)
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    accumulator = accumulator * a_s * b_s
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + offs_cm[:, None] * 0 + stride_cn * offs_cn[None, :]
-    # See block-FP8 kernel above: BLOCK_SIZE_M lanes alias the same C row;
-    # mask so only lane 0 stores to avoid hardware-undefined duplicate writes on XPU.
-    tl.store(c_ptrs, c, mask=(offs_cm == 0)[:, None])
-@triton_op("finegrained_fp8::w8a8_block_fp8_matmul_batched", mutates_args=())
-def _w8a8_block_fp8_matmul_batched(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    expert_ids: torch.Tensor,
-    block_size: list[int],
-) -> torch.Tensor:
-    """Block-scale batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T, with fused act quant.
-    A:  (S, K) raw bf16/fp16 activations
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
-    """
-    assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    assert A.shape[1] == B.shape[2], (
-        f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
-    )
-    S, K = A.shape
-    E, N, _ = B.shape
-    assert len(block_size) == 2, (
-        f"block_size must be [block_n, block_k], got {block_size}"
-    )
-    block_n, block_k = block_size[0], block_size[1]
-    # MoE expert dimensions must be block-aligned; non-aligned N/K is not supported.
-    assert N % block_n == 0, f"N ({N}) must be divisible by block_n ({block_n})"
-    assert K % block_k == 0, f"K ({K}) must be divisible by block_k ({block_k})"
-    assert Bs.ndim == 3, (
-        f"Bs must be 3D (E, N//block_n, K//block_k), got ndim={Bs.ndim}"
-    )
-    assert Bs.shape == (E, N // block_n, K // block_k), (
-        f"Bs shape {tuple(Bs.shape)} != expected ({E}, {N // block_n}, {K // block_k})"
-    )
-    C = A.new_empty(S, N)
-    # Adaptive BLOCK_SIZE_M: smallest power-of-2 >= M, floored at 16, capped at 128.
-    # Matches the WGMMA tile to the actual row count — smaller tiles use less
-    # register pressure and a better-matched FP8 WGMMA instruction, improving
-    # both accuracy and performance for small M (decode).
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
-    grid = (S, triton.cdiv(N, block_n))
-    with device_context(A.device):
-        wrap_triton(w8a8_block_fp8_matmul_batched_kernel)[grid](
-            A,
-            B,
-            C,
-            Bs,
-            expert_ids,
-            S,
-            N,
-            K,
-            A.stride(0),
-            A.stride(1),
-            B.stride(0),
-            B.stride(2),
-            B.stride(1),
-            C.stride(0),
-            C.stride(1),
-            Bs.stride(0),
-            Bs.stride(2),
-            Bs.stride(1),
-            BLOCK_SIZE_N=block_n,
-            BLOCK_SIZE_K=block_k,
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-        )
-    return C
-@triton_op("finegrained_fp8::w8a8_tensor_fp8_matmul_batched", mutates_args=())
-def _w8a8_tensor_fp8_matmul_batched(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    expert_ids: torch.Tensor,
-) -> torch.Tensor:
-    """Tensor-scale batched FP8 matmul: C[s] = A[s] @ B[expert_ids[s]].T, with fused act quant.
-    A:  (S, K) raw bf16/fp16 activations
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
-    """
-    assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    assert A.shape[1] == B.shape[2], (
-        f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
-    )
-    S, K = A.shape
-    E, N, _ = B.shape
-    # Normalize Bs to (E, 1, 1)
-    if Bs.ndim == 1:
-        assert Bs.shape[0] == E, f"Bs shape {tuple(Bs.shape)} != expected ({E},)"
-        Bs = Bs.reshape(E, 1, 1)
-    else:
-        assert Bs.shape == (E, 1, 1), (
-            f"Bs shape {tuple(Bs.shape)} != expected ({E}, 1, 1)"
-        )
-    BLOCK_SIZE_N = 128
-    BLOCK_SIZE_K = 128
-    C = A.new_empty(S, N)
-    qA, As = fp8_act_quant(A, K)
-    grid = (S, triton.cdiv(N, BLOCK_SIZE_N))
-    # Adaptive BLOCK_SIZE_M: smallest power-of-2 >= M, floored at 16, capped at 128.
-    # Matches the WGMMA tile to the actual row count — smaller tiles use less
-    # register pressure and a better-matched FP8 WGMMA instruction, improving
-    # both accuracy and performance for small M (decode).
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
-    grid = (S, triton.cdiv(N, BLOCK_SIZE_N))
-    with device_context(A.device):
-        wrap_triton(w8a8_tensor_fp8_matmul_batched_kernel)[grid](
-            qA,
-            B,
-            C,
-            As,
-            Bs,
-            expert_ids,
-            S,
-            N,
-            K,
-            qA.stride(0),
-            qA.stride(1),
-            B.stride(0),
-            B.stride(2),
-            B.stride(1),
-            C.stride(0),
-            C.stride(1),
-            As.stride(0),
-            Bs.stride(0),
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-        )
-    return C
-def w8a8_block_fp8_matmul_batched(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    expert_ids: torch.Tensor,
-    block_size: list[int],
-) -> torch.Tensor:
-    """Block-scale batched FP8 matmul with fused activation quantization.
-    A:  (S, K) raw activations, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
-    """
-    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_batched(
-        A, B, Bs, expert_ids, block_size
-    )
-def w8a8_tensor_fp8_matmul_batched(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    expert_ids: torch.Tensor,
-) -> torch.Tensor:
-    """Tensor-scale batched FP8 matmul with fused activation quantization.
-    A:  (S, K) raw activations, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
-    """
-    return torch.ops.finegrained_fp8.w8a8_tensor_fp8_matmul_batched(
-        A, B, Bs, expert_ids
-    )
-def w8a8_fp8_matmul_batched(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    expert_ids: torch.Tensor,
-    block_size: list[int] | None,
-) -> torch.Tensor:
-    """Unified batched W8A8 FP8 matmul dispatcher.
-    Dispatch rules:
-    - tensor mode when ``block_size is None``
-    - tensor mode when ``block_size == [N, K]``
-    - otherwise block mode
-    Returns:
-        Output tensor ``[S, N]`` in the same dtype as ``A``.
-    """
-    if block_size is None or (
-        block_size[0] == B.size(1) and block_size[1] == B.size(2)
-    ):
-        return w8a8_tensor_fp8_matmul_batched(A, B, Bs, expert_ids)
-    return w8a8_block_fp8_matmul_batched(A, B, Bs, expert_ids, block_size)

build/torch-xpu/finegrained_fp8/__init__.py DELETED Viewed

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

build/torch-xpu/grouped.py DELETED Viewed

@@ -1,477 +0,0 @@
-# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-import triton
-import triton.language as tl
-from .act_quant import fp8_act_quant
-from torch.library import triton_op, wrap_triton
-from .utils import device_context
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
-    key=["N", "K", "BLOCK_SIZE_M"],
-)
-@triton.jit
-def w8a8_block_fp8_matmul_grouped_kernel(
-    A,  # (S, K)  raw BF16/FP16 activations, sorted/grouped by expert id
-    B,  # (E, N, K) FP8 weight matrices
-    C,  # (S, N)  output
-    Bs,  # (E, N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales
-    Offsets,  # (E,) int32 — cumulative row-end per expert
-    TileOffsets,  # (E,) int32 — cumulative tile-end per expert
-    # Shape
-    S,
-    N,
-    K,
-    # Strides
-    stride_am,
-    stride_ak,
-    stride_be,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_bs_e,
-    stride_bs_k,
-    stride_bs_n,
-    # Meta-parameters
-    NUM_EXPERTS: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
-    NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
-):
-    """Block-scale grouped FP8 expert matmul kernel.
-    Tokens are assumed sorted by expert. The kernel maps each M-tile to its
-    owning expert via ``TileOffsets`` and applies fused activation quantization.
-    """
-    pid_m = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    # Exit early for programs beyond the actual tile count.
-    total_tiles = tl.load(TileOffsets + NUM_EXPERTS - 1)
-    if pid_m >= total_tiles:
-        return
-    # Binary search in TileOffsets to find the owning expert.
-    # Finds the smallest e such that TileOffsets[e] > pid_m (upper_bound semantics),
-    # which is the expert whose tile range contains pid_m.
-    # O(log2(NUM_EXPERTS)) loads instead of the O(NUM_EXPERTS) linear scan.
-    # NUM_EXPERTS_BIT_LENGTH is ceil(log2(E))+1 for powers-of-two, giving one
-    # harmless extra iteration when lo==hi; it's a compile-time constant so the
-    # loop is fully unrolled by the compiler.
-    lo = 0
-    hi = NUM_EXPERTS
-    for _ in tl.static_range(NUM_EXPERTS_BIT_LENGTH):
-        mid = (lo + hi) >> 1
-        mid_val = tl.load(TileOffsets + mid)
-        is_left = mid_val <= pid_m
-        lo = tl.where(is_left, mid + 1, lo)
-        hi = tl.where(is_left, hi, mid)
-    # Cast expert_id to int64 to prevent int32 overflow when computing
-    # expert_id * stride_be (e.g. 255 * 9_437_184 > 2^31 for 256 experts of
-    # 3072×3072 FP8 weights).
-    expert_id = lo.to(tl.int64)
-    prev_eid = tl.maximum(expert_id - 1, 0)
-    expert_start = tl.where(expert_id == 0, 0, tl.load(Offsets + prev_eid))
-    expert_end = tl.load(Offsets + expert_id)
-    M_expert = expert_end - expert_start
-    expert_tile_start = tl.where(expert_id == 0, 0, tl.load(TileOffsets + prev_eid))
-    local_tile = pid_m - expert_tile_start
-    m_off = local_tile * BLOCK_SIZE_M
-    offs_am = m_off + tl.arange(0, BLOCK_SIZE_M)
-    row_mask = offs_am < M_expert
-    offs_global_m = expert_start + offs_am
-    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + offs_global_m[:, None] * stride_am + offs_k[None, :] * stride_ak
-    b_ptrs = (
-        B
-        + expert_id * stride_be
-        + offs_k[:, None] * stride_bk
-        + offs_bn[None, :] * stride_bn
-    )
-    bs_ptrs = Bs + expert_id * stride_bs_e + pid_n * stride_bs_n
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        # ---- fused fp8_act_quant ----
-        a_raw = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0).to(tl.float32)
-        a_s = tl.max(tl.abs(a_raw), axis=1) / 448.0
-        a = (a_raw / tl.maximum(a_s[:, None], 1e-12)).to(tl.float8e4nv)
-        # ---- matmul ----
-        b = tl.load(b_ptrs)
-        b_s = tl.load(bs_ptrs + k * stride_bs_k)
-        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    c_ptrs = C + stride_cm * offs_global_m[:, None] + stride_cn * offs_bn[None, :]
-    c_mask = row_mask[:, None]
-    tl.store(c_ptrs, c, mask=c_mask)
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4, 5]
-    ],
-    key=["N", "K", "BLOCK_SIZE_M"],
-)
-@triton.jit
-def w8a8_tensor_fp8_matmul_grouped_kernel(
-    A,  # (S, K) raw BF16/FP16 activations, sorted/grouped by expert idc
-    B,  # (E, N, K) FP8 weight matrices
-    C,  # (S, N) output
-    As,  # (S, 1) activation scales
-    Bs,  # (E, 1, 1) per-tensor weight scales
-    Offsets,
-    TileOffsets,
-    S,
-    N,
-    K,
-    stride_am,
-    stride_ak,
-    stride_be,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_as_m,
-    stride_bs_e,
-    NUM_EXPERTS: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    BLOCK_SIZE_M: tl.constexpr,
-    NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
-):
-    """Tensor-scale grouped FP8 expert matmul kernel.
-    Uses grouped expert scheduling with pre-quantized activations plus
-    per-token activation scales and per-expert tensor weight scales.
-    """
-    pid_m = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    total_tiles = tl.load(TileOffsets + NUM_EXPERTS - 1)
-    if pid_m >= total_tiles:
-        return
-    lo = 0
-    hi = NUM_EXPERTS
-    for _ in tl.static_range(NUM_EXPERTS_BIT_LENGTH):
-        mid = (lo + hi) >> 1
-        mid_val = tl.load(TileOffsets + mid)
-        is_left = mid_val <= pid_m
-        lo = tl.where(is_left, mid + 1, lo)
-        hi = tl.where(is_left, hi, mid)
-    expert_id = lo.to(tl.int64)
-    prev_eid = tl.maximum(expert_id - 1, 0)
-    expert_start = tl.where(expert_id == 0, 0, tl.load(Offsets + prev_eid))
-    expert_end = tl.load(Offsets + expert_id)
-    M_expert = expert_end - expert_start
-    expert_tile_start = tl.where(expert_id == 0, 0, tl.load(TileOffsets + prev_eid))
-    local_tile = pid_m - expert_tile_start
-    m_off = local_tile * BLOCK_SIZE_M
-    offs_am = m_off + tl.arange(0, BLOCK_SIZE_M)
-    row_mask = offs_am < M_expert
-    offs_global_m = expert_start + offs_am
-    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + offs_global_m[:, None] * stride_am + offs_k[None, :] * stride_ak
-    b_ptrs = (
-        B
-        + expert_id * stride_be
-        + offs_k[:, None] * stride_bk
-        + offs_bn[None, :] * stride_bn
-    )
-    a_s = tl.load(As + offs_global_m * stride_as_m, mask=row_mask, other=0.0)
-    b_s = tl.load(Bs + expert_id * stride_bs_e)
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        a = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0)
-        b = tl.load(b_ptrs)
-        accumulator += tl.dot(a, b)
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    accumulator = accumulator * a_s[:, None] * b_s
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    c_ptrs = C + stride_cm * offs_global_m[:, None] + stride_cn * offs_bn[None, :]
-    c_mask = row_mask[:, None]
-    tl.store(c_ptrs, c, mask=c_mask)
-@triton_op("finegrained_fp8::w8a8_block_fp8_matmul_grouped", mutates_args=())
-def _w8a8_block_fp8_matmul_grouped(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    offsets: torch.Tensor,
-    tokens_per_expert: torch.Tensor,
-    block_size: list[int],
-) -> torch.Tensor:
-    """Block-scale grouped FP8 matmul: C = A @ B.T per expert, with fused act quant.
-    A:  (S, K) raw bf16/fp16 activations, sorted by expert
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
-    """
-    assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    assert A.shape[1] == B.shape[2], (
-        f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
-    )
-    S, K = A.shape
-    E, N, _ = B.shape
-    assert len(block_size) == 2, (
-        f"block_size must be [block_n, block_k], got {block_size}"
-    )
-    block_n, block_k = block_size[0], block_size[1]
-    # MoE expert dimensions must be block-aligned; non-aligned N/K is not supported.
-    assert N % block_n == 0, f"N ({N}) must be divisible by block_n ({block_n})"
-    assert K % block_k == 0, f"K ({K}) must be divisible by block_k ({block_k})"
-    assert Bs.ndim == 3, (
-        f"Bs must be 3D (E, N//block_n, K//block_k), got ndim={Bs.ndim}"
-    )
-    assert Bs.shape == (E, N // block_n, K // block_k), (
-        f"Bs shape {tuple(Bs.shape)} != expected ({E}, {N // block_n}, {K // block_k})"
-    )
-    C = A.new_empty(S, N)
-    # Adaptive BLOCK_SIZE_M: match tile to average tokens per expert.
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
-    tiles_per_expert = (tokens_per_expert + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M
-    tile_offsets = torch.cumsum(tiles_per_expert, dim=0).to(torch.int32)
-    # Upper bound on M-tiles: sum_e ceil(M_e / BLOCK_M) <= ceil(S / BLOCK_M) + E.
-    # Programs beyond the real tile count exit immediately via the early-return
-    # guard inside the kernel. This is faster than syncing for the exact count
-    # and keeps the grid size data-independent (cuda-graph / torch.compile safe).
-    max_M_tiles = triton.cdiv(S, BLOCK_SIZE_M) + E
-    grid = (max_M_tiles, triton.cdiv(N, block_n))
-    with device_context(A.device):
-        wrap_triton(w8a8_block_fp8_matmul_grouped_kernel)[grid](
-            A,
-            B,
-            C,
-            Bs,
-            offsets,
-            tile_offsets,
-            S,
-            N,
-            K,
-            A.stride(0),
-            A.stride(1),
-            B.stride(0),
-            B.stride(2),
-            B.stride(1),
-            C.stride(0),
-            C.stride(1),
-            Bs.stride(0),
-            Bs.stride(2),
-            Bs.stride(1),
-            # Meta-parameters
-            NUM_EXPERTS=E,
-            BLOCK_SIZE_N=block_n,
-            BLOCK_SIZE_K=block_k,
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-            NUM_EXPERTS_BIT_LENGTH=E.bit_length(),
-        )
-    return C
-@triton_op("finegrained_fp8::w8a8_tensor_fp8_matmul_grouped", mutates_args=())
-def _w8a8_tensor_fp8_matmul_grouped(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    offsets: torch.Tensor,
-    tokens_per_expert: torch.Tensor,
-) -> torch.Tensor:
-    """Tensor-scale grouped FP8 matmul: C = A @ B.T per expert, with fused act quant.
-    A:  (S, K) raw bf16/fp16 activations, sorted by expert
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
-    """
-    assert A.ndim == 2, f"A must be 2D (S, K), got ndim={A.ndim}"
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 3, f"B must be 3D (E, N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    assert A.shape[1] == B.shape[2], (
-        f"K mismatch: A has K={A.shape[1]}, B has K={B.shape[2]}"
-    )
-    S, K = A.shape
-    E, N, _ = B.shape
-    # Normalize Bs to (E, 1, 1)
-    if Bs.ndim == 1:
-        assert Bs.shape[0] == E, f"Bs shape {tuple(Bs.shape)} != expected ({E},)"
-        Bs = Bs.reshape(E, 1, 1)
-    else:
-        assert Bs.shape == (E, 1, 1), (
-            f"Bs shape {tuple(Bs.shape)} != expected ({E}, 1, 1)"
-        )
-    BLOCK_SIZE_N = 128
-    BLOCK_SIZE_K = 128
-    C = A.new_empty(S, N)
-    qA, As = fp8_act_quant(A, K)
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
-    tiles_per_expert = (tokens_per_expert + BLOCK_SIZE_M - 1) // BLOCK_SIZE_M
-    tile_offsets = torch.cumsum(tiles_per_expert, dim=0).to(torch.int32)
-    # Upper bound on M-tiles: sum_e ceil(M_e / BLOCK_M) <= ceil(S / BLOCK_M) + E.
-    # Programs beyond the real tile count exit immediately via the early-return
-    # guard inside the kernel. This is faster than syncing for the exact count
-    # and keeps the grid size data-independent (cuda-graph / torch.compile safe).
-    max_M_tiles = triton.cdiv(S, BLOCK_SIZE_M) + E
-    grid = (max_M_tiles, triton.cdiv(N, BLOCK_SIZE_N))
-    with device_context(A.device):
-        wrap_triton(w8a8_tensor_fp8_matmul_grouped_kernel)[grid](
-            qA,
-            B,
-            C,
-            As,
-            Bs,
-            offsets,
-            tile_offsets,
-            S,
-            N,
-            K,
-            qA.stride(0),
-            qA.stride(1),
-            B.stride(0),
-            B.stride(2),
-            B.stride(1),
-            C.stride(0),
-            C.stride(1),
-            As.stride(0),
-            Bs.stride(0),
-            # Meta-parameters
-            NUM_EXPERTS=E,
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-            NUM_EXPERTS_BIT_LENGTH=E.bit_length(),
-        )
-    return C
-def w8a8_block_fp8_matmul_grouped(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    offsets: torch.Tensor,
-    tokens_per_expert: torch.Tensor,
-    block_size: list[int],
-) -> torch.Tensor:
-    """Block-scale grouped FP8 matmul with fused activation quantization.
-    A:  (S, K) raw activations sorted by expert, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E, N // block_n, K // block_k) per-block weight scales
-    """
-    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_grouped(
-        A, B, Bs, offsets, tokens_per_expert, block_size
-    )
-def w8a8_tensor_fp8_matmul_grouped(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    offsets: torch.Tensor,
-    tokens_per_expert: torch.Tensor,
-) -> torch.Tensor:
-    """Tensor-scale grouped FP8 matmul with fused activation quantization.
-    A:  (S, K) raw activations sorted by expert, bf16/fp16/fp32
-    B:  (E, N, K) FP8 expert weights
-    Bs: (E,) or (E, 1, 1) per-expert weight scales
-    """
-    return torch.ops.finegrained_fp8.w8a8_tensor_fp8_matmul_grouped(
-        A, B, Bs, offsets, tokens_per_expert
-    )
-def w8a8_fp8_matmul_grouped(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    Bs: torch.Tensor,
-    offsets: torch.Tensor,
-    tokens_per_expert: torch.Tensor,
-    block_size: list[int] | None,
-) -> torch.Tensor:
-    """Unified grouped W8A8 FP8 matmul dispatcher.
-    Dispatch rules:
-    - tensor mode when ``block_size is None``
-    - tensor mode when ``block_size == [N, K]``
-    - otherwise block mode
-    Returns:
-        Output tensor ``[S, N]`` in the same dtype as ``A``, in expert-sorted order.
-    """
-    if block_size is None or (
-        block_size[0] == B.size(1) and block_size[1] == B.size(2)
-    ):
-        return w8a8_tensor_fp8_matmul_grouped(A, B, Bs, offsets, tokens_per_expert)
-    return w8a8_block_fp8_matmul_grouped(
-        A, B, Bs, offsets, tokens_per_expert, block_size
-    )

build/torch-xpu/matmul.py DELETED Viewed

@@ -1,411 +0,0 @@
-# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-import torch
-import triton
-import triton.language as tl
-from torch.library import triton_op, wrap_triton
-from .utils import device_context
-# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4]
-    ],
-    key=["N", "K", "BLOCK_SIZE_M"],
-)
-@triton.jit
-def w8a8_block_fp8_matmul_kernel(
-    # Pointers to inputs and output
-    A,
-    B,
-    C,
-    As,
-    Bs,
-    # Shape for matmul
-    M,
-    N,
-    K,
-    stride_am,
-    stride_ak,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_as_m,
-    stride_as_k,
-    stride_bs_k,
-    stride_bs_n,
-    # Meta-parameters
-    BLOCK_SIZE_M: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    GROUP_SIZE_M: tl.constexpr,
-):
-    """Block-scale FP8 GEMM kernel.
-    Computes ``C = A @ B.T`` with block-wise activation/weight scales.
-    Uses a 2D grid with swizzle for L2 cache locality on B tiles.
-    """
-    pid_m = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-    pid_m, pid_n = tl.swizzle2d(pid_m, pid_n, num_pid_m, num_pid_n, GROUP_SIZE_M)
-    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
-    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
-    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
-    as_ptrs = As + offs_am * stride_as_m
-    offs_bsn = offs_bn // BLOCK_SIZE_N
-    bs_ptrs = Bs + offs_bsn * stride_bs_n
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        k_remaining = K - k * BLOCK_SIZE_K
-        a = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
-        b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
-        a_s = tl.load(as_ptrs + k * stride_as_k)
-        b_s = tl.load(bs_ptrs + k * stride_bs_k)
-        accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-    tl.store(c_ptrs, c, mask=c_mask)
-@triton.autotune(
-    configs=[
-        triton.Config({}, num_warps=w, num_stages=s)
-        for w in [2, 4, 8, 16]
-        for s in [2, 3, 4]
-    ],
-    key=["N", "K", "BLOCK_SIZE_M"],
-)
-@triton.jit
-def w8a8_tensor_fp8_matmul_kernel(
-    A,
-    B,
-    C,
-    As,
-    Bs,
-    M,
-    N,
-    K,
-    stride_am,
-    stride_ak,
-    stride_bk,
-    stride_bn,
-    stride_cm,
-    stride_cn,
-    stride_as_m,
-    BLOCK_SIZE_M: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
-    GROUP_SIZE_M: tl.constexpr,
-):
-    """Tensor-scale FP8 GEMM kernel.
-    Computes ``C = A @ B.T`` with one activation scale per row and one
-    weight scale for the full matrix.
-    Uses a 2D grid with swizzle for L2 cache locality on B tiles.
-    """
-    pid_m = tl.program_id(axis=0)
-    pid_n = tl.program_id(axis=1)
-    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
-    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
-    pid_m, pid_n = tl.swizzle2d(pid_m, pid_n, num_pid_m, num_pid_n, GROUP_SIZE_M)
-    offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
-    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
-    offs_k = tl.arange(0, BLOCK_SIZE_K)
-    a_ptrs = A + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
-    b_ptrs = B + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn
-    a_s = tl.load(As + offs_am * stride_as_m)
-    b_s = tl.load(Bs)
-    # Accumulate raw dot products, apply scales once after the loop.
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
-    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
-        k_remaining = K - k * BLOCK_SIZE_K
-        a = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0)
-        b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
-        accumulator += tl.dot(a, b)
-        a_ptrs += BLOCK_SIZE_K * stride_ak
-        b_ptrs += BLOCK_SIZE_K * stride_bk
-    accumulator = accumulator * a_s[:, None] * b_s
-    if C.dtype.element_ty == tl.bfloat16:
-        c = accumulator.to(tl.bfloat16)
-    elif C.dtype.element_ty == tl.float16:
-        c = accumulator.to(tl.float16)
-    else:
-        c = accumulator.to(tl.float32)
-    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
-    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = C + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
-    tl.store(c_ptrs, c, mask=c_mask)
-@triton_op("finegrained_fp8::w8a8_block_fp8_matmul", mutates_args=())
-def _w8a8_block_fp8_matmul(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    As: torch.Tensor,
-    Bs: torch.Tensor,
-    block_size: list[int],
-    output_dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """Block-scale FP8 matmul: C = A @ B.T with per-block scales.
-    As: (M, K // block_k) — per-token-group activation scales
-    Bs: (N // block_n, K // block_k) — per-block weight scales
-    """
-    assert len(block_size) == 2, (
-        f"block_size must be [block_n, block_k], got {block_size}"
-    )
-    block_n, block_k = block_size[0], block_size[1]
-    assert A.shape[-1] == B.shape[-1], (
-        f"K mismatch: A has K={A.shape[-1]}, B has K={B.shape[-1]}"
-    )
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 2, f"B must be 2D (N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    N, K = B.shape
-    M = A.numel() // A.shape[-1]
-    assert As.ndim >= 2, f"As must be at least 2D, got ndim={As.ndim}"
-    assert As.shape[-1] == triton.cdiv(K, block_k), (
-        f"As last dim {As.shape[-1]} != expected {triton.cdiv(K, block_k)} (cdiv(K={K}, block_k={block_k}))"
-    )
-    assert Bs.ndim == 2, f"Bs must be 2D (N//block_n, K//block_k), got ndim={Bs.ndim}"
-    assert Bs.shape == (triton.cdiv(N, block_n), triton.cdiv(K, block_k)), (
-        f"Bs shape {tuple(Bs.shape)} != expected ({triton.cdiv(N, block_n)}, {triton.cdiv(K, block_k)})"
-    )
-    BLOCK_SIZE_K = block_k
-    BLOCK_SIZE_N = block_n
-    C_shape = A.shape[:-1] + (N,)
-    C = A.new_empty(C_shape, dtype=output_dtype)
-    # Adaptive BLOCK_SIZE_M: smallest power-of-2 >= M, floored at 16, capped at 128.
-    # Matches the WGMMA tile to the actual row count — smaller tiles use less
-    # register pressure and a better-matched FP8 WGMMA instruction, improving
-    # both accuracy and performance for small M (decode).
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2(M), 16), 128)
-    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
-    with device_context(A.device):
-        wrap_triton(w8a8_block_fp8_matmul_kernel)[grid](
-            A,
-            B,
-            C,
-            As,
-            Bs,
-            M,
-            N,
-            K,
-            A.stride(-2),
-            A.stride(-1),
-            B.stride(1),
-            B.stride(0),
-            C.stride(-2),
-            C.stride(-1),
-            As.stride(-2),
-            As.stride(-1),
-            Bs.stride(1),
-            Bs.stride(0),
-            # Meta-parameters
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
-            GROUP_SIZE_M=8,
-        )
-    return C
-@triton_op("finegrained_fp8::w8a8_tensor_fp8_matmul", mutates_args=())
-def _w8a8_tensor_fp8_matmul(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    As: torch.Tensor,
-    Bs: torch.Tensor,
-    output_dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """Tensor-scale FP8 matmul: C = A @ B.T with per-row / per-tensor scales.
-    As: scalar, (M,), or (M, 1) — per-row activation scales
-    Bs: scalar, (1,), or (1, 1) — single weight scale
-    """
-    assert A.shape[-1] == B.shape[-1], (
-        f"K mismatch: A has K={A.shape[-1]}, B has K={B.shape[-1]}"
-    )
-    assert A.is_contiguous(), "A must be contiguous"
-    assert B.ndim == 2, f"B must be 2D (N, K), got ndim={B.ndim}"
-    assert B.is_contiguous(), "B must be contiguous"
-    N, K = B.shape
-    M = A.numel() // A.shape[-1]
-    # Normalize As to (M,)
-    if As.numel() == 1:
-        As = As.reshape(1).expand(M).contiguous()
-    elif As.ndim == 2:
-        As = As.reshape(M)
-    assert As.ndim == 1 and As.shape[0] == M, (
-        f"As must be scalar, (M,), or (M,1) with M={M}, got {tuple(As.shape)}"
-    )
-    # Normalize Bs to (1,)
-    assert Bs.numel() == 1, f"Bs must be scalar or (1,), got {tuple(Bs.shape)}"
-    Bs = Bs.reshape(1)
-    BLOCK_SIZE_N = 128
-    BLOCK_SIZE_K = 128
-    C_shape = A.shape[:-1] + (N,)
-    C = A.new_empty(C_shape, dtype=output_dtype)
-    BLOCK_SIZE_M = min(max(triton.next_power_of_2(M), 16), 128)
-    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
-    with device_context(A.device):
-        wrap_triton(w8a8_tensor_fp8_matmul_kernel)[grid](
-            A,
-            B,
-            C,
-            As,
-            Bs,
-            M,
-            N,
-            K,
-            A.stride(-2),
-            A.stride(-1),
-            B.stride(1),
-            B.stride(0),
-            C.stride(-2),
-            C.stride(-1),
-            As.stride(0),
-            # Meta-parameters
-            BLOCK_SIZE_M=BLOCK_SIZE_M,
-            BLOCK_SIZE_N=BLOCK_SIZE_N,
-            BLOCK_SIZE_K=BLOCK_SIZE_K,
-            GROUP_SIZE_M=8,
-        )
-    return C
-def w8a8_block_fp8_matmul(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    As: torch.Tensor,
-    Bs: torch.Tensor,
-    block_size: list[int],
-    output_dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """Block-wise W8A8 FP8 matrix multiplication.
-    Computes ``C = A @ B.T`` where both operands are pre-quantized to
-    ``float8_e4m3fn`` with per-block scales, and accumulates in float32
-    before casting to ``output_dtype``.
-    Args:
-        A: Quantized activation tensor ``[M, K]`` in ``float8_e4m3fn``.
-        B: Quantized weight tensor ``[N, K]`` in ``float8_e4m3fn``.
-        As: Per-token-group activation scales ``[M, K // block_size[1]]``.
-        Bs: Per-block weight scales ``[N // block_size[0], K // block_size[1]]``.
-        block_size: ``[block_n, block_k]`` quantization block dimensions, e.g. ``[128, 128]``.
-        output_dtype: dtype of the returned tensor (default: ``torch.float32``).
-    Returns:
-        Output tensor ``[M, N]`` in ``output_dtype``.
-    """
-    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul(
-        A, B, As, Bs, block_size, output_dtype
-    )
-def w8a8_tensor_fp8_matmul(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    As: torch.Tensor,
-    Bs: torch.Tensor,
-    output_dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """Tensor-scale W8A8 FP8 matrix multiplication.
-    Computes ``C = A @ B.T`` in tensor-scale mode using pre-quantized FP8
-    activations/weights and tensor scales.
-    Args:
-        A: Quantized activation tensor ``[M, K]`` in ``float8_e4m3fn``.
-        B: Quantized weight tensor ``[N, K]`` in ``float8_e4m3fn``.
-        As: Per-row activation scales ``[M]``.
-        Bs: Single weight scale, scalar or ``[1]``.
-        output_dtype: dtype of the returned tensor.
-    Returns:
-        Output tensor ``[M, N]`` in ``output_dtype``.
-    """
-    return torch.ops.finegrained_fp8.w8a8_tensor_fp8_matmul(A, B, As, Bs, output_dtype)
-def w8a8_fp8_matmul(
-    A: torch.Tensor,
-    B: torch.Tensor,
-    As: torch.Tensor,
-    Bs: torch.Tensor,
-    block_size: list[int] | None,
-    output_dtype: torch.dtype = torch.float32,
-) -> torch.Tensor:
-    """Unified W8A8 FP8 matmul dispatcher.
-    Dispatch rules:
-    - tensor mode when ``block_size is None``
-    - tensor mode when ``block_size == [N, K]``
-    - otherwise block mode
-    Returns:
-        Output tensor ``[M, N]`` in ``output_dtype``.
-    """
-    if block_size is None or (
-        block_size[0] == B.size(0) and block_size[1] == B.size(1)
-    ):
-        return w8a8_tensor_fp8_matmul(A, B, As, Bs, output_dtype)
-    return w8a8_block_fp8_matmul(A, B, As, Bs, block_size, output_dtype)

build/torch-xpu/metadata.json DELETED Viewed

@@ -1,10 +0,0 @@
-{
-  "name": "finegrained-fp8",
-  "id": "_finegrained_fp8_xpu_7c5619e",
-  "version": 1,
-  "license": "Apache-2.0",
-  "python-depends": [],
-  "backend": {
-    "type": "xpu"
-  }
-}

build/torch-xpu/utils.py DELETED Viewed

@@ -1,13 +0,0 @@
-import torch
-from contextlib import contextmanager
-@contextmanager
-def device_context(device: torch.device):
-    """Context manager that sets the active device for any backend (cuda, xpu, etc.)."""
-    backend = getattr(torch, device.type, None)
-    if backend is not None and hasattr(backend, "device"):
-        with backend.device(device):
-            yield
-    else:
-        yield