Build uploaded using `kernels`.

build/torch-rocm/_ops.py

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

build/torch-rocm/act_quant.py

@@ -28,8 +28,7 @@ def _fp8_act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr):
     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)
+    y = (x / s).to(y_ptr.dtype.element_ty)
     tl.store(y_ptr + offs, y)
     tl.store(s_ptr + pid, s)
 

build/torch-rocm/batched.py

@@ -23,14 +23,12 @@ 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
+    Bs,  # (E, N // block_n, K // block_k) weight scales
     ExpertIds,  # (S,) — which expert each batch element routes to
     # Shape
+    S,
     N,
     K,
-    # Block size for block-wise quantization
-    group_n,
-    group_k,
     # Per-row strides
     stride_ak,
     stride_bk,
@@ -41,50 +39,50 @@ def w8a8_block_fp8_matmul_batched_kernel(
     # Batch / expert strides
     stride_Ab,  # stride between rows in A (one token per program)
     stride_Eb,  # stride between experts in B
-    stride_Cb,
+    stride_Cb,  # stride between rows in C (one token per program)
     stride_Esb,  # stride between experts in Bs
     # Meta-parameters
+    block_n: tl.constexpr,
+    block_k: tl.constexpr,
     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)
+    batch_id = tl.program_id(axis=0)
+    pid_n = 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)
+    # 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_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.
+    offs_bn = pid_n * block_n + tl.arange(0, block_n)
+    offs_k = tl.arange(0, block_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)
 
-    offs_bsn = offs_bn // group_n
+    offs_bsn = offs_bn // block_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)
+    accumulator = tl.zeros((BLOCK_SIZE_M, block_n), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, block_k)):
+        # ---- fused fp8_act_quant ----
+        a_raw = tl.load(a_ptrs).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 = tl.load(b_ptrs)
+        k_start = k * block_k
+        offs_ks = k_start // block_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
+        a_ptrs += block_k * stride_ak
+        b_ptrs += block_k * stride_bk
 
     if C.dtype.element_ty == tl.bfloat16:
         c = accumulator.to(tl.bfloat16)
@@ -93,12 +91,10 @@ def w8a8_block_fp8_matmul_batched_kernel(
     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)
+    offs_cn = pid_n * block_n + tl.arange(0, block_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)
+    tl.store(c_ptrs, c)
 
 
 @triton_op("finegrained_fp8::w8a8_block_fp8_matmul_batched", mutates_args=())
@@ -118,52 +114,81 @@ def _w8a8_block_fp8_matmul_batched(
     Activation quantization (``act_quant``) is fused into the matmul loop.
     """
     assert A.ndim == 2, "A must be (S, K)"
-    assert A.is_contiguous()
+    assert A.is_contiguous(), "A must be contiguous"
 
     assert B.ndim == 3, "B must be (E, N, K)"
-    assert B.is_contiguous()
+    assert B.is_contiguous(), "B must be contiguous"
 
     assert A.shape[1] == B.shape[2], "K dimension mismatch between A and B"
-    assert expert_ids.is_contiguous()
-    assert Bs.is_contiguous()
+    assert expert_ids.is_contiguous(), "expert_ids must be contiguous"
+    assert Bs.is_contiguous(), "Bs must be contiguous"
+
+    S, K = A.shape
+    E, N, K = B.shape
 
     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)
+    # if we have per-tensor quantization, we use 128x128 block size for tiled matmul multiplication
+    if block_n == N and block_k == K:
+        block_n, block_k = 128, 128
 
-    # 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)
+    # we specifically require block-aligned shapes to keep the kernel simpler
+    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})"
+
+    # For per-tensor scales, expand to block-scale shape with strides (0, 0).
+    # This is a zero-copy view; all loads for a given expert hit the same value.
+    if Bs.ndim == 1:
+        Bs = Bs.reshape(E, 1, 1).expand(E, N // block_n, K // block_k)
+    elif Bs.ndim == 3 and Bs.shape[0] == E and Bs.shape[1] == 1 and Bs.shape[2] == 1:
+        Bs = Bs.expand(E, N // block_n, K // block_k)
+    else:
+        assert Bs.ndim == 3, (
+            "Bs must be either (E,) / (E,1,1) for per-tensor scales or (E,N//block_n,K//block_k)"
+        )
+        assert Bs.shape[0] == E, (
+            f"Bs expert dim mismatch: expected {E}, got {Bs.shape[0]}"
+        )
+        assert Bs.shape[1] == N // block_n, (
+            f"Bs N-block dim mismatch: expected {N // block_n}, got {Bs.shape[1]}"
+        )
+        assert Bs.shape[2] == K // block_k, (
+            f"Bs K-block dim mismatch: expected {K // block_k}, got {Bs.shape[2]}"
+        )
 
-    grid = (triton.cdiv(N, block_n), S)
+    C = A.new_empty(S, N)
+    grid = (S, triton.cdiv(N, block_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)
     wrap_triton(w8a8_block_fp8_matmul_batched_kernel)[grid](
         A,
         B,
         C,
         Bs,
         expert_ids,
+        S,
         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
+        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
+        A.stride(0),  # stride_Ab
+        B.stride(0),  # stride_Eb
+        C.stride(0),  # stride_Cb
         Bs.stride(0),  # stride_Esb
+        # Meta-parameters
+        block_n=block_n,
+        block_k=block_k,
         BLOCK_SIZE_M=BLOCK_SIZE_M,
-        BLOCK_SIZE_N=block_n,
-        BLOCK_SIZE_K=block_k,
     )
 
     return C
@@ -193,4 +218,6 @@ def w8a8_block_fp8_matmul_batched(
     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)
+    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_batched(
+        A, B, Bs, expert_ids, block_size
+    )

build/torch-rocm/grouped.py

@@ -23,16 +23,13 @@ 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
+    Bs,  # (E, N // block_n, K // block_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,
@@ -45,15 +42,13 @@ def w8a8_block_fp8_grouped_mm_kernel(
     stride_Bsk,
     stride_Bsn,
     # Meta-parameters
-    NUM_EXPERTS: tl.constexpr,
+    block_n: tl.constexpr,
+    block_k: tl.constexpr,
     BLOCK_SIZE_M: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
+    NUM_EXPERTS: 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
+    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)
@@ -75,7 +70,11 @@ def w8a8_block_fp8_grouped_mm_kernel(
         is_left = mid_val <= pid_m
         lo = tl.where(is_left, mid + 1, lo)
         hi = tl.where(is_left, hi, mid)
-    expert_id = lo
+
+    # 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 = lo.to(tl.int64)
 
     prev_eid = tl.maximum(expert_id - 1, 0)
 
@@ -91,44 +90,33 @@ def w8a8_block_fp8_grouped_mm_kernel(
     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_bn = pid_n * block_n + tl.arange(0, block_n)
+    offs_k = tl.arange(0, block_k)
 
-    offs_am_safe = offs_global_m % S
-
-    a_ptrs = A + offs_am_safe[:, None] * stride_am + offs_k[None, :] * stride_ak
+    a_ptrs = A + offs_global_m[:, 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_bn[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)
+    offs_bsn = offs_bn // block_n
+    Bs_ptrs = Bs + expert_id * stride_Esb + offs_bsn * stride_Bsn
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, block_n), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, block_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  # 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 = tl.load(b_ptrs)
+        k_start = k * block_k
+        offs_ks = k_start // block_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
+        a_ptrs += block_k * stride_ak
+        b_ptrs += block_k * stride_bk
 
     if C.dtype.element_ty == tl.bfloat16:
         c = accumulator.to(tl.bfloat16)
@@ -138,7 +126,7 @@ def w8a8_block_fp8_grouped_mm_kernel(
         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)
+    c_mask = row_mask[:, None]
     tl.store(c_ptrs, c, mask=c_mask)
 
 
@@ -160,23 +148,51 @@ def _w8a8_block_fp8_matmul_grouped(
     ``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 A.is_contiguous(), "A must be contiguous"
 
     assert B.ndim == 3, "B must be (E, N, K)"
-    assert B.is_contiguous()
+    assert B.is_contiguous(), "B must be contiguous"
 
-    assert tokens_per_expert.is_contiguous()
-    assert offsets.is_contiguous()
-    assert Bs.is_contiguous()
+    assert A.shape[1] == B.shape[2], "K dimension mismatch between A and B"
+    assert tokens_per_expert.is_contiguous(), "tokens_per_expert must be contiguous"
+    assert offsets.is_contiguous(), "offsets must be contiguous"
+    assert Bs.is_contiguous(), "Bs must be contiguous"
+
+    S, K = A.shape
+    E, N, K = B.shape
 
     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
+    if block_n == N and block_k == K:
+        block_n, block_k = 128, 128
+
+    # we specifically require block-aligned shapes to keep the kernel simpler
+    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})"
+
+    # For per-tensor scales, expand to block-scale shape with strides (0, 0).
+    # This is a zero-copy view; all loads for a given expert hit the same value.
+    if Bs.ndim == 1:
+        Bs = Bs.reshape(E, 1, 1).expand(E, N // block_n, K // block_k)
+    elif Bs.ndim == 3 and Bs.shape[0] == E and Bs.shape[1] == 1 and Bs.shape[2] == 1:
+        Bs = Bs.expand(E, N // block_n, K // block_k)
+    else:
+        assert Bs.ndim == 3, (
+            "Bs must be either (E,) / (E,1,1) for per-tensor scales or (E,N//block_n,K//block_k)"
+        )
+        assert Bs.shape[0] == E, (
+            f"Bs expert dim mismatch: expected {E}, got {Bs.shape[0]}"
+        )
+        assert Bs.shape[1] == N // block_n, (
+            f"Bs N-block dim mismatch: expected {N // block_n}, got {Bs.shape[1]}"
+        )
+        assert Bs.shape[2] == K // block_k, (
+            f"Bs K-block dim mismatch: expected {K // block_k}, got {Bs.shape[2]}"
+        )
+
     C = A.new_empty(S, N)
 
     # Adaptive BLOCK_SIZE_M: match tile to average tokens per expert.
@@ -189,7 +205,7 @@ def _w8a8_block_fp8_matmul_grouped(
     # 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),)
+    grid = (max_M_tiles, triton.cdiv(N, block_n))
     wrap_triton(w8a8_block_fp8_grouped_mm_kernel)[grid](
         A,
         B,
@@ -200,8 +216,6 @@ def _w8a8_block_fp8_matmul_grouped(
         S,
         N,
         K,
-        block_n,
-        block_k,
         A.stride(0),  # stride_am
         A.stride(1),  # stride_ak
         B.stride(0),  # stride_Eb
@@ -212,10 +226,11 @@ def _w8a8_block_fp8_matmul_grouped(
         Bs.stride(0),  # stride_Esb
         Bs.stride(2),  # stride_Bsk
         Bs.stride(1),  # stride_Bsn
-        NUM_EXPERTS=E,
+        # Meta-parameters
+        block_n=block_n,
+        block_k=block_k,
         BLOCK_SIZE_M=BLOCK_SIZE_M,
-        BLOCK_SIZE_N=block_n,
-        BLOCK_SIZE_K=block_k,
+        NUM_EXPERTS=E,
     )
 
     return C

build/torch-rocm/{metadata-rocm.json → metadata.json}

@@ -1,4 +1,5 @@
 {
   "version": 1,
+  "license": "Apache-2.0",
   "python-depends": []
 }

build/torch-xpu/_ops.py

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

build/torch-xpu/act_quant.py

@@ -28,8 +28,7 @@ def _fp8_act_quant_kernel(x_ptr, y_ptr, s_ptr, BLOCK_SIZE: tl.constexpr):
     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)
+    y = (x / s).to(y_ptr.dtype.element_ty)
     tl.store(y_ptr + offs, y)
     tl.store(s_ptr + pid, s)
 

build/torch-xpu/batched.py

@@ -23,14 +23,12 @@ 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
+    Bs,  # (E, N // block_n, K // block_k) weight scales
     ExpertIds,  # (S,) — which expert each batch element routes to
     # Shape
+    S,
     N,
     K,
-    # Block size for block-wise quantization
-    group_n,
-    group_k,
     # Per-row strides
     stride_ak,
     stride_bk,
@@ -41,50 +39,50 @@ def w8a8_block_fp8_matmul_batched_kernel(
     # Batch / expert strides
     stride_Ab,  # stride between rows in A (one token per program)
     stride_Eb,  # stride between experts in B
-    stride_Cb,
+    stride_Cb,  # stride between rows in C (one token per program)
     stride_Esb,  # stride between experts in Bs
     # Meta-parameters
+    block_n: tl.constexpr,
+    block_k: tl.constexpr,
     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)
+    batch_id = tl.program_id(axis=0)
+    pid_n = 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)
+    # 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_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.
+    offs_bn = pid_n * block_n + tl.arange(0, block_n)
+    offs_k = tl.arange(0, block_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)
 
-    offs_bsn = offs_bn // group_n
+    offs_bsn = offs_bn // block_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)
+    accumulator = tl.zeros((BLOCK_SIZE_M, block_n), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, block_k)):
+        # ---- fused fp8_act_quant ----
+        a_raw = tl.load(a_ptrs).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 = tl.load(b_ptrs)
+        k_start = k * block_k
+        offs_ks = k_start // block_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
+        a_ptrs += block_k * stride_ak
+        b_ptrs += block_k * stride_bk
 
     if C.dtype.element_ty == tl.bfloat16:
         c = accumulator.to(tl.bfloat16)
@@ -93,12 +91,10 @@ def w8a8_block_fp8_matmul_batched_kernel(
     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)
+    offs_cn = pid_n * block_n + tl.arange(0, block_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)
+    tl.store(c_ptrs, c)
 
 
 @triton_op("finegrained_fp8::w8a8_block_fp8_matmul_batched", mutates_args=())
@@ -118,52 +114,81 @@ def _w8a8_block_fp8_matmul_batched(
     Activation quantization (``act_quant``) is fused into the matmul loop.
     """
     assert A.ndim == 2, "A must be (S, K)"
-    assert A.is_contiguous()
+    assert A.is_contiguous(), "A must be contiguous"
 
     assert B.ndim == 3, "B must be (E, N, K)"
-    assert B.is_contiguous()
+    assert B.is_contiguous(), "B must be contiguous"
 
     assert A.shape[1] == B.shape[2], "K dimension mismatch between A and B"
-    assert expert_ids.is_contiguous()
-    assert Bs.is_contiguous()
+    assert expert_ids.is_contiguous(), "expert_ids must be contiguous"
+    assert Bs.is_contiguous(), "Bs must be contiguous"
+
+    S, K = A.shape
+    E, N, K = B.shape
 
     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)
+    # if we have per-tensor quantization, we use 128x128 block size for tiled matmul multiplication
+    if block_n == N and block_k == K:
+        block_n, block_k = 128, 128
 
-    # 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)
+    # we specifically require block-aligned shapes to keep the kernel simpler
+    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})"
+
+    # For per-tensor scales, expand to block-scale shape with strides (0, 0).
+    # This is a zero-copy view; all loads for a given expert hit the same value.
+    if Bs.ndim == 1:
+        Bs = Bs.reshape(E, 1, 1).expand(E, N // block_n, K // block_k)
+    elif Bs.ndim == 3 and Bs.shape[0] == E and Bs.shape[1] == 1 and Bs.shape[2] == 1:
+        Bs = Bs.expand(E, N // block_n, K // block_k)
+    else:
+        assert Bs.ndim == 3, (
+            "Bs must be either (E,) / (E,1,1) for per-tensor scales or (E,N//block_n,K//block_k)"
+        )
+        assert Bs.shape[0] == E, (
+            f"Bs expert dim mismatch: expected {E}, got {Bs.shape[0]}"
+        )
+        assert Bs.shape[1] == N // block_n, (
+            f"Bs N-block dim mismatch: expected {N // block_n}, got {Bs.shape[1]}"
+        )
+        assert Bs.shape[2] == K // block_k, (
+            f"Bs K-block dim mismatch: expected {K // block_k}, got {Bs.shape[2]}"
+        )
 
-    grid = (triton.cdiv(N, block_n), S)
+    C = A.new_empty(S, N)
+    grid = (S, triton.cdiv(N, block_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)
     wrap_triton(w8a8_block_fp8_matmul_batched_kernel)[grid](
         A,
         B,
         C,
         Bs,
         expert_ids,
+        S,
         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
+        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
+        A.stride(0),  # stride_Ab
+        B.stride(0),  # stride_Eb
+        C.stride(0),  # stride_Cb
         Bs.stride(0),  # stride_Esb
+        # Meta-parameters
+        block_n=block_n,
+        block_k=block_k,
         BLOCK_SIZE_M=BLOCK_SIZE_M,
-        BLOCK_SIZE_N=block_n,
-        BLOCK_SIZE_K=block_k,
     )
 
     return C
@@ -193,4 +218,6 @@ def w8a8_block_fp8_matmul_batched(
     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)
+    return torch.ops.finegrained_fp8.w8a8_block_fp8_matmul_batched(
+        A, B, Bs, expert_ids, block_size
+    )

build/torch-xpu/grouped.py

@@ -23,16 +23,13 @@ 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
+    Bs,  # (E, N // block_n, K // block_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,
@@ -45,15 +42,13 @@ def w8a8_block_fp8_grouped_mm_kernel(
     stride_Bsk,
     stride_Bsn,
     # Meta-parameters
-    NUM_EXPERTS: tl.constexpr,
+    block_n: tl.constexpr,
+    block_k: tl.constexpr,
     BLOCK_SIZE_M: tl.constexpr,
-    BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_K: tl.constexpr,
+    NUM_EXPERTS: 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
+    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)
@@ -75,7 +70,11 @@ def w8a8_block_fp8_grouped_mm_kernel(
         is_left = mid_val <= pid_m
         lo = tl.where(is_left, mid + 1, lo)
         hi = tl.where(is_left, hi, mid)
-    expert_id = lo
+
+    # 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 = lo.to(tl.int64)
 
     prev_eid = tl.maximum(expert_id - 1, 0)
 
@@ -91,44 +90,33 @@ def w8a8_block_fp8_grouped_mm_kernel(
     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_bn = pid_n * block_n + tl.arange(0, block_n)
+    offs_k = tl.arange(0, block_k)
 
-    offs_am_safe = offs_global_m % S
-
-    a_ptrs = A + offs_am_safe[:, None] * stride_am + offs_k[None, :] * stride_ak
+    a_ptrs = A + offs_global_m[:, 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_bn[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)
+    offs_bsn = offs_bn // block_n
+    Bs_ptrs = Bs + expert_id * stride_Esb + offs_bsn * stride_Bsn
+
+    accumulator = tl.zeros((BLOCK_SIZE_M, block_n), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, block_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  # 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 = tl.load(b_ptrs)
+        k_start = k * block_k
+        offs_ks = k_start // block_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
+        a_ptrs += block_k * stride_ak
+        b_ptrs += block_k * stride_bk
 
     if C.dtype.element_ty == tl.bfloat16:
         c = accumulator.to(tl.bfloat16)
@@ -138,7 +126,7 @@ def w8a8_block_fp8_grouped_mm_kernel(
         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)
+    c_mask = row_mask[:, None]
     tl.store(c_ptrs, c, mask=c_mask)
 
 
@@ -160,23 +148,51 @@ def _w8a8_block_fp8_matmul_grouped(
     ``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 A.is_contiguous(), "A must be contiguous"
 
     assert B.ndim == 3, "B must be (E, N, K)"
-    assert B.is_contiguous()
+    assert B.is_contiguous(), "B must be contiguous"
 
-    assert tokens_per_expert.is_contiguous()
-    assert offsets.is_contiguous()
-    assert Bs.is_contiguous()
+    assert A.shape[1] == B.shape[2], "K dimension mismatch between A and B"
+    assert tokens_per_expert.is_contiguous(), "tokens_per_expert must be contiguous"
+    assert offsets.is_contiguous(), "offsets must be contiguous"
+    assert Bs.is_contiguous(), "Bs must be contiguous"
+
+    S, K = A.shape
+    E, N, K = B.shape
 
     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
+    if block_n == N and block_k == K:
+        block_n, block_k = 128, 128
+
+    # we specifically require block-aligned shapes to keep the kernel simpler
+    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})"
+
+    # For per-tensor scales, expand to block-scale shape with strides (0, 0).
+    # This is a zero-copy view; all loads for a given expert hit the same value.
+    if Bs.ndim == 1:
+        Bs = Bs.reshape(E, 1, 1).expand(E, N // block_n, K // block_k)
+    elif Bs.ndim == 3 and Bs.shape[0] == E and Bs.shape[1] == 1 and Bs.shape[2] == 1:
+        Bs = Bs.expand(E, N // block_n, K // block_k)
+    else:
+        assert Bs.ndim == 3, (
+            "Bs must be either (E,) / (E,1,1) for per-tensor scales or (E,N//block_n,K//block_k)"
+        )
+        assert Bs.shape[0] == E, (
+            f"Bs expert dim mismatch: expected {E}, got {Bs.shape[0]}"
+        )
+        assert Bs.shape[1] == N // block_n, (
+            f"Bs N-block dim mismatch: expected {N // block_n}, got {Bs.shape[1]}"
+        )
+        assert Bs.shape[2] == K // block_k, (
+            f"Bs K-block dim mismatch: expected {K // block_k}, got {Bs.shape[2]}"
+        )
+
     C = A.new_empty(S, N)
 
     # Adaptive BLOCK_SIZE_M: match tile to average tokens per expert.
@@ -189,7 +205,7 @@ def _w8a8_block_fp8_matmul_grouped(
     # 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),)
+    grid = (max_M_tiles, triton.cdiv(N, block_n))
     wrap_triton(w8a8_block_fp8_grouped_mm_kernel)[grid](
         A,
         B,
@@ -200,8 +216,6 @@ def _w8a8_block_fp8_matmul_grouped(
         S,
         N,
         K,
-        block_n,
-        block_k,
         A.stride(0),  # stride_am
         A.stride(1),  # stride_ak
         B.stride(0),  # stride_Eb
@@ -212,10 +226,11 @@ def _w8a8_block_fp8_matmul_grouped(
         Bs.stride(0),  # stride_Esb
         Bs.stride(2),  # stride_Bsk
         Bs.stride(1),  # stride_Bsn
-        NUM_EXPERTS=E,
+        # Meta-parameters
+        block_n=block_n,
+        block_k=block_k,
         BLOCK_SIZE_M=BLOCK_SIZE_M,
-        BLOCK_SIZE_N=block_n,
-        BLOCK_SIZE_K=block_k,
+        NUM_EXPERTS=E,
     )
 
     return C

build/torch-xpu/{metadata-xpu.json → metadata.json}

@@ -1,4 +1,5 @@
 {
   "version": 1,
+  "license": "Apache-2.0",
   "python-depends": []
 }