Build uploaded using `kernels`.

build/torch-cuda/__init__.py

@@ -0,0 +1,11 @@
+from .act_quant import fp8_act_quant
+from .batched import w8a8_block_fp8_matmul_batched
+from .grouped import w8a8_block_fp8_matmul_grouped
+from .matmul import w8a8_block_fp8_matmul
+
+__all__ = [
+    "fp8_act_quant",
+    "w8a8_block_fp8_matmul",
+    "w8a8_block_fp8_matmul_batched",
+    "w8a8_block_fp8_matmul_grouped",
+]

build/torch-cuda/_ops.py

@@ -0,0 +1,8 @@
+import torch
+ops = torch.ops._finegrained_fp8_2c10f72
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_finegrained_fp8_2c10f72::{op_name}"

build/torch-cuda/act_quant.py

@@ -0,0 +1,70 @@
+# 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
+
+
+_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
+    y = y.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)
+    s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
+
+    grid = (triton.cdiv(x.numel(), block_size),)
+    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

@@ -0,0 +1,196 @@
+# 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
+
+
+@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 // group_n, K // group_k) weight scales
+    ExpertIds,  # (S,) — which expert each batch element routes to
+    # Shape
+    N,
+    K,
+    # Block size for block-wise quantization
+    group_n,
+    group_k,
+    # Per-row strides
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cn,
+    stride_Bs_k,
+    stride_Bs_n,
+    # Batch / expert strides
+    stride_Ab,  # stride between rows in A (one token per program)
+    stride_Eb,  # stride between experts in B
+    stride_Cb,
+    stride_Esb,  # stride between experts in Bs
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    pid_n = tl.program_id(axis=0)
+    batch_id = tl.program_id(axis=1)
+
+    # Advance base pointers to this token's activation row and its expert's
+    # weight / scale slice. No pre-gather of weights needed (like in non-fp8 impls)
+    expert_id = tl.load(ExpertIds + batch_id)
+    A = A + batch_id * stride_Ab
+    B = B + expert_id * stride_Eb
+    C = C + batch_id * stride_Cb
+    Bs = Bs + expert_id * stride_Esb
+
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # M=1: broadcast the single activation row to BLOCK_SIZE_M identical rows
+    # so tl.dot gets the required (BLOCK_SIZE_M, BLOCK_SIZE_K) shape.
+    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)
+
+    offs_bsn = offs_bn // group_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)):
+        # ---- fused act_quant (replaces: a = tl.load(a_ptrs); a_s = tl.load(As_ptrs)) ----
+        a_raw = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0).to(tl.float32)
+        a_s = tl.max(tl.abs(a_raw)) / 448.0  # per-block scale (scalar for M=1)
+        a = (a_raw / tl.maximum(a_s, 1e-12)).to(tl.float8e4nv)
+        # ---- same as baseline from here ----
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        k_start = k * BLOCK_SIZE_K
+        offs_ks = k_start // group_k
+        b_s = tl.load(Bs_ptrs + offs_ks * 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)
+
+    # Only write row 0 (M=1); the broadcast rows are discarded.
+    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, :]
+    c_mask = (offs_cm[:, None] < 1) & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+
+@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] | None,
+) -> torch.Tensor:
+    """Batched FP8 block-wise matmul with fused activation quantization.
+
+    Mirrors ``_batched_linear`` for FP8 weights: A is the raw (BF16/FP16)
+    activation matrix, B / Bs are the stacked expert weights / scales.
+    The kernel looks up ``expert_ids[batch_id]`` to address the correct expert
+    slice of B directly — no (S, N, K) weight gather is needed.
+    Activation quantization (``act_quant``) is fused into the matmul loop.
+    """
+    assert A.ndim == 2, "A must be (S, K)"
+    assert A.is_contiguous()
+
+    assert B.ndim == 3, "B must be (E, N, K)"
+    assert B.is_contiguous()
+
+    assert A.shape[1] == B.shape[2], "K dimension mismatch between A and B"
+    assert expert_ids.is_contiguous()
+    assert Bs.is_contiguous()
+
+    if block_size is None:
+        block_n, block_k = 128, 128
+    else:
+        assert len(block_size) == 2
+        block_n, block_k = block_size[0], block_size[1]
+
+    S, K = A.shape
+    E, N, _ = B.shape
+    C = A.new_empty(S, N)
+
+    # Adaptive BLOCK_SIZE_M: match the tile to the average tokens per expert
+    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
+
+    grid = (triton.cdiv(N, block_n), S)
+    wrap_triton(w8a8_block_fp8_matmul_batched_kernel)[grid](
+        A,
+        B,
+        C,
+        Bs,
+        expert_ids,
+        N,
+        K,
+        block_n,
+        block_k,
+        A.stride(1),   # stride_ak
+        B.stride(2),   # stride_bk
+        B.stride(1),   # stride_bn
+        C.stride(1),   # stride_cn
+        Bs.stride(2),  # stride_Bs_k
+        Bs.stride(1),  # stride_Bs_n
+        A.stride(0),   # stride_Ab
+        B.stride(0),   # stride_Eb
+        C.stride(0),   # stride_Cb
+        Bs.stride(0),  # stride_Esb
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=block_n,
+        BLOCK_SIZE_K=block_k,
+    )
+
+    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] | None,
+) -> torch.Tensor:
+    """Batched W8A8 FP8 matmul for MoE expert dispatch with fused activation quantization.
+
+    Each token in ``A`` is routed to one expert via ``expert_ids``. The kernel
+    quantizes ``A`` to FP8 on-the-fly (fused ``act_quant``), reads the correct
+    expert weight slice of ``B`` directly using ``expert_ids[batch_id]`` — no
+    ``(S, N, K)`` weight gather needed — and accumulates in float32.
+
+    Args:
+        A: Raw activation matrix ``[S, K]`` in bf16/fp16/fp32.
+        B: Stacked expert weight tensor ``[E, N, K]`` in ``float8_e4m3fn``.
+        Bs: Stacked expert weight scales ``[E, N // block_size[0], K // block_size[1]]``.
+        expert_ids: Expert index per token ``[S]``, values in ``[0, E)``.
+        block_size: ``[block_n, block_k]`` quantization block dimensions, e.g. ``[128, 128]``.
+
+    Returns:
+        Output tensor ``[S, N]`` in the same dtype as ``A``.
+    """
+    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_batched(A, B, Bs, expert_ids, block_size)

build/torch-cuda/finegrained_fp8/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+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-cuda/grouped.py

@@ -0,0 +1,252 @@
+# 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
+
+
+@triton.jit
+def w8a8_block_fp8_grouped_mm_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 // group_n, K // group_k) weight scales
+    Offsets,  # (E,) int32 — cumulative row-end per expert
+    TileOffsets,  # (E,) int32 — cumulative tile-end per expert
+    # Shape
+    S,
+    N,
+    K,
+    # Block size for block-wise quantization
+    group_n,
+    group_k,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_Eb,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_Esb,
+    stride_Bsk,
+    stride_Bsn,
+    # Meta-parameters
+    NUM_EXPERTS: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    pid_m = pid // num_pid_n
+    pid_n = pid % num_pid_n
+
+    # 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)
+    expert_id = lo
+
+    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_bn_safe = offs_bn % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+
+    offs_am_safe = offs_global_m % S
+
+    a_ptrs = A + offs_am_safe[:, None] * stride_am + offs_k[None, :] * stride_ak
+    b_ptrs = (
+        B
+        + expert_id * stride_Eb
+        + offs_k[:, None] * stride_bk
+        + offs_bn_safe[None, :] * stride_bn
+    )
+    offs_bsn_safe = offs_bn_safe // group_n
+    Bs_ptrs = Bs + expert_id * stride_Esb + offs_bsn_safe * stride_Bsn
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        # ---- fused act_quant (replaces: a = tl.load(a_ptrs); a_s = tl.load(As_ptrs)) ----
+        a_raw = tl.load(
+            a_ptrs,
+            mask=row_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
+            other=0.0,
+        ).to(tl.float32)
+        a_s = tl.max(tl.abs(a_raw), axis=1) / 448.0  # per-row scale  (BLOCK_SIZE_M,)
+        # clamp denominator so masked all-zero rows don't produce NaN
+        # (their a_s multiplier is 0 anyway, so the output row is correct)
+        a = (a_raw / tl.maximum(a_s[:, None], 1e-12)).to(tl.float8e4nv)
+        # ---- same as baseline from here ----
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        k_start = k * BLOCK_SIZE_K
+        offs_ks = k_start // group_k
+        b_s = tl.load(Bs_ptrs + offs_ks * stride_Bsk)
+
+        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] & (offs_bn[None, :] < N)
+    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] | None,
+) -> torch.Tensor:
+    """Grouped FP8 block-wise matmul with fused activation quantization.
+
+    Mirrors ``_grouped_linear`` / ``_grouped_mm`` for FP8 weights: A is the
+    raw (BF16/FP16) activation matrix sorted by expert, B / Bs are the stacked
+    expert weights / scales.  Activation quantization (``act_quant``) is fused
+    into the matmul loop.  ``tokens_per_expert`` is needed (in addition to
+    ``offsets``) to build the per-expert tile schedule inside the kernel.
+    """
+    assert A.ndim == 2, "A must be (S, K)"
+    assert A.is_contiguous()
+
+    assert B.ndim == 3, "B must be (E, N, K)"
+    assert B.is_contiguous()
+
+    assert tokens_per_expert.is_contiguous()
+    assert offsets.is_contiguous()
+    assert Bs.is_contiguous()
+
+    if block_size is None:
+        block_n, block_k = 128, 128
+    else:
+        assert len(block_size) == 2
+        block_n, block_k = block_size[0], block_size[1]
+
+    S, K = A.shape
+    E, N, _ = B.shape
+    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.
+    # Using a static upper bound keeps the grid size data-independent, which is
+    # required for cuda-graph compatibility.  Programs beyond the real tile count
+    # exit immediately via the early-return guard inside the kernel.
+    max_M_tiles = triton.cdiv(S, BLOCK_SIZE_M) + E
+
+    grid = (max_M_tiles * triton.cdiv(N, block_n),)
+    wrap_triton(w8a8_block_fp8_grouped_mm_kernel)[grid](
+        A,
+        B,
+        C,
+        Bs,
+        offsets,
+        tile_offsets,
+        S,
+        N,
+        K,
+        block_n,
+        block_k,
+        A.stride(0),  # stride_am
+        A.stride(1),  # stride_ak
+        B.stride(0),  # stride_Eb
+        B.stride(2),  # stride_bk
+        B.stride(1),  # stride_bn
+        C.stride(0),  # stride_cm
+        C.stride(1),  # stride_cn
+        Bs.stride(0),  # stride_Esb
+        Bs.stride(2),  # stride_Bsk
+        Bs.stride(1),  # stride_Bsn
+        NUM_EXPERTS=E,
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=block_n,
+        BLOCK_SIZE_K=block_k,
+    )
+
+    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] | None,
+) -> torch.Tensor:
+    """Grouped W8A8 FP8 matmul for MoE expert dispatch with fused activation quantization.
+
+    Tokens in ``A`` must be pre-sorted by expert id. The kernel quantizes ``A``
+    to FP8 on-the-fly (fused ``act_quant``), uses a static over-provisioned grid
+    for CUDA-graph compatibility, and resolves each tile's expert via an O(log E)
+    binary search over ``offsets``.
+
+    Args:
+        A: Raw activation matrix ``[S, K]`` sorted by expert, in bf16/fp16/fp32.
+        B: Stacked expert weight tensor ``[E, N, K]`` in ``float8_e4m3fn``.
+        Bs: Stacked expert weight scales ``[E, N // block_size[0], K // block_size[1]]``.
+        offsets: Cumulative token counts per expert ``[E]`` (i.e. ``cumsum(tokens_per_expert)``).
+        tokens_per_expert: Number of tokens routed to each expert ``[E]``.
+        block_size: ``[block_n, block_k]`` quantization block dimensions, e.g. ``[128, 128]``.
+
+    Returns:
+        Output tensor ``[S, N]`` in the same dtype as ``A``, in expert-sorted order.
+    """
+    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_grouped(
+        A, B, Bs, offsets, tokens_per_expert, block_size
+    )

build/torch-cuda/matmul.py

@@ -0,0 +1,232 @@
+# 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
+
+
+# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py
+@triton.jit
+def w8a8_block_fp8_matmul_kernel(
+    # Pointers to inputs and output
+    A,
+    B,
+    C,
+    As,
+    Bs,
+    # Shape for matmul
+    M,
+    N,
+    K,
+    # Block size for block-wise quantization
+    group_n,
+    group_k,
+    # Stride for inputs and output
+    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,
+):
+    """Triton-accelerated function used to perform linear operations (dot
+    product) on input tensors `A` and `B` with block-wise quantization, and
+    store the result in output tensor `C`.
+    """
+
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + (pid % group_size_m)
+    pid_n = (pid % num_pid_in_group) // 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 // group_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)):
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        k_start = k * BLOCK_SIZE_K
+        offs_ks = k_start // group_k
+        a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
+        b_s = tl.load(Bs_ptrs + offs_ks * 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_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] | None,
+    output_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """This function performs matrix multiplication with block-wise
+    quantization.
+    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
+    The output is returned in the specified `output_dtype`.
+    Args:
+        A: The input tensor, e.g., activation.
+        B: The input tensor, e.g., weight.
+        As: The per-token-group quantization scale for `A`.
+        Bs: The per-block quantization scale for `B`.
+        block_size: The block size for per-block quantization. It should
+        be 2-dim, e.g., [128, 128].
+        output_dtype: The dtype of the returned tensor.
+    Returns:
+        torch.Tensor: The result of matmul.
+    """
+    if block_size is None:
+        block_n, block_k = 128, 128
+    else:
+        assert len(block_size) == 2
+        block_n, block_k = block_size[0], block_size[1]
+
+    # if we have per-tensor quantization, we use 128x128 block size for tiled matmul multiplication
+    if block_n == B.shape[-2] and block_k == B.shape[-1]:
+        block_n = 128
+        block_k = 128
+
+    assert A.shape[-1] == B.shape[-1]
+    assert A.is_contiguous()
+
+    assert B.ndim == 2
+    assert B.is_contiguous()
+
+    N, K = B.shape
+    M = A.numel() // A.shape[-1]
+
+    # For per-tensor scales (scalar), expand to block-scale shape with strides (0, 0).
+    # This is a zero-copy view; all loads inside the kernel hit the same cached value.
+    if As.numel() == 1:
+        As = As.reshape(1, 1).expand(M, triton.cdiv(K, block_k))
+    else:
+        assert A.shape[:-1] == As.shape[:-1]
+        assert triton.cdiv(K, block_k) == As.shape[-1]
+    if Bs.numel() == 1:
+        Bs = Bs.reshape(1, 1).expand(triton.cdiv(N, block_n), triton.cdiv(K, block_k))
+    else:
+        assert Bs.ndim == 2
+        assert triton.cdiv(N, block_n) == Bs.shape[0], f"{N}, {block_n}, {Bs.shape}"
+        assert triton.cdiv(K, block_k) == Bs.shape[1], f"{K}, {block_k}, {Bs.shape}"
+
+    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)
+    BLOCK_SIZE_K = block_k
+    BLOCK_SIZE_N = block_n
+
+    grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),)
+    wrap_triton(w8a8_block_fp8_matmul_kernel)[grid](
+        A,
+        B,
+        C,
+        As,
+        Bs,
+        M,
+        N,
+        K,
+        block_n,
+        block_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),
+        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] | None,
+    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)

build/torch-cuda/metadata-cuda.json

@@ -0,0 +1,4 @@
+{
+  "version": 1,
+  "python-depends": []
+}

build/torch-rocm/__init__.py

@@ -0,0 +1,11 @@
+from .act_quant import fp8_act_quant
+from .batched import w8a8_block_fp8_matmul_batched
+from .grouped import w8a8_block_fp8_matmul_grouped
+from .matmul import w8a8_block_fp8_matmul
+
+__all__ = [
+    "fp8_act_quant",
+    "w8a8_block_fp8_matmul",
+    "w8a8_block_fp8_matmul_batched",
+    "w8a8_block_fp8_matmul_grouped",
+]

build/torch-rocm/_ops.py

@@ -0,0 +1,8 @@
+import torch
+ops = torch.ops._finegrained_fp8_2c10f72
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_finegrained_fp8_2c10f72::{op_name}"

build/torch-rocm/act_quant.py

@@ -0,0 +1,70 @@
+# 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
+
+
+_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
+    y = y.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)
+    s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
+
+    grid = (triton.cdiv(x.numel(), block_size),)
+    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

@@ -0,0 +1,196 @@
+# 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
+
+
+@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 // group_n, K // group_k) weight scales
+    ExpertIds,  # (S,) — which expert each batch element routes to
+    # Shape
+    N,
+    K,
+    # Block size for block-wise quantization
+    group_n,
+    group_k,
+    # Per-row strides
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cn,
+    stride_Bs_k,
+    stride_Bs_n,
+    # Batch / expert strides
+    stride_Ab,  # stride between rows in A (one token per program)
+    stride_Eb,  # stride between experts in B
+    stride_Cb,
+    stride_Esb,  # stride between experts in Bs
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    pid_n = tl.program_id(axis=0)
+    batch_id = tl.program_id(axis=1)
+
+    # Advance base pointers to this token's activation row and its expert's
+    # weight / scale slice. No pre-gather of weights needed (like in non-fp8 impls)
+    expert_id = tl.load(ExpertIds + batch_id)
+    A = A + batch_id * stride_Ab
+    B = B + expert_id * stride_Eb
+    C = C + batch_id * stride_Cb
+    Bs = Bs + expert_id * stride_Esb
+
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # M=1: broadcast the single activation row to BLOCK_SIZE_M identical rows
+    # so tl.dot gets the required (BLOCK_SIZE_M, BLOCK_SIZE_K) shape.
+    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)
+
+    offs_bsn = offs_bn // group_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)):
+        # ---- fused act_quant (replaces: a = tl.load(a_ptrs); a_s = tl.load(As_ptrs)) ----
+        a_raw = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0).to(tl.float32)
+        a_s = tl.max(tl.abs(a_raw)) / 448.0  # per-block scale (scalar for M=1)
+        a = (a_raw / tl.maximum(a_s, 1e-12)).to(tl.float8e4nv)
+        # ---- same as baseline from here ----
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        k_start = k * BLOCK_SIZE_K
+        offs_ks = k_start // group_k
+        b_s = tl.load(Bs_ptrs + offs_ks * 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)
+
+    # Only write row 0 (M=1); the broadcast rows are discarded.
+    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, :]
+    c_mask = (offs_cm[:, None] < 1) & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+
+@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] | None,
+) -> torch.Tensor:
+    """Batched FP8 block-wise matmul with fused activation quantization.
+
+    Mirrors ``_batched_linear`` for FP8 weights: A is the raw (BF16/FP16)
+    activation matrix, B / Bs are the stacked expert weights / scales.
+    The kernel looks up ``expert_ids[batch_id]`` to address the correct expert
+    slice of B directly — no (S, N, K) weight gather is needed.
+    Activation quantization (``act_quant``) is fused into the matmul loop.
+    """
+    assert A.ndim == 2, "A must be (S, K)"
+    assert A.is_contiguous()
+
+    assert B.ndim == 3, "B must be (E, N, K)"
+    assert B.is_contiguous()
+
+    assert A.shape[1] == B.shape[2], "K dimension mismatch between A and B"
+    assert expert_ids.is_contiguous()
+    assert Bs.is_contiguous()
+
+    if block_size is None:
+        block_n, block_k = 128, 128
+    else:
+        assert len(block_size) == 2
+        block_n, block_k = block_size[0], block_size[1]
+
+    S, K = A.shape
+    E, N, _ = B.shape
+    C = A.new_empty(S, N)
+
+    # Adaptive BLOCK_SIZE_M: match the tile to the average tokens per expert
+    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
+
+    grid = (triton.cdiv(N, block_n), S)
+    wrap_triton(w8a8_block_fp8_matmul_batched_kernel)[grid](
+        A,
+        B,
+        C,
+        Bs,
+        expert_ids,
+        N,
+        K,
+        block_n,
+        block_k,
+        A.stride(1),   # stride_ak
+        B.stride(2),   # stride_bk
+        B.stride(1),   # stride_bn
+        C.stride(1),   # stride_cn
+        Bs.stride(2),  # stride_Bs_k
+        Bs.stride(1),  # stride_Bs_n
+        A.stride(0),   # stride_Ab
+        B.stride(0),   # stride_Eb
+        C.stride(0),   # stride_Cb
+        Bs.stride(0),  # stride_Esb
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=block_n,
+        BLOCK_SIZE_K=block_k,
+    )
+
+    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] | None,
+) -> torch.Tensor:
+    """Batched W8A8 FP8 matmul for MoE expert dispatch with fused activation quantization.
+
+    Each token in ``A`` is routed to one expert via ``expert_ids``. The kernel
+    quantizes ``A`` to FP8 on-the-fly (fused ``act_quant``), reads the correct
+    expert weight slice of ``B`` directly using ``expert_ids[batch_id]`` — no
+    ``(S, N, K)`` weight gather needed — and accumulates in float32.
+
+    Args:
+        A: Raw activation matrix ``[S, K]`` in bf16/fp16/fp32.
+        B: Stacked expert weight tensor ``[E, N, K]`` in ``float8_e4m3fn``.
+        Bs: Stacked expert weight scales ``[E, N // block_size[0], K // block_size[1]]``.
+        expert_ids: Expert index per token ``[S]``, values in ``[0, E)``.
+        block_size: ``[block_n, block_k]`` quantization block dimensions, e.g. ``[128, 128]``.
+
+    Returns:
+        Output tensor ``[S, N]`` in the same dtype as ``A``.
+    """
+    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_batched(A, B, Bs, expert_ids, block_size)

build/torch-rocm/finegrained_fp8/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+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

@@ -0,0 +1,252 @@
+# 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
+
+
+@triton.jit
+def w8a8_block_fp8_grouped_mm_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 // group_n, K // group_k) weight scales
+    Offsets,  # (E,) int32 — cumulative row-end per expert
+    TileOffsets,  # (E,) int32 — cumulative tile-end per expert
+    # Shape
+    S,
+    N,
+    K,
+    # Block size for block-wise quantization
+    group_n,
+    group_k,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_Eb,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_Esb,
+    stride_Bsk,
+    stride_Bsn,
+    # Meta-parameters
+    NUM_EXPERTS: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    pid_m = pid // num_pid_n
+    pid_n = pid % num_pid_n
+
+    # 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)
+    expert_id = lo
+
+    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_bn_safe = offs_bn % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+
+    offs_am_safe = offs_global_m % S
+
+    a_ptrs = A + offs_am_safe[:, None] * stride_am + offs_k[None, :] * stride_ak
+    b_ptrs = (
+        B
+        + expert_id * stride_Eb
+        + offs_k[:, None] * stride_bk
+        + offs_bn_safe[None, :] * stride_bn
+    )
+    offs_bsn_safe = offs_bn_safe // group_n
+    Bs_ptrs = Bs + expert_id * stride_Esb + offs_bsn_safe * stride_Bsn
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        # ---- fused act_quant (replaces: a = tl.load(a_ptrs); a_s = tl.load(As_ptrs)) ----
+        a_raw = tl.load(
+            a_ptrs,
+            mask=row_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
+            other=0.0,
+        ).to(tl.float32)
+        a_s = tl.max(tl.abs(a_raw), axis=1) / 448.0  # per-row scale  (BLOCK_SIZE_M,)
+        # clamp denominator so masked all-zero rows don't produce NaN
+        # (their a_s multiplier is 0 anyway, so the output row is correct)
+        a = (a_raw / tl.maximum(a_s[:, None], 1e-12)).to(tl.float8e4nv)
+        # ---- same as baseline from here ----
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        k_start = k * BLOCK_SIZE_K
+        offs_ks = k_start // group_k
+        b_s = tl.load(Bs_ptrs + offs_ks * stride_Bsk)
+
+        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] & (offs_bn[None, :] < N)
+    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] | None,
+) -> torch.Tensor:
+    """Grouped FP8 block-wise matmul with fused activation quantization.
+
+    Mirrors ``_grouped_linear`` / ``_grouped_mm`` for FP8 weights: A is the
+    raw (BF16/FP16) activation matrix sorted by expert, B / Bs are the stacked
+    expert weights / scales.  Activation quantization (``act_quant``) is fused
+    into the matmul loop.  ``tokens_per_expert`` is needed (in addition to
+    ``offsets``) to build the per-expert tile schedule inside the kernel.
+    """
+    assert A.ndim == 2, "A must be (S, K)"
+    assert A.is_contiguous()
+
+    assert B.ndim == 3, "B must be (E, N, K)"
+    assert B.is_contiguous()
+
+    assert tokens_per_expert.is_contiguous()
+    assert offsets.is_contiguous()
+    assert Bs.is_contiguous()
+
+    if block_size is None:
+        block_n, block_k = 128, 128
+    else:
+        assert len(block_size) == 2
+        block_n, block_k = block_size[0], block_size[1]
+
+    S, K = A.shape
+    E, N, _ = B.shape
+    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.
+    # Using a static upper bound keeps the grid size data-independent, which is
+    # required for cuda-graph compatibility.  Programs beyond the real tile count
+    # exit immediately via the early-return guard inside the kernel.
+    max_M_tiles = triton.cdiv(S, BLOCK_SIZE_M) + E
+
+    grid = (max_M_tiles * triton.cdiv(N, block_n),)
+    wrap_triton(w8a8_block_fp8_grouped_mm_kernel)[grid](
+        A,
+        B,
+        C,
+        Bs,
+        offsets,
+        tile_offsets,
+        S,
+        N,
+        K,
+        block_n,
+        block_k,
+        A.stride(0),  # stride_am
+        A.stride(1),  # stride_ak
+        B.stride(0),  # stride_Eb
+        B.stride(2),  # stride_bk
+        B.stride(1),  # stride_bn
+        C.stride(0),  # stride_cm
+        C.stride(1),  # stride_cn
+        Bs.stride(0),  # stride_Esb
+        Bs.stride(2),  # stride_Bsk
+        Bs.stride(1),  # stride_Bsn
+        NUM_EXPERTS=E,
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=block_n,
+        BLOCK_SIZE_K=block_k,
+    )
+
+    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] | None,
+) -> torch.Tensor:
+    """Grouped W8A8 FP8 matmul for MoE expert dispatch with fused activation quantization.
+
+    Tokens in ``A`` must be pre-sorted by expert id. The kernel quantizes ``A``
+    to FP8 on-the-fly (fused ``act_quant``), uses a static over-provisioned grid
+    for CUDA-graph compatibility, and resolves each tile's expert via an O(log E)
+    binary search over ``offsets``.
+
+    Args:
+        A: Raw activation matrix ``[S, K]`` sorted by expert, in bf16/fp16/fp32.
+        B: Stacked expert weight tensor ``[E, N, K]`` in ``float8_e4m3fn``.
+        Bs: Stacked expert weight scales ``[E, N // block_size[0], K // block_size[1]]``.
+        offsets: Cumulative token counts per expert ``[E]`` (i.e. ``cumsum(tokens_per_expert)``).
+        tokens_per_expert: Number of tokens routed to each expert ``[E]``.
+        block_size: ``[block_n, block_k]`` quantization block dimensions, e.g. ``[128, 128]``.
+
+    Returns:
+        Output tensor ``[S, N]`` in the same dtype as ``A``, in expert-sorted order.
+    """
+    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_grouped(
+        A, B, Bs, offsets, tokens_per_expert, block_size
+    )

build/torch-rocm/matmul.py

@@ -0,0 +1,232 @@
+# 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
+
+
+# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py
+@triton.jit
+def w8a8_block_fp8_matmul_kernel(
+    # Pointers to inputs and output
+    A,
+    B,
+    C,
+    As,
+    Bs,
+    # Shape for matmul
+    M,
+    N,
+    K,
+    # Block size for block-wise quantization
+    group_n,
+    group_k,
+    # Stride for inputs and output
+    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,
+):
+    """Triton-accelerated function used to perform linear operations (dot
+    product) on input tensors `A` and `B` with block-wise quantization, and
+    store the result in output tensor `C`.
+    """
+
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + (pid % group_size_m)
+    pid_n = (pid % num_pid_in_group) // 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 // group_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)):
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        k_start = k * BLOCK_SIZE_K
+        offs_ks = k_start // group_k
+        a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
+        b_s = tl.load(Bs_ptrs + offs_ks * 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_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] | None,
+    output_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """This function performs matrix multiplication with block-wise
+    quantization.
+    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
+    The output is returned in the specified `output_dtype`.
+    Args:
+        A: The input tensor, e.g., activation.
+        B: The input tensor, e.g., weight.
+        As: The per-token-group quantization scale for `A`.
+        Bs: The per-block quantization scale for `B`.
+        block_size: The block size for per-block quantization. It should
+        be 2-dim, e.g., [128, 128].
+        output_dtype: The dtype of the returned tensor.
+    Returns:
+        torch.Tensor: The result of matmul.
+    """
+    if block_size is None:
+        block_n, block_k = 128, 128
+    else:
+        assert len(block_size) == 2
+        block_n, block_k = block_size[0], block_size[1]
+
+    # if we have per-tensor quantization, we use 128x128 block size for tiled matmul multiplication
+    if block_n == B.shape[-2] and block_k == B.shape[-1]:
+        block_n = 128
+        block_k = 128
+
+    assert A.shape[-1] == B.shape[-1]
+    assert A.is_contiguous()
+
+    assert B.ndim == 2
+    assert B.is_contiguous()
+
+    N, K = B.shape
+    M = A.numel() // A.shape[-1]
+
+    # For per-tensor scales (scalar), expand to block-scale shape with strides (0, 0).
+    # This is a zero-copy view; all loads inside the kernel hit the same cached value.
+    if As.numel() == 1:
+        As = As.reshape(1, 1).expand(M, triton.cdiv(K, block_k))
+    else:
+        assert A.shape[:-1] == As.shape[:-1]
+        assert triton.cdiv(K, block_k) == As.shape[-1]
+    if Bs.numel() == 1:
+        Bs = Bs.reshape(1, 1).expand(triton.cdiv(N, block_n), triton.cdiv(K, block_k))
+    else:
+        assert Bs.ndim == 2
+        assert triton.cdiv(N, block_n) == Bs.shape[0], f"{N}, {block_n}, {Bs.shape}"
+        assert triton.cdiv(K, block_k) == Bs.shape[1], f"{K}, {block_k}, {Bs.shape}"
+
+    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)
+    BLOCK_SIZE_K = block_k
+    BLOCK_SIZE_N = block_n
+
+    grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),)
+    wrap_triton(w8a8_block_fp8_matmul_kernel)[grid](
+        A,
+        B,
+        C,
+        As,
+        Bs,
+        M,
+        N,
+        K,
+        block_n,
+        block_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),
+        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] | None,
+    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)

build/torch-rocm/metadata-rocm.json

@@ -0,0 +1,4 @@
+{
+  "version": 1,
+  "python-depends": []
+}

build/torch-xpu/__init__.py

@@ -0,0 +1,11 @@
+from .act_quant import fp8_act_quant
+from .batched import w8a8_block_fp8_matmul_batched
+from .grouped import w8a8_block_fp8_matmul_grouped
+from .matmul import w8a8_block_fp8_matmul
+
+__all__ = [
+    "fp8_act_quant",
+    "w8a8_block_fp8_matmul",
+    "w8a8_block_fp8_matmul_batched",
+    "w8a8_block_fp8_matmul_grouped",
+]

build/torch-xpu/_ops.py

@@ -0,0 +1,8 @@
+import torch
+ops = torch.ops._finegrained_fp8_2c10f72
+
+def add_op_namespace_prefix(op_name: str):
+    """
+    Prefix op by namespace.
+    """
+    return f"_finegrained_fp8_2c10f72::{op_name}"

build/torch-xpu/act_quant.py

@@ -0,0 +1,70 @@
+# 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
+
+
+_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
+    y = y.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)
+    s = x.new_empty(*x.size()[:-1], x.size(-1) // block_size, dtype=torch.float32)
+
+    grid = (triton.cdiv(x.numel(), block_size),)
+    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

@@ -0,0 +1,196 @@
+# 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
+
+
+@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 // group_n, K // group_k) weight scales
+    ExpertIds,  # (S,) — which expert each batch element routes to
+    # Shape
+    N,
+    K,
+    # Block size for block-wise quantization
+    group_n,
+    group_k,
+    # Per-row strides
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cn,
+    stride_Bs_k,
+    stride_Bs_n,
+    # Batch / expert strides
+    stride_Ab,  # stride between rows in A (one token per program)
+    stride_Eb,  # stride between experts in B
+    stride_Cb,
+    stride_Esb,  # stride between experts in Bs
+    # Meta-parameters
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    pid_n = tl.program_id(axis=0)
+    batch_id = tl.program_id(axis=1)
+
+    # Advance base pointers to this token's activation row and its expert's
+    # weight / scale slice. No pre-gather of weights needed (like in non-fp8 impls)
+    expert_id = tl.load(ExpertIds + batch_id)
+    A = A + batch_id * stride_Ab
+    B = B + expert_id * stride_Eb
+    C = C + batch_id * stride_Cb
+    Bs = Bs + expert_id * stride_Esb
+
+    offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+    # M=1: broadcast the single activation row to BLOCK_SIZE_M identical rows
+    # so tl.dot gets the required (BLOCK_SIZE_M, BLOCK_SIZE_K) shape.
+    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)
+
+    offs_bsn = offs_bn // group_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)):
+        # ---- fused act_quant (replaces: a = tl.load(a_ptrs); a_s = tl.load(As_ptrs)) ----
+        a_raw = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0).to(tl.float32)
+        a_s = tl.max(tl.abs(a_raw)) / 448.0  # per-block scale (scalar for M=1)
+        a = (a_raw / tl.maximum(a_s, 1e-12)).to(tl.float8e4nv)
+        # ---- same as baseline from here ----
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        k_start = k * BLOCK_SIZE_K
+        offs_ks = k_start // group_k
+        b_s = tl.load(Bs_ptrs + offs_ks * 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)
+
+    # Only write row 0 (M=1); the broadcast rows are discarded.
+    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, :]
+    c_mask = (offs_cm[:, None] < 1) & (offs_cn[None, :] < N)
+    tl.store(c_ptrs, c, mask=c_mask)
+
+
+@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] | None,
+) -> torch.Tensor:
+    """Batched FP8 block-wise matmul with fused activation quantization.
+
+    Mirrors ``_batched_linear`` for FP8 weights: A is the raw (BF16/FP16)
+    activation matrix, B / Bs are the stacked expert weights / scales.
+    The kernel looks up ``expert_ids[batch_id]`` to address the correct expert
+    slice of B directly — no (S, N, K) weight gather is needed.
+    Activation quantization (``act_quant``) is fused into the matmul loop.
+    """
+    assert A.ndim == 2, "A must be (S, K)"
+    assert A.is_contiguous()
+
+    assert B.ndim == 3, "B must be (E, N, K)"
+    assert B.is_contiguous()
+
+    assert A.shape[1] == B.shape[2], "K dimension mismatch between A and B"
+    assert expert_ids.is_contiguous()
+    assert Bs.is_contiguous()
+
+    if block_size is None:
+        block_n, block_k = 128, 128
+    else:
+        assert len(block_size) == 2
+        block_n, block_k = block_size[0], block_size[1]
+
+    S, K = A.shape
+    E, N, _ = B.shape
+    C = A.new_empty(S, N)
+
+    # Adaptive BLOCK_SIZE_M: match the tile to the average tokens per expert
+    BLOCK_SIZE_M = min(max(triton.next_power_of_2((S + E - 1) // E), 16), 128)
+
+    grid = (triton.cdiv(N, block_n), S)
+    wrap_triton(w8a8_block_fp8_matmul_batched_kernel)[grid](
+        A,
+        B,
+        C,
+        Bs,
+        expert_ids,
+        N,
+        K,
+        block_n,
+        block_k,
+        A.stride(1),   # stride_ak
+        B.stride(2),   # stride_bk
+        B.stride(1),   # stride_bn
+        C.stride(1),   # stride_cn
+        Bs.stride(2),  # stride_Bs_k
+        Bs.stride(1),  # stride_Bs_n
+        A.stride(0),   # stride_Ab
+        B.stride(0),   # stride_Eb
+        C.stride(0),   # stride_Cb
+        Bs.stride(0),  # stride_Esb
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=block_n,
+        BLOCK_SIZE_K=block_k,
+    )
+
+    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] | None,
+) -> torch.Tensor:
+    """Batched W8A8 FP8 matmul for MoE expert dispatch with fused activation quantization.
+
+    Each token in ``A`` is routed to one expert via ``expert_ids``. The kernel
+    quantizes ``A`` to FP8 on-the-fly (fused ``act_quant``), reads the correct
+    expert weight slice of ``B`` directly using ``expert_ids[batch_id]`` — no
+    ``(S, N, K)`` weight gather needed — and accumulates in float32.
+
+    Args:
+        A: Raw activation matrix ``[S, K]`` in bf16/fp16/fp32.
+        B: Stacked expert weight tensor ``[E, N, K]`` in ``float8_e4m3fn``.
+        Bs: Stacked expert weight scales ``[E, N // block_size[0], K // block_size[1]]``.
+        expert_ids: Expert index per token ``[S]``, values in ``[0, E)``.
+        block_size: ``[block_n, block_k]`` quantization block dimensions, e.g. ``[128, 128]``.
+
+    Returns:
+        Output tensor ``[S, N]`` in the same dtype as ``A``.
+    """
+    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_batched(A, B, Bs, expert_ids, block_size)

build/torch-xpu/finegrained_fp8/__init__.py

@@ -0,0 +1,26 @@
+import ctypes
+import sys
+
+import importlib
+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

@@ -0,0 +1,252 @@
+# 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
+
+
+@triton.jit
+def w8a8_block_fp8_grouped_mm_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 // group_n, K // group_k) weight scales
+    Offsets,  # (E,) int32 — cumulative row-end per expert
+    TileOffsets,  # (E,) int32 — cumulative tile-end per expert
+    # Shape
+    S,
+    N,
+    K,
+    # Block size for block-wise quantization
+    group_n,
+    group_k,
+    # Strides
+    stride_am,
+    stride_ak,
+    stride_Eb,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    stride_Esb,
+    stride_Bsk,
+    stride_Bsn,
+    # Meta-parameters
+    NUM_EXPERTS: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_N: tl.constexpr,
+    BLOCK_SIZE_K: tl.constexpr,
+):
+    pid = tl.program_id(axis=0)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    pid_m = pid // num_pid_n
+    pid_n = pid % num_pid_n
+
+    # 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)
+    expert_id = lo
+
+    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_bn_safe = offs_bn % N
+    offs_k = tl.arange(0, BLOCK_SIZE_K)
+
+    offs_am_safe = offs_global_m % S
+
+    a_ptrs = A + offs_am_safe[:, None] * stride_am + offs_k[None, :] * stride_ak
+    b_ptrs = (
+        B
+        + expert_id * stride_Eb
+        + offs_k[:, None] * stride_bk
+        + offs_bn_safe[None, :] * stride_bn
+    )
+    offs_bsn_safe = offs_bn_safe // group_n
+    Bs_ptrs = Bs + expert_id * stride_Esb + offs_bsn_safe * stride_Bsn
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
+        # ---- fused act_quant (replaces: a = tl.load(a_ptrs); a_s = tl.load(As_ptrs)) ----
+        a_raw = tl.load(
+            a_ptrs,
+            mask=row_mask[:, None] & (offs_k[None, :] < K - k * BLOCK_SIZE_K),
+            other=0.0,
+        ).to(tl.float32)
+        a_s = tl.max(tl.abs(a_raw), axis=1) / 448.0  # per-row scale  (BLOCK_SIZE_M,)
+        # clamp denominator so masked all-zero rows don't produce NaN
+        # (their a_s multiplier is 0 anyway, so the output row is correct)
+        a = (a_raw / tl.maximum(a_s[:, None], 1e-12)).to(tl.float8e4nv)
+        # ---- same as baseline from here ----
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        k_start = k * BLOCK_SIZE_K
+        offs_ks = k_start // group_k
+        b_s = tl.load(Bs_ptrs + offs_ks * stride_Bsk)
+
+        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] & (offs_bn[None, :] < N)
+    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] | None,
+) -> torch.Tensor:
+    """Grouped FP8 block-wise matmul with fused activation quantization.
+
+    Mirrors ``_grouped_linear`` / ``_grouped_mm`` for FP8 weights: A is the
+    raw (BF16/FP16) activation matrix sorted by expert, B / Bs are the stacked
+    expert weights / scales.  Activation quantization (``act_quant``) is fused
+    into the matmul loop.  ``tokens_per_expert`` is needed (in addition to
+    ``offsets``) to build the per-expert tile schedule inside the kernel.
+    """
+    assert A.ndim == 2, "A must be (S, K)"
+    assert A.is_contiguous()
+
+    assert B.ndim == 3, "B must be (E, N, K)"
+    assert B.is_contiguous()
+
+    assert tokens_per_expert.is_contiguous()
+    assert offsets.is_contiguous()
+    assert Bs.is_contiguous()
+
+    if block_size is None:
+        block_n, block_k = 128, 128
+    else:
+        assert len(block_size) == 2
+        block_n, block_k = block_size[0], block_size[1]
+
+    S, K = A.shape
+    E, N, _ = B.shape
+    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.
+    # Using a static upper bound keeps the grid size data-independent, which is
+    # required for cuda-graph compatibility.  Programs beyond the real tile count
+    # exit immediately via the early-return guard inside the kernel.
+    max_M_tiles = triton.cdiv(S, BLOCK_SIZE_M) + E
+
+    grid = (max_M_tiles * triton.cdiv(N, block_n),)
+    wrap_triton(w8a8_block_fp8_grouped_mm_kernel)[grid](
+        A,
+        B,
+        C,
+        Bs,
+        offsets,
+        tile_offsets,
+        S,
+        N,
+        K,
+        block_n,
+        block_k,
+        A.stride(0),  # stride_am
+        A.stride(1),  # stride_ak
+        B.stride(0),  # stride_Eb
+        B.stride(2),  # stride_bk
+        B.stride(1),  # stride_bn
+        C.stride(0),  # stride_cm
+        C.stride(1),  # stride_cn
+        Bs.stride(0),  # stride_Esb
+        Bs.stride(2),  # stride_Bsk
+        Bs.stride(1),  # stride_Bsn
+        NUM_EXPERTS=E,
+        BLOCK_SIZE_M=BLOCK_SIZE_M,
+        BLOCK_SIZE_N=block_n,
+        BLOCK_SIZE_K=block_k,
+    )
+
+    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] | None,
+) -> torch.Tensor:
+    """Grouped W8A8 FP8 matmul for MoE expert dispatch with fused activation quantization.
+
+    Tokens in ``A`` must be pre-sorted by expert id. The kernel quantizes ``A``
+    to FP8 on-the-fly (fused ``act_quant``), uses a static over-provisioned grid
+    for CUDA-graph compatibility, and resolves each tile's expert via an O(log E)
+    binary search over ``offsets``.
+
+    Args:
+        A: Raw activation matrix ``[S, K]`` sorted by expert, in bf16/fp16/fp32.
+        B: Stacked expert weight tensor ``[E, N, K]`` in ``float8_e4m3fn``.
+        Bs: Stacked expert weight scales ``[E, N // block_size[0], K // block_size[1]]``.
+        offsets: Cumulative token counts per expert ``[E]`` (i.e. ``cumsum(tokens_per_expert)``).
+        tokens_per_expert: Number of tokens routed to each expert ``[E]``.
+        block_size: ``[block_n, block_k]`` quantization block dimensions, e.g. ``[128, 128]``.
+
+    Returns:
+        Output tensor ``[S, N]`` in the same dtype as ``A``, in expert-sorted order.
+    """
+    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_grouped(
+        A, B, Bs, offsets, tokens_per_expert, block_size
+    )

build/torch-xpu/matmul.py

@@ -0,0 +1,232 @@
+# 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
+
+
+# Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/layers/quantization/fp8_kernel.py
+@triton.jit
+def w8a8_block_fp8_matmul_kernel(
+    # Pointers to inputs and output
+    A,
+    B,
+    C,
+    As,
+    Bs,
+    # Shape for matmul
+    M,
+    N,
+    K,
+    # Block size for block-wise quantization
+    group_n,
+    group_k,
+    # Stride for inputs and output
+    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,
+):
+    """Triton-accelerated function used to perform linear operations (dot
+    product) on input tensors `A` and `B` with block-wise quantization, and
+    store the result in output tensor `C`.
+    """
+
+    pid = tl.program_id(axis=0)
+    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
+    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
+    num_pid_in_group = GROUP_SIZE_M * num_pid_n
+    group_id = pid // num_pid_in_group
+    first_pid_m = group_id * GROUP_SIZE_M
+    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
+    pid_m = first_pid_m + (pid % group_size_m)
+    pid_n = (pid % num_pid_in_group) // 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 // group_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)):
+        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
+        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
+
+        k_start = k * BLOCK_SIZE_K
+        offs_ks = k_start // group_k
+        a_s = tl.load(As_ptrs + offs_ks * stride_As_k)
+        b_s = tl.load(Bs_ptrs + offs_ks * 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_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] | None,
+    output_dtype: torch.dtype = torch.float32,
+) -> torch.Tensor:
+    """This function performs matrix multiplication with block-wise
+    quantization.
+    It takes two input tensors `A` and `B` with scales `As` and `Bs`.
+    The output is returned in the specified `output_dtype`.
+    Args:
+        A: The input tensor, e.g., activation.
+        B: The input tensor, e.g., weight.
+        As: The per-token-group quantization scale for `A`.
+        Bs: The per-block quantization scale for `B`.
+        block_size: The block size for per-block quantization. It should
+        be 2-dim, e.g., [128, 128].
+        output_dtype: The dtype of the returned tensor.
+    Returns:
+        torch.Tensor: The result of matmul.
+    """
+    if block_size is None:
+        block_n, block_k = 128, 128
+    else:
+        assert len(block_size) == 2
+        block_n, block_k = block_size[0], block_size[1]
+
+    # if we have per-tensor quantization, we use 128x128 block size for tiled matmul multiplication
+    if block_n == B.shape[-2] and block_k == B.shape[-1]:
+        block_n = 128
+        block_k = 128
+
+    assert A.shape[-1] == B.shape[-1]
+    assert A.is_contiguous()
+
+    assert B.ndim == 2
+    assert B.is_contiguous()
+
+    N, K = B.shape
+    M = A.numel() // A.shape[-1]
+
+    # For per-tensor scales (scalar), expand to block-scale shape with strides (0, 0).
+    # This is a zero-copy view; all loads inside the kernel hit the same cached value.
+    if As.numel() == 1:
+        As = As.reshape(1, 1).expand(M, triton.cdiv(K, block_k))
+    else:
+        assert A.shape[:-1] == As.shape[:-1]
+        assert triton.cdiv(K, block_k) == As.shape[-1]
+    if Bs.numel() == 1:
+        Bs = Bs.reshape(1, 1).expand(triton.cdiv(N, block_n), triton.cdiv(K, block_k))
+    else:
+        assert Bs.ndim == 2
+        assert triton.cdiv(N, block_n) == Bs.shape[0], f"{N}, {block_n}, {Bs.shape}"
+        assert triton.cdiv(K, block_k) == Bs.shape[1], f"{K}, {block_k}, {Bs.shape}"
+
+    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)
+    BLOCK_SIZE_K = block_k
+    BLOCK_SIZE_N = block_n
+
+    grid = (triton.cdiv(M, BLOCK_SIZE_M) * triton.cdiv(N, BLOCK_SIZE_N),)
+    wrap_triton(w8a8_block_fp8_matmul_kernel)[grid](
+        A,
+        B,
+        C,
+        As,
+        Bs,
+        M,
+        N,
+        K,
+        block_n,
+        block_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),
+        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] | None,
+    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)

build/torch-xpu/metadata-xpu.json

@@ -0,0 +1,4 @@
+{
+  "version": 1,
+  "python-depends": []
+}