Kernels:

kernels-community
/

finegrained-fp8

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on Jun 17

Commit

682a416

verified ·

1 Parent(s): 9f00a55

Uploaded using `kernel-builder`.

Browse files

Files changed (9) hide show

build/torch-rocm/__init__.py +14 -0
build/torch-rocm/_ops.py +38 -0
build/torch-rocm/batched.py +773 -0
build/torch-rocm/finegrained_fp8/__init__.py +26 -0
build/torch-rocm/grouped.py +951 -0
build/torch-rocm/matmul.py +941 -0
build/torch-rocm/metadata.json +22 -0
build/torch-rocm/metadata.json.sigstore +1 -0
build/torch-rocm/utils.py +269 -0

build/torch-rocm/__init__.py ADDED Viewed

	@@ -0,0 +1,14 @@

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

build/torch-rocm/_ops.py ADDED Viewed

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

+import torch
+def get_backend() -> str:
+    """Detect the backend by inspecting torch."""
+    import torch
+    if hasattr(torch, "neuron"):
+        # Needs to be sorted before specific Torch builds, since Neuron
+        # extension can be loaded into e.g. CUDA Torch builds.
+        return "neuron"
+    elif torch.version.cuda is not None:
+        return "cuda"
+    elif torch.version.hip is not None:
+        return "rocm"
+    elif torch.backends.mps.is_available():
+        return "metal"
+    elif hasattr(torch.version, "xpu") and torch.version.xpu is not None:
+        return "xpu"
+    else:
+        return "cpu"
+def _find_ops_name() -> str:
+    kernel_name = "finegrained_fp8"
+    unique_id = "846165b"
+    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/batched.py ADDED Viewed

	@@ -0,0 +1,773 @@

+# 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 ._ops import add_op_namespace_prefix, ops
+from .utils import (
+    MX_SCALE_GROUP_K,
+    NIBBLES_PER_BYTE,
+    decode_ue8m0_scale,
+    device_context,
+    mx_act_quant_inline,
+    fp8_act_quant,
+    fp8_act_quant_inline,
+    get_accelerator_autotuning_configs,
+    is_mxfp4,
+    is_mxfp8,
+    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"],
+)
+@triton.jit
+def w8a8_block_dynamic_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,
+    # 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,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    NUM_EXPERTS: 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)
+    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)):
+        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(with_block_sizes=True),
+    key=["N", "K"],
+)
+@triton.jit
+def w8a8_tensor_dynamic_fp8_matmul_batched_kernel(
+    A,  # (S, K) pre-quantized FP8 activations
+    B,  # (E, N, K) FP8 weight matrices
+    C,  # (S, N) output
+    As,  # (S,) per-token activation scales
+    Bs,  # (E, 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_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_as_m,
+    stride_bs_e,
+    stride_eid,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    NUM_EXPERTS: 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)
+    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
+    _store_row(C, accumulator, pid_n, stride_cn, BLOCK_SIZE_M, BLOCK_SIZE_N)
+@triton.autotune(
+    configs=get_accelerator_autotuning_configs(with_block_sizes=True),
+    key=["N", "K"],
+)
+@triton.jit
+def w4a8_mx_dynamic_fp4_matmul_batched_kernel(
+    A,  # (S, K) raw BF16/FP16 activations
+    B,  # (E, N, K // 2) packed FP4 (E2M1) expert weights as int8
+    C,  # (S, N) output
+    Bs,  # (E, 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,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    NIBBLES_PER_BYTE: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    NUM_EXPERTS: tl.constexpr,
+):
+    """Batched MXFP4 (W4A8) expert matmul with fused activation quant.
+    Each program handles one routed token row and one N-tile, looks up the
+    owning expert from ``ExpertIds``, quantizes ``A`` to FP8 per K-group inline
+    (UE8M0 scale), then ``tl.dot_scaled`` against packed FP4 weights.
+    """
+    _, 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_k_byte = tl.arange(0, BLOCK_SIZE_K // NIBBLES_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_k_byte[:, 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 = mx_act_quant_inline(
+            a_raw, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K
+        )
+        b = tl.load(b_ptrs).to(tl.uint8)
+        b_s = tl.load(bs_ptrs).to(tl.uint8)
+        accumulator += tl.dot_scaled(a, a_scale, "e4m3", b, b_s, "e2m1")
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += (BLOCK_SIZE_K // NIBBLES_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.autotune(
+    configs=get_accelerator_autotuning_configs(with_block_sizes=True),
+    key=["N", "K"],
+)
+@triton.jit
+def w8a8_mx_dynamic_fp8_matmul_batched_kernel(
+    A,  # (S, K) raw BF16/FP16 activations
+    B,  # (E, N, K) E4M3 expert weights (unpacked)
+    C,  # (S, N) output
+    Bs,  # (E, 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,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+    NUM_EXPERTS: tl.constexpr,
+):
+    """Batched MXFP8 (W8A8) expert matmul with fused activation quant.
+    Each program handles one routed token row and one N-tile, looks up the
+    owning expert from ``ExpertIds``, quantizes ``A`` to E4M3 per K-group inline
+    (UE8M0 scale), then ``tl.dot_scaled`` against E4M3 weights. Mirrors the
+    batched MXFP4 kernel but with unpacked weights and ``"e4m3"`` on both operands.
+    """
+    _, 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_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_k[:, 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 = mx_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)
+        accumulator += tl.dot_scaled(a, a_scale, "e4m3", b, b_s, "e4m3")
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * 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:  (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})"
+    )
+    # Decode handles one routed row per program; BLOCK_SIZE_M > 1 would just
+    # duplicate the same row computation and keep one row on store.
+    BLOCK_SIZE_M = 1
+    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,
+            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),
+            expert_ids.stride(0),
+            BLOCK_SIZE_N=block_n,
+            BLOCK_SIZE_K=block_k,
+            BLOCK_SIZE_M=BLOCK_SIZE_M,
+            NUM_EXPERTS=E,
+        )
+    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:  (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)"
+        )
+    # Decode handles one routed row per program; BLOCK_SIZE_M > 1 would just
+    # duplicate the same row computation and keep one row on store.
+    BLOCK_SIZE_M = 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,
+            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),
+            expert_ids.stride(0),
+            BLOCK_SIZE_M=BLOCK_SIZE_M,
+            NUM_EXPERTS=E,
+        )
+    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:  (E, N, K) FP8 expert weights
+    Bs: (E, 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:  (E, N, K) FP8 expert weights
+    Bs: (E,) or (E, 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("w4a8_mx_dynamic_fp4_matmul_batched"), mutates_args=()
+)
+def _w4a8_mx_dynamic_fp4_matmul_batched(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    expert_ids: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Batched MXFP4 (W4A8) matmul with fused activation quant.
+    A:  (S, K) raw activations, bf16/fp16/fp32 (quantized inline to FP8)
+    B:  (E, N, K // 2) packed FP4 (E2M1) expert weights, two codes per int8
+    Bs: (E, N, K // 32) UE8M0 weight scales
+    expert_ids: (S,) which expert each routed row uses
+    BLOCK_SIZE_N and BLOCK_SIZE_K are autotuned.
+    """
+    assert A.ndim == 2 and B.ndim == 3 and Bs.ndim == 3
+    assert expert_ids.ndim == 1
+    assert B.dtype == torch.int8, f"B must be int8 (packed FP4), got {B.dtype}"
+    assert Bs.dtype == torch.float8_e8m0fnu, (
+        f"Bs must be float8_e8m0fnu, got {Bs.dtype}"
+    )
+    S, K = A.shape
+    E, N, K_half = B.shape
+    assert expert_ids.shape[0] == S
+    assert K == NIBBLES_PER_BYTE * K_half, (
+        f"K (={K}) must equal {NIBBLES_PER_BYTE} * B.shape[2] (={K_half})"
+    )
+    assert K % MX_SCALE_GROUP_K == 0, (
+        f"K (={K}) must be a multiple of {MX_SCALE_GROUP_K}"
+    )
+    assert Bs.shape == (E, N, K // MX_SCALE_GROUP_K), (
+        f"Bs shape {tuple(Bs.shape)} != ({E}, {N}, {K // MX_SCALE_GROUP_K})"
+    )
+    # Decode handles one routed row per program; BLOCK_SIZE_M > 1 would just
+    # duplicate the same row computation and keep one row on store.
+    BLOCK_SIZE_M = 1
+    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(w4a8_mx_dynamic_fp4_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=BLOCK_SIZE_M,
+            NIBBLES_PER_BYTE=NIBBLES_PER_BYTE,
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+            NUM_EXPERTS=E,
+        )
+    return C
+def w4a8_mx_dynamic_fp4_matmul_batched(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    expert_ids: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Batched MXFP4 (W4A8) matmul with fused activation quant. Tile
+    shape autotuned; FP4 scale granularity is fixed at 32."""
+    return ops.w4a8_mx_dynamic_fp4_matmul_batched(A, B, Bs, expert_ids, output_dtype)
+@triton_op(
+    add_op_namespace_prefix("w8a8_mx_dynamic_fp8_matmul_batched"), mutates_args=()
+)
+def _w8a8_mx_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:
+    """Batched MXFP8 (W8A8) matmul: ``C[s] = A[s] @ B[expert_ids[s]].T`` (E4M3 ×
+    E4M3, UE8M0 group-32 scales) with fused activation quant.
+    A:  (S, K) raw activations, bf16/fp16/fp32 (quantized inline to E4M3)
+    B:  (E, N, K) E4M3 expert weights (unpacked)
+    Bs: (E, N, K // 32) UE8M0 weight scales
+    expert_ids: (S,) which expert each routed row uses
+    BLOCK_SIZE_N and BLOCK_SIZE_K are autotuned.
+    """
+    assert A.ndim == 2 and B.ndim == 3 and Bs.ndim == 3
+    assert expert_ids.ndim == 1
+    assert B.dtype == torch.float8_e4m3fn, (
+        f"B must be float8_e4m3fn (E4M3 weights), got {B.dtype}"
+    )
+    assert Bs.dtype == torch.float8_e8m0fnu, (
+        f"Bs must be float8_e8m0fnu, got {Bs.dtype}"
+    )
+    S, K = A.shape
+    E, N, K_b = B.shape
+    assert expert_ids.shape[0] == S
+    assert K == K_b, f"K mismatch: A has K={K}, B has K={K_b}"
+    assert K % MX_SCALE_GROUP_K == 0, (
+        f"K (={K}) must be a multiple of {MX_SCALE_GROUP_K}"
+    )
+    assert Bs.shape == (E, N, K // MX_SCALE_GROUP_K), (
+        f"Bs shape {tuple(Bs.shape)} != ({E}, {N}, {K // MX_SCALE_GROUP_K})"
+    )
+    # Decode handles one routed row per program; BLOCK_SIZE_M > 1 would just
+    # duplicate the same row computation and keep one row on store.
+    BLOCK_SIZE_M = 1
+    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(w8a8_mx_dynamic_fp8_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=BLOCK_SIZE_M,
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+            NUM_EXPERTS=E,
+        )
+    return C
+def w8a8_mx_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:
+    """Batched MXFP8 (W8A8) matmul with fused activation quant (E4M3 × E4M3,
+    UE8M0 group-32) via the MX scaled MMA. Tile shape autotuned."""
+    return ops.w8a8_mx_dynamic_fp8_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``:
+    - ``B.dtype == int8`` (packed FP4) → ``w4a8_mx_dynamic_fp4_matmul_batched``
+      (``block_size`` is ignored; FP4 tile shape is autotuned).
+    - ``B.dtype == float8_e4m3fn`` with UE8M0 group-32 ``Bs`` (shape ``[E, N, K//32]``)
+      → ``w8a8_mx_dynamic_fp8_matmul_batched`` (MX scaled MMA; ``block_size`` ignored).
+    - ``block_size`` None or full ``[N, K]`` → ``w8a8_tensor_dynamic_fp8_matmul_batched``.
+    - otherwise → ``w8a8_block_dynamic_fp8_matmul_batched``.
+    """
+    if is_mxfp4(B, Bs):
+        return w4a8_mx_dynamic_fp4_matmul_batched(A, B, Bs, expert_ids, output_dtype)
+    if is_mxfp8(B, Bs):
+        return w8a8_mx_dynamic_fp8_matmul_batched(A, B, Bs, expert_ids, output_dtype)
+    if block_size is None or (
+        block_size[0] == B.size(1) and block_size[1] == B.size(2)
+    ):
+        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-rocm/finegrained_fp8/__init__.py ADDED Viewed

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

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

build/torch-rocm/grouped.py ADDED Viewed

	@@ -0,0 +1,951 @@

+# 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 ._ops import add_op_namespace_prefix, ops
+from .utils import (
+    MX_SCALE_GROUP_K,
+    NIBBLES_PER_BYTE,
+    adaptive_block_size_m,
+    device_context,
+    mx_act_quant_inline,
+    fp8_act_quant,
+    fp8_act_quant_inline,
+    get_accelerator_autotuning_configs,
+    is_mxfp4,
+    is_mxfp8,
+    ue8m0_as_uint8,
+    decode_ue8m0_scale,
+)
+def grouped_tile_layout(
+    tokens_per_expert: torch.Tensor,
+    block_size_m: int,
+    S: int,
+    E: int,
+) -> tuple[torch.Tensor, int]:
+    """Compute the M-tile layout for the grouped kernels.
+    Returns ``(tile_offsets, max_m_tiles)``:
+    - ``tile_offsets``: int32 (E,) 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) + E
+    return tile_offsets, max_m_tiles
+@triton.jit
+def _grouped_tile_setup(
+    pid_m,
+    pid_n,
+    Offsets,
+    TileOffsets,
+    stride_offs,
+    stride_tile,
+    NUM_EXPERTS: tl.constexpr,
+    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(E))+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 E
+    # × 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,  # (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,
+    stride_offs,
+    stride_tile,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    NUM_EXPERTS: 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)
+    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 = (
+        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)):
+        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
+    )
+@triton.autotune(
+    configs=get_accelerator_autotuning_configs(with_block_sizes=True),
+    key=["N", "K", "BLOCK_SIZE_M"],
+)
+@triton.jit
+def w8a8_tensor_dynamic_fp8_matmul_grouped_kernel(
+    A,  # (S, K) pre-quantized FP8 activations, sorted/grouped by expert id
+    B,  # (E, N, K) FP8 weight matrices
+    C,  # (S, N) output
+    As,  # (S,) per-token activation scales
+    Bs,  # (E, 1, 1) per-tensor 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_as_m,
+    stride_bs_e,
+    stride_offs,
+    stride_tile,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    NUM_EXPERTS: 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) * 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 = (
+        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
+    _store_tile(
+        C, accumulator, offs_global_m, offs_bn, row_mask, stride_cm, stride_cn
+    )
+@triton.autotune(
+    configs=get_accelerator_autotuning_configs(with_block_sizes=True),
+    key=["N", "K", "BLOCK_SIZE_M"],
+)
+@triton.jit
+def w4a8_mx_dynamic_fp4_matmul_grouped_kernel(
+    A,  # (S, K) raw BF16/FP16 activations, sorted by expert id
+    B,  # (E, N, K // 2) packed FP4 (E2M1) expert weights as int8
+    C,  # (S, N) output
+    Bs,  # (E, N, K // SCALE_GROUP_K) UE8M0 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,
+    stride_offs,
+    stride_tile,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    NUM_EXPERTS: tl.constexpr,
+    NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
+    NIBBLES_PER_BYTE: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+):
+    """Grouped MXFP4 (W4A8) expert matmul with fused activation quant.
+    Tokens are assumed sorted by expert. The kernel maps each M-tile to its
+    owning expert via ``TileOffsets``, quantizes ``A`` to FP8 per K-group inline
+    (UE8M0 scale), then ``tl.dot_scaled`` against packed FP4 weights.
+    """
+    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_k_byte = tl.arange(0, BLOCK_SIZE_K // NIBBLES_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_k_byte[:, 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 = mx_act_quant_inline(
+            a_raw, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K
+        )
+        b = tl.load(b_ptrs).to(tl.uint8)
+        b_s = tl.load(bs_ptrs).to(tl.uint8)
+        accumulator += tl.dot_scaled(a, a_scale, "e4m3", b, b_s, "e2m1")
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += (BLOCK_SIZE_K // NIBBLES_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.autotune(
+    configs=get_accelerator_autotuning_configs(with_block_sizes=True),
+    key=["N", "K", "BLOCK_SIZE_M"],
+)
+@triton.jit
+def w8a8_mx_dynamic_fp8_matmul_grouped_kernel(
+    A,  # (S, K) raw BF16/FP16 activations, sorted by expert id
+    B,  # (E, N, K) E4M3 expert weights (unpacked)
+    C,  # (S, N) output
+    Bs,  # (E, N, K // SCALE_GROUP_K) UE8M0 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,
+    stride_offs,
+    stride_tile,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    NUM_EXPERTS: tl.constexpr,
+    NUM_EXPERTS_BIT_LENGTH: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+):
+    """Grouped MXFP8 (W8A8) expert matmul with fused activation quant.
+    Tokens are assumed sorted by expert. The kernel maps each M-tile to its
+    owning expert via ``TileOffsets``, quantizes ``A`` to E4M3 per K-group inline
+    (UE8M0 scale), then ``tl.dot_scaled`` against E4M3 weights. Mirrors the
+    grouped MXFP4 kernel but with unpacked weights and ``"e4m3"`` on both operands.
+    """
+    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_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_k[:, 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 = mx_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)
+        accumulator += tl.dot_scaled(a, a_scale, "e4m3", b, b_s, "e4m3")
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * 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:  (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})"
+    )
+    Bs = ue8m0_as_uint8(Bs)
+    C = A.new_empty(S, N, dtype=output_dtype)
+    BLOCK_SIZE_M = adaptive_block_size_m((S + E - 1) // E)
+    tile_offsets, max_m_tiles = grouped_tile_layout(
+        tokens_per_expert, BLOCK_SIZE_M, S, E
+    )
+    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,
+            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),
+            offsets.stride(0),
+            tile_offsets.stride(0),
+            # 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(
+    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:  (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)"
+        )
+    qA, As = fp8_act_quant(A, K)
+    C = A.new_empty(S, N, dtype=output_dtype)
+    BLOCK_SIZE_M = adaptive_block_size_m((S + E - 1) // E)
+    tile_offsets, max_m_tiles = grouped_tile_layout(
+        tokens_per_expert, BLOCK_SIZE_M, S, E
+    )
+    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,
+            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),
+            offsets.stride(0),
+            tile_offsets.stride(0),
+            # Meta-parameters (BLOCK_SIZE_N, BLOCK_SIZE_K come from autotune Config)
+            NUM_EXPERTS=E,
+            BLOCK_SIZE_M=BLOCK_SIZE_M,
+            NUM_EXPERTS_BIT_LENGTH=E.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:  (E, N, K) FP8 expert weights
+    Bs: (E, 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:  (E, N, K) FP8 expert weights
+    Bs: (E,) or (E, 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("w4a8_mx_dynamic_fp4_matmul_grouped"), mutates_args=()
+)
+def _w4a8_mx_dynamic_fp4_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 MXFP4 (W4A8) matmul with fused activation quant.
+    A:  (S, K) raw activations, bf16/fp16/fp32, expert-sorted (quantized inline)
+    B:  (E, N, K // 2) packed FP4 (E2M1) expert weights, two codes per int8
+    Bs: (E, N, K // 32) UE8M0 weight scales
+    offsets: (E,) — exclusive prefix of expert token counts (cumsum)
+    tokens_per_expert: (E,) — per-expert row count
+    BLOCK_SIZE_N and BLOCK_SIZE_K are autotuned.
+    """
+    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 == torch.int8, f"B must be int8 (packed FP4), got {B.dtype}"
+    assert Bs.dtype == torch.float8_e8m0fnu, (
+        f"Bs must be float8_e8m0fnu, got {Bs.dtype}"
+    )
+    S, K = A.shape
+    E, N, K_half = B.shape
+    assert offsets.shape[0] == E and tokens_per_expert.shape[0] == E
+    assert K == NIBBLES_PER_BYTE * K_half, (
+        f"K (={K}) must equal {NIBBLES_PER_BYTE} * B.shape[2] (={K_half})"
+    )
+    assert K % MX_SCALE_GROUP_K == 0, (
+        f"K (={K}) must be a multiple of {MX_SCALE_GROUP_K}"
+    )
+    assert Bs.shape == (E, N, K // MX_SCALE_GROUP_K), (
+        f"Bs shape {tuple(Bs.shape)} != ({E}, {N}, {K // MX_SCALE_GROUP_K})"
+    )
+    bs_u8 = ue8m0_as_uint8(Bs)
+    C = A.new_empty((S, N), dtype=output_dtype)
+    BLOCK_SIZE_M = adaptive_block_size_m((S + E - 1) // E)
+    tile_offsets, max_m_tiles = grouped_tile_layout(
+        tokens_per_expert, BLOCK_SIZE_M, S, E
+    )
+    def grid(META):
+        return (max_m_tiles, triton.cdiv(N, META["BLOCK_SIZE_N"]))
+    with device_context(A.device):
+        wrap_triton(w4a8_mx_dynamic_fp4_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=E,
+            BLOCK_SIZE_M=BLOCK_SIZE_M,
+            NUM_EXPERTS_BIT_LENGTH=E.bit_length(),
+            NIBBLES_PER_BYTE=NIBBLES_PER_BYTE,
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+        )
+    return C
+def w4a8_mx_dynamic_fp4_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 MXFP4 (W4A8) matmul with fused activation quant. Per-expert
+    ``C[s] = A[s] @ B[e].T`` over contiguous, expert-sorted rows. Tile shape
+    autotuned; FP4 scale granularity is fixed at 32."""
+    return ops.w4a8_mx_dynamic_fp4_matmul_grouped(
+        A, B, Bs, offsets, tokens_per_expert, output_dtype
+    )
+@triton_op(
+    add_op_namespace_prefix("w8a8_mx_dynamic_fp8_matmul_grouped"), mutates_args=()
+)
+def _w8a8_mx_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:
+    """Grouped MXFP8 (W8A8) matmul: ``C[s] = A[s] @ B[e].T`` (E4M3 × E4M3, UE8M0
+    group-32 scales) with fused activation quant.
+    A:  (S, K) raw activations, bf16/fp16/fp32, expert-sorted (quantized inline)
+    B:  (E, N, K) E4M3 expert weights (unpacked)
+    Bs: (E, N, K // 32) UE8M0 weight scales
+    offsets: (E,) — exclusive prefix of expert token counts (cumsum)
+    tokens_per_expert: (E,) — per-expert row count
+    BLOCK_SIZE_N and BLOCK_SIZE_K are autotuned; MX scale granularity is fixed at 32.
+    """
+    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 == torch.float8_e4m3fn, (
+        f"B must be float8_e4m3fn (E4M3 weights), got {B.dtype}"
+    )
+    assert Bs.dtype == torch.float8_e8m0fnu, (
+        f"Bs must be float8_e8m0fnu, got {Bs.dtype}"
+    )
+    S, K = A.shape
+    E, N, K_b = B.shape
+    assert offsets.shape[0] == E and tokens_per_expert.shape[0] == E
+    assert K == K_b, f"K mismatch: A has K={K}, B has K={K_b}"
+    assert K % MX_SCALE_GROUP_K == 0, (
+        f"K (={K}) must be a multiple of {MX_SCALE_GROUP_K}"
+    )
+    assert Bs.shape == (E, N, K // MX_SCALE_GROUP_K), (
+        f"Bs shape {tuple(Bs.shape)} != ({E}, {N}, {K // MX_SCALE_GROUP_K})"
+    )
+    bs_u8 = ue8m0_as_uint8(Bs)
+    C = A.new_empty((S, N), dtype=output_dtype)
+    BLOCK_SIZE_M = adaptive_block_size_m((S + E - 1) // E)
+    tile_offsets, max_m_tiles = grouped_tile_layout(
+        tokens_per_expert, BLOCK_SIZE_M, S, E
+    )
+    def grid(META):
+        return (max_m_tiles, triton.cdiv(N, META["BLOCK_SIZE_N"]))
+    with device_context(A.device):
+        wrap_triton(w8a8_mx_dynamic_fp8_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=E,
+            BLOCK_SIZE_M=BLOCK_SIZE_M,
+            NUM_EXPERTS_BIT_LENGTH=E.bit_length(),
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+        )
+    return C
+def w8a8_mx_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:
+    """Grouped MXFP8 (W8A8) matmul with fused activation quant. Per-expert
+    ``C[s] = A[s] @ B[e].T`` over contiguous, expert-sorted rows via the MX
+    scaled MMA (E4M3 × E4M3, UE8M0 group-32). Tile shape autotuned."""
+    return ops.w8a8_mx_dynamic_fp8_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``:
+    - ``B.dtype == int8`` (packed FP4) → ``w4a8_mx_dynamic_fp4_matmul_grouped``
+      (``block_size`` is ignored; FP4 tile shape is autotuned).
+    - ``B.dtype == float8_e4m3fn`` with UE8M0 group-32 ``Bs`` (shape ``[E, N, K//32]``)
+      → ``w8a8_mx_dynamic_fp8_matmul_grouped`` (MX scaled MMA; ``block_size`` ignored).
+    - ``block_size`` None or full ``[N, K]`` → ``w8a8_tensor_dynamic_fp8_matmul_grouped``.
+    - otherwise → ``w8a8_block_dynamic_fp8_matmul_grouped``.
+    """
+    if is_mxfp4(B, Bs):
+        return w4a8_mx_dynamic_fp4_matmul_grouped(
+            A, B, Bs, offsets, tokens_per_expert, output_dtype
+        )
+    if is_mxfp8(B, Bs):
+        return w8a8_mx_dynamic_fp8_matmul_grouped(
+            A, B, Bs, offsets, tokens_per_expert, output_dtype
+        )
+    if block_size is None or (
+        block_size[0] == B.size(1) and block_size[1] == B.size(2)
+    ):
+        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-rocm/matmul.py ADDED Viewed

	@@ -0,0 +1,941 @@

+# 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 ._ops import add_op_namespace_prefix, ops
+from .utils import (
+    NIBBLES_PER_BYTE,
+    MX_SCALE_GROUP_K,
+    adaptive_block_size_m,
+    decode_ue8m0_scale,
+    device_context,
+    fp8_act_quant,
+    fp8_act_quant_inline,
+    get_accelerator_autotuning_configs,
+    is_mxfp4,
+    is_mxfp8,
+    mx_act_quant_inline,
+    ue8m0_as_uint8,
+)
+@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_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    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(with_block_sizes=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_bn,
+    stride_cm,
+    stride_cn,
+    stride_as_m,
+    # Meta-parameters
+    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, 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
+    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
+    _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,
+    )
+@triton.autotune(
+    configs=get_accelerator_autotuning_configs(with_block_sizes=True),
+    key=["N", "K", "BLOCK_SIZE_M"],
+)
+@triton.jit
+def w4a8_mx_dynamic_fp4_matmul_kernel(
+    A,  # (M, K) raw BF16/FP16 activations
+    B,  # (N, K // 2) packed FP4 (E2M1) weights as int8
+    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,
+    NIBBLES_PER_BYTE: tl.constexpr,
+    SCALE_GROUP_K: tl.constexpr,
+):
+    """MXFP4 (W4A8) GEMM kernel with fused activation quantization.
+    Computes ``C = A @ B.T`` with bf16/fp16 ``A`` quantized to FP8 (E4M3) per
+    K-group of ``SCALE_GROUP_K`` elements inline (UE8M0 scale), and packed FP4
+    (E2M1) weights with their own UE8M0 scales. 2D grid with swizzle for L2
+    cache locality on B tiles, ``tl.dot_scaled`` for the scaled MMA.
+    """
+    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_k_byte = tl.arange(0, BLOCK_SIZE_K // NIBBLES_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_k_byte[:, 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 = mx_act_quant_inline(
+            a_raw, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K
+        )
+        b = tl.load(
+            b_ptrs, mask=offs_k_byte[:, None] < k_remaining // NIBBLES_PER_BYTE, other=0
+        ).to(tl.uint8)
+        b_s = tl.load(
+            bs_ptrs, mask=offs_sf[None, :] < k_remaining // SCALE_GROUP_K, other=0
+        ).to(tl.uint8)
+        accumulator += tl.dot_scaled(a, a_scale, "e4m3", b, b_s, "e2m1")
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += (BLOCK_SIZE_K // NIBBLES_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.autotune(
+    configs=get_accelerator_autotuning_configs(with_block_sizes=True),
+    key=["N", "K", "BLOCK_SIZE_M"],
+)
+@triton.jit
+def w8a8_mx_dynamic_fp8_matmul_kernel(
+    A,  # (M, K) raw BF16/FP16 activations
+    B,  # (N, K) E4M3 weights (unpacked, one byte per value)
+    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,
+    SCALE_GROUP_K: tl.constexpr,
+):
+    """MXFP8 block-scale GEMM kernel with fused activation quantization.
+    Computes ``C = A @ B.T`` with bf16/fp16 ``A`` quantized to FP8 (E4M3) per
+    K-group of ``SCALE_GROUP_K`` elements inline (UE8M0 scale), against E4M3
+    weights with their own UE8M0 K-group scales. Mirrors the W4A8 FP4 kernel but
+    keeps weights unpacked (one E4M3 byte per value); ``tl.dot_scaled`` drives
+    the MX scaled MMA with ``"e4m3"`` on both operands.
+    """
+    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_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_k[:, 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 = mx_act_quant_inline(
+            a_raw, BLOCK_SIZE_M, BLOCK_SIZE_K, SCALE_GROUP_K
+        )
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
+        b_s = tl.load(
+            bs_ptrs, mask=offs_sf[None, :] < k_remaining // SCALE_GROUP_K, other=0
+        ).to(tl.uint8)
+        accumulator += tl.dot_scaled(a, a_scale, "e4m3", b, b_s, "e4m3")
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * 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}"
+    )
+    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 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,
+            K,
+            A.stride(-2),
+            A.stride(-1),
+            B.stride(1),
+            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=8,
+        )
+    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,
+            As,
+            Bs,
+            M,
+            N,
+            K,
+            A.stride(-2),
+            A.stride(-1),
+            B.stride(1),
+            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=8,
+        )
+    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=8,
+        )
+    return C
+@triton_op(add_op_namespace_prefix("w4a8_mx_dynamic_fp4_matmul"), mutates_args=())
+def _w4a8_mx_dynamic_fp4_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """MXFP4 (W4A8) matmul: ``C = A @ B.T`` with fused activation quant.
+    A:  (M, K) raw activations, bf16/fp16/fp32 (quantized inline to FP8)
+    B:  (N, K // 2) packed FP4 (E2M1) weights, two codes per int8
+    Bs: (N, K // 32) UE8M0 weight scales
+    On CUDA and other backends, BLOCK_SIZE_N and BLOCK_SIZE_K are autotuned.
+    On XPU they are fixed to 128x128 and only num_warps/num_stages are autotuned.
+    FP4 scale granularity is fixed at 32, so tile shape is purely a perf knob.
+    """
+    assert A.ndim == 2 and B.ndim == 2 and Bs.ndim == 2
+    assert B.dtype == torch.int8, f"B must be int8 (packed FP4), 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"
+    M, K = A.shape
+    N, K_half = B.shape
+    assert K == NIBBLES_PER_BYTE * K_half, (
+        f"K (={K}) must equal {NIBBLES_PER_BYTE} * B.shape[1] (={K_half})"
+    )
+    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})"
+    )
+    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(w4a8_mx_dynamic_fp4_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=8,
+            NIBBLES_PER_BYTE=NIBBLES_PER_BYTE,
+            SCALE_GROUP_K=MX_SCALE_GROUP_K,
+        )
+    return C
+@triton_op(add_op_namespace_prefix("w8a8_mx_dynamic_fp8_matmul"), mutates_args=())
+def _w8a8_mx_dynamic_fp8_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """MXFP8 W8A8 matmul: ``C = A @ B.T`` (E4M3 × E4M3, UE8M0 group-32 scales).
+    A:  (M, K) raw activations, bf16/fp16/fp32 (quantized inline to E4M3,
+        MX group-32 UE8M0 scales)
+    B:  (N, K) E4M3 weights (unpacked)
+    Bs: (N, K // 32) UE8M0 weight scales
+    Tile shape (BLOCK_SIZE_N, BLOCK_SIZE_K) is autotuned; MX scale granularity is
+    fixed at 32 (the MX-format spec), so tile shape is purely a perf knob.
+    """
+    assert A.ndim == 2 and B.ndim == 2 and Bs.ndim == 2
+    assert B.dtype == torch.float8_e4m3fn, (
+        f"B must be float8_e4m3fn (E4M3 weights), 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"
+    M, K = A.shape
+    N, K_b = B.shape
+    assert K == K_b, f"K mismatch: A has K={K}, B has K={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})"
+    )
+    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(w8a8_mx_dynamic_fp8_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=8,
+            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 w4a8_mx_dynamic_fp4_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """MXFP4 (W4A8) matmul with fused activation quantization.
+    Computes ``C = A @ B.T`` with bf16/fp16/fp32 ``A`` quantized to FP8 (E4M3)
+    inline and packed FP4 (E2M1) weights; both scales are UE8M0 at K-group
+    granularity 32 (the MX-format spec). Tile shape (BLOCK_SIZE_N, BLOCK_SIZE_K)
+    is autotuned — FP4 has no caller-controlled quantization block.
+    Args:
+        A: Raw activations ``[M, K]`` in bf16/fp16/fp32.
+        B: Packed FP4 weights ``[N, K // 2]`` (``int8``, two codes per byte).
+        Bs: UE8M0 weight scales ``[N, K // 32]``.
+        output_dtype: dtype of the returned tensor (default ``bfloat16``).
+    """
+    return ops.w4a8_mx_dynamic_fp4_matmul(A, B, Bs, output_dtype)
+def w8a8_mx_dynamic_fp8_matmul(
+    A: torch.Tensor,
+    B: torch.Tensor,
+    Bs: torch.Tensor,
+    output_dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """MXFP8 W8A8 matmul (E4M3 × E4M3, UE8M0 group-32 scales) with fused act quant.
+    Computes ``C = A @ B.T`` with bf16/fp16/fp32 ``A`` quantized inline to E4M3
+    against per-row, per-32-K-group UE8M0 scales, and E4M3 weights with their own
+    UE8M0 K-group scales. Both scales feed ``tl.dot_scaled`` (the MX scaled MMA),
+    unlike ``w8a8_block_dynamic_fp8_matmul`` which applies 128×128 block scales in
+    software. Tile shape is autotuned; MX scale granularity is fixed at 32.
+    Args:
+        A: Raw activations ``[M, K]`` in bf16/fp16/fp32.
+        B: E4M3 weights ``[N, K]`` (``float8_e4m3fn``, unpacked).
+        Bs: UE8M0 weight scales ``[N, K // 32]`` (``float8_e8m0fnu``).
+        output_dtype: dtype of the returned tensor (default ``bfloat16``).
+    """
+    return ops.w8a8_mx_dynamic_fp8_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``:
+    - ``B.dtype == int8`` (packed FP4) → ``w4a8_mx_dynamic_fp4_matmul``
+      (``block_size`` is ignored; FP4 scale granularity is fixed at 32 and
+      tile sizes are autotuned).
+    - ``B.dtype == float8_e4m3fn`` with UE8M0 group-32 ``Bs`` (shape ``[N, K//32]``)
+      → ``w8a8_mx_dynamic_fp8_matmul`` (MX scaled MMA; ``block_size`` ignored).
+    - ``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 activation_scale is not None:
+        if B.dtype == torch.int8:
+            raise NotImplementedError(
+                "static activation_scale is not supported on the FP4 path"
+            )
+        if block_size is None or (
+            block_size[0] == B.size(0) and block_size[1] == B.size(1)
+        ):
+            raise NotImplementedError(
+                "static activation_scale requires block-wise weights, "
+                "not tensor-mode (block_size None or full [N, K])"
+            )
+        return w8a8_block_static_fp8_matmul(
+            A, B, Bs, activation_scale, block_size, output_dtype
+        )
+    if is_mxfp4(B, Bs):
+        return w4a8_mx_dynamic_fp4_matmul(A, B, Bs, output_dtype)
+    if is_mxfp8(B, Bs):
+        return w8a8_mx_dynamic_fp8_matmul(A, B, Bs, output_dtype)
+    if block_size is None or (
+        block_size[0] == B.size(0) and block_size[1] == B.size(1)
+    ):
+        return w8a8_tensor_dynamic_fp8_matmul(A, B, Bs, output_dtype)
+    return w8a8_block_dynamic_fp8_matmul(A, B, Bs, block_size, output_dtype)

build/torch-rocm/metadata.json ADDED Viewed

	@@ -0,0 +1,22 @@

+{
+  "name": "finegrained-fp8",
+  "id": "_finegrained_fp8_rocm_846165b",
+  "version": 3,
+  "license": "Apache-2.0",
+  "python-depends": [],
+  "backend": {
+    "type": "rocm"
+  },
+  "digest": {
+    "algorithm": "sha256",
+    "files": {
+      "__init__.py": "jOglX0eto6QNL9D8T9DzKIb3XYUF5vZAjvEL32MCAHY=",
+      "_ops.py": "2e6XSotI60yh7+FU07X7EGCLbnBQRgMN9Kh06vIw9CY=",
+      "batched.py": "/ZuVQYNILMt6KBM321fgXxhFRKQM9tMWvvfAEKZY42U=",
+      "finegrained_fp8/__init__.py": "DFYPlrhXwYjEqCl/8n0SmWGZV8NFml5DPhMjKfv98GY=",
+      "grouped.py": "EiotInBQUPtX35Ms+eEZ6b7YyWcxhWTWGq3EiMN0fMs=",
+      "matmul.py": "AfyGST+Lrr8AjWqcIx4txrekpKH4iwyS8177TiyeW/M=",
+      "utils.py": "pJ6NJrMMz/mSHXZakTmcXW/ZgQ91YnphRk8PVa8yj+U="
+    }
+  }
+}

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

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtKgAwIBAgIUePTBhgrVPrG0jyrZOpoWjJWMrN4wCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjE3MTMxNDA0WhcNMjYwNjE3MTMyNDA0WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAE0VUnz693BcejXkNTO62Z4hPYtmgg22tjv3G+w7jq79WyMCmcSVaofu2ZJlOMsA9v36EETYobwETef0B9lLMqAKOCBfEwggXtMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUWg0Gm4lpyM98cT8KBL1cWUOoUKAwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoODQ2MTY1YjExZDU3MTA1ODI5MmQ1NDRiNDIyMzQ2MzE3ZGUxMDJkZjATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoODQ2MTY1YjExZDU3MTA1ODI5MmQ1NDRiNDIyMzQ2MzE3ZGUxMDJkZjAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoODQ2MTY1YjExZDU3MTA1ODI5MmQ1NDRiNDIyMzQ2MzE3ZGUxMDJkZjAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDg0NjE2NWIxMWQ1NzEwNTgyOTJkNTQ0YjQyMjM0NjMxN2RlMTAyZGYwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjc2OTEyMTUxMzkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBiwYKKwYBBAHWeQIEAgR9BHsAeQB3AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABntW3q3sAAAQDAEgwRgIhAJbEJ3yIOVCYyNmuAJfsYvWyfk+1pDFQY8XlT3ChiK4jAiEA3eQjC77zGyucbZa6WiNYOz2RO+RYirxBVNgK4f5YQ38wCgYIKoZIzj0EAwMDZwAwZAIwG8x9gZM5WmgUIeDX2l+fLc8BMs/EZFaKD/MTa15tB3hpVYaqTyom3dQSIglzETtxAjAbp0rCeLzhA5nmE+ctKvOwCvxYpD06P51lfdqcUHrgSoK0hNCk26zd8j2S5sAkITM="},"tlogEntries":[{"logIndex":"1851564527","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1781702045","inclusionPromise":{"signedEntryTimestamp":"MEUCIQCXnosi1Fkyem6EcQ7sHvaUqAmtQuF+6Rqe5C3zf6nAzgIgWFY8FhIFCZtR/KGwJjUXFmLG/q6h2SxlBmN4MD2SHUE="},"inclusionProof":{"logIndex":"1729660265","rootHash":"bA+ZSMR5ebA+bxzst9Muh7lqhA3+EZn+j04BxfvmN9w=","treeSize":"1729660347","hashes":["G0kerrCVnmrmJIucj6z35eeb9nQSM/bHQgcR4XUI+h0=","lF6pG0uLVoW8P9IFRM9FuVBL878Tb4de5y8Nos8MY1U=","0us52j044UZxn5OwNPtbzHCxxWeaCW5w/tP1ey0bbdg=","zrWsV+4wyv1S7Qp+yH/FBzQ1C+W8SYogESc7f1QzoH0=","1E3pfOGXxKWEmoIzMvVEkmCUtzhElCK1v5f/rXUgKto=","inihgnOSx1WccBIuIvfshXFGfocE8mH19j08bmnPgMQ=","Xiy1HnHrEBvLAb8qevuD1odTsmP2TEbHNBPhdHPs738=","G75SHivXVw4HNLVI6MwIJ8rOpL3N5b3UTo+0mwcasrU=","/MUwk7ZfPoiJVazjBloofRJ6hS2B8Az3apT2GBKtJ3g=","DT5JAn4B3Dm8gmFcDKslYJrA/3Q9NNFCX4dpHJUGPEU=","cRNcuWZEC/h+ACT/IKS0cgSD29B0pvP9++uNHzrSlj8=","4up1iqXVaWkl4o3JcV/XPRm3hVSZfjTk1Iaqp+cfi/o=","M9+xE/67JCQM3UvBz/mjGSEhOSJ8QCD3Xi7mH8EKFEk=","mqM7J+i75IpuD09QejiUBjeH85AZa+fSm4RXXFmj3lI=","72FC5FYLhxB4a4iiC956o0B/fT54ip4R41vsw2QBKtU=","lYGQ9ibwC8+smMkPQ6TchJm3H9Nc/aTYLfdRacGFChw=","daxmZaajRpZV+JxHiOYZhJBiSKN5ucqjh2WnGbHhirw=","DOCeoSMovIvLExkhIvisow9AuNXgeWs4ECkyR6EcqYU="],"checkpoint":{"envelope":"rekor.sigstore.dev - 1193050959916656506\n1729660347\nbA+ZSMR5ebA+bxzst9Muh7lqhA3+EZn+j04BxfvmN9w=\n\n— rekor.sigstore.dev wNI9ajBFAiEAuxuJjQEyMmOUsx8f2UggYVM4d36K6WVhM8FqnsaH+0wCIE0MIegCiydXCBDLCY7uPBMp08uzZiISGY+h/oc+gsp/\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiJiMWZhMGY0YWIyY2ZhNzJkNzc4YzUzOTMxNTNlYzQ1ZGUwZDMzNjFhMDFiNTQ1NjNlNjQyZTkxNWFlYjU4NjZmIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJRy9lenNqWnRlTWhBY1U5TmhMWURNUDFrcllDd0t6bkF3LzVaaUZnbm5VbkFpRUEyMnJhcXZzMm92OGtIcnJJa2JteDhCSnVPU0g4dmh0WmlQMnRBTUdSa0NzPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5STFowRjNTVUpCWjBsVlpWQlVRbWhuY2xaUWNrY3dhbmx5V2s5d2IxZHFTbGROY2s0MGQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxUlROTlZFMTRUa1JCTUZkb1kwNU5hbGwzVG1wRk0wMVVUWGxPUkVFd1YycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVV3VmxWdWVqWTVNMEpqWldwWWEwNVVUell5V2pSb1VGbDBiV2RuTWpKMGFuWXpSeXNLZHpkcWNUYzVWM2xOUTIxalUxWmhiMloxTWxwS2JFOU5jMEU1ZGpNMlJVVlVXVzlpZDBWVVpXWXdRamxzVEUxeFFVdFBRMEptUlhkbloxaDBUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZYWnpCSENtMDBiSEI1VFRrNFkxUTRTMEpNTVdOWFZVOXZWVXRCZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVBSRkV5VFZSWk1WbHFSWGhhUkZVelRWUkJNVTlFU1RWTmJWRXhDazVFVW1sT1JFbDVUWHBSTWsxNlJUTmFSMVY0VFVSS2ExcHFRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDlFVVRKTlZGa3hXV3BGZUZwRVZUTk5WRUV4VDBSSk5VMXRVVEZPUkZKcFRrUkplVTE2VVRJS1RYcEZNMXBIVlhoTlJFcHJXbXBCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDlFVVRKTlZGa3hXV3BGZUZwRVZUTk5WRUV4VDBSSk5VMXRVVEVLVGtSU2FVNUVTWGxOZWxFeVRYcEZNMXBIVlhoTlJFcHJXbXBCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUm5NRTVxUlRJS1RsZEplRTFYVVRGT2VrVjNUbFJuZVU5VVNtdE9WRkV3V1dwUmVVMXFUVEJPYWsxNFRqSlNiRTFVUVhsYVIxbDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1NeVQxUkZlVTFVVlhoTmVtdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwZDFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT1VKSWMwRUtaVkZDTTBGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTBWek54TTNOQlFVRlJSQXBCUldkM1VtZEphRUZLWWtWS00zbEpUMVpEV1hsT2JYVkJTbVp6V1haWGVXWnJLekZ3UkVaUldUaFliRlF6UTJocFN6UnFRV2xGUVRObFVXcEROemQ2Q2tkNWRXTmlXbUUyVjJsT1dVOTZNbEpQSzFKWmFYSjRRbFpPWjBzMFpqVlpVVE00ZDBObldVbExiMXBKZW1vd1JVRjNUVVJhZDBGM1drRkpkMGM0ZURrS1oxcE5OVmR0WjFWSlpVUllNbXdyWmt4ak9FSk5jeTlGV2taaFMwUXZUVlJoTVRWMFFqTm9jRlpaWVhGVWVXOXRNMlJSVTBsbmJIcEZWSFI0UVdwQllncHdNSEpEWlV4NmFFRTFibTFGSzJOMFMzWlBkME4yZUZsd1JEQTJVRFV4Ykdaa2NXTlZTSEpuVTI5TE1HaE9RMnN5Tm5wa09Hb3lVelZ6UVd0SlZFMDlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyDADAgEAMIICvwYJKoZIhvcNAQcCoIICsDCCAqwCAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQgirsf3iN2RuZt8kLZKD5VUGqJ0tF1saNu4OqNKmEEAQkCFEvPI1QROe8iGgplF9//L1TeLChVGA8yMDI2MDYxNzEzMTQwNFowAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdowggHWAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNjE3MTMxNDA0WjAvBgkqhkiG9w0BCQQxIgQg4zyqLLBaOA3qBIJ3rbimqn9CbsEwnKlOhCtR7uq4B7cwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGYwZAIwT9UpLg/DWWo/eugeIjzk42GOVBb1Gb6MAuIheAGnW2fNlhwC+NrHfPPdktaiStmtAjBmoKRxlPXAqSOmvHg1DYzFdMGSGdm2JJ1wfV3/oYNMuymIIXPWookWoIWZQHgKKzE="}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"sfoPSrLPpy13jFOTFT7EXeDTNhoBtUVj5kLpFa61hm8="},"signature":"MEUCIG/ezsjZteMhAcU9NhLYDMP1krYCwKznAw/5ZiFgnnUnAiEA22raqvs2ov8kHrrIkbmx8BJuOSH8vhtZiP2tAMGRkCs="}}

build/torch-rocm/utils.py ADDED Viewed

	@@ -0,0 +1,269 @@

+# 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):
+    """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
+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 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 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(*, with_block_sizes: 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, with_block_sizes)``, then crossed
+    into the config list.
+    ``with_block_sizes=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. ``with_block_sizes=False`` emits a single empty meta-dict (block-scaled
+    kernels take the tile from the caller's ``block_size``).
+    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"
+    if with_block_sizes:
+        num_stages = [2, 3]
+        num_warps = [8, 16] if is_xpu else [4, 8, 16]
+        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:
+        num_stages = [2, 3, 4]
+        num_warps = [8, 16] if is_xpu else [2, 4, 8, 16]
+        blocks = [{}]
+    return [
+        triton.Config(b, num_warps=w, num_stages=s)
+        for b in blocks
+        for w in num_warps
+        for s in num_stages
+    ]
+# ── 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 mx_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
+# ── 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)