diff --git a/python/freetoken/kernel/nvfp4_moe.py b/python/freetoken/kernel/nvfp4_moe.py index a61afda56..0aa494222 100644 --- a/python/freetoken/kernel/nvfp4_moe.py +++ b/python/freetoken/kernel/nvfp4_moe.py @@ -8,6 +8,7 @@ import functools import re +from dataclasses import dataclass, replace from typing import Any, Dict, Tuple import torch @@ -174,12 +175,51 @@ def _tensor_rho(name: str, major: int, minor: int) -> float: 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)} + + +@dataclass(frozen=True) +class _PrefillTuning: + """Triton prefill GEMM tuning of one GPU class.""" + + steps: Tuple[float, float, float] # tokens per expert where BLOCK_SIZE_M becomes 32, 64, 128 + tiles: Dict[int, Tuple[int, int, int]] + pair_decode_below: float = 0.0 # tokens per expert under which FAST_DECODE stays 0 + small_tile_decode: int = 2 # FAST_DECODE of BLOCK_SIZE_M <= 32 + large_tile_decode: int = 1 # FAST_DECODE of BLOCK_SIZE_M >= 64 + scale_words: bool = False # SCALE_WORDS on BLOCK_SIZE_M 64 word-decode tiles + long_k_warps: int = 4 # num_warps of SCALE_WORDS GEMMs with K >= _LONG_K + + +_INF = float("inf") +_LONG_K = 2048 +_TUNING_H100 = _PrefillTuning((12, 24, 64), _PREFILL_TILES_H100) +_TUNING_RTX_PRO = _PrefillTuning((16, 36, 150), _PREFILL_TILES_RTX_PRO) +_TUNING_GEFORCE = _PrefillTuning((16, 150, _INF), _PREFILL_TILES_GEFORCE) +# RTX 4090 A/B: up to 32 tokens per expert the pair decode streams the weights 2-18% faster, from 51 words win 11-14%, also on 32-row tiles +_SM89_DECODE = dict(pair_decode_below=48, small_tile_decode=1) +# H100 A/B: bytes beat words at BLOCK_SIZE_M 64 / 128 by 1-17% per layer (24-31% on Gemma-4-26B-A4B's K = 704 down GEMM) +_SM90_DECODE = dict(small_tile_decode=2, large_tile_decode=2) +# RTX 4090 / 3080 Ti layer A/B: scale words + 8 warps on the long-K GEMM give 1.03-1.09x / 1.05-1.19x; with them sm_86 gains from 64-row tiles at 100 tokens per expert, sm_89 does not +_GEFORCE_SCALE_WORDS = dict(scale_words=True, long_k_warps=8) +# (rho class, compute capability) -> tuning, from sweeps on H100, RTX PRO 6000 and RTX 5090 / 4090 / 3080 Ti; (class, None) covers the rest of the class. +# The RTX tiles come from cc 8.6 / 8.9 / 12.0 sweeps (mma.sync, 99 KB shared memory); other GPUs keep the H100 tiles. +_PREFILL_TUNING = { + ("wide", (9, 0)): replace(_TUNING_H100, **_SM90_DECODE), + ("wide", None): _TUNING_H100, + ("pro", (8, 6)): _TUNING_RTX_PRO, + ("pro", (8, 9)): replace(_TUNING_RTX_PRO, **_SM89_DECODE), + ("pro", (9, 0)): replace(_TUNING_RTX_PRO, tiles=_PREFILL_TILES_H100, **_SM90_DECODE), + ("pro", (12, 0)): _TUNING_RTX_PRO, + ("pro", None): replace(_TUNING_RTX_PRO, tiles=_PREFILL_TILES_H100), + ("geforce", (8, 6)): replace(_TUNING_GEFORCE, steps=(16, 100, _INF), **_GEFORCE_SCALE_WORDS), + ("geforce", (8, 9)): replace(_TUNING_GEFORCE, **_SM89_DECODE, **_GEFORCE_SCALE_WORDS), + ("geforce", (12, 0)): _TUNING_GEFORCE, + ("geforce", None): replace(_TUNING_GEFORCE, tiles=_PREFILL_TILES_H100), +} # steps while the Gluon kernel runs BLOCK_SIZE_M >= 64, from an RTX PRO 6000 sweep; classes without an entry keep the Triton steps _PREFILL_BLOCK_M_STEPS_GLUON = ((3.0, (16, 36, 96)),) @@ -213,19 +253,29 @@ def _gluon_prefill(device_index: int): @functools.lru_cache(maxsize=None) -def _prefill_tuning(device_index: int) -> Tuple[Tuple[float, float, float], Dict[int, Tuple[int, int, int]]]: +def _device_rho(device_index: int) -> float: + props = torch.cuda.get_device_properties(device_index) + return _tensor_rho(props.name, props.major, props.minor) + + +@functools.lru_cache(maxsize=None) +def _prefill_tuning(device_index: int) -> _PrefillTuning: 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) + return _TUNING_H100 + rho = _device_rho(device_index) + rho_class = "wide" if rho >= 8.0 else "pro" if rho >= 3.0 else "geforce" + cc = torch.cuda.get_device_capability(device_index) + return _PREFILL_TUNING.get((rho_class, cc), _PREFILL_TUNING[(rho_class, None)]) + + +@functools.lru_cache(maxsize=None) +def _prefill_block_m_steps(device_index: int) -> Tuple[float, float, float]: + steps = _prefill_tuning(device_index).steps if _gluon_prefill(device_index) is not None: + rho = _device_rho(device_index) steps = next((s for min_rho, s in _PREFILL_BLOCK_M_STEPS_GLUON if rho >= min_rho), steps) - # 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 + return steps @functools.lru_cache(maxsize=None) @@ -241,31 +291,29 @@ def _prefill_use_tma(packed: torch.Tensor, scale: torch.Tensor, K: int, device_i and packed.data_ptr() % 16 == 0 and scale.data_ptr() % 16 == 0) -@functools.lru_cache(maxsize=None) -def _device_cc(device_index: int) -> Tuple[int, int]: - return torch.cuda.get_device_capability(device_index) +def _prefill_use_scale_words(scale: torch.Tensor, K: int, cfg: Dict[str, Any], device_index: int) -> bool: + # K // 16 scale bytes per row split into whole int32 words, and a K step reads at least one word + return (_prefill_tuning(device_index).scale_words and cfg["BLOCK_SIZE_M"] == 64 and cfg["BLOCK_SIZE_KB"] >= 32 + and K % 64 == 0 and scale.is_contiguous() and scale.data_ptr() % 4 == 0) def _prefill_fast_decode(block_m: int, per_expert: float, device_index: int) -> int: """FAST_DECODE of the prefill kernel: 0 keeps the fp16 pair decode (ROCm: the fast decode is PTX inline asm), 1 words, 2 bytes.""" if is_rocm(): return 0 - if _device_cc(device_index) == (8, 9): - # RTX 4090 A/B: up to 32 tokens per expert the pair decode streams the weights 2-18% faster, from 51 words win 11-14% - return 0 if per_expert < 48 else 1 - if _device_cc(device_index) == (9, 0): - # H100 A/B: bytes beat words at BLOCK_SIZE_M 64 / 128 by 1-17% per layer (24-31% on Gemma-4-26B-A4B's K = 704 down GEMM) - return 2 + tuning = _prefill_tuning(device_index) + if per_expert < tuning.pair_decode_below: + return 0 # the word decode halves the loop but stages the decoded tile through shared memory; small tiles keep the byte form - return 2 if block_m <= 32 else 1 + return tuning.small_tile_decode if block_m <= 32 else tuning.large_tile_decode 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) + to32, to64, to128 = _prefill_block_m_steps(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] + block_n, block_kb, stages = _prefill_tuning(device_index).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, FAST_DECODE=_prefill_fast_decode(block_m, per_expert, device_index)) @@ -308,6 +356,11 @@ def prefill_gemm( use_tma = _prefill_use_tma(packed, scale, K, a.device.index) # the 16-byte scale box costs registers that the BLOCK_SIZE_M >= 64 tiles do not have tma_scales = cfg["BLOCK_SIZE_M"] <= 32 + scale_words = not use_tma and fast == 1 and _prefill_use_scale_words(scale, K, cfg, a.device.index) + if scale_words: + scale = scale.view(torch.int32) + if K >= _LONG_K: + cfg["num_warps"] = _prefill_tuning(a.device.index).long_k_warps w_desc = s_desc = None if use_tma: rows = packed.shape[0] * N @@ -332,5 +385,6 @@ def prefill_gemm( USE_TMA=use_tma, TMA_SCALES=tma_scales, FAST_DECODE=fast, + SCALE_WORDS=scale_words, **cfg, ) diff --git a/python/freetoken/kernel/triton/nvfp4_fused_moe.py b/python/freetoken/kernel/triton/nvfp4_fused_moe.py index 0fa4d4aa9..61e867b16 100644 --- a/python/freetoken/kernel/triton/nvfp4_fused_moe.py +++ b/python/freetoken/kernel/triton/nvfp4_fused_moe.py @@ -385,6 +385,7 @@ def _prefill_nvfp4_moe_kernel( USE_TMA: tl.constexpr = False, # needs K % 256 == 0 TMA_SCALES: tl.constexpr = False, # with USE_TMA: scales by TMA too, else pointer loads FAST_DECODE: tl.constexpr = 0, # 0: fp16 pair decode, 1: packed_ptr / w_desc are int32 views of 8 codes, 2: bytes + SCALE_WORDS: tl.constexpr = False, # with FAST_DECODE == 1 and pointer loads: scale_ptr is an int32 view of 4 scales ): # 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. @@ -420,7 +421,13 @@ def _prefill_nvfp4_moe_kernel( p_ptrs = packed_ptr + slot * stride_pe + offs_bn[:, None] * stride_pn + offs_kw[None, :] * stride_pkb else: 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 + tl.static_assert(not SCALE_WORDS or (WORDS and not USE_TMA and BLOCK_SIZE_KB >= 32)) + SW: tl.constexpr = BLOCK_SIZE_KB // 32 + if SCALE_WORDS: + offs_sw = tl.arange(0, SW) + s_ptrs = scale_ptr + slot * stride_se + offs_bn[:, None] * stride_sn + offs_sw[None, :] * stride_sblk + else: + s_ptrs = scale_ptr + slot * stride_se + offs_bn[:, None] * stride_sn + offs_sb[None, :] * stride_sblk accumulator = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32) K_BYTES = K // 2 @@ -450,7 +457,16 @@ def _prefill_nvfp4_moe_kernel( else: byte_mask = kb * BLOCK_SIZE_KB + offs_kb < K_BYTES v = tl.load(p_ptrs, mask=byte_mask[None, :], other=0).to(tl.int32) - if e4m3_native_cx(): + if SCALE_WORDS: + # 4-byte loads can go through the cp.async pipeline, byte loads cannot; spread each word over the 8 weight words it scales + sw = tl.load(s_ptrs, mask=(kb * SW + offs_sw < K_BYTES // 32)[None, :], other=0) + sw = tl.reshape(tl.broadcast_to(sw[:, :, None], (BLOCK_SIZE_N, SW, KW // SW)), (BLOCK_SIZE_N, KW)) + s = ((sw >> ((offs_kw // 2 % 4) * 8)[None, :]) & 0xFF).to(tl.uint8) + if e4m3_native_cx(): + s = s.to(tl.float8e4nv, bitcast=True).to(tl.float32) + else: + s = e4m3_u8_to_f32(s) + elif e4m3_native_cx(): s = tl.load(s_ptrs, mask=sblk_mask[None, :], other=0.0).to(tl.float32) else: s = e4m3_u8_to_f32(tl.load(s_ptrs, mask=sblk_mask[None, :], other=0)) @@ -459,8 +475,9 @@ def _prefill_nvfp4_moe_kernel( s = (s * 3.3230699894622794e35).to(tl.bfloat16).to(tl.uint16, bitcast=True).to(tl.uint32) s = s | (s << 16) if WORDS: - # one scale per 16 values, two words - s = tl.reshape(tl.broadcast_to(s[:, :, None], (BLOCK_SIZE_N, SB, 2)), (BLOCK_SIZE_N, KW)) + if not SCALE_WORDS: + # one scale per 16 values, two words + s = tl.reshape(tl.broadcast_to(s[:, :, None], (BLOCK_SIZE_N, SB, 2)), (BLOCK_SIZE_N, KW)) w = _e2m1x8_bf16(v.to(tl.uint32, bitcast=True), s, BLOCK_SIZE_N, KW) else: # one scale per 16 values, eight bytes @@ -481,7 +498,7 @@ def _prefill_nvfp4_moe_kernel( if not USE_TMA: p_ptrs += (KW if WORDS else BLOCK_SIZE_KB) * stride_pkb if not (USE_TMA and TMA_SCALES): - s_ptrs += (BLOCK_SIZE_KB // 8) * stride_sblk + s_ptrs += (SW if SCALE_WORDS else BLOCK_SIZE_KB // 8) * stride_sblk g = tl.load(global_ptr + slot * stride_ge + offs_bn * stride_gn).to(tl.float32) if FAST_DECODE: