diff --git a/python/freetoken/kernel/gluon/__init__.py b/python/freetoken/kernel/gluon/__init__.py new file mode 100644 index 000000000..eaf4365ed --- /dev/null +++ b/python/freetoken/kernel/gluon/__init__.py @@ -0,0 +1,35 @@ +"""Gluon versions of Triton kernels for specific GPUs, written against the Gluon API of Triton :data:`TRITON_VERSION`. + +Callers route to a Gluon kernel only when :func:`available` is true and the kernel supports the device and inputs; otherwise the Triton kernel runs. +""" + +from __future__ import annotations + +import functools + +from freetoken.kernel.backend import is_rocm +from freetoken.kernel.utils import _env_enabled +from freetoken.utils import init_logger + +logger = init_logger(__name__) + +DISABLE_GLUON_ENV = "FREETOKEN_DISABLE_GLUON" +# Gluon is experimental and its API changes between Triton releases +TRITON_VERSION = (3, 8) + + +@functools.cache +def available() -> bool: + """True on CUDA with Triton :data:`TRITON_VERSION`, unless FREETOKEN_DISABLE_GLUON is set; call it only for a GPU that has a Gluon kernel, as another Triton version logs a warning.""" + if _env_enabled(DISABLE_GLUON_ENV) or is_rocm(): + return False + import triton + from packaging.version import Version + + if Version(triton.__version__).release[:2] == TRITON_VERSION: + return True + logger.warning( + "Gluon kernels need Triton %s, found %s; this GPU runs the slower Triton kernels instead (%s=1 silences this)", + ".".join(map(str, TRITON_VERSION)), triton.__version__, DISABLE_GLUON_ENV, + ) + return False diff --git a/python/freetoken/kernel/gluon/nvfp4/moe/sm120.py b/python/freetoken/kernel/gluon/nvfp4/moe/sm120.py new file mode 100644 index 000000000..0455b237a --- /dev/null +++ b/python/freetoken/kernel/gluon/nvfp4/moe/sm120.py @@ -0,0 +1,354 @@ +"""NVFP4 W4A16 MoE prefill GEMM for sm_120: a loader warp (TMA weights, cp.async scales and tokens), decode warps and mma.sync warps. + +Same contract as the Triton prefill GEMM; :mod:`freetoken.kernel.nvfp4_moe` picks it for BLOCK_SIZE_M 64 and 128 and sets its constexprs. +""" + +import torch +import triton +from triton.experimental import gluon +from triton.experimental.gluon import language as gl +from triton.experimental.gluon.language.nvidia.ampere import async_copy, mma_v2 +from triton.experimental.gluon.language.nvidia.hopper import fence_async_shared, mbarrier, tma +from triton.experimental.gluon.nvidia.hopper import TensorDescriptor + +from ..utils import e2m1x8_bf16x2, scale_bf16x2 + + +@gluon.jit +def _tile(t, num_pid_m, num_pid_n, GROUP_SIZE_M: gl.constexpr): + num_pid_in_group = GROUP_SIZE_M * num_pid_n + first_pid_m = (t // num_pid_in_group) * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((t % num_pid_in_group) % group_size_m) + pid_n = (t % num_pid_in_group) // group_size_m + return pid_m, pid_n + + +@gluon.jit +def _load_partition(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles, num_pid_m, num_valid, + stride_am, ks4, n_rows, top_k: gl.constexpr, K_TAIL: gl.constexpr, GROUP_SIZE_M: gl.constexpr, + PERSISTENT: gl.constexpr): + BN: gl.constexpr = bufs[1].type.shape[1] + SPC: gl.constexpr = 256 // (bufs[1].type.shape[2] * 2) + if PERSISTENT: + # it / jt count chunks and scale blocks across tiles, so the rings keep their phase from one tile to the next + it = 0 + jt = 0 + for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)): + pid_m, pid_n = _tile(t, num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M) + _load_tile(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, pid_m, pid_n, N, nk, it, jt, num_valid, + stride_am, ks4, n_rows, top_k, K_TAIL) + it += nk + jt += gl.cdiv(nk, SPC) + else: + pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M) + _load_tile(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, pid_m, pid_n, N, nk, 0, 0, num_valid, + stride_am, ks4, n_rows, top_k, K_TAIL) + + +@gluon.jit +def _load_tile(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, pid_m, pid_n, N, nk, it, jt, num_valid, + stride_am, ks4, n_rows, top_k: gl.constexpr, K_TAIL: gl.constexpr): + x_bufs, w_bufs, s_bufs, d_bufs = bufs + x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready = bars + KS: gl.constexpr = w_bufs.type.shape[2] // 8 + XS: gl.constexpr = x_bufs.type.shape[0] + BM: gl.constexpr = x_bufs.type.shape[1] // KS + BK: gl.constexpr = KS * 16 + WS: gl.constexpr = w_bufs.type.shape[0] + KB: gl.constexpr = w_bufs.type.shape[2] + SS: gl.constexpr = s_bufs.type.shape[0] + BN: gl.constexpr = s_bufs.type.shape[1] + SPC: gl.constexpr = 256 // BK + + # 16 bytes per cp.async; adjacent lanes read one token's contiguous K + NWL: gl.constexpr = gl.num_warps() + xl: gl.constexpr = gl.BlockedLayout([1, 1, 8], [KS, 16 // KS, 2], [1, NWL, 1], [2, 0, 1]) + offs_s = gl.arange(0, KS, layout=gl.SliceLayout(1, gl.SliceLayout(2, xl))) + offs_c = gl.arange(0, 16, layout=gl.SliceLayout(0, gl.SliceLayout(1, xl))) + cols = gl.expand_dims(gl.expand_dims(offs_s * 16, 1), 2) + gl.expand_dims(gl.expand_dims(offs_c, 0), 1) + + # scale rows are K / 16 bytes: 16-byte copies when K % 256 == 0, else 4-byte copies zeroed past the row end + SV: gl.constexpr = 4 if K_TAIL else 16 + sll: gl.constexpr = gl.BlockedLayout([1, SV], [32 * SV // 16, 16 // SV], [NWL, 1], [1, 0]) + s_cols = gl.arange(0, 16, layout=gl.SliceLayout(0, sll)) + # both forms equal K // 16; the second tells the compiler the row is a multiple of 16 bytes + row_len = ks4 * 4 if K_TAIL else (ks4 // 4) * 16 + + w_row = gl.load(eid_ptr + pid_m) * N + pid_n * BN + offs_m = pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, gl.SliceLayout(2, xl))) + tok = gl.load(ids_ptr + offs_m) + tok_mask = gl.expand_dims(gl.expand_dims(tok < num_valid, 0), 2) + rows = gl.expand_dims(gl.expand_dims((tok // top_k).to(gl.int64) * stride_am, 0), 2) + x_ptrs = a_ptr + rows + cols + s_rows = w_row + gl.arange(0, BN, layout=gl.SliceLayout(1, sll)) + s_ptrs = scale_ptr + s_rows.to(gl.int64)[:, None] * row_len + s_cols[None, :] + s_row_ok = (s_rows < n_rows)[:, None] + for kk in range(nk): + if kk % SPC == 0: + jl = kk // SPC + j = jt + jl + ss = j % SS + mbarrier.wait(s_empty.index(ss), (j // SS & 1) ^ 1) + async_copy.async_load(s_bufs.index(ss), s_ptrs + jl * 16, mask=s_row_ok & (jl * 16 + s_cols < row_len)[None, :]) + async_copy.mbarrier_arrive(s_ready.index(ss)) + mbarrier.arrive(s_ready.index(ss)) + i = it + kk + ws = i % WS + mbarrier.wait(w_empty.index(ws), (i // WS & 1) ^ 1) + mbarrier.expect(w_ready.index(ws), w_desc.block_type.nbytes) + tma.async_load(w_desc, [w_row, kk * KB], w_ready.index(ws), w_bufs.index(ws)) + xs = i % XS + mbarrier.wait(x_empty.index(xs), (i // XS & 1) ^ 1) + async_copy.async_load(x_bufs.index(xs).reshape([KS, BM, 16]), x_ptrs + kk * BK, mask=tok_mask) + async_copy.mbarrier_arrive(x_ready.index(xs)) + mbarrier.arrive(x_ready.index(xs)) + + +@gluon.jit +def _decode_partition(a_ptr, ids_ptr, eid_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles, num_pid_m, num_valid, + stride_am, ks4, n_rows, top_k: gl.constexpr, K_TAIL: gl.constexpr, GROUP_SIZE_M: gl.constexpr, + PERSISTENT: gl.constexpr): + SPC: gl.constexpr = 256 // (bufs[1].type.shape[2] * 2) + if PERSISTENT: + it = 0 + jt = 0 + for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)): + _decode_tile(bufs, bars, nk, it, jt, K_TAIL) + it += nk + jt += gl.cdiv(nk, SPC) + else: + _decode_tile(bufs, bars, nk, 0, 0, K_TAIL) + + +@gluon.jit +def _decode_tile(bufs, bars, nk, it, jt, K_TAIL: gl.constexpr): + x_bufs, w_bufs, s_bufs, d_bufs = bufs + x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready = bars + WS: gl.constexpr = w_bufs.type.shape[0] + BN: gl.constexpr = w_bufs.type.shape[1] + KB: gl.constexpr = w_bufs.type.shape[2] + BK: gl.constexpr = KB * 2 + SS: gl.constexpr = s_bufs.type.shape[0] + SB: gl.constexpr = BK // 16 + DS: gl.constexpr = d_bufs.type.shape[0] + SPC: gl.constexpr = 256 // BK + NW: gl.constexpr = gl.num_warps() + + # one thread per row, so the k16 slice stores of a warp hit 32 consecutive rows + ql: gl.constexpr = gl.BlockedLayout([1, 2 * SB], [32, 1], [NW, 1], [1, 0]) + q3l: gl.constexpr = gl.BlockedLayout([1, SB, 2], [32, 1, 1], [NW, 1, 1], [2, 1, 0]) + qsl: gl.constexpr = gl.SliceLayout(2, q3l) + w32_layout: gl.constexpr = gl.NVMMASharedLayout(swizzle_byte_width=w_bufs.type.layout.swizzle_byte_width, + element_bitwidth=32, rank=2) + d32_layout: gl.constexpr = gl.NVMMASharedLayout(swizzle_byte_width=32, element_bitwidth=32, rank=2) + w32 = w_bufs._reinterpret(gl.int32, [WS, BN, KB // 4], w32_layout) + d32 = d_bufs._reinterpret(gl.int32, [DS, SB * BN, 8], d32_layout) + + # with K % 256 != 0 the last scale block covers fewer than SPC chunks + for jb in range(gl.cdiv(nk, SPC)): + j = jt + jb + ss = j % SS + mbarrier.wait(s_ready.index(ss), j // SS & 1) + for c in gl.static_range(SPC): + kk = jb * SPC + c + if not K_TAIL or kk < nk: + i = it + kk + ws = i % WS + mbarrier.wait(w_ready.index(ws), i // WS & 1) + q = w32.index(ws).load(ql) + s = s_bufs.index(ss).slice(c * SB, SB, dim=1).load(qsl) + fence_async_shared() + mbarrier.arrive(w_empty.index(ws)) + q = gl.convert_layout(gl.reshape(q, [BN, SB, 2]), q3l, assert_trivial=True).to(gl.uint32, bitcast=True) + w0, w1, w2, w3 = e2m1x8_bf16x2(q, gl.expand_dims(scale_bf16x2(s), 2)) + # join appends the new axis last, so this flattens to w0, w1, w2, w3; then k16 slice major for the MMA warps + w = gl.reshape(gl.join(gl.join(w0, w2), gl.join(w1, w3)), [BN, SB, 8]) + w = gl.permute(w, (1, 0, 2)).to(gl.int32, bitcast=True) + + ds = i % DS + # the MMA warps release x and d slots together on x_empty + mbarrier.wait(x_empty.index(ds), (i // DS & 1) ^ 1) + d32.index(ds).reshape([SB, BN, 8]).store(w) + mbarrier.arrive(d_ready.index(ds)) + fence_async_shared() + mbarrier.arrive(s_empty.index(ss)) + + +@gluon.jit +def _mma_chunks(bufs, bars, acc, it, nk, al: gl.constexpr, bl: gl.constexpr, BN: gl.constexpr, BM: gl.constexpr): + x_bufs, w_bufs, s_bufs, d_bufs = bufs + x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready = bars + KS: gl.constexpr = w_bufs.type.shape[2] // 8 + XS: gl.constexpr = x_bufs.type.shape[0] + for kk in range(nk): + i = it + kk + xs = i % XS + mbarrier.wait(d_ready.index(xs), i // XS & 1) + mbarrier.wait(x_ready.index(xs), i // XS & 1) + for k16 in gl.static_range(KS): + a = d_bufs.index(xs).slice(k16 * BN, BN).load(al) + b = x_bufs.index(xs).slice(k16 * BM, BM).permute((1, 0)).load(bl) + acc = mma_v2(a, b, acc) + mbarrier.arrive(x_empty.index(xs)) + return acc + + +@gluon.jit +def _mma_partition(bufs, bars, ids_ptr, eid_ptr, global_ptr, c_ptr, topk_weights_ptr, N, nk, num_tiles, num_pid_m, + num_valid, stride_ge, stride_cm, stride_tw, mma: gl.constexpr, BN: gl.constexpr, BM: gl.constexpr, + MUL_ROUTED_WEIGHT: gl.constexpr, GROUP_SIZE_M: gl.constexpr, PERSISTENT: gl.constexpr): + al: gl.constexpr = gl.DotOperandLayout(0, mma, 2) + bl: gl.constexpr = gl.DotOperandLayout(1, mma, 2) + if PERSISTENT: + it = 0 + for t in range(gl.program_id(0), num_tiles, gl.num_programs(0)): + acc = _mma_chunks(bufs, bars, gl.zeros([BN, BM], gl.float32, mma), it, nk, al, bl, BN, BM) + it += nk + # the epilogue runs here while the loader and decode warps already fill the rings for the next tile + # stored straight from the MMA layout, as a convert_layout needs scratch shared memory on top of the rings + pid_m, pid_n = _tile(t, num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M) + slot = gl.load(eid_ptr + pid_m) + offs_n = pid_n * BN + gl.arange(0, BN, layout=gl.SliceLayout(1, mma)) + offs_t = pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, mma)) + _epilogue(acc, ids_ptr, global_ptr, c_ptr, topk_weights_ptr, slot, offs_n, offs_t, N, num_valid, + stride_ge, stride_cm, stride_tw, MUL_ROUTED_WEIGHT) + else: + return _mma_chunks(bufs, bars, gl.zeros([BN, BM], gl.float32, mma), 0, nk, al, bl, BN, BM) + + +@gluon.jit +def _epilogue(acc, ids_ptr, global_ptr, c_ptr, topk_weights_ptr, slot, offs_n, offs_t, N, num_valid, + stride_ge, stride_cm, stride_tw, MUL_ROUTED_WEIGHT: gl.constexpr): + route = gl.load(ids_ptr + offs_t) + route_mask = route < num_valid + n_mask = offs_n < N + g = gl.load(global_ptr + slot.to(gl.int64) * stride_ge + offs_n, mask=n_mask, other=0.0).to(gl.float32) + # 256 undoes the 2^-8 that the decode leaves in every weight + acc = acc * (g * 256.0)[:, None] + if MUL_ROUTED_WEIGHT: + wt = gl.load(topk_weights_ptr + route * stride_tw, mask=route_mask, other=0.0) + acc = acc * wt[None, :] + c_ptrs = c_ptr + route.to(gl.int64)[None, :] * stride_cm + offs_n[:, None] + gl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty), mask=route_mask[None, :] & n_mask[:, None]) + + +SCALE_SMEM_LAYOUT = gl.constexpr(gl.NVMMASharedLayout(swizzle_byte_width=0, element_bitwidth=8, rank=2)) + + +@gluon.jit +def _prefill_nvfp4_moe_kernel( + a_ptr, w_desc, scale_ptr, global_ptr, c_ptr, topk_weights_ptr, + sorted_token_ids_ptr, expert_ids_ptr, num_tokens_post_padded_ptr, + N, K, EM, num_valid_tokens, ks4, n_rows, + stride_am, stride_ge, stride_cm, stride_tw, + BLOCK_SIZE_M: gl.constexpr, + BLOCK_SIZE_N: gl.constexpr, + BLOCK_SIZE_K: gl.constexpr, + GROUP_SIZE_M: gl.constexpr, + STAGES: gl.constexpr, + W_STAGES: gl.constexpr, + PERSISTENT: gl.constexpr, + K_TAIL: gl.constexpr, + DECODE_WARPS: gl.constexpr, + MMA_WARPS_N: gl.constexpr, + LOAD_REGS: gl.constexpr, + DECODE_REGS: gl.constexpr, + MUL_ROUTED_WEIGHT: gl.constexpr, + top_k: gl.constexpr, +): + BM: gl.constexpr = BLOCK_SIZE_M + BN: gl.constexpr = BLOCK_SIZE_N + BK: gl.constexpr = BLOCK_SIZE_K + live_m = gl.cdiv(gl.load(num_tokens_post_padded_ptr), BM) + if PERSISTENT: + num_pid_m = live_m + else: + # the grid covers the most m-blocks moe_align can produce; each CTA takes the tile of its index + num_pid_m = gl.cdiv(EM, BM) + pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M) + if pid_m >= live_m: + return + num_tiles = num_pid_m * gl.cdiv(N, BN) + nk = K // BK + + # x and the decoded weights are stored as k16 slices, [stage, (k // 16) * rows, 16], and share STAGES slots + x_layout: gl.constexpr = gl.NVMMASharedLayout.get_default_for([BM, 16], gl.bfloat16) + d_layout: gl.constexpr = gl.NVMMASharedLayout.get_default_for([BN, 16], gl.bfloat16) + x_bufs = gl.allocate_shared_memory(gl.bfloat16, [STAGES, (BK // 16) * BM, 16], x_layout) + w_bufs = gl.allocate_shared_memory(gl.uint8, [W_STAGES, BN, BK // 2], w_desc.layout) + s_bufs = gl.allocate_shared_memory(gl.uint8, [2, BN, 16], SCALE_SMEM_LAYOUT) + d_bufs = gl.allocate_shared_memory(gl.bfloat16, [STAGES, (BK // 16) * BN, 16], d_layout) + bufs = (x_bufs, w_bufs, s_bufs, d_bufs) + + x_ready = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout()) + x_empty = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout()) + w_ready = gl.allocate_shared_memory(gl.int64, [W_STAGES, 1], mbarrier.MBarrierLayout()) + w_empty = gl.allocate_shared_memory(gl.int64, [W_STAGES, 1], mbarrier.MBarrierLayout()) + s_ready = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout()) + s_empty = gl.allocate_shared_memory(gl.int64, [2, 1], mbarrier.MBarrierLayout()) + d_ready = gl.allocate_shared_memory(gl.int64, [STAGES, 1], mbarrier.MBarrierLayout()) + for i in gl.static_range(STAGES): + mbarrier.init(x_ready.index(i), count=1) + mbarrier.init(x_empty.index(i), count=1) + mbarrier.init(d_ready.index(i), count=1) + for i in gl.static_range(W_STAGES): + mbarrier.init(w_ready.index(i), count=1) + mbarrier.init(w_empty.index(i), count=1) + for i in gl.static_range(2): + mbarrier.init(s_ready.index(i), count=1) + mbarrier.init(s_empty.index(i), count=1) + bars = (x_ready, x_empty, w_ready, w_empty, s_ready, s_empty, d_ready) + + NW: gl.constexpr = gl.num_warps() + mma: gl.constexpr = gl.NVMMADistributedLayout(version=[2, 0], warps_per_cta=[MMA_WARPS_N, NW // MMA_WARPS_N], + instr_shape=[16, 8]) + # tuples go inline: a tuple stored in a variable turns its constexprs into runtime values + acc = gl.warp_specialize([ + (_mma_partition, (bufs, bars, sorted_token_ids_ptr, expert_ids_ptr, global_ptr, c_ptr, topk_weights_ptr, N, nk, + num_tiles, num_pid_m, num_valid_tokens, stride_ge, stride_cm, stride_tw, mma, BN, BM, + MUL_ROUTED_WEIGHT, GROUP_SIZE_M, PERSISTENT)), + (_load_partition, (a_ptr, sorted_token_ids_ptr, expert_ids_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles, + num_pid_m, num_valid_tokens, stride_am, ks4, n_rows, top_k, K_TAIL, GROUP_SIZE_M, PERSISTENT)), + (_decode_partition, (a_ptr, sorted_token_ids_ptr, expert_ids_ptr, w_desc, scale_ptr, bufs, bars, N, nk, num_tiles, + num_pid_m, num_valid_tokens, stride_am, ks4, n_rows, top_k, K_TAIL, GROUP_SIZE_M, PERSISTENT)), + ], [1, DECODE_WARPS], [LOAD_REGS, DECODE_REGS]) + + if not PERSISTENT: + # a coalesced store layout; the rings are idle by now, so the conversion's scratch fits + out_l: gl.constexpr = gl.BlockedLayout([8, 1], [BN // 8, 32 * 8 // BN], [1, NW], [0, 1]) + pid_m, pid_n = _tile(gl.program_id(0), num_pid_m, gl.cdiv(N, BN), GROUP_SIZE_M) + slot = gl.load(expert_ids_ptr + pid_m) + offs_n = pid_n * BN + gl.arange(0, BN, layout=gl.SliceLayout(1, out_l)) + offs_t = pid_m * BM + gl.arange(0, BM, layout=gl.SliceLayout(0, out_l)) + _epilogue(gl.convert_layout(acc, out_l), sorted_token_ids_ptr, global_ptr, c_ptr, topk_weights_ptr, slot, offs_n, + offs_t, N, num_valid_tokens, stride_ge, stride_cm, stride_tw, MUL_ROUTED_WEIGHT) + + +def supports(a: torch.Tensor, packed: torch.Tensor, scale: torch.Tensor, c: torch.Tensor) -> bool: + """Whether :func:`prefill_gemm` can run this GEMM on an sm_120 GPU.""" + K = packed.shape[2] * 2 + # the TMA weight loads and 16-byte cp.async need contiguous, 16-byte aligned rows + return (a.dtype == torch.bfloat16 and K % 64 == 0 + and packed.is_contiguous() and scale.is_contiguous() and c.is_contiguous() + and packed.data_ptr() % 16 == 0 and scale.data_ptr() % 16 == 0 + and a.stride(1) == 1 and a.stride(0) % 8 == 0 and a.data_ptr() % 16 == 0) + + +def prefill_gemm(a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded, + num_valid_tokens, top_k, mul_routed_weight, cfg): + """Launch the kernel with the tile / pipeline constexprs in ``cfg``; K must be a multiple of cfg["BLOCK_SIZE_K"].""" + S, N, KB = packed.shape + K = KB * 2 + BM, BN, BK = cfg["BLOCK_SIZE_M"], cfg["BLOCK_SIZE_N"], cfg["BLOCK_SIZE_K"] + w_layout = gl.NVMMASharedLayout.get_default_for([BN, BK // 2], gl.uint8) + w_desc = TensorDescriptor.from_tensor(packed.view(S * N, KB), [BN, BK // 2], w_layout) + max_tiles = triton.cdiv(sorted_ids.shape[0], BM) * triton.cdiv(N, BN) + if cfg["PERSISTENT"]: + max_tiles = min(max_tiles, torch.cuda.get_device_properties(a.device).multi_processor_count) + _prefill_nvfp4_moe_kernel[(max_tiles,)]( + a, w_desc, scale.view(torch.uint8), glob, c, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded, + N, K, sorted_ids.shape[0], num_valid_tokens, K // 64, S * N, + a.stride(0), glob.stride(0), c.stride(1), topk_weights_flat.stride(0), + K_TAIL=K % 256 != 0, MUL_ROUTED_WEIGHT=mul_routed_weight, top_k=top_k, **cfg, + ) diff --git a/python/freetoken/kernel/gluon/nvfp4/utils.py b/python/freetoken/kernel/gluon/nvfp4/utils.py new file mode 100644 index 000000000..430c146d4 --- /dev/null +++ b/python/freetoken/kernel/gluon/nvfp4/utils.py @@ -0,0 +1,43 @@ +"""NVFP4 helpers shared by the Gluon kernels.""" + +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +@gluon.jit +def bf16x2_mul(a, b): + # sm_90+; sm_80-89 need fma.rn.bf16x2 with a -0.0 addend + return gl.inline_asm_elementwise("mul.rn.bf16x2 $0, $1, $2;", "=r,r,r", [a, b], gl.uint32, True, 1) + + +@gluon.jit +def prmt(a, b, sel: gl.constexpr): + return gl.inline_asm_elementwise(f"prmt.b32 $0, $1, $2, {sel};", "=r,r,r", [a, b], gl.uint32, True, 1) + + +@gluon.jit +def scale_bf16x2(s): + """e4m3 scale bytes -> bf16x2 words holding scale * 2^118 twice.""" + # 2^118 * the 2^-126 of e2m1x8_bf16x2 leaves 2^-8 for the caller to undo; every step is exact in bf16 + s = (s.to(gl.float8e4nv, bitcast=True).to(gl.float32) * 3.3230699894622794e35).to(gl.bfloat16) + s = s.to(gl.uint16, bitcast=True).to(gl.uint32) + return s | (s << 16) + + +@gluon.jit +def e2m1x8_bf16x2(q, s2): + """8 e2m1 codes in a uint32 -> 4 bf16x2 words in k order, each value times the bf16x2 scale s2 and 2^-126.""" + # each 16-bit half's top nibble lands in a bf16 as sign + exponent/mantissa bits, a value of e2m1 * 2^-126 + o0 = (q & 0x80008000) | ((q >> 6) & 0x01C001C0) + q1 = q << 4 + o1 = (q1 & 0x80008000) | ((q1 >> 6) & 0x01C001C0) + q2 = q << 8 + o2 = (q2 & 0x80008000) | ((q2 >> 6) & 0x01C001C0) + q3 = q << 12 + o3 = (q3 & 0x80008000) | ((q3 >> 6) & 0x01C001C0) + # o0..o3 hold the k pairs (3, 7), (2, 6), (1, 5), (0, 4) + w0 = bf16x2_mul(prmt(o3, o2, 0x5410), s2) + w1 = bf16x2_mul(prmt(o1, o0, 0x5410), s2) + w2 = bf16x2_mul(prmt(o3, o2, 0x7632), s2) + w3 = bf16x2_mul(prmt(o1, o0, 0x7632), s2) + return w0, w1, w2, w3 diff --git a/python/freetoken/kernel/nvfp4_moe.py b/python/freetoken/kernel/nvfp4_moe.py new file mode 100644 index 000000000..a61afda56 --- /dev/null +++ b/python/freetoken/kernel/nvfp4_moe.py @@ -0,0 +1,336 @@ +"""Kernel dispatch for the inline-dequant NVFP4 MoE GEMMs: which kernel runs, with which tiles and switches, and its launch. + +The kernels are templates (:mod:`freetoken.kernel.triton.nvfp4_fused_moe`, :mod:`freetoken.kernel.gluon.nvfp4.moe.sm120`); +:mod:`freetoken.moe.fused_nvfp4` strings the GEMMs into the MoE layer. +""" + +from __future__ import annotations + +import functools +import re +from typing import Any, Dict, Tuple + +import torch +import triton +import triton.language as tl +from triton.tools.tensor_descriptor import TensorDescriptor + +from freetoken.kernel.backend import is_rocm +from freetoken.kernel.triton.e4m3_compat import e4m3_kernel_view, e4m3_native +from freetoken.kernel.triton.nvfp4_fused_moe import ( + _decode_nvfp4_marlin_kernel, + _decode_nvfp4_moe_kernel, + _e2m1_lut, + _prefill_nvfp4_moe_kernel, +) +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. +# These drive the original LUT-gather decode (decode_gemm), kept only for A/B. +_DECODE_BLOCK_N = 64 +_DECODE_BLOCK_KB = 128 +_DECODE_WARPS = 4 + +# Marlin-style decode: every K-loop load must keep the word tile's register layout, so BLOCK_N * BLOCK_KW <= 128 * num_warps and BLOCK_KW <= 32 at 4 warps. +# Long K (gate_up) uses the narrow N tile for more programs at M=1; short K (down) uses the wide one. +_DECODE_MARLIN_BLOCK_N = 16 +_DECODE_MARLIN_BLOCK_KW = 16 +_DECODE_MARLIN_UNROLL = 2 +_DECODE_MARLIN_WARPS = 4 +_DECODE_MARLIN_DEEPK_BLOCK_N = 8 +_DECODE_MARLIN_DEEPK_BLOCK_KW = 32 +_DECODE_MARLIN_DEEPK_UNROLL = 2 +_DECODE_MARLIN_DEEPK_THRESHOLD = 2048 + + +def _tl_dtype(dt: torch.dtype): + if dt == torch.bfloat16: + return tl.bfloat16 + if dt == torch.float16: + return tl.float16 + return tl.float32 + + +def decode_gemm( + a: torch.Tensor, + packed: torch.Tensor, + scale: torch.Tensor, + glob: torch.Tensor, + c: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + mul_routed_weight: bool, + a_row_is_route: bool, +) -> None: + M, top_k = topk_ids.shape + N = packed.shape[1] + K = packed.shape[2] * 2 + scale = e4m3_kernel_view(scale) + total_routes = M * top_k + grid = (total_routes, triton.cdiv(N, _DECODE_BLOCK_N)) + _decode_nvfp4_moe_kernel[grid]( + a, packed, scale, glob, c, topk_weights, topk_ids, + _e2m1_lut(a.device.index), + total_routes, N, K, + a.stride(0), a.stride(1), + packed.stride(0), packed.stride(1), packed.stride(2), + scale.stride(0), scale.stride(1), scale.stride(2), + glob.stride(0), glob.stride(1), + c.stride(0), c.stride(1), c.stride(2), + topk_weights.stride(0), topk_weights.stride(1), + topk_ids.stride(0), topk_ids.stride(1), + BLOCK_SIZE_N=_DECODE_BLOCK_N, + BLOCK_SIZE_KB=_DECODE_BLOCK_KB, + TOP_K=top_k, + A_ROW_IS_ROUTE=a_row_is_route, + MUL_ROUTED_WEIGHT=mul_routed_weight, + compute_type=_tl_dtype(c.dtype), + num_warps=_DECODE_WARPS, + ) + + +def decode_gemm_marlin( + a: torch.Tensor, + packed: torch.Tensor, + scale: torch.Tensor, + glob: torch.Tensor, + c: torch.Tensor, + topk_weights: torch.Tensor, + topk_ids: torch.Tensor, + mul_routed_weight: bool, + a_row_is_route: bool, +) -> None: + """Marlin-style decode GEMV: int32 wide loads + deferred reduction + (:func:`_decode_nvfp4_marlin_kernel`). ``packed`` is the uint8 ``[S, N, K//2]`` bank; + it is reinterpreted as int32 ``[S, N, K//8]`` (contiguous, K%8==0 for NVFP4).""" + M, top_k = topk_ids.shape + N = packed.shape[1] + K = packed.shape[2] * 2 + packed_i32 = packed.view(torch.int32) # [S, N, K // 8] + scale = e4m3_kernel_view(scale) + total_routes = M * top_k + deep_k = K >= _DECODE_MARLIN_DEEPK_THRESHOLD + block_n = _DECODE_MARLIN_DEEPK_BLOCK_N if deep_k else _DECODE_MARLIN_BLOCK_N + block_kw = _DECODE_MARLIN_DEEPK_BLOCK_KW if deep_k else _DECODE_MARLIN_BLOCK_KW + unroll = _DECODE_MARLIN_DEEPK_UNROLL if deep_k else _DECODE_MARLIN_UNROLL + # int64 activation loads need contiguous 16-bit rows that start 16-byte aligned + a_wide = ( + a.dtype in (torch.bfloat16, torch.float16) + and a.stride(1) == 1 and a.stride(0) % 8 == 0 and a.data_ptr() % 16 == 0 + ) + grid = (total_routes, triton.cdiv(N, block_n)) + _decode_nvfp4_marlin_kernel[grid]( + a, packed_i32, scale, glob, c, topk_weights, topk_ids, + total_routes, N, K, + a.stride(0), a.stride(1), + packed_i32.stride(0), packed_i32.stride(1), packed_i32.stride(2), + scale.stride(0), scale.stride(1), scale.stride(2), + glob.stride(0), glob.stride(1), + c.stride(0), c.stride(1), c.stride(2), + topk_weights.stride(0), topk_weights.stride(1), + topk_ids.stride(0), topk_ids.stride(1), + BLOCK_SIZE_N=block_n, + BLOCK_SIZE_KW=block_kw, + UNROLL=unroll, + TOP_K=top_k, + A_ROW_IS_ROUTE=a_row_is_route, + A_WIDE=a_wide, + MUL_ROUTED_WEIGHT=mul_routed_weight, + compute_type=_tl_dtype(c.dtype), + num_warps=_DECODE_MARLIN_WARPS, + ) + + +# 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)} +# 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)),) + + +# Gluon prefill GEMM tiles: BLOCK_SIZE_N -> (BLOCK_SIZE_K, setmaxnreg of the loader / decode warps) +_GLUON_PREFILL_TILES = {128: (64, 96, 96), 256: (32, 88, 72)} + + +def _gluon_prefill_config(N: int, block_m: int, mul_routed_weight: bool) -> Dict[str, Any]: + # BLOCK_SIZE_M 64 tiles are short (few tokens per expert, often short K): persistent CTAs overlap one tile's epilogue and ring fill with the next tile, and a third x / decoded-weight stage absorbs DRAM latency + persistent = block_m == 64 + # a 256-row tile halves the token loads per weight byte, but a partial last tile still decodes all 256 rows + # the persistent epilogue keeps one 64-bit address per element in registers, which spills with the routed-weight loads + block_n = 256 if N % 256 == 0 and not (persistent and mul_routed_weight) else 128 + block_k, load_regs, decode_regs = _GLUON_PREFILL_TILES[block_n] + return dict(BLOCK_SIZE_M=block_m, BLOCK_SIZE_N=block_n, BLOCK_SIZE_K=block_k, GROUP_SIZE_M=8, + STAGES=3 if persistent else 2, W_STAGES=4, PERSISTENT=persistent, DECODE_WARPS=4, MMA_WARPS_N=4, + LOAD_REGS=load_regs, DECODE_REGS=decode_regs, num_warps=8) + + +@functools.lru_cache(maxsize=None) +def _gluon_prefill(device_index: int): + """The Gluon module whose prefill GEMM replaces the Triton kernel on this GPU, or None.""" + from freetoken.kernel import gluon + + if torch.cuda.get_device_capability(device_index) != (12, 0) or not gluon.available(): + return None + from freetoken.kernel.gluon.nvfp4.moe import sm120 + + return sm120 + + +@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) + if _gluon_prefill(device_index) is not None: + 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 + + +@functools.lru_cache(maxsize=None) +def _tma_device(device_index: int) -> bool: + return not is_rocm() and torch.cuda.get_device_capability(device_index) == (12, 0) + + +def _prefill_use_tma(packed: torch.Tensor, scale: torch.Tensor, K: int, device_index: int) -> bool: + # a TMA box row is >= 16 bytes and every global stride a multiple of 16 bytes, so K // 16 scale bytes need K % 256 == 0; + # the TMA scale path reads the scales as fp8, so the forced e4m3 emulation stays on pointer loads + return (_tma_device(device_index) and e4m3_native() and K % 256 == 0 + and packed.is_contiguous() and scale.is_contiguous() + 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_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 + # 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 + + +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, + FAST_DECODE=_prefill_fast_decode(block_m, per_expert, device_index)) + + +def prefill_gemm( + a: torch.Tensor, + packed: torch.Tensor, + scale: torch.Tensor, + glob: torch.Tensor, + c: torch.Tensor, + topk_weights_flat: torch.Tensor, + sorted_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + num_valid_tokens: int, + kernel_top_k: int, + mul_routed_weight: bool, + cfg: Dict[str, Any], +) -> None: + gluon = _gluon_prefill(a.device.index) + if gluon is not None and cfg["BLOCK_SIZE_M"] >= 64 and gluon.supports(a, packed, scale, c): + gluon.prefill_gemm(a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded, + num_valid_tokens, kernel_top_k, mul_routed_weight, + _gluon_prefill_config(packed.shape[1], cfg["BLOCK_SIZE_M"], mul_routed_weight)) + return + N = packed.shape[1] + K = packed.shape[2] * 2 + EM = sorted_ids.shape[0] + scale = e4m3_kernel_view(scale) + grid = lambda META: ( # noqa: E731 + triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), + ) + cfg = dict(cfg) + fast = cfg.pop("FAST_DECODE", 0) + if K % 8: + fast = 0 + if fast == 1: + packed = packed.view(torch.int32) + 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 + w_desc = s_desc = None + if use_tma: + rows = packed.shape[0] * N + box_k = cfg["BLOCK_SIZE_KB"] // 4 if fast == 1 else cfg["BLOCK_SIZE_KB"] + w_desc = TensorDescriptor.from_tensor(packed.view(rows, packed.shape[2]), [cfg["BLOCK_SIZE_N"], box_k]) + if tma_scales: + s_desc = TensorDescriptor.from_tensor(scale.view(rows, K // 16), [cfg["BLOCK_SIZE_N"], 16]) + _prefill_nvfp4_moe_kernel[grid]( + a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids, + num_tokens_post_padded, + w_desc, s_desc, + N, K, EM, num_valid_tokens, + a.stride(0), a.stride(1), + packed.stride(0), packed.stride(1), packed.stride(2), + scale.stride(0), scale.stride(1), scale.stride(2), + glob.stride(0), glob.stride(1), + c.stride(1), c.stride(2), + topk_weights_flat.stride(0), + MUL_ROUTED_WEIGHT=mul_routed_weight, + top_k=kernel_top_k, + compute_type=_tl_dtype(c.dtype), + USE_TMA=use_tma, + TMA_SCALES=tma_scales, + FAST_DECODE=fast, + **cfg, + ) diff --git a/python/freetoken/kernel/triton/nvfp4_fused_moe.py b/python/freetoken/kernel/triton/nvfp4_fused_moe.py index ca9b9dbe9..0fa4d4aa9 100644 --- a/python/freetoken/kernel/triton/nvfp4_fused_moe.py +++ b/python/freetoken/kernel/triton/nvfp4_fused_moe.py @@ -28,6 +28,8 @@ import torch import triton import triton.language as tl +from triton.language import target_info +from triton.runtime.jit import constexpr_function from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, e4m3_u8_to_f32 from freetoken.kernel.triton.nvfp4_linear import _nvfp4_pair_f32 @@ -299,6 +301,57 @@ def _e2m1_f16_x2p14(code): return bits.to(tl.uint16).to(tl.float16, bitcast=True) +@constexpr_function +def _bf16x2_mul_native(): + return target_info.cuda_capability_geq(9, 0) + + +@triton.jit +def _bf16x2_scale(w, s2): + if _bf16x2_mul_native(): + return tl.inline_asm_elementwise("mul.rn.bf16x2 $0, $1, $2;", "=r,r,r", [w, s2], tl.uint32, True, 1) + else: + # mul.bf16x2 is sm_90+; fma with a -0.0 addend gives the same product + return tl.inline_asm_elementwise("{ .reg .b32 z; mov.b32 z, 0x80008000; fma.rn.bf16x2 $0, $1, $2, z; }", + "=r,r,r", [w, s2], tl.uint32, True, 1) + + +@triton.jit +def _bf16x2_halves(w): + lo = (w & 0xFFFF).to(tl.uint16).to(tl.bfloat16, bitcast=True) + hi = (w >> 16).to(tl.uint16).to(tl.bfloat16, bitcast=True) + return lo, hi + + +@triton.jit +def _e2m1x8_bf16(q, s2, BLOCK_N: tl.constexpr, BLOCK_W: tl.constexpr): + """[BLOCK_N, BLOCK_W] uint32 words of 8 e2m1 codes -> [BLOCK_N, 8 * BLOCK_W] bf16 in k order, times s2 * 2^-126.""" + # each 16-bit half's top nibble becomes a bf16 sign + exponent/mantissa: e2m1 * 2^-126, exact incl. e == 0 + o0 = _bf16x2_scale((q & 0x80008000) | ((q >> 6) & 0x01C001C0), s2) + q1 = q << 4 + o1 = _bf16x2_scale((q1 & 0x80008000) | ((q1 >> 6) & 0x01C001C0), s2) + q2 = q << 8 + o2 = _bf16x2_scale((q2 & 0x80008000) | ((q2 >> 6) & 0x01C001C0), s2) + q3 = q << 12 + o3 = _bf16x2_scale((q3 & 0x80008000) | ((q3 >> 6) & 0x01C001C0), s2) + # o0..o3 hold the k pairs (3, 7), (2, 6), (1, 5), (0, 4); join appends the new axis last + k0, k4 = _bf16x2_halves(o3) + k1, k5 = _bf16x2_halves(o2) + k2, k6 = _bf16x2_halves(o1) + k3, k7 = _bf16x2_halves(o0) + w = tl.join(tl.join(tl.join(k0, k4), tl.join(k2, k6)), tl.join(tl.join(k1, k5), tl.join(k3, k7))) + return tl.reshape(w, (BLOCK_N, 8 * BLOCK_W)) + + +@triton.jit +def _e2m1x2_bf16(v, s2, BLOCK_N: tl.constexpr, BLOCK_KB: tl.constexpr): + """[BLOCK_N, BLOCK_KB] bytes of 2 e2m1 codes -> [BLOCK_N, 2 * BLOCK_KB] bf16 in k order, times s2 * 2^-126.""" + # low nibble -> low bf16, high nibble -> high bf16: sign to bit 15 / 31, exponent and mantissa to bits 8..6 / 24..22 + pair = ((v & 0x8) << 12) | ((v & 0x7) << 6) | ((v & 0x80) << 24) | ((v & 0x70) << 18) + lo, hi = _bf16x2_halves(_bf16x2_scale(pair.to(tl.uint32, bitcast=True), s2)) + return tl.reshape(tl.join(lo, hi), (BLOCK_N, 2 * BLOCK_KB)) + + @triton.jit def _prefill_nvfp4_moe_kernel( a_ptr, # [M, K] activations @@ -310,6 +363,8 @@ def _prefill_nvfp4_moe_kernel( sorted_token_ids_ptr, expert_ids_ptr, # cache slot per M-block num_tokens_post_padded_ptr, + w_desc, # TMA descriptor over packed as [S * N, K // 2], or None + s_desc, # TMA descriptor over scale as [S * N, K // 16], or None N, K, EM, @@ -327,6 +382,9 @@ def _prefill_nvfp4_moe_kernel( MUL_ROUTED_WEIGHT: tl.constexpr, top_k: tl.constexpr, compute_type: tl.constexpr, + 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 ): # 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. @@ -352,38 +410,82 @@ def _prefill_nvfp4_moe_kernel( offs_kb = tl.arange(0, BLOCK_SIZE_KB) offs_k = tl.arange(0, 2 * BLOCK_SIZE_KB) offs_sb = tl.arange(0, BLOCK_SIZE_KB // 8) + KW: tl.constexpr = BLOCK_SIZE_KB // 4 + offs_kw = tl.arange(0, KW) 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) - p_ptrs = packed_ptr + slot * stride_pe + offs_bn[:, None] * stride_pn + offs_kb[None, :] * stride_pkb + WORDS: tl.constexpr = FAST_DECODE == 1 + if WORDS: + 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 accumulator = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_M), dtype=tl.float32) K_BYTES = K // 2 + # rows past N belong to the next slot; their outputs are masked on store + w_row = (slot * N + pid_n * BLOCK_SIZE_N).to(tl.int32) + SB: tl.constexpr = BLOCK_SIZE_KB // 8 for kb in range(0, tl.cdiv(K_BYTES, BLOCK_SIZE_KB)): - 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(): + if USE_TMA and WORDS: + v = w_desc.load([w_row, kb * KW]) + elif USE_TMA: + v = w_desc.load([w_row, kb * BLOCK_SIZE_KB]).to(tl.int32) + if USE_TMA and TMA_SCALES: + # a TMA box row needs >= 16 bytes, so load the 16 scales that hold this step and pick its SB of them + s16 = s_desc.load([w_row, (kb * SB) // 16 * 16]).to(tl.uint8, bitcast=True) + s3 = tl.reshape(s16, (BLOCK_SIZE_N, 16 // SB, SB)) + pick = tl.arange(0, 16 // SB) == (kb * SB) % 16 // SB + s = tl.sum(tl.where(pick[None, :, None], s3, 0), axis=1).to(tl.uint8) + s = s.to(tl.float8e4nv, bitcast=True).to(tl.float32) + elif USE_TMA: + sblk_mask = kb * SB + offs_sb < K_BYTES // 8 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)) - # 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)) + sblk_mask = kb * SB + offs_sb < K_BYTES // 8 + if WORDS: + v = tl.load(p_ptrs, mask=(kb * KW + offs_kw < K_BYTES // 4)[None, :], other=0) + 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(): + 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)) + if FAST_DECODE: + # 2^118 here and the decode's 2^-126 leave 2^-8, undone in the epilogue; every step is exact in bf16 + 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)) + w = _e2m1x8_bf16(v.to(tl.uint32, bitcast=True), s, BLOCK_SIZE_N, KW) + else: + # one scale per 16 values, eight bytes + s = tl.reshape(tl.broadcast_to(s[:, :, None], (BLOCK_SIZE_N, SB, 8)), (BLOCK_SIZE_N, BLOCK_SIZE_KB)) + w = _e2m1x2_bf16(v, s, BLOCK_SIZE_N, BLOCK_SIZE_KB) + else: + # 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 + 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 g = tl.load(global_ptr + slot * stride_ge + offs_bn * stride_gn).to(tl.float32) + if FAST_DECODE: + g = g * 256.0 accumulator = accumulator * g[:, None] if MUL_ROUTED_WEIGHT: diff --git a/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py index 9202f5f38..197ef5472 100644 --- a/python/freetoken/moe/fused_nvfp4.py +++ b/python/freetoken/moe/fused_nvfp4.py @@ -7,144 +7,11 @@ from __future__ import annotations -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, - _decode_nvfp4_moe_kernel, - _e2m1_lut, - _prefill_nvfp4_moe_kernel, -) +from freetoken.kernel import moe_sum_reduce_triton, nvfp4_moe 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. -# These drive the original LUT-gather decode (_decode_gemm), kept only for A/B. -_DECODE_BLOCK_N = 64 -_DECODE_BLOCK_KB = 128 -_DECODE_WARPS = 4 - -# Marlin-style decode: every K-loop load must keep the word tile's register layout, so BLOCK_N * BLOCK_KW <= 128 * num_warps and BLOCK_KW <= 32 at 4 warps. -# Long K (gate_up) uses the narrow N tile for more programs at M=1; short K (down) uses the wide one. -_DECODE_MARLIN_BLOCK_N = 16 -_DECODE_MARLIN_BLOCK_KW = 16 -_DECODE_MARLIN_UNROLL = 2 -_DECODE_MARLIN_WARPS = 4 -_DECODE_MARLIN_DEEPK_BLOCK_N = 8 -_DECODE_MARLIN_DEEPK_BLOCK_KW = 32 -_DECODE_MARLIN_DEEPK_UNROLL = 2 -_DECODE_MARLIN_DEEPK_THRESHOLD = 2048 - - -def _tl_dtype(dt: torch.dtype): - if dt == torch.bfloat16: - return tl.bfloat16 - if dt == torch.float16: - return tl.float16 - return tl.float32 - - -def _decode_gemm( - a: torch.Tensor, - packed: torch.Tensor, - scale: torch.Tensor, - glob: torch.Tensor, - c: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - mul_routed_weight: bool, - a_row_is_route: bool, -) -> None: - M, top_k = topk_ids.shape - N = packed.shape[1] - K = packed.shape[2] * 2 - scale = e4m3_kernel_view(scale) - total_routes = M * top_k - grid = (total_routes, triton.cdiv(N, _DECODE_BLOCK_N)) - _decode_nvfp4_moe_kernel[grid]( - a, packed, scale, glob, c, topk_weights, topk_ids, - _e2m1_lut(a.device.index), - total_routes, N, K, - a.stride(0), a.stride(1), - packed.stride(0), packed.stride(1), packed.stride(2), - scale.stride(0), scale.stride(1), scale.stride(2), - glob.stride(0), glob.stride(1), - c.stride(0), c.stride(1), c.stride(2), - topk_weights.stride(0), topk_weights.stride(1), - topk_ids.stride(0), topk_ids.stride(1), - BLOCK_SIZE_N=_DECODE_BLOCK_N, - BLOCK_SIZE_KB=_DECODE_BLOCK_KB, - TOP_K=top_k, - A_ROW_IS_ROUTE=a_row_is_route, - MUL_ROUTED_WEIGHT=mul_routed_weight, - compute_type=_tl_dtype(c.dtype), - num_warps=_DECODE_WARPS, - ) - - -def _decode_gemm_marlin( - a: torch.Tensor, - packed: torch.Tensor, - scale: torch.Tensor, - glob: torch.Tensor, - c: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - mul_routed_weight: bool, - a_row_is_route: bool, -) -> None: - """Marlin-style decode GEMV: int32 wide loads + deferred reduction - (:func:`_decode_nvfp4_marlin_kernel`). ``packed`` is the uint8 ``[S, N, K//2]`` bank; - it is reinterpreted as int32 ``[S, N, K//8]`` (contiguous, K%8==0 for NVFP4).""" - M, top_k = topk_ids.shape - N = packed.shape[1] - K = packed.shape[2] * 2 - packed_i32 = packed.view(torch.int32) # [S, N, K // 8] - scale = e4m3_kernel_view(scale) - total_routes = M * top_k - deep_k = K >= _DECODE_MARLIN_DEEPK_THRESHOLD - block_n = _DECODE_MARLIN_DEEPK_BLOCK_N if deep_k else _DECODE_MARLIN_BLOCK_N - block_kw = _DECODE_MARLIN_DEEPK_BLOCK_KW if deep_k else _DECODE_MARLIN_BLOCK_KW - unroll = _DECODE_MARLIN_DEEPK_UNROLL if deep_k else _DECODE_MARLIN_UNROLL - # int64 activation loads need contiguous 16-bit rows that start 16-byte aligned - a_wide = ( - a.dtype in (torch.bfloat16, torch.float16) - and a.stride(1) == 1 and a.stride(0) % 8 == 0 and a.data_ptr() % 16 == 0 - ) - grid = (total_routes, triton.cdiv(N, block_n)) - _decode_nvfp4_marlin_kernel[grid]( - a, packed_i32, scale, glob, c, topk_weights, topk_ids, - total_routes, N, K, - a.stride(0), a.stride(1), - packed_i32.stride(0), packed_i32.stride(1), packed_i32.stride(2), - scale.stride(0), scale.stride(1), scale.stride(2), - glob.stride(0), glob.stride(1), - c.stride(0), c.stride(1), c.stride(2), - topk_weights.stride(0), topk_weights.stride(1), - topk_ids.stride(0), topk_ids.stride(1), - BLOCK_SIZE_N=block_n, - BLOCK_SIZE_KW=block_kw, - UNROLL=unroll, - TOP_K=top_k, - A_ROW_IS_ROUTE=a_row_is_route, - A_WIDE=a_wide, - MUL_ROUTED_WEIGHT=mul_routed_weight, - compute_type=_tl_dtype(c.dtype), - num_warps=_DECODE_MARLIN_WARPS, - ) def _fused_experts_decode_nvfp4( @@ -164,8 +31,8 @@ def _fused_experts_decode_nvfp4( act_limit: float = 7.0, ) -> torch.Tensor: """Shared decode body (gemm1 -> act -> gemm2 -> sum-reduce); ``gemm_fn`` is either - the marlin-style int32 GEMV (:func:`_decode_gemm_marlin`) or the original LUT-gather - GEMV (:func:`_decode_gemm`), both with the same calling convention.""" + the marlin-style int32 GEMV (:func:`freetoken.kernel.nvfp4_moe.decode_gemm_marlin`) or the original LUT-gather + GEMV (:func:`freetoken.kernel.nvfp4_moe.decode_gemm`), both with the same calling convention.""" M, H = hidden_states.shape top_k = topk_ids.shape[1] two_i = gate_up_packed.shape[1] @@ -206,7 +73,7 @@ def fused_experts_decode_nvfp4_marlin( ) -> torch.Tensor: """Decode inline-NVFP4 MoE using the Marlin-style int32 wide-load GEMV.""" return _fused_experts_decode_nvfp4( - _decode_gemm_marlin, + nvfp4_moe.decode_gemm_marlin, hidden_states, gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global, topk_weights, topk_ids, activation, apply_router_weight_on_input, @@ -232,7 +99,7 @@ def fused_experts_decode_nvfp4_serial( """Original LUT-gather decode (one program per route, full K reduction). Retained for A/B benchmarking against the marlin decode path; not on the production decode path.""" return _fused_experts_decode_nvfp4( - _decode_gemm, + nvfp4_moe.decode_gemm, hidden_states, gate_up_packed, gate_up_scale, gate_up_global, down_packed, down_scale, down_global, topk_weights, topk_ids, activation, apply_router_weight_on_input, @@ -240,107 +107,6 @@ def fused_experts_decode_nvfp4_serial( ) -# 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( - a: torch.Tensor, - packed: torch.Tensor, - scale: torch.Tensor, - glob: torch.Tensor, - c: torch.Tensor, - topk_weights_flat: torch.Tensor, - sorted_ids: torch.Tensor, - expert_ids: torch.Tensor, - num_tokens_post_padded: torch.Tensor, - num_valid_tokens: int, - kernel_top_k: int, - mul_routed_weight: bool, - cfg: Dict[str, Any], -) -> None: - N = packed.shape[1] - K = packed.shape[2] * 2 - EM = sorted_ids.shape[0] - scale = e4m3_kernel_view(scale) - grid = lambda META: ( # noqa: E731 - triton.cdiv(EM, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]), - ) - _prefill_nvfp4_moe_kernel[grid]( - a, packed, scale, glob, c, topk_weights_flat, sorted_ids, expert_ids, - num_tokens_post_padded, - N, K, EM, num_valid_tokens, - a.stride(0), a.stride(1), - packed.stride(0), packed.stride(1), packed.stride(2), - scale.stride(0), scale.stride(1), scale.stride(2), - glob.stride(0), glob.stride(1), - c.stride(1), c.stride(2), - topk_weights_flat.stride(0), - MUL_ROUTED_WEIGHT=mul_routed_weight, - top_k=kernel_top_k, - compute_type=_tl_dtype(c.dtype), - **cfg, - ) - - def fused_experts_nvfp4( hidden_states: torch.Tensor, gate_up_packed: torch.Tensor, @@ -365,14 +131,14 @@ 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, hidden_states.device.index) + cfg = nvfp4_moe.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() num_valid = topk_ids.numel() ic1 = torch.empty((M, top_k, two_i), device=dev, dtype=dt) - _prefill_gemm( + nvfp4_moe.prefill_gemm( hidden_states, gate_up_packed, gate_up_scale, gate_up_global, ic1, tw, sorted_ids, expert_ids, ntpp, num_valid, top_k, apply_router_weight_on_input, cfg, @@ -380,7 +146,7 @@ def fused_experts_nvfp4( ic2 = torch.empty((M * top_k, inter), device=dev, dtype=dt) gated_act_and_mul(activation, ic1.view(-1, two_i), ic2, alpha=act_alpha, limit=act_limit) ic3 = torch.empty((M, top_k, H), device=dev, dtype=dt) - _prefill_gemm( + nvfp4_moe.prefill_gemm( ic2, down_packed, down_scale, down_global, ic3, tw, sorted_ids, expert_ids, ntpp, num_valid, 1, not apply_router_weight_on_input, cfg, diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py index dbf06bcc6..29c042ba4 100644 --- a/tests/moe/test_nvfp4_backends.py +++ b/tests/moe/test_nvfp4_backends.py @@ -10,6 +10,7 @@ baseline-kernel equality guard. This is the path used on sm_120 + CUDA 12.x. - Marlin (sm_80..sm_99, e.g. H100): prefill + decode + overlap. - b12x (sm_120 + CUDA>=13): pack + fused decode, skipped where the kernel cannot run. + - Gluon (sm_120): the prefill GEMMs it takes over from the Triton kernel, against that kernel. Banks are built the way the loader builds them: the bound kernel object packs native pieces. """ @@ -366,6 +367,52 @@ def test_triton_swigluoai_matches_dequant_reference(): _assert_close(out, ref) +@cuda +@pytest.mark.skipif(not torch.cuda.is_available() or torch.cuda.get_device_capability() != (12, 0), reason="needs sm_120") +@pytest.mark.parametrize("hidden_size,inter", [(256, 128), (384, 192)]) +@pytest.mark.parametrize("block_m", [64, 128]) +def test_triton_prefill_routes_to_gluon_on_sm120(monkeypatch, hidden_size, inter, block_m): + """Both prefill GEMMs go to the Gluon kernel and match the Triton kernel: persistent (BLOCK_SIZE_M 64) and not (128), 256- and 128-row tiles, K % 256 == 0 and not.""" + import freetoken.moe.fused_nvfp4 as fused_nvfp4 + from freetoken.kernel import gluon, nvfp4_moe + + if not gluon.available(): + pytest.skip(f"Gluon is off or needs Triton {'.'.join(map(str, gluon.TRITON_VERSION))}") + from freetoken.kernel.gluon.nvfp4.moe import sm120 + + # the GPU class picks BLOCK_SIZE_M from tokens per expert; pin it so every sm_120 class runs both kernel modes + config = nvfp4_moe.prefill_config + monkeypatch.setattr(nvfp4_moe, "prefill_config", lambda *args: {**config(*args), "BLOCK_SIZE_M": block_m}) + + device = torch.device("cuda") + g = torch.Generator(device=device).manual_seed(0) + + def bank(n, k): + packed = torch.randint(0, 256, (E, n, k // 2), dtype=torch.uint8, device=device, generator=g) + scale = (torch.rand(E, n, k // 16, device=device, generator=g) + 0.25).to(torch.float8_e4m3fn) + return packed, scale, (torch.rand(E, n, device=device, generator=g) + 0.5).to(torch.float16) + + banks = (*bank(2 * inter, hidden_size), *bank(hidden_size, inter)) + M = 1024 + hidden = torch.randn(M, hidden_size, dtype=torch.bfloat16, device=device, generator=g) / 4 + topk_ids = torch.randint(0, E, (M, TOPK), dtype=torch.int32, device=device, generator=g) + topk_weights = torch.rand(M, TOPK, device=device, generator=g) + + calls = [] + gluon_gemm = sm120.prefill_gemm + + def counted(*args): + calls.append(args) + return gluon_gemm(*args) + + monkeypatch.setattr(sm120, "prefill_gemm", counted) + out = fused_nvfp4.fused_experts_nvfp4(hidden, *banks, topk_weights, topk_ids, E) + assert len(calls) == 2 + monkeypatch.setattr(nvfp4_moe, "_gluon_prefill", lambda device_index: None) + ref = fused_nvfp4.fused_experts_nvfp4(hidden, *banks, topk_weights, topk_ids, E) + torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2 * float(ref.abs().max())) + + def _triton_cache(device, *, cache_size=S, prefill_overlap=False): """Native 6-bank NVFP4 cache (no repack), consumed directly by the Triton kernels. The banks are not transformed, so ``sources`` doubles as the dequant reference."""