perf(kernel): rewrite the triton nvfp4 moe prefill gemm

python/freetoken/kernel/triton/fp8_blockscale_moe.py

@@ -128,8 +128,7 @@ def fused_experts_decode_fp8_blockscale(
 
 
 # ======================================================================================
-# Prefill: M-tiled inline-dequant grouped GEMM (moe_align_block_size), mirrors
-# nvfp4_fused_moe._prefill_nvfp4_moe_kernel with fp8*bf16-block-scale dequant.
+# Prefill: M-tiled inline-dequant grouped GEMM (moe_align_block_size) with fp8*bf16-block-scale dequant.
 # BLOCK_SIZE_N == BLOCK_SIZE_K == 128 so each tile maps to one weight scale block.
 # ======================================================================================
 @triton.jit

python/freetoken/kernel/triton/nvfp4_fused_moe.py

@@ -292,18 +292,24 @@ def _decode_nvfp4_marlin_kernel(
     tl.store(c_ptrs, accumulator.to(compute_type), mask=(route_id < total_routes) & n_mask)
 
 
+@triton.jit
+def _e2m1_f16_x2p14(code):
+    """e2m1 code (low 4 bits) -> fp16 value * 2^-14: magnitude bits go to fp16 [11:9], the sign to [15]; fp16 subnormals line up the e == 0 codes."""
+    bits = ((code & 7) << 9) | ((code & 8) << 12)
+    return bits.to(tl.uint16).to(tl.float16, bitcast=True)
+
+
 @triton.jit
 def _prefill_nvfp4_moe_kernel(
     a_ptr,             # [M, K] activations
     packed_ptr,        # [S, N, K // 2] uint8
-    scale_ptr,         # [S, N, K // 16] fp8-e4m3
+    scale_ptr,         # [S, N, K // 16] fp8-e4m3 (uint8 view before sm_89)
     global_ptr,        # [S, N] fp16
     c_ptr,             # [num_valid_tokens, N] output (flat over M*top_k)
     topk_weights_ptr,
     sorted_token_ids_ptr,
     expert_ids_ptr,    # cache slot per M-block
     num_tokens_post_padded_ptr,
-    lut_ptr,
     N,
     K,
     EM,
@@ -314,14 +320,16 @@ def _prefill_nvfp4_moe_kernel(
     stride_ge, stride_gn,
     stride_cm, stride_cn,
     stride_tw,
-    BLOCK_SIZE_M: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,   # tokens per block, the moe_align padding
     BLOCK_SIZE_N: tl.constexpr,
-    BLOCK_SIZE_KB: tl.constexpr,
+    BLOCK_SIZE_KB: tl.constexpr,  # packed bytes per K step (2 values each)
     GROUP_SIZE_M: tl.constexpr,
     MUL_ROUTED_WEIGHT: tl.constexpr,
     top_k: tl.constexpr,
     compute_type: tl.constexpr,
 ):
+    # out^T = W @ X^T: wgmma reads its A operand (the decoded weights) from registers, while B must be in shared memory.
+    # The token block is the MMA N dimension, so a 16-token block still fills the MMA.
     pid = tl.program_id(axis=0)
     num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M)
     num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
@@ -342,51 +350,50 @@ def _prefill_nvfp4_moe_kernel(
 
     offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
     offs_kb = tl.arange(0, BLOCK_SIZE_KB)
-    a_ptrs_lo = a_ptr + (offs_token[:, None] // top_k * stride_am + (2 * offs_kb)[None, :] * stride_ak)
-    a_ptrs_hi = a_ptr + (offs_token[:, None] // top_k * stride_am + (2 * offs_kb + 1)[None, :] * stride_ak)
+    offs_k = tl.arange(0, 2 * BLOCK_SIZE_KB)
+    offs_sb = tl.arange(0, BLOCK_SIZE_KB // 8)
+    xt_ptrs = a_ptr + (offs_token[None, :] // top_k * stride_am + offs_k[:, None] * stride_ak)
 
     slot = tl.load(expert_ids_ptr + pid_m).to(tl.int64)
-    packed_base = packed_ptr + slot * stride_pe + offs_bn[None, :] * stride_pn
-    scale_base = scale_ptr + slot * stride_se + offs_bn[None, :] * stride_sn
+    p_ptrs = packed_ptr + slot * stride_pe + offs_bn[:, None] * stride_pn + offs_kb[None, :] * stride_pkb
+    s_ptrs = scale_ptr + slot * stride_se + offs_bn[:, None] * stride_sn + offs_sb[None, :] * stride_sblk
 
-    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
+    accumulator = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32)
     K_BYTES = K // 2
     for kb in range(0, tl.cdiv(K_BYTES, BLOCK_SIZE_KB)):
-        byte_idx = kb * BLOCK_SIZE_KB + offs_kb
-        byte_mask = byte_idx < K_BYTES
-
-        p_ptrs = packed_base + byte_idx[:, None] * stride_pkb
-        bytes_ = tl.load(p_ptrs, mask=byte_mask[:, None], other=0).to(tl.int32)
-        lo = bytes_ & 0xF
-        hi = (bytes_ >> 4) & 0xF
-        sblk = byte_idx // 8
-        s_ptrs = scale_base + sblk[:, None] * stride_sblk
+        byte_mask = kb * BLOCK_SIZE_KB + offs_kb < K_BYTES
+        sblk_mask = kb * (BLOCK_SIZE_KB // 8) + offs_sb < K_BYTES // 8
+        k_mask = kb * 2 * BLOCK_SIZE_KB + offs_k < K
+        v = tl.load(p_ptrs, mask=byte_mask[None, :], other=0).to(tl.int32)
         if e4m3_native_cx():
-            scale = tl.load(s_ptrs, mask=byte_mask[:, None], other=0.0).to(tl.float32)
+            s = tl.load(s_ptrs, mask=sblk_mask[None, :], other=0.0).to(tl.float32)
         else:
-            scale = e4m3_u8_to_f32(tl.load(s_ptrs, mask=byte_mask[:, None], other=0))
-        b_lo = tl.load(lut_ptr + lo) * scale  # [BLOCK_KB, BLOCK_N]
-        b_hi = tl.load(lut_ptr + hi) * scale
-
-        a_lo = tl.load(a_ptrs_lo, mask=token_mask[:, None] & byte_mask[None, :], other=0.0)
-        a_hi = tl.load(a_ptrs_hi, mask=token_mask[:, None] & byte_mask[None, :], other=0.0)
-        accumulator += tl.dot(a_lo, b_lo.to(a_lo.dtype))
-        accumulator += tl.dot(a_hi, b_hi.to(a_hi.dtype))
-
-        a_ptrs_lo += BLOCK_SIZE_KB * 2 * stride_ak
-        a_ptrs_hi += BLOCK_SIZE_KB * 2 * stride_ak
+            s = e4m3_u8_to_f32(tl.load(s_ptrs, mask=sblk_mask[None, :], other=0))
+        # one scale per 16 values; the 2^14 undoes the decode's 2^-14
+        s = tl.reshape(tl.broadcast_to((s * 16384.0)[:, :, None], (BLOCK_SIZE_N, BLOCK_SIZE_KB // 8, 8)), (BLOCK_SIZE_N, BLOCK_SIZE_KB))
+        s = s.to(tl.bfloat16)
+        # exact in bf16: e2m1 has <= 2 significant bits and e4m3 <= 4
+        w_lo = _e2m1_f16_x2p14(v & 0xF).to(tl.bfloat16) * s
+        w_hi = _e2m1_f16_x2p14(v >> 4).to(tl.bfloat16) * s
+        # interleave back to k = 2 * byte + {0, 1} so one dot covers the K step
+        w = tl.reshape(tl.join(w_lo, w_hi), (BLOCK_SIZE_N, 2 * BLOCK_SIZE_KB))
+        xt = tl.load(xt_ptrs, mask=k_mask[:, None] & token_mask[None, :], other=0.0)
+        accumulator = tl.dot(w.to(xt.dtype), xt, accumulator)
+        xt_ptrs += 2 * BLOCK_SIZE_KB * stride_ak
+        p_ptrs += BLOCK_SIZE_KB * stride_pkb
+        s_ptrs += (BLOCK_SIZE_KB // 8) * stride_sblk
 
     g = tl.load(global_ptr + slot * stride_ge + offs_bn * stride_gn).to(tl.float32)
-    accumulator = accumulator * g[None, :]
+    accumulator = accumulator * g[:, None]
 
     if MUL_ROUTED_WEIGHT:
         moe_weight = tl.load(topk_weights_ptr + offs_token * stride_tw, mask=token_mask, other=0)
-        accumulator = accumulator * moe_weight[:, None]
+        accumulator = accumulator * moe_weight[None, :]
 
     accumulator = accumulator.to(compute_type)
     offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
-    c_ptrs = c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_cn[None, :]
-    c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
+    c_ptrs = c_ptr + stride_cm * offs_token[None, :] + stride_cn * offs_cn[:, None]
+    c_mask = token_mask[None, :] & (offs_cn[:, None] < N)
     tl.store(c_ptrs, accumulator, mask=c_mask)
 
 

python/freetoken/moe/fused_nvfp4.py

@@ -7,13 +7,16 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict
+import functools
+import re
+from typing import Any, Dict, Tuple
 
 import torch
 import triton
 import triton.language as tl
 
 from freetoken.kernel import moe_sum_reduce_triton
+from freetoken.kernel.backend import is_rocm
 from freetoken.kernel.triton.e4m3_compat import e4m3_kernel_view
 from freetoken.kernel.triton.nvfp4_fused_moe import (
     _decode_nvfp4_marlin_kernel,
@@ -23,6 +26,9 @@
 )
 from freetoken.layers import gated_act_and_mul
 from freetoken.moe.fused import moe_align_block_size
+from freetoken.utils import init_logger
+
+logger = init_logger(__name__)
 
 # Decode is captured into a CUDA graph, so the config must be fixed (no triton.autotune,
 # which benchmarks at run time). Tuned offline against the NVFP4 decode kernels.
@@ -234,15 +240,66 @@ def fused_experts_decode_nvfp4_serial(
     )
 
 
-def _prefill_config(M: int) -> Dict[str, int]:
-    # ``BLOCK_SIZE_M`` is coupled to host-side ``moe_align_block_size`` (token padding),
-    # so it cannot be picked by triton.autotune; these were chosen by an offline sweep
-    # over (BLOCK_M, BLOCK_N, BLOCK_KB, num_warps, num_stages) for the MiniMax-M2 shapes.
-    if M <= 64:
-        return dict(BLOCK_SIZE_M=16, BLOCK_SIZE_N=64, BLOCK_SIZE_KB=32,
-                    GROUP_SIZE_M=1, num_warps=8, num_stages=4)
-    return dict(BLOCK_SIZE_M=32, BLOCK_SIZE_N=64, BLOCK_SIZE_KB=32,
-                GROUP_SIZE_M=8, num_warps=8, num_stages=4)
+# rho = bf16 tensor-core FLOP / fp32 ALU FLOP per SM per clock, with fp32 accumulate (the only accumulator for bf16 MMA).
+# Each M block decodes its weight tile once on the ALUs, and the MMA also runs on the padding up to BLOCK_SIZE_M.
+# High rho makes decode the cost, so big blocks win; low rho makes padded MMA the cost, so small blocks win.
+# Spec, FLOP/clk/SM: GeForce RTX 30/40/50 512 / 256 = 2, as GeForce runs fp32 accumulate at half rate.
+# RTX A, RTX Ada Generation, RTX PRO Blackwell, A10, A40, L4, L40S 1024 / 256 = 4; A100 2048 / 128 = 16; H100 4096 / 256 = 16; H20 148 / 44 TFLOP/s = 3.4.
+# Measured (cuBLAS bf16 / Triton fp32 FMA): RTX PRO 6000 4.5, RTX 5090 2.3, RTX 4090 1.9, RTX 3080 Ti 1.6.
+_PRO_GPU_RE = re.compile(
+    r"RTX\s*A\d{3,4}"                        # RTX A400 .. A6000
+    r"|RTX\s*\d{3,4}E?\s*(SFF\s*)?-?\s*Ada"  # RTX 500 .. 6000 Ada Generation
+    r"|RTX\s*PRO\s*\d{3,4}"                  # RTX PRO 500 .. 6000 Blackwell
+    r"|RTX\s*\d{4}D\b"                       # RTX 6000D
+    r"|Quadro|TITAN"
+    r"|^(NVIDIA\s+)?(GRID\s+)?(A2|A10[GM]?|A16|A40|L2|L20|L4|L40S?)\b",
+    re.I,
+)
+# also matches names without "GeForce", e.g. "NVIDIA RTX 4090"
+_GEFORCE_GPU_RE = re.compile(r"GeForce|RTX\s*[2-6]0[5-9]0(?!\d)", re.I)
+
+
+def _tensor_rho(name: str, major: int, minor: int) -> float:
+    if major not in (8, 12) or (major, minor) == (8, 0):
+        return 3.4 if re.search(r"\bH20\b", name) else 16.0
+    if _PRO_GPU_RE.search(name):
+        return 4.0
+    if _GEFORCE_GPU_RE.search(name):
+        return 2.0
+    logger.warning("unknown GPU %r: NVFP4 MoE prefill assumes full-rate fp32-accumulate tensor cores", name)
+    return 4.0
+
+
+# (min rho, tokens per expert where BLOCK_SIZE_M becomes 32, 64, 128), from sweeps on H100, RTX PRO 6000 and RTX 5090 / 4090 / 3080 Ti
+_PREFILL_BLOCK_M_STEPS = ((8.0, (12, 24, 64)), (3.0, (16, 36, 150)), (0.0, (16, 150, float("inf"))))
+# BLOCK_SIZE_M -> (BLOCK_SIZE_N, BLOCK_SIZE_KB, num_stages)
+_PREFILL_TILES_H100 = {16: (64, 64, 3), 32: (64, 32, 3), 64: (64, 32, 3), 128: (64, 32, 3)}
+_PREFILL_TILES_RTX_PRO = {16: (64, 64, 4), 32: (128, 64, 3), 64: (128, 64, 2), 128: (64, 32, 2)}
+_PREFILL_TILES_GEFORCE = {16: (64, 64, 4), 32: (128, 64, 2), 64: (128, 32, 2)}
+
+
+@functools.lru_cache(maxsize=None)
+def _prefill_tuning(device_index: int) -> Tuple[Tuple[float, float, float], Dict[int, Tuple[int, int, int]]]:
+    if is_rocm():
+        # TODO: tune for AMD GPUs; they keep the H100 tuning
+        return _PREFILL_BLOCK_M_STEPS[0][1], _PREFILL_TILES_H100
+    props = torch.cuda.get_device_properties(device_index)
+    rho = _tensor_rho(props.name, props.major, props.minor)
+    steps = next(s for min_rho, s in _PREFILL_BLOCK_M_STEPS if rho >= min_rho)
+    # the RTX tables come from cc 8.6 / 8.9 / 12.0 sweeps (mma.sync, 99 KB shared memory); other GPUs keep the H100 one
+    if (props.major, props.minor) not in ((8, 6), (8, 9), (12, 0)):
+        return steps, _PREFILL_TILES_H100
+    return steps, _PREFILL_TILES_RTX_PRO if rho >= 3.0 else _PREFILL_TILES_GEFORCE
+
+
+def _prefill_config(M: int, top_k: int, num_experts: int, device_index: int) -> Dict[str, int]:
+    # moe_align_block_size pads each expert's token count to a multiple of BLOCK_SIZE_M, so triton.autotune cannot pick it.
+    (to32, to64, to128), tiles = _prefill_tuning(device_index)
+    per_expert = M * top_k / num_experts
+    block_m = 128 if per_expert >= to128 else 64 if per_expert >= to64 else 32 if per_expert >= to32 else 16
+    block_n, block_kb, stages = tiles[block_m]
+    return dict(BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=block_n, BLOCK_SIZE_KB=block_kb,
+                GROUP_SIZE_M=8, num_warps=4, num_stages=stages)
 
 
 def _prefill_gemm(
@@ -270,7 +327,6 @@ def _prefill_gemm(
     _prefill_nvfp4_moe_kernel[grid](
         a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids,
         num_tokens_post_padded,
-        _e2m1_lut(a.device.index),
         N, K, EM, num_valid_tokens,
         a.stride(0), a.stride(1),
         packed.stride(0), packed.stride(1), packed.stride(2),
@@ -309,7 +365,7 @@ def fused_experts_nvfp4(
     two_i = gate_up_packed.shape[1]
     inter = two_i // 2
     dev, dt = hidden_states.device, hidden_states.dtype
-    cfg = _prefill_config(M)
+    cfg = _prefill_config(M, top_k, num_experts, hidden_states.device.index)
 
     sorted_ids, expert_ids, ntpp = moe_align_block_size(topk_ids, cfg["BLOCK_SIZE_M"], num_experts)
     tw = topk_weights.reshape(-1).contiguous()