Kernels:

kernels-community
/

finegrained-fp8

Trusted publisher

Kernel card Files Files and versions

xet

Community

kernels-bot commited on Jun 17

Commit

099314d

verified ·

1 Parent(s): 0c9fbb9

Uploaded using `kernel-builder`.

Browse files

Files changed (27) hide show

build/torch-cuda/__init__.py +8 -26
build/torch-cuda/_ops.py +1 -1
build/torch-cuda/act_quant.py +0 -73
build/torch-cuda/batched.py +492 -117
build/torch-cuda/grouped.py +638 -164
build/torch-cuda/matmul.py +692 -162
build/torch-cuda/metadata.json +14 -2
build/torch-cuda/metadata.json.sigstore +1 -0
build/torch-cuda/utils.py +257 -1
build/torch-rocm/__init__.py +0 -32
build/torch-rocm/_ops.py +0 -38
build/torch-rocm/act_quant.py +0 -73
build/torch-rocm/batched.py +0 -390
build/torch-rocm/finegrained_fp8/__init__.py +0 -26
build/torch-rocm/grouped.py +0 -477
build/torch-rocm/matmul.py +0 -411
build/torch-rocm/metadata.json +0 -10
build/torch-rocm/utils.py +0 -13
build/torch-xpu/__init__.py +0 -32
build/torch-xpu/_ops.py +0 -38
build/torch-xpu/act_quant.py +0 -73
build/torch-xpu/batched.py +0 -398
build/torch-xpu/finegrained_fp8/__init__.py +0 -26
build/torch-xpu/grouped.py +0 -477
build/torch-xpu/matmul.py +0 -411
build/torch-xpu/metadata.json +0 -10
build/torch-xpu/utils.py +0 -13

build/torch-cuda/__init__.py CHANGED Viewed

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

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

build/torch-cuda/_ops.py CHANGED Viewed

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

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

build/torch-cuda/act_quant.py DELETED Viewed

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

build/torch-cuda/batched.py CHANGED Viewed

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

 import torch
 import triton
 import triton.language as tl
 from torch.library import triton_op, wrap_triton
+from ._ops import add_op_namespace_prefix, ops
+from .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_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)
     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_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)
     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.
         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.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.
             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,
             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-cuda/grouped.py CHANGED Viewed

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

 import torch
 import triton
 import triton.language as tl
 from torch.library import triton_op, wrap_triton
+from ._ops import add_op_namespace_prefix, ops
+from .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
     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.
     pid_m = tl.program_id(axis=0)
     pid_n = tl.program_id(axis=1)
+    total_tiles = tl.load(TileOffsets + (NUM_EXPERTS - 1) * stride_tile)
     if pid_m >= total_tiles:
         return
+    expert_id, offs_global_m, row_mask, offs_bn, offs_k = _grouped_tile_setup(
+        pid_m,
+        pid_n,
+        Offsets,
+        TileOffsets,
+        stride_offs,
+        stride_tile,
+        NUM_EXPERTS,
+        NUM_EXPERTS_BIT_LENGTH,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+        BLOCK_SIZE_K,
+    )
     a_ptrs = A + offs_global_m[:, None] * stride_am + offs_k[None, :] * stride_ak
     b_ptrs = (
     accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
     for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a_raw = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0).to(tl.float32)
+        a, a_s = fp8_act_quant_inline(a_raw)
         b = tl.load(b_ptrs)
+        b_s = decode_ue8m0_scale(tl.load(bs_ptrs))
         accumulator += tl.dot(a, b) * a_s[:, None] * b_s[None, :]
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
+        bs_ptrs += stride_bs_k
+    _store_tile(
+        C, accumulator, offs_global_m, offs_bn, row_mask, stride_cm, stride_cn
+    )
 @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_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.
     pid_m = tl.program_id(axis=0)
     pid_n = tl.program_id(axis=1)
+    total_tiles = tl.load(TileOffsets + (NUM_EXPERTS - 1) * stride_tile)
     if pid_m >= total_tiles:
         return
+    expert_id, offs_global_m, row_mask, offs_bn, offs_k = _grouped_tile_setup(
+        pid_m,
+        pid_n,
+        Offsets,
+        TileOffsets,
+        stride_offs,
+        stride_tile,
+        NUM_EXPERTS,
+        NUM_EXPERTS_BIT_LENGTH,
+        BLOCK_SIZE_M,
+        BLOCK_SIZE_N,
+        BLOCK_SIZE_K,
+    )
     a_ptrs = A + offs_global_m[:, None] * stride_am + offs_k[None, :] * stride_ak
     b_ptrs = (
     for _ in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
         a = tl.load(a_ptrs, mask=row_mask[:, None], other=0.0)
         b = tl.load(b_ptrs)
         accumulator += tl.dot(a, b)
         a_ptrs += BLOCK_SIZE_K * stride_ak
         b_ptrs += BLOCK_SIZE_K * stride_bk
     accumulator = accumulator * a_s[:, None] * b_s
+    _store_tile(
+        C, accumulator, offs_global_m, offs_bn, row_mask, stride_cm, stride_cn
+    )
+@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.
         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.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,
     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.
             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,
             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-cuda/matmul.py CHANGED Viewed

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

 import triton.language as tl
 from torch.library import triton_op, wrap_triton
+from ._ops import add_op_namespace_prefix, ops
+from .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_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_cm,
     stride_cn,
     stride_as_m,
+    # Meta-parameters
     BLOCK_SIZE_M: tl.constexpr,
     BLOCK_SIZE_N: tl.constexpr,
     BLOCK_SIZE_K: tl.constexpr,
     weight scale for the full matrix.
     Uses a 2D grid with swizzle for L2 cache locality on B tiles.
     """
+    pid_m, pid_n, offs_am, offs_bn, offs_k = _swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M
+    )
     a_ptrs = A + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
     b_ptrs = B + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn
     accumulator = accumulator * a_s[:, None] * b_s
+    _store_masked(
+        C, accumulator, pid_m, pid_n, M, N, stride_cm, stride_cn,
+        BLOCK_SIZE_M, BLOCK_SIZE_N,
+    )
+@triton.autotune(
+    configs=get_accelerator_autotuning_configs(),
+    key=["N", "K", "BLOCK_SIZE_M"],
+)
+@triton.jit
+def w8a8_block_static_fp8_matmul_kernel(
+    A,  # (M, K) raw BF16/FP16 activations
+    B,  # (N, K) FP8 weights
+    C,  # (M, N) output
+    As,  # scalar — static per-tensor activation scale (calibration-time)
+    Bs,  # (N // BLOCK_SIZE_N, K // BLOCK_SIZE_K) weight scales (fp32 or uint8/UE8M0)
+    # Shape
+    M,
+    N,
+    K,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_bs_k,
+    stride_bs_n,
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+    GROUP_SIZE_M: tl.constexpr,
+):
+    """Block-scale FP8 GEMM with static (per-tensor) activation scale.
+    ``A`` is raw bf16/fp16; the kernel divides by the scalar ``As`` and casts
+    to FP8 inline. Per-block weight scales apply per-K-tile during
+    accumulation; the scalar activation scale factors out of the loop and
+    is applied once at the end.
+    """
+    pid_m, pid_n, offs_am, offs_bn, offs_k = _swizzle_offsets(
+        M, N, BLOCK_SIZE_M, BLOCK_SIZE_N, BLOCK_SIZE_K, GROUP_SIZE_M
+    )
+    a_ptrs = A + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak)
+    b_ptrs = B + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn)
+    offs_bsn = offs_bn // BLOCK_SIZE_N
+    bs_ptrs = Bs + offs_bsn * stride_bs_n
+    a_s_static = tl.load(As)
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        k_remaining = K - k * BLOCK_SIZE_K
+        a_raw = tl.load(a_ptrs, mask=offs_k[None, :] < k_remaining, other=0.0).to(
+            tl.float32
+        )
+        a = (a_raw / a_s_static).to(tl.float8e4nv)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < k_remaining, other=0.0)
+        b_s = decode_ue8m0_scale(tl.load(bs_ptrs))
+        accumulator += tl.dot(a, b) * b_s[None, :]
+        a_ptrs += BLOCK_SIZE_K * stride_ak
+        b_ptrs += BLOCK_SIZE_K * stride_bk
+        bs_ptrs += stride_bs_k
+    accumulator = accumulator * a_s_static
+    _store_masked(
+        C, accumulator, pid_m, pid_n, M, N, stride_cm, stride_cn,
+        BLOCK_SIZE_M, BLOCK_SIZE_N,
+    )
+@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}"
     N, K = B.shape
     M = A.numel() // A.shape[-1]
     assert Bs.ndim == 2, f"Bs must be 2D (N//block_n, K//block_k), got ndim={Bs.ndim}"
     assert Bs.shape == (triton.cdiv(N, block_n), triton.cdiv(K, block_k)), (
         f"Bs shape {tuple(Bs.shape)} != expected ({triton.cdiv(N, block_n)}, {triton.cdiv(K, block_k)})"
     BLOCK_SIZE_K = block_k
     BLOCK_SIZE_N = block_n
+    Bs = ue8m0_as_uint8(Bs)
     C_shape = A.shape[:-1] + (N,)
+    BLOCK_SIZE_M = adaptive_block_size_m(M)
     C = A.new_empty(C_shape, dtype=output_dtype)
     grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
     with device_context(A.device):
+        wrap_triton(w8a8_block_dynamic_fp8_matmul_kernel)[grid](
             A,
             B,
             C,
             Bs,
             M,
             N,
             B.stride(0),
             C.stride(-2),
             C.stride(-1),
             Bs.stride(1),
             Bs.stride(0),
             # Meta-parameters
     return C
+@triton_op(add_op_namespace_prefix("w8a8_block_static_fp8_matmul"), mutates_args=())
+def _w8a8_block_static_fp8_matmul(
     A: torch.Tensor,
     B: torch.Tensor,
     Bs: torch.Tensor,
+    As: torch.Tensor,
+    block_size: list[int],
+    output_dtype: torch.dtype | None = None,
 ) -> torch.Tensor:
+    """Block-scale FP8 matmul with static (per-tensor) activation quantization.
+    A:  (..., K) raw bf16/fp16 activations — quantized to FP8 inline against ``As``
+    B:  (N, K) FP8 weights
+    Bs: (N // block_n, K // block_k) per-block weight scales
+    As: scalar / (1,) — per-tensor static activation scale
     """
+    assert len(block_size) == 2, (
+        f"block_size must be [block_n, block_k], got {block_size}"
+    )
+    block_n, block_k = block_size[0], block_size[1]
+    assert B.dtype != torch.int8, (
+        "static activation quant is not supported on the FP4 path"
+    )
+    assert not (block_n == B.size(0) and block_k == B.size(1)), (
+        "static activation quant requires block-wise weights, not tensor-mode"
+    )
     assert A.shape[-1] == B.shape[-1], (
         f"K mismatch: A has K={A.shape[-1]}, B has K={B.shape[-1]}"
     )
     assert A.is_contiguous(), "A must be contiguous"
     assert B.ndim == 2, f"B must be 2D (N, K), got ndim={B.ndim}"
     assert B.is_contiguous(), "B must be contiguous"
+    assert As.numel() == 1, f"As must be scalar or (1,), got {tuple(As.shape)}"
     N, K = B.shape
     M = A.numel() // A.shape[-1]
+    assert Bs.ndim == 2, f"Bs must be 2D (N//block_n, K//block_k), got ndim={Bs.ndim}"
+    assert Bs.shape == (triton.cdiv(N, block_n), triton.cdiv(K, block_k)), (
+        f"Bs shape {tuple(Bs.shape)} != expected ({triton.cdiv(N, block_n)}, {triton.cdiv(K, block_k)})"
     )
+    BLOCK_SIZE_K = block_k
+    BLOCK_SIZE_N = block_n
+    BLOCK_SIZE_M = adaptive_block_size_m(M)
+    grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N))
+    Bs = ue8m0_as_uint8(Bs)
     C_shape = A.shape[:-1] + (N,)
     C = A.new_empty(C_shape, dtype=output_dtype)
+    As = As.reshape(1).to(torch.float32)
     with device_context(A.device):
+        wrap_triton(w8a8_block_static_fp8_matmul_kernel)[grid](
             A,
             B,
             C,
             B.stride(0),
             C.stride(-2),
             C.stride(-1),
+            Bs.stride(1),
+            Bs.stride(0),
             # Meta-parameters
             BLOCK_SIZE_M=BLOCK_SIZE_M,
             BLOCK_SIZE_N=BLOCK_SIZE_N,
     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-cuda/metadata.json CHANGED Viewed

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

 {
   "name": "finegrained-fp8",
+  "id": "_finegrained_fp8_cuda_846165b",
+  "version": 3,
   "license": "Apache-2.0",
   "python-depends": [],
   "backend": {
     "type": "cuda"
+  },
+  "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-cuda/metadata.json.sigstore ADDED Viewed

	@@ -0,0 +1 @@

+ {"mediaType":"application/vnd.dev.sigstore.bundle.v0.3+json","verificationMaterial":{"certificate":{"rawBytes":"MIIHSzCCBtGgAwIBAgIUK7eu4ilTcA9MKi79s+pe6U8MDSAwCgYIKoZIzj0EAwMwNzEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MR4wHAYDVQQDExVzaWdzdG9yZS1pbnRlcm1lZGlhdGUwHhcNMjYwNjE3MTMxMjA3WhcNMjYwNjE3MTMyMjA3WjAAMFkwEwYHKoZIzj0CAQYIKoZIzj0DAQcDQgAECZZ+JP1Y7BHCrd++8XaWUPtIBgQzWsQXm4M11h7CRoh50JcA1HGMeBp9PXfC1tklfJm0yLlkcFuO412eiXFNjqOCBfAwggXsMA4GA1UdDwEB/wQEAwIHgDATBgNVHSUEDDAKBggrBgEFBQcDAzAdBgNVHQ4EFgQUraCbypzc0xBBe6P5b4UxBwI6smwwHwYDVR0jBBgwFoAU39Ppz1YkEZb5qNjpKFWixi4YZD8wawYDVR0RAQH/BGEwX4ZdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDkGCisGAQQBg78wAQEEK2h0dHBzOi8vdG9rZW4uYWN0aW9ucy5naXRodWJ1c2VyY29udGVudC5jb20wHwYKKwYBBAGDvzABAgQRd29ya2Zsb3dfZGlzcGF0Y2gwNgYKKwYBBAGDvzABAwQoODQ2MTY1YjExZDU3MTA1ODI5MmQ1NDRiNDIyMzQ2MzE3ZGUxMDJkZjATBgorBgEEAYO/MAEEBAVCdWlsZDArBgorBgEEAYO/MAEFBB1odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eTAdBgorBgEEAYO/MAEGBA9yZWZzL2hlYWRzL21haW4wOwYKKwYBBAGDvzABCAQtDCtodHRwczovL3Rva2VuLmFjdGlvbnMuZ2l0aHVidXNlcmNvbnRlbnQuY29tMG0GCisGAQQBg78wAQkEXwxdaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5Ly5naXRodWIvd29ya2Zsb3dzL2J1aWxkLnlhbWxAcmVmcy9oZWFkcy9tYWluMDgGCisGAQQBg78wAQoEKgwoODQ2MTY1YjExZDU3MTA1ODI5MmQ1NDRiNDIyMzQ2MzE3ZGUxMDJkZjAbBgorBgEEAYO/MAELBA0MC3NlbGYtaG9zdGVkMEAGCisGAQQBg78wAQwEMgwwaHR0cHM6Ly9naXRodWIuY29tL2h1Z2dpbmdmYWNlL2tlcm5lbHMtY29tbXVuaXR5MDgGCisGAQQBg78wAQ0EKgwoODQ2MTY1YjExZDU3MTA1ODI5MmQ1NDRiNDIyMzQ2MzE3ZGUxMDJkZjAfBgorBgEEAYO/MAEOBBEMD3JlZnMvaGVhZHMvbWFpbjAaBgorBgEEAYO/MAEPBAwMCjEwNzE0NzU1MjkwLgYKKwYBBAGDvzABEAQgDB5odHRwczovL2dpdGh1Yi5jb20vaHVnZ2luZ2ZhY2UwGAYKKwYBBAGDvzABEQQKDAgyNTcyMDc0MzBtBgorBgEEAYO/MAESBF8MXWh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS8uZ2l0aHViL3dvcmtmbG93cy9idWlsZC55YW1sQHJlZnMvaGVhZHMvbWFpbjA4BgorBgEEAYO/MAETBCoMKDg0NjE2NWIxMWQ1NzEwNTgyOTJkNTQ0YjQyMjM0NjMxN2RlMTAyZGYwIQYKKwYBBAGDvzABFAQTDBF3b3JrZmxvd19kaXNwYXRjaDBkBgorBgEEAYO/MAEVBFYMVGh0dHBzOi8vZ2l0aHViLmNvbS9odWdnaW5nZmFjZS9rZXJuZWxzLWNvbW11bml0eS9hY3Rpb25zL3J1bnMvMjc2OTEyMTUxMzkvYXR0ZW1wdHMvMTAWBgorBgEEAYO/MAEWBAgMBnB1YmxpYzBGBgorBgEEAYO/MAEYBDgMNnJlcG86aHVnZ2luZ2ZhY2Uva2VybmVscy1jb21tdW5pdHk6cmVmOnJlZnMvaGVhZHMvbWFpbjCBigYKKwYBBAHWeQIEAgR8BHoAeAB2AN09MGrGxxEyYxkeHJlnNwKiSl643jyt/4eKcoAvKe6OAAABntW14zAAAAQDAEcwRQIhANR0FINBlVxQ+UrWqe1Gm1asRyym4uVoCvI6Xi68HdS8AiAwqtFKjJoz4ZpWp1/pr2GAJF2bAAQ9Nny7688Fb3xtbjAKBggqhkjOPQQDAwNoADBlAjEA9duWR/CLVkZJwUaI2+LNZjLBAym9nsaLgByw++Htoke2IXx15s4jD568fDI8lWhjAjBdtOz6szSSLeZFRnSyD90l7ssm9INZR1+HC5iczBnJmcyhSa1sXznWwpkaU/j1FYw="},"tlogEntries":[{"logIndex":"1851547484","logId":{"keyId":"wNI9atQGlz+VWfO6LRygH4QUfY/8W4RFwiT5i5WRgB0="},"kindVersion":{"kind":"hashedrekord","version":"0.0.1"},"integratedTime":"1781701928","inclusionPromise":{"signedEntryTimestamp":"MEYCIQDl/9r4CLLrS4QQ3pgoERXkqgOu8J175XMWRogZd0pD9QIhAKKqKcFs/urgUIojMcTzma/v0miHksvtf/e0CYAcTyf/"},"inclusionProof":{"logIndex":"1729643222","rootHash":"d2YR8PsiavxlU1mWQTuknMnLG0hB94xQrxkWQH4MAWg=","treeSize":"1729643237","hashes":["WOkbi8t5T96oGLPXvMJLBG6tjgE+FbwH5WD/DlXW7V8=","96lbMQvpGQmDasvoGfaSrU6aALFCpwNgHChWFhXWNQw=","scSXRqlAwZCM9FW1ZTArFT+WdafQWCPkvi63WjCwGSI=","2JH/NjO8pW7Tl4rmHFgPPciYyS7eHC5p6lCAWN4Hwgc=","Ucph7+nD7Xsl136N9pouoYNdji3ZbYWg+jl98wduUpc=","2kfcco1m25Ift/pgK5YZ3iVPOrz9r83x+S2FJjhnOqg=","1lDJ0gtYDj4AY6Y0aeoHLvrfFCnfZT39PmXC6qldZ/A=","N3X9cecXj3ztkQd5ni4RsrA6ynSQZ/BYoisbueVosuY=","WOT6l6C/2O+HgFq4iZSilKHsx7x+Vnf9VvTqGgb/odM=","v8dESBp721LFExEAv/0tR1e6Rh1pBFHHAvDk7Oc7U1E=","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\n1729643237\nd2YR8PsiavxlU1mWQTuknMnLG0hB94xQrxkWQH4MAWg=\n\n— rekor.sigstore.dev wNI9ajBFAiABOnfelsKpgYmcBv15SCy1mX6GPGnZa0WFylcmDc26vAIhAMMMND1TTD6sbL4B3XGxjHQOG5oLWNpXx2vrB+FABV43\n"}},"canonicalizedBody":"eyJhcGlWZXJzaW9uIjoiMC4wLjEiLCJraW5kIjoiaGFzaGVkcmVrb3JkIiwic3BlYyI6eyJkYXRhIjp7Imhhc2giOnsiYWxnb3JpdGhtIjoic2hhMjU2IiwidmFsdWUiOiI2NmRiYTFlNTBjZmE4ZGMxZmFhMDk1ZGNjNzNkZGQ3ZWYxODYzYThlNGIwMzAwODQzMGZjMDZkZDQ0YTQ3NzNlIn19LCJzaWduYXR1cmUiOnsiY29udGVudCI6Ik1FVUNJUURCYkI3UXRoa1FvL2JOUWY3dVlzakhkV0FzZjV2TkRJaUFLVFB1UTAzL2xRSWdEWjFyVGJpWW9pQXFlMUhqUWlwcU0ra3QzTzN4M1JDcEZEU3FnT2pSdlVjPSIsInB1YmxpY0tleSI6eyJjb250ZW50IjoiTFMwdExTMUNSVWRKVGlCRFJWSlVTVVpKUTBGVVJTMHRMUzB0Q2sxSlNVaFRla05EUW5SSFowRjNTVUpCWjBsVlN6ZGxkVFJwYkZSalFUbE5TMmszT1hNcmNHVTJWVGhOUkZOQmQwTm5XVWxMYjFwSmVtb3dSVUYzVFhjS1RucEZWazFDVFVkQk1WVkZRMmhOVFdNeWJHNWpNMUoyWTIxVmRWcEhWakpOVWpSM1NFRlpSRlpSVVVSRmVGWjZZVmRrZW1SSE9YbGFVekZ3WW01U2JBcGpiVEZzV2tkc2FHUkhWWGRJYUdOT1RXcFpkMDVxUlROTlZFMTRUV3BCTTFkb1kwNU5hbGwzVG1wRk0wMVVUWGxOYWtFelYycEJRVTFHYTNkRmQxbElDa3R2V2tsNmFqQkRRVkZaU1V0dldrbDZhakJFUVZGalJGRm5RVVZEV2xvclNsQXhXVGRDU0VOeVpDc3JPRmhoVjFWUWRFbENaMUY2VjNOUldHMDBUVEVLTVdnM1ExSnZhRFV3U21OQk1VaEhUV1ZDY0RsUVdHWkRNWFJyYkdaS2JUQjVUR3hyWTBaMVR6UXhNbVZwV0VaT2FuRlBRMEptUVhkbloxaHpUVUUwUndwQk1WVmtSSGRGUWk5M1VVVkJkMGxJWjBSQlZFSm5UbFpJVTFWRlJFUkJTMEpuWjNKQ1owVkdRbEZqUkVGNlFXUkNaMDVXU0ZFMFJVWm5VVlZ5WVVOaUNubHdlbU13ZUVKQ1pUWlFOV0kwVlhoQ2QwazJjMjEzZDBoM1dVUldVakJxUWtKbmQwWnZRVlV6T1ZCd2VqRlphMFZhWWpWeFRtcHdTMFpYYVhocE5Ga0tXa1E0ZDJGM1dVUldVakJTUVZGSUwwSkhSWGRZTkZwa1lVaFNNR05JVFRaTWVUbHVZVmhTYjJSWFNYVlpNamwwVERKb01Wb3laSEJpYldSdFdWZE9iQXBNTW5Sc1kyMDFiR0pJVFhSWk1qbDBZbGhXZFdGWVVqVk1lVFZ1WVZoU2IyUlhTWFprTWpsNVlUSmFjMkl6WkhwTU1rb3hZVmQ0YTB4dWJHaGlWM2hCQ21OdFZtMWplVGx2V2xkR2EyTjVPWFJaVjJ4MVRVUnJSME5wYzBkQlVWRkNaemM0ZDBGUlJVVkxNbWd3WkVoQ2VrOXBPSFprUnpseVdsYzBkVmxYVGpBS1lWYzVkV041Tlc1aFdGSnZaRmRLTVdNeVZubFpNamwxWkVkV2RXUkROV3BpTWpCM1NIZFpTMHQzV1VKQ1FVZEVkbnBCUWtGblVWSmtNamw1WVRKYWN3cGlNMlJtV2tkc2VtTkhSakJaTW1kM1RtZFpTMHQzV1VKQ1FVZEVkbnBCUWtGM1VXOVBSRkV5VFZSWk1WbHFSWGhhUkZVelRWUkJNVTlFU1RWTmJWRXhDazVFVW1sT1JFbDVUWHBSTWsxNlJUTmFSMVY0VFVSS2ExcHFRVlJDWjI5eVFtZEZSVUZaVHk5TlFVVkZRa0ZXUTJSWGJITmFSRUZ5UW1kdmNrSm5SVVVLUVZsUEwwMUJSVVpDUWpGdlpGZGtibUZYTlc1YWJVWnFXbE01Y2xwWVNuVmFWM2g2VEZkT2RtSlhNVEZpYld3d1pWUkJaRUpuYjNKQ1owVkZRVmxQTHdwTlFVVkhRa0U1ZVZwWFducE1NbWhzV1ZkU2Vrd3lNV2hoVnpSM1QzZFpTMHQzV1VKQ1FVZEVkbnBCUWtOQlVYUkVRM1J2WkVoU2QyTjZiM1pNTTFKMkNtRXlWblZNYlVacVpFZHNkbUp1VFhWYU1td3dZVWhXYVdSWVRteGpiVTUyWW01U2JHSnVVWFZaTWpsMFRVY3dSME5wYzBkQlVWRkNaemM0ZDBGUmEwVUtXSGQ0WkdGSVVqQmpTRTAyVEhrNWJtRllVbTlrVjBsMVdUSTVkRXd5YURGYU1tUndZbTFrYlZsWFRteE1NblJzWTIwMWJHSklUWFJaTWpsMFlsaFdkUXBoV0ZJMVRIazFibUZZVW05a1YwbDJaREk1ZVdFeVduTmlNMlI2VERKS01XRlhlR3RNYm14b1lsZDRRV050Vm0xamVUbHZXbGRHYTJONU9YUlpWMngxQ2sxRVowZERhWE5IUVZGUlFtYzNPSGRCVVc5RlMyZDNiMDlFVVRKTlZGa3hXV3BGZUZwRVZUTk5WRUV4VDBSSk5VMXRVVEZPUkZKcFRrUkplVTE2VVRJS1RYcEZNMXBIVlhoTlJFcHJXbXBCWWtKbmIzSkNaMFZGUVZsUEwwMUJSVXhDUVRCTlF6Tk9iR0pIV1hSaFJ6bDZaRWRXYTAxRlFVZERhWE5IUVZGUlFncG5OemgzUVZGM1JVMW5kM2RoU0ZJd1kwaE5Oa3g1T1c1aFdGSnZaRmRKZFZreU9YUk1NbWd4V2pKa2NHSnRaRzFaVjA1c1RESjBiR050Tld4aVNFMTBDbGt5T1hSaVdGWjFZVmhTTlUxRVowZERhWE5IUVZGUlFtYzNPSGRCVVRCRlMyZDNiMDlFVVRKTlZGa3hXV3BGZUZwRVZUTk5WRUV4VDBSSk5VMXRVVEVLVGtSU2FVNUVTWGxOZWxFeVRYcEZNMXBIVlhoTlJFcHJXbXBCWmtKbmIzSkNaMFZGUVZsUEwwMUJSVTlDUWtWTlJETktiRnB1VFhaaFIxWm9Xa2hOZGdwaVYwWndZbXBCWVVKbmIzSkNaMFZGUVZsUEwwMUJSVkJDUVhkTlEycEZkMDU2UlRCT2VsVXhUV3ByZDB4bldVdExkMWxDUWtGSFJIWjZRVUpGUVZGbkNrUkNOVzlrU0ZKM1kzcHZka3d5WkhCa1IyZ3hXV2sxYW1JeU1IWmhTRlp1V2pKc2RWb3lXbWhaTWxWM1IwRlpTMHQzV1VKQ1FVZEVkbnBCUWtWUlVVc0tSRUZuZVU1VVkzbE5SR013VFhwQ2RFSm5iM0pDWjBWRlFWbFBMMDFCUlZOQ1JqaE5XRmRvTUdSSVFucFBhVGgyV2pKc01HRklWbWxNYlU1MllsTTVid3BrVjJSdVlWYzFibHB0Um1wYVV6bHlXbGhLZFZwWGVIcE1WMDUyWWxjeE1XSnRiREJsVXpoMVdqSnNNR0ZJVm1sTU0yUjJZMjEwYldKSE9UTmplVGxwQ21SWGJITmFRelUxV1ZjeGMxRklTbXhhYmsxMllVZFdhRnBJVFhaaVYwWndZbXBCTkVKbmIzSkNaMFZGUVZsUEwwMUJSVlJDUTI5TlMwUm5NRTVxUlRJS1RsZEplRTFYVVRGT2VrVjNUbFJuZVU5VVNtdE9WRkV3V1dwUmVVMXFUVEJPYWsxNFRqSlNiRTFVUVhsYVIxbDNTVkZaUzB0M1dVSkNRVWRFZG5wQlFncEdRVkZVUkVKR00ySXpTbkphYlhoMlpERTVhMkZZVG5kWldGSnFZVVJDYTBKbmIzSkNaMFZGUVZsUEwwMUJSVlpDUmxsTlZrZG9NR1JJUW5wUGFUaDJDbG95YkRCaFNGWnBURzFPZG1KVE9XOWtWMlJ1WVZjMWJscHRSbXBhVXpseVdsaEtkVnBYZUhwTVYwNTJZbGN4TVdKdGJEQmxVemxvV1ROU2NHSXlOWG9LVEROS01XSnVUWFpOYW1NeVQxUkZlVTFVVlhoTmVtdDJXVmhTTUZwWE1YZGtTRTEyVFZSQlYwSm5iM0pDWjBWRlFWbFBMMDFCUlZkQ1FXZE5RbTVDTVFwWmJYaHdXWHBDUjBKbmIzSkNaMFZGUVZsUEwwMUJSVmxDUkdkTlRtNUtiR05IT0RaaFNGWnVXakpzZFZveVdtaFpNbFYyWVRKV2VXSnRWbk5qZVRGcUNtSXlNWFJrVnpWd1pFaHJObU50Vm0xUGJrcHNXbTVOZG1GSFZtaGFTRTEyWWxkR2NHSnFRMEpwWjFsTFMzZFpRa0pCU0ZkbFVVbEZRV2RTT0VKSWIwRUtaVUZDTWtGT01EbE5SM0pIZUhoRmVWbDRhMlZJU214dVRuZExhVk5zTmpRemFubDBMelJsUzJOdlFYWkxaVFpQUVVGQlFtNTBWekUwZWtGQlFVRlJSQXBCUldOM1VsRkphRUZPVWpCR1NVNUNiRlo0VVN0VmNsZHhaVEZIYlRGaGMxSjVlVzAwZFZadlEzWkpObGhwTmpoSVpGTTRRV2xCZDNGMFJrdHFTbTk2Q2pSYWNGZHdNUzl3Y2pKSFFVcEdNbUpCUVZFNVRtNTVOelk0T0VaaU0zaDBZbXBCUzBKblozRm9hMnBQVUZGUlJFRjNUbTlCUkVKc1FXcEZRVGxrZFZjS1VpOURURlpyV2twM1ZXRkpNaXRNVGxwcVRFSkJlVzA1Ym5OaFRHZENlWGNySzBoMGIydGxNa2xZZURFMWN6UnFSRFUyT0daRVNUaHNWMmhxUVdwQ1pBcDBUM28yYzNwVFUweGxXa1pTYmxONVJEa3diRGR6YzIwNVNVNWFVakVyU0VNMWFXTjZRbTVLYldONWFGTmhNWE5ZZW01WGQzQnJZVlV2YWpGR1dYYzlDaTB0TFMwdFJVNUVJRU5GVWxSSlJrbERRVlJGTFMwdExTMEsifX19fQ=="}],"timestampVerificationData":{"rfc3161Timestamps":[{"signedTimestamp":"MIICyTADAgEAMIICwAYJKoZIhvcNAQcCoIICsTCCAq0CAQMxDTALBglghkgBZQMEAgEwgbcGCyqGSIb3DQEJEAEEoIGnBIGkMIGhAgEBBgkrBgEEAYO/MAIwMTANBglghkgBZQMEAgEFAAQg95ymQflF1KOugV3rCEAcia0It+U6GO/o+inRA6Ukw/MCFH/a74pw2WWfo0u5xXACrtDVUl1qGA8yMDI2MDYxNzEzMTIwN1owAwIBAaAypDAwLjEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MRUwEwYDVQQDEwxzaWdzdG9yZS10c2GgADGCAdswggHXAgEBMFEwOTEVMBMGA1UEChMMc2lnc3RvcmUuZGV2MSAwHgYDVQQDExdzaWdzdG9yZS10c2Etc2VsZnNpZ25lZAIUOhNULwyQYe68wUMvy4qOiyojiwwwCwYJYIZIAWUDBAIBoIH8MBoGCSqGSIb3DQEJAzENBgsqhkiG9w0BCRABBDAcBgkqhkiG9w0BCQUxDxcNMjYwNjE3MTMxMjA3WjAvBgkqhkiG9w0BCQQxIgQgMSUrKLihMpIxLMfbbVMWPYDLZVdNCYeXctPeWo8abpYwgY4GCyqGSIb3DQEJEAIvMX8wfTB7MHkEIIX5J7wHq2LKw7RDVsEO/IGyxog/2nq55thw2dE6zQW3MFUwPaQ7MDkxFTATBgNVBAoTDHNpZ3N0b3JlLmRldjEgMB4GA1UEAxMXc2lnc3RvcmUtdHNhLXNlbGZzaWduZWQCFDoTVC8MkGHuvMFDL8uKjosqI4sMMAoGCCqGSM49BAMCBGcwZQIxALro8l8xdiK/G4V4W4HOiShvXk+S4Tm2a9UbLWfAdViCoS6oGyXhq40R01O+dYsb9AIwOaeIGFnr6SpMU4s117Qf6v8XdnkgsyROipTnBQ97AKa4/1/EV73fmQzxk5M8VDtM"}]}},"messageSignature":{"messageDigest":{"algorithm":"SHA2_256","digest":"Ztuh5Qz6jcH6oJXcxz3dfvGGOo5LAwCEMPwG3USkdz4="},"signature":"MEUCIQDBbB7QthkQo/bNQf7uYsjHdWAsf5vNDIiAKTPuQ03/lQIgDZ1rTbiYoiAqe1HjQipqM+kt3O3x3RCpFDSqgOjRvUc="}}

build/torch-cuda/utils.py CHANGED Viewed

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

+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import functools
 from contextlib import contextmanager
+import torch
+import triton
+import triton.language as tl
+from torch.library import triton_op, wrap_triton
+from ._ops import add_op_namespace_prefix, ops
+# ── Format constants ──────────────────────────────────────────────────────────
+# FP8 (E4M3) is the main format for weights and activations;
+FP8_DTYPE = torch.float8_e4m3fn
+# FP4 (E2M1) packs two 4-bit nibbles per byte. MX formats (MXFP4 weights, MXFP8
+# E4M3 weights/activations) share one UE8M0 scale per 32-element K-group — the OCP
+# MX block size, consumed by ``tl.dot_scaled``. Format constants, not tunables.
+NIBBLES_PER_BYTE = 2
+MX_SCALE_GROUP_K = 32
+# ── Host-side helpers ─────────────────────────────────────────────────────────
 @contextmanager
 def device_context(device: torch.device):
             yield
     else:
         yield
+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)

build/torch-rocm/__init__.py DELETED Viewed

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

build/torch-rocm/_ops.py DELETED Viewed

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

build/torch-rocm/act_quant.py DELETED Viewed

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

build/torch-rocm/batched.py DELETED Viewed

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

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

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

build/torch-rocm/grouped.py DELETED Viewed

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

build/torch-rocm/matmul.py DELETED Viewed

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

build/torch-rocm/metadata.json DELETED Viewed

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

build/torch-rocm/utils.py DELETED Viewed

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

build/torch-xpu/__init__.py DELETED Viewed

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

build/torch-xpu/_ops.py DELETED Viewed

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

build/torch-xpu/act_quant.py DELETED Viewed

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

build/torch-xpu/batched.py DELETED Viewed

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

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

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

build/torch-xpu/grouped.py DELETED Viewed

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

build/torch-xpu/matmul.py DELETED Viewed

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

build/torch-xpu/metadata.json DELETED Viewed

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

build/torch-xpu/utils.py DELETED Viewed

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