From 6b4086ca33cccdd7aca7ae70ea0633cd0bdfa55b Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Fri, 16 Oct 2026 19:55:40 +0000 Subject: [PATCH 1/2] perf(kernel): rewrite the triton nvfp4 moe prefill gemm --- .../kernel/triton/fp8_blockscale_moe.py | 3 +- .../kernel/triton/nvfp4_fused_moe.py | 75 ++++++++++--------- python/freetoken/moe/fused_nvfp4.py | 19 ++--- 3 files changed, 50 insertions(+), 47 deletions(-) diff --git a/python/freetoken/kernel/triton/fp8_blockscale_moe.py b/python/freetoken/kernel/triton/fp8_blockscale_moe.py index 1db54a8072..429134d23f 100644 --- a/python/freetoken/kernel/triton/fp8_blockscale_moe.py +++ b/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 diff --git a/python/freetoken/kernel/triton/nvfp4_fused_moe.py b/python/freetoken/kernel/triton/nvfp4_fused_moe.py index b1de443787..ca9b9dbe9d 100644 --- a/python/freetoken/kernel/triton/nvfp4_fused_moe.py +++ b/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) diff --git a/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py index 375d2a547a..d557cec755 100644 --- a/python/freetoken/moe/fused_nvfp4.py +++ b/python/freetoken/moe/fused_nvfp4.py @@ -234,15 +234,13 @@ 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) +def _prefill_config(M: int, top_k: int, num_experts: 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. + # BLOCK_SIZE_M follows the mean tokens per expert; the thresholds and the other fields come from an H100 sweep. + per_expert = M * top_k / num_experts + block_m = 128 if per_expert >= 64 else 64 if per_expert >= 24 else 32 if per_expert >= 12 else 16 + return dict(BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=64, BLOCK_SIZE_KB=64 if block_m == 16 else 32, + GROUP_SIZE_M=8, num_warps=4, num_stages=3) def _prefill_gemm( @@ -270,7 +268,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 +306,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) sorted_ids, expert_ids, ntpp = moe_align_block_size(topk_ids, cfg["BLOCK_SIZE_M"], num_experts) tw = topk_weights.reshape(-1).contiguous() From c4acce345dc711f42c487a94450b9179171a003b Mon Sep 17 00:00:00 2001 From: Xiaoze Fan Date: Sat, 17 Oct 2026 02:26:30 +0000 Subject: [PATCH 2/2] perf(moe): tune nvfp4 prefill tiles per gpu class GeForce runs fp32-accumulate tensor MMA at half rate, so RTX parts want smaller BLOCK_SIZE_M than H100 and their own BLOCK_N / BLOCK_KB / stages. Tables come from tile sweeps on RTX PRO 6000, RTX 5090, RTX 4090 and RTX 3080 Ti; H100-class GPUs and AMD keep the previous tuning. --- python/freetoken/moe/fused_nvfp4.py | 73 ++++++++++++++++++++++++++--- 1 file changed, 66 insertions(+), 7 deletions(-) diff --git a/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py index d557cec755..9202f5f38a 100644 --- a/python/freetoken/moe/fused_nvfp4.py +++ b/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,13 +240,66 @@ def fused_experts_decode_nvfp4_serial( ) -def _prefill_config(M: int, top_k: int, num_experts: int) -> Dict[str, int]: +# 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. - # BLOCK_SIZE_M follows the mean tokens per expert; the thresholds and the other fields come from an H100 sweep. + (to32, to64, to128), tiles = _prefill_tuning(device_index) per_expert = M * top_k / num_experts - block_m = 128 if per_expert >= 64 else 64 if per_expert >= 24 else 32 if per_expert >= 12 else 16 - return dict(BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=64, BLOCK_SIZE_KB=64 if block_m == 16 else 32, - GROUP_SIZE_M=8, num_warps=4, num_stages=3) + 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( @@ -306,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, top_k, num_experts) + 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()