perf(moe): dequant nvfp4 prefill experts by bit placement and load activations contiguously

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)

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:

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)
 

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,
     )
 

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)
 
 

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