diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 14a7acd835..4b77cf4e23 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -109,6 +109,16 @@ def e4m3_u8_to_f16_x128(v): return h.to(tl.float16, bitcast=True) * 32768.0 +@jit +def e4m3_to_f16_x128(raw): + """An e4m3 value as a branched kernel loads it (fp8 on sm_89+, its uint8 bits + before) -> fp16 pre-scaled by 128, the same value on both.""" + if e4m3_native_cx(): + return raw.to(tl.float16) * 128.0 + else: + return e4m3_u8_to_f16_x128(raw) + + @jit def e4m3_f32_to_u8(x): """Encode an fp32 value that already sits on the e4m3 grid (see :func:`round_e4m3`, |x| <= 448) diff --git a/python/freetoken/kernel/triton/nvfp4_fused_moe.py b/python/freetoken/kernel/triton/nvfp4_fused_moe.py index b1de443787..f3625f065c 100644 --- a/python/freetoken/kernel/triton/nvfp4_fused_moe.py +++ b/python/freetoken/kernel/triton/nvfp4_fused_moe.py @@ -19,6 +19,10 @@ MiniMax-M2 decode toward the RTX 5090 read-bandwidth ceiling (~87%); the residual gap is FP4 dequant ALU, which a swizzled-layout tensor-core path (marlin / flashinfer b12x) closes but those need sm_80-99 / CUDA>=13 respectively. + +Prefill (:func:`_prefill_nvfp4_moe_kernel`) is tensor-core bound, so it dequantizes with +``nvfp4_linear.py``'s fp16 bit placement rather than the table, and loads its activations as +one contiguous tile, split into the even and odd k the two nibbles of a byte multiply. """ from __future__ import annotations @@ -29,8 +33,8 @@ import triton import triton.language as tl -from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, e4m3_u8_to_f32 -from freetoken.kernel.triton.nvfp4_linear import _nvfp4_pair_f32 +from freetoken.kernel.triton.e4m3_compat import e4m3_native_cx, e4m3_to_f16_x128, e4m3_u8_to_f32 +from freetoken.kernel.triton.nvfp4_linear import _nvfp4_pair_f16, _nvfp4_pair_f32 _E2M1_VALUES = [ 0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, @@ -45,6 +49,13 @@ def _e2m1_lut(device_index: int) -> torch.Tensor: ) +@triton.jit +def _e2m1_byte_f16_x2pow_neg14(packed): + """The two e2m1 codes of each byte (int32, low nibble first) -> (lo, hi) fp16 = value * 2^-14: + ``_nvfp4_pair_f16`` once the high nibble also sits at bits [19:16].""" + return _nvfp4_pair_f16(packed | (packed << 12)) + + @triton.jit def _decode_nvfp4_moe_kernel( a_ptr, # [M, K] activations (compute dtype) @@ -303,7 +314,6 @@ def _prefill_nvfp4_moe_kernel( sorted_token_ids_ptr, expert_ids_ptr, # cache slot per M-block num_tokens_post_padded_ptr, - lut_ptr, N, K, EM, @@ -321,6 +331,7 @@ def _prefill_nvfp4_moe_kernel( MUL_ROUTED_WEIGHT: tl.constexpr, top_k: tl.constexpr, compute_type: tl.constexpr, + EVEN_K: tl.constexpr, ): pid = tl.program_id(axis=0) num_pid_m = tl.cdiv(EM, BLOCK_SIZE_M) @@ -342,8 +353,8 @@ 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) + a_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 @@ -353,30 +364,32 @@ def _prefill_nvfp4_moe_kernel( 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 - if e4m3_native_cx(): - scale = tl.load(s_ptrs, mask=byte_mask[:, None], other=0.0).to(tl.float32) + s_ptrs = scale_base + (byte_idx // 8)[:, None] * stride_sblk + if EVEN_K: + bytes_ = tl.load(p_ptrs).to(tl.int32) + raw_scale = tl.load(s_ptrs) 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 - - g = tl.load(global_ptr + slot * stride_ge + offs_bn * stride_gn).to(tl.float32) + byte_mask = byte_idx < K_BYTES + bytes_ = tl.load(p_ptrs, mask=byte_mask[:, None], other=0).to(tl.int32) + raw_scale = tl.load(s_ptrs, mask=byte_mask[:, None], other=0.0) + # [BLOCK_KB, BLOCK_N]; every e2m1 * e4m3 product is exact in fp16 and in the dot dtype + scale = e4m3_to_f16_x128(raw_scale) + lo, hi = _e2m1_byte_f16_x2pow_neg14(bytes_) + b_lo = (lo * scale).to(a_ptr.dtype.element_ty) + b_hi = (hi * scale).to(a_ptr.dtype.element_ty) + if EVEN_K: + a = tl.load(a_ptrs, mask=token_mask[:, None], other=0.0) + else: + k_mask = kb * 2 * BLOCK_SIZE_KB + offs_k < K + a = tl.load(a_ptrs, mask=token_mask[:, None] & k_mask[None, :], other=0.0) + a_lo, a_hi = tl.split(tl.reshape(a, (BLOCK_SIZE_M, BLOCK_SIZE_KB, 2))) + accumulator += tl.dot(a_lo, b_lo) + accumulator += tl.dot(a_hi, b_hi) + a_ptrs += BLOCK_SIZE_KB * 2 * stride_ak + + # the placement's 2^-14 and the scale's 2^7 leave the products 2^-7 short + g = tl.load(global_ptr + slot * stride_ge + offs_bn * stride_gn).to(tl.float32) * 128.0 accumulator = accumulator * g[None, :] if MUL_ROUTED_WEIGHT: diff --git a/python/freetoken/kernel/triton/nvfp4_linear.py b/python/freetoken/kernel/triton/nvfp4_linear.py index 53abdedc71..943a6e50e8 100644 --- a/python/freetoken/kernel/triton/nvfp4_linear.py +++ b/python/freetoken/kernel/triton/nvfp4_linear.py @@ -43,7 +43,7 @@ from freetoken.kernel.triton.e4m3_compat import ( e4m3_kernel_view, e4m3_native_cx, - e4m3_u8_to_f16_x128, + e4m3_to_f16_x128, e4m3_u8_to_f32, ) @@ -426,19 +426,13 @@ def _nvfp4_gemm_kernel( for t in range(tiles_per): if EVEN_K: word = tl.load(b_ptrs, mask=n_mask[:, None], other=0) - if e4m3_native_cx(): - s128 = tl.load(s_ptrs, mask=n_mask[:, None], other=0.0).to(tl.float16) * 128.0 - else: - s128 = e4m3_u8_to_f16_x128(tl.load(s_ptrs, mask=n_mask[:, None], other=0)) + s128 = e4m3_to_f16_x128(tl.load(s_ptrs, mask=n_mask[:, None], other=0.0)) a_tile = tl.load(a_ptrs, mask=m_mask[:, None], other=0.0) else: kw_ok = (tile0 + t) * BLOCK_KW + offs_kw < K_WORDS word = tl.load(b_ptrs, mask=n_mask[:, None] & kw_ok[None, :], other=0) s_mask = n_mask[:, None] & kw_ok[None, :] - if e4m3_native_cx(): - s128 = tl.load(s_ptrs, mask=s_mask, other=0.0).to(tl.float16) * 128.0 - else: - s128 = e4m3_u8_to_f16_x128(tl.load(s_ptrs, mask=s_mask, other=0)) + s128 = e4m3_to_f16_x128(tl.load(s_ptrs, mask=s_mask, other=0.0)) ak_ok = (tile0 + t) * BLOCK_KW * 8 + offs_ak < K a_tile = tl.load(a_ptrs, mask=m_mask[:, None] & ak_ok[None, :], other=0.0) diff --git a/python/freetoken/moe/fused_nvfp4.py b/python/freetoken/moe/fused_nvfp4.py index 375d2a547a..e8ad477021 100644 --- a/python/freetoken/moe/fused_nvfp4.py +++ b/python/freetoken/moe/fused_nvfp4.py @@ -270,7 +270,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), @@ -281,6 +280,7 @@ def _prefill_gemm( MUL_ROUTED_WEIGHT=mul_routed_weight, top_k=kernel_top_k, compute_type=_tl_dtype(c.dtype), + EVEN_K=(K // 2) % cfg["BLOCK_SIZE_KB"] == 0, **cfg, ) diff --git a/tests/kernels/test_e4m3_compat.py b/tests/kernels/test_e4m3_compat.py index fa8a1f1d32..0911fc1f04 100644 --- a/tests/kernels/test_e4m3_compat.py +++ b/tests/kernels/test_e4m3_compat.py @@ -251,6 +251,22 @@ def _emit_all(path: str) -> None: out["nv_moe_prefill"] = f32(fused_experts_nvfp4( hid4, gup, gus4, gug, dnp, dns4, dng, tw, tids, S)) + # a K no tile divides takes the masked loads, whose fill value must suit an fp8 pointer + Kr, H5, I5 = 272, 272, 80 + pkr = torch.randint(0, 256, (N, Kr // 2), dtype=torch.uint8, device=dev) + scr = (torch.rand(N, Kr // 16, device=dev) * 2 + 0.1).to(FP8) + out["nv_gemm_inkernel_ragged_k"] = f32(nvfp4_dense_linear( + torch.randn(8, Kr, dtype=torch.bfloat16, device=dev), pkr, scr, g)) + gup5 = torch.randint(0, 256, (S, 2 * I5, H5 // 2), dtype=torch.uint8, device=dev) + gus5 = (torch.rand(S, 2 * I5, H5 // 16, device=dev) + 0.1).to(FP8) + gug5 = (torch.rand(S, 2 * I5, device=dev) * 0.05 + 0.01).to(torch.float16) + dnp5 = torch.randint(0, 256, (S, H5, I5 // 2), dtype=torch.uint8, device=dev) + dns5 = (torch.rand(S, H5, I5 // 16, device=dev) + 0.1).to(FP8) + dng5 = (torch.rand(S, H5, device=dev) * 0.05 + 0.01).to(torch.float16) + out["nv_moe_prefill_ragged_k"] = f32(fused_experts_nvfp4( + torch.randn(2, H5, dtype=torch.bfloat16, device=dev), + gup5, gus5, gug5, dnp5, dns5, dng5, tw, tids, S)) + torch.save(out, path) diff --git a/tests/moe/test_nvfp4_backends.py b/tests/moe/test_nvfp4_backends.py index 0546fbcd6f..7bdbc95681 100644 --- a/tests/moe/test_nvfp4_backends.py +++ b/tests/moe/test_nvfp4_backends.py @@ -46,7 +46,9 @@ def _dequant_ref(packed: torch.Tensor, scale: torch.Tensor, row_global: torch.Te return w * s * row_global.float().unsqueeze(1) -def _make_native_sources(device: torch.device, seed: int = 0) -> dict[str, list[torch.Tensor]]: +def _make_native_sources( + device: torch.device, seed: int = 0, hidden: int = H, inter: int = I +) -> dict[str, list[torch.Tensor]]: """Random ModelOpt-style banks, CPU pinned, with one expert whose w1/w3 globals differ. One flat ``[L*E, ...]`` RNG draw (so seeding is unaffected) split into L @@ -61,15 +63,15 @@ def rand_u8(*shape): def rand_scale(*shape): return (torch.rand(*shape, generator=g) * 1.5 + 0.25).to(torch.float8_e4m3fn) - gate_up_global = torch.full((total, 2 * I), 1.0, dtype=torch.float16) - gate_up_global[:, I:] = 0.5 # w3 global != w1 global: exercises the alpha fold - down_global = torch.full((total, H), 0.75, dtype=torch.float16) + gate_up_global = torch.full((total, 2 * inter), 1.0, dtype=torch.float16) + gate_up_global[:, inter:] = 0.5 # w3 global != w1 global: exercises the alpha fold + down_global = torch.full((total, hidden), 0.75, dtype=torch.float16) flat = { - "gate_up_packed": rand_u8(total, 2 * I, H // 2), - "gate_up_scale": rand_scale(total, 2 * I, H // 16), + "gate_up_packed": rand_u8(total, 2 * inter, hidden // 2), + "gate_up_scale": rand_scale(total, 2 * inter, hidden // 16), "gate_up_global": gate_up_global, - "down_packed": rand_u8(total, H, I // 2), - "down_scale": rand_scale(total, H, I // 16), + "down_packed": rand_u8(total, hidden, inter // 2), + "down_scale": rand_scale(total, hidden, inter // 16), "down_global": down_global, } return {name: list(t.pin_memory().split(E)) for name, t in flat.items()} @@ -84,8 +86,9 @@ def _assert_close(out: torch.Tensor, ref: torch.Tensor) -> None: def _swigluoai_ref(h: torch.Tensor, alpha: float = 1.702, limit: float = 7.0) -> torch.Tensor: """MiniMax-M3 / gpt-oss clamped swiglu over UNINTERLEAVED [gate; up] halves.""" - gate = h[:I].clamp(max=limit) - up = h[I:].clamp(-limit, limit) + gate, up = h.chunk(2) + gate = gate.clamp(max=limit) + up = up.clamp(-limit, limit) return gate * torch.sigmoid(gate * alpha) * (up + 1.0) @@ -110,7 +113,8 @@ def _ref_moe(sources, layer_id, hidden, topk_weights, topk_ids, activation="silu if activation == "swigluoai": act = _swigluoai_ref(h) else: - act = torch.nn.functional.silu(h[:I]) * h[I:] + gate, up = h.chunk(2) + act = torch.nn.functional.silu(gate) * up out[t] += float(topk_weights[t, j]) * (dn @ act) return out @@ -325,6 +329,74 @@ def test_triton_overlap_prefill_matches_dequant_reference(): cache.release_prefill_layer(layer_id) +@cuda +def test_triton_prefill_dequant_reads_every_code_and_scale_exactly(): + """The prefill GEMM decodes e2m1 by placing its bits in fp16 rather than through a table: + every code must read 2^-14 of its value, and every product with an e4m3 block scale must be + exact in bf16 -- the property that keeps the GEMM bit-identical to the table decode.""" + import triton + import triton.language as tl + + from freetoken.kernel.triton.e4m3_compat import e4m3_kernel_view, e4m3_to_f16_x128 + from freetoken.kernel.triton.nvfp4_fused_moe import _e2m1_byte_f16_x2pow_neg14 + + @triton.jit + def dequant(bytes_ptr, scales_ptr, placed_ptr, weights_ptr, BLOCK: tl.constexpr): + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + lo, hi = _e2m1_byte_f16_x2pow_neg14(tl.load(bytes_ptr + offs)) + scale = e4m3_to_f16_x128(tl.load(scales_ptr + offs)) + tl.store(placed_ptr + 2 * offs, lo) + tl.store(placed_ptr + 2 * offs + 1, hi) + tl.store(weights_ptr + 2 * offs, (lo * scale).to(tl.bfloat16)) + tl.store(weights_ptr + 2 * offs + 1, (hi * scale).to(tl.bfloat16)) + + device = torch.device("cuda") + # every byte against every scale: each (low code, high code) pair meets all 256 scale codes + packed = torch.arange(256, dtype=torch.int32, device=device).repeat_interleave(256) + scales = torch.arange(256, dtype=torch.uint8, device=device).repeat(256).view(torch.float8_e4m3fn) + placed = torch.empty(2 * packed.numel(), dtype=torch.float16, device=device) + weights = torch.empty(2 * packed.numel(), dtype=torch.bfloat16, device=device) + dequant[(packed.numel() // 1024,)](packed, e4m3_kernel_view(scales), placed, weights, BLOCK=1024) + + codes = torch.stack([packed & 0xF, packed >> 4], dim=1).reshape(-1) + value = _E2M1.double().to(device)[codes.long()] + assert torch.equal(placed.double(), value * 2.0**-14) + assert torch.equal(placed.signbit(), value.signbit()) + scale = scales.double().repeat_interleave(2) + finite = ~scale.isnan() + assert int(finite.sum()) == 2 * 256 * 254 + assert torch.equal(weights.double()[finite], (value * scale * 2.0**-7)[finite]) + + +@cuda +@pytest.mark.parametrize( + ("hidden", "inter", "tokens"), + [(H, I, 96), (272, 80, 96)], + ids=["tiled-k", "k-not-a-tile-multiple"], +) +def test_triton_prefill_large_batch_matches_dequant_reference(hidden, inter, tokens): + """A batch past the small-batch tile (> 64 tokens) through the Triton prefill grouped GEMM, + at a K the byte tile divides and at one it does not (272 and 80 leave 8 bytes over).""" + from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4 + + device = torch.device("cuda") + sources = _make_native_sources(device, seed=13, hidden=hidden, inter=inter) + torch.manual_seed(14) + x = torch.randn(tokens, hidden, dtype=torch.bfloat16, device=device) / 4 + topk_ids = torch.randint(0, E, (tokens, TOPK), dtype=torch.int32, device=device) + topk_weights = torch.rand(tokens, TOPK, dtype=torch.float32, device=device) + banks = [ + sources[name][0].to(device) + for name in ( + "gate_up_packed", "gate_up_scale", "gate_up_global", + "down_packed", "down_scale", "down_global", + ) + ] + ref = _ref_moe(sources, 0, x, topk_weights, topk_ids) + out = fused_experts_nvfp4(x, *banks, topk_weights, topk_ids, E, "silu", False) + _assert_close(out, ref) + + @cuda def test_triton_swigluoai_matches_dequant_reference(): """MiniMax-M3's swigluoai routed experts through the Triton prefill grouped GEMM