perf(moe): prefill nvfp4 W4A4 moe kernel on sm_120

python/freetoken/kernel/gluon/nvfp4/utils.py

@@ -2,6 +2,11 @@
 
 from triton.experimental import gluon
 from triton.experimental.gluon import language as gl
+from triton.experimental.gluon.language import _core as ttgl
+from triton.experimental.gluon.language._layouts import (
+    DotOperandLayout,
+    NVMMADistributedLayout,
+)
 
 
 @gluon.jit
@@ -41,3 +46,37 @@ def e2m1x8_bf16x2(q, s2):
     w2 = bf16x2_mul(prmt(o3, o2, 0x7632), s2)
     w3 = bf16x2_mul(prmt(o1, o0, 0x7632), s2)
     return w0, w1, w2, w3
+
+
+def scale_layout(op_idx: int, rows: int, kgroups: int, mma_warps: list[int]) -> gl.DistributedLinearLayout:
+    """Layout of a [rows, K / 16] e4m3 scale tile for operand op_idx of an NVFP4 mma_v2; mirrors Triton's getSM120DotScaledScaleLayout."""
+    def bits(n):
+        return n.bit_length() - 1
+
+    w0, w1 = mma_warps
+    reg = [[0, 1 << i] for i in range(bits(kgroups))]
+    if op_idx == 0:
+        lane = [[8, 0], [0, 0], [1, 0], [2, 0], [4, 0]]
+        warp = [[0, 0]] * bits(w1) + [[16 << i, 0] for i in range(bits(w0))]
+        cover = 16 * w0
+    else:
+        lane = [[0, 0], [0, 0], [1, 0], [2, 0], [4, 0]]
+        warp = [[8 << i, 0] for i in range(bits(w1))] + [[0, 0]] * bits(w0)
+        cover = 8 * w1
+    reg += [[cover << i, 0] for i in range(max(0, bits(rows) - bits(cover)))]
+    # bases past the tile's rows broadcast
+    lane = [[b[0] if b[0] < rows else 0, b[1]] for b in lane]
+    warp = [[b[0] if b[0] < rows else 0, b[1]] for b in warp]
+    return gl.DistributedLinearLayout(reg, lane, warp, [], [rows, kgroups])
+
+
+@ttgl.builtin
+def mma_nvf4(a, a_scale, b, b_scale, acc, _semantic=None):
+    """acc += (a * a_scale) @ (b * b_scale) on sm_120 (mma kind::mxf4nvf4, scale_vec::4X).
+
+    a [M, K/2] and b [K/2, N] are packed e2m1 in DotOperandLayout(k_width=4) of acc's mma_v2 layout; the e4m3 scales are [M, K/16] / [N, K/16] in scale_layout.
+    """
+    assert isinstance(acc.type.layout, NVMMADistributedLayout) and acc.type.layout.version == [2, 0]
+    assert isinstance(a.type.layout, DotOperandLayout) and isinstance(b.type.layout, DotOperandLayout)
+    out = _semantic.dot_scaled(a, a_scale, "e2m1", b, b_scale, "e2m1", acc, False, True, True, ttgl.float32)
+    return ttgl.tensor(out.handle, acc.type)

python/freetoken/kernel/nvfp4_moe.py

@@ -1,6 +1,6 @@
-"""Kernel dispatch for the inline-dequant NVFP4 MoE GEMMs: which kernel runs, with which tiles and switches, and its launch.
+"""Kernel dispatch for the 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`);
+The kernels are templates (:mod:`freetoken.kernel.triton.nvfp4_fused_moe`, :mod:`freetoken.kernel.gluon.nvfp4.moe.sm120`, :mod:`freetoken.kernel.gluon.nvfp4.moe.sm120_w4a4`);
 :mod:`freetoken.moe.fused_nvfp4` strings the GEMMs into the MoE layer.
 """
 
@@ -388,3 +388,73 @@ def prefill_gemm(
         SCALE_WORDS=scale_words,
         **cfg,
     )
+
+
+@functools.lru_cache(maxsize=None)
+def _w4a4_prefill(device_index: int):
+    """The Gluon module of the W4A4 prefill GEMMs 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_w4a4
+
+    return sm120_w4a4
+
+
+def w4a4_available(device_index: int) -> bool:
+    return _w4a4_prefill(device_index) is not None
+
+
+# Below this many tokens per expert both prefills stream the weights at the same speed (RTX PRO 6000 and RTX 5090 sweeps over the
+# NVFP4 MoE shapes); W4A16 keeps the exact activations there
+_W4A4_MIN_TOKENS_PER_EXPERT = 4
+# (min rho, tokens per expert where BLOCK_SIZE_M becomes 128), from the same sweeps
+_W4A4_BLOCK_M_STEPS = ((3.0, 128), (0.0, 64))
+
+
+@functools.lru_cache(maxsize=None)
+def _w4a4_device(device_index: int) -> Tuple[float, int]:
+    props = torch.cuda.get_device_properties(device_index)
+    rho = _tensor_rho(props.name, props.major, props.minor)
+    return next(t for min_rho, t in _W4A4_BLOCK_M_STEPS if rho >= min_rho), props.L2_cache_size
+
+
+def w4a4_prefill_config(M: int, top_k: int, num_experts: int, H: int, I: int, activation: str, dtype: torch.dtype,
+                        device_index: int) -> Dict[str, Any] | None:
+    """Tiles of the W4A4 prefill GEMMs, or None where the W4A16 prefill runs instead."""
+    w4a4 = _w4a4_prefill(device_index)
+    if w4a4 is None or dtype != torch.bfloat16 or activation not in w4a4.ACTIVATIONS or not w4a4.supports(H, I):
+        return None
+    to128, l2_bytes = _w4a4_device(device_index)
+    per_expert = M * top_k / num_experts
+    if per_expert < _W4A4_MIN_TOKENS_PER_EXPERT:
+        return None
+    block_m = 128 if per_expert >= to128 else 64
+    common = dict(BLOCK_SIZE_M=block_m, GROUP_SIZE_M=8, MMA_WARPS_N=4, LOAD_WARPS=4, num_warps=8)
+    # the gated tile holds BN / 2 gate rows and their up rows, so I must split into whole half tiles
+    gate_up = dict(common, BLOCK_SIZE_N=256 if I % 128 == 0 else 128, BLOCK_SIZE_K=128 if H % 128 == 0 else 64,
+                   N_GROUP=0, STAGES=3, LOAD_REGS=88)
+    if I < 1024:
+        # short K: the persistent kernel overlaps a tile's atomics with the next mainloop; N strips keep the [M, H] output
+        # in L2 once it outgrows half of it
+        down = dict(common, BLOCK_SIZE_N=256, BLOCK_SIZE_K=64, STAGES=4, N_GROUP=2 if M * H * 2 > l2_bytes // 2 else 0,
+                    LOAD_REGS=56, EPI_WARPS=4, EPI_REGS=72, PERSISTENT=True)
+    else:
+        down = dict(common, BLOCK_SIZE_N=256 if H % 256 == 0 else 128, BLOCK_SIZE_K=128 if I % 128 == 0 else 64,
+                    N_GROUP=0, STAGES=3, LOAD_REGS=88, PERSISTENT=False)
+    return dict(BLOCK_SIZE_M=block_m, gate_up=gate_up, down=down)
+
+
+def w4a4_gate_up_gemm(xq, xs, g_in, packed, scale, glob, q, qs, g_out, sorted_ids, expert_ids, num_tokens_post_padded,
+                      num_valid_tokens, top_k, activation, act_alpha, act_limit, cfg):
+    _w4a4_prefill(xq.device.index).gate_up_gemm(
+        xq, xs, g_in, packed, scale, glob, q, qs, g_out, sorted_ids, expert_ids, num_tokens_post_padded,
+        num_valid_tokens, top_k, activation, act_alpha, act_limit, cfg)
+
+
+def w4a4_down_gemm(xq, xs, g_in, packed, scale, glob, out, topk_weights_flat, sorted_ids, expert_ids,
+                   num_tokens_post_padded, num_valid_tokens, top_k, cfg):
+    _w4a4_prefill(xq.device.index).down_gemm(
+        xq, xs, g_in, packed, scale, glob, out, topk_weights_flat, sorted_ids, expert_ids, num_tokens_post_padded,
+        num_valid_tokens, top_k, cfg)

python/freetoken/kernel/triton/nvfp4_act_quant.py

@@ -0,0 +1,42 @@
+"""NVFP4 activation quantization under a static global scale: x ~= e2m1 * e4m3 block scale (per 16) * g."""
+
+from __future__ import annotations
+
+import torch
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _e2m1x2(hi, lo):
+    # sm_100+ / sm_120a; cvt packs its first source into the high nibble
+    return tl.inline_asm_elementwise(
+        "{ .reg .b8 t; cvt.rn.satfinite.e2m1x2.f32 t, $1, $2; cvt.u16.u8 $0, t; }",
+        "=h,r,r", [hi, lo], dtype=tl.uint16, is_pure=True, pack=1)
+
+
+@triton.jit
+def _quant_kernel(x_ptr, q_ptr, s_ptr, g_ptr, K, stride_x, BLOCK_K: tl.constexpr):
+    row = tl.program_id(0).to(tl.int64)
+    offs = tl.arange(0, BLOCK_K)
+    x = tl.load(x_ptr + row * stride_x + offs, mask=offs < K, other=0.0).to(tl.float32)
+    g = tl.load(g_ptr)
+    xb = tl.reshape(x, [BLOCK_K // 16, 16])
+    s = (tl.max(tl.abs(xb), 1) / (6.0 * g)).to(tl.float8e4nv)
+    sf = s.to(tl.float32) * g
+    inv = tl.where(sf > 0, 1.0 / sf, 0.0)
+    lo, hi = tl.split(tl.reshape(xb * inv[:, None], [BLOCK_K // 2, 2]))
+    oq = tl.arange(0, BLOCK_K // 2)
+    tl.store(q_ptr + row * (K // 2) + oq, _e2m1x2(hi, lo).to(tl.uint8), mask=oq < K // 2)
+    os_ = tl.arange(0, BLOCK_K // 16)
+    tl.store(s_ptr + row * (K // 16) + os_, s.to(tl.uint8, bitcast=True), mask=os_ < K // 16)
+
+
+def quantize_nvfp4(x: torch.Tensor, g: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """x [R, K] -> (codes [R, K / 2] uint8, low nibble first; block scales [R, K / 16] e4m3) under the fp32 global g [1]; K % 16 == 0."""
+    R, K = x.shape
+    q = torch.empty(R, K // 2, dtype=torch.uint8, device=x.device)
+    s = torch.empty(R, K // 16, dtype=torch.uint8, device=x.device)
+    if R:
+        _quant_kernel[(R,)](x, q, s, g, K, x.stride(0), BLOCK_K=triton.next_power_of_2(K), num_warps=4)
+    return q, s.view(torch.float8_e4m3fn)

python/freetoken/layers/quantization/moe/nvfp4.py

@@ -78,6 +78,56 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
         return fused_experts_decode_nvfp4_marlin(x, *banks, topk_weights, topk_ids, layer.activation, layer.apply_router_weight_on_input, alpha, limit)
 
 
+class W4a4Nvfp4MoEKernel(TritonNvfp4MoEKernel):
+    """The Triton kernels' banks and decode, with the prefill GEMMs on NVFP4 activations (W4A4) where the GPU and the batch favour it.
+
+    Needs a checkpoint whose experts declare static activation scales; the GPU-resident alphas carry them per expert (gate_up_alpha: the larger of gate's and up's input_scale).
+    """
+
+    name = "w4a4"
+
+    def unusable_reason(self, cfg: MoEConfig) -> str | None:
+        reason = super().unusable_reason(cfg)
+        if reason:
+            return reason
+        if cfg.scheme is None or not cfg.scheme.has("input_scale"):
+            return "the experts declare no static activation scale (W4A16 checkpoint)"
+        if cfg.apply_router_weight_on_input:
+            return "the routing weight scales the input, which is quantized once per token"
+        return None
+
+    def worth_it(self, cfg: MoEConfig) -> bool:
+        from freetoken.kernel import nvfp4_moe
+
+        return not backend.is_rocm() and torch.cuda.is_available() and nvfp4_moe.w4a4_available(torch.cuda.current_device())
+
+    def layout(self, cfg: MoEConfig) -> dict[str, BankSpec]:
+        return {
+            **super().layout(cfg),
+            "gate_up_alpha": BankSpec((), torch.float32, resident=True),
+            "down_alpha": BankSpec((), torch.float32, resident=True),
+        }
+
+    def pack(self, pieces, cfg: MoEConfig, out):
+        super().pack(pieces, cfg, out)
+        if not all(f"{p}_input" in pieces for p in ("gate", "up", "down")):
+            raise ValueError("w4a4 nvfp4 MoE kernel: the checkpoint has no expert input_scale tensors")
+        gate, up, down = (pieces[f"{p}_input"].reshape(len(pieces[f"{p}_input"]), -1)[:, 0].float() for p in ("gate", "up", "down"))
+        return {"gate_up_alpha": torch.maximum(gate, up), "down_alpha": down}
+
+    def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefill: bool):
+        if not is_prefill:
+            return super().apply(layer, x, topk_weights, topk_ids, view, is_prefill=False)
+        from freetoken.moe.fused_nvfp4 import fused_experts_nvfp4_w4a4
+
+        t = view.tensors
+        assert view.alphas is not None
+        return fused_experts_nvfp4_w4a4(
+            x, t["gate_up"], t["gate_up_scale"], t["gate_up_global"], t["down"], t["down_scale"], t["down_global"],
+            view.alphas[0], view.alphas[1], topk_weights, topk_ids, view.n, layer.activation, float(layer.alpha), limit_or_inf(layer),
+        )
+
+
 # ---------------------------------------------------------------------------
 # Marlin: vLLM's fused Marlin MoE (W4A16, dequant in kernel) over Marlin-tiled banks.
 # The Marlin kernels come from vLLM (not sgl-kernel) deliberately: vLLM ships the
@@ -570,4 +620,4 @@ def apply(self, layer, x, topk_weights, topk_ids, view: ExpertView, *, is_prefil
 
 @register_method(QuantKind.NVFP4, LayerKind.MOE)
 class Nvfp4MoEMethod(MoEMethod):
-    candidates = (TritonNvfp4MoEKernel, MarlinNvfp4MoEKernel, B12xNvfp4MoEKernel)
+    candidates = (W4a4Nvfp4MoEKernel, TritonNvfp4MoEKernel, MarlinNvfp4MoEKernel, B12xNvfp4MoEKernel)

python/freetoken/models/nvfp4_banks.py

@@ -39,6 +39,28 @@ def _ingest_global(spec: "Nvfp4ExpertSourceSpec", tensor: torch.Tensor) -> torch
     return tensor.to(torch.float16)
 
 
+def _activation_scale():
+    """The checkpoint tensor of the NVFP4 static activation scale under the installed dialect, or None."""
+    from freetoken.layers.quantization import QuantKind, get_quant_config
+
+    try:
+        quant = get_quant_config()
+    except RuntimeError:
+        return None
+    if quant is None or QuantKind.NVFP4 not in quant.STORAGE:
+        return None
+    return quant.stored_tensors(QuantKind.NVFP4).get("input_scale")
+
+
+def _ingest(spec: "Nvfp4ExpertSourceSpec", act, role: str, tensor: torch.Tensor) -> torch.Tensor:
+    if role.endswith("_global"):
+        return _ingest_global(spec, tensor)
+    if role.endswith("_input"):
+        tensor = tensor.float()
+        return 1.0 / tensor if act.reciprocal else tensor
+    return tensor
+
+
 def _num_moe_layers(config) -> int:
     value = getattr(config, "num_moe_layers", None)
     if value is not None:
@@ -77,7 +99,8 @@ def iter_nvfp4_expert_pieces(
 ):
     """One piece per routed expert: ``gate`` / ``up`` / ``down`` codes plus their ``_scale``
     (fp8 block scales) and ``_global`` (the per-tensor scale, reciprocal for quant-side dialects,
-    fp16) companions, straight from the safetensors shards.
+    fp16) companions, straight from the safetensors shards. When every expert also stores the
+    dialect's static activation scale, the pieces carry it as fp32 ``_input`` (dequant side).
 
     Serial reads walk the shards in order; ``parallel`` uses the chunked O_DIRECT reader. Either
     way tensors of one expert may span shards, so they are grouped by (layer, expert) as they land.
@@ -92,6 +115,7 @@ def iter_nvfp4_expert_pieces(
     weight_map = safetensors_weight_map(folder)
 
     keep = None if layers is None else set(layers)
+    act = _activation_scale()
     wanted: dict[str, tuple[int, int, str]] = {}
     for name in weight_map:
         match = spec.key_pattern.match(name)
@@ -107,9 +131,20 @@ def iter_nvfp4_expert_pieces(
         if kind not in ("weight", "weight_scale", "weight_scale_2"):
             raise ValueError(f"{spec.desc}: unknown NVFP4 expert tensor kind {kind!r}")
         wanted[name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + _kind_suffix(kind))
-    expected = (_num_moe_layers(config) if keep is None else len(keep)) * config.num_experts * 9
-    if len(wanted) != expected:
-        raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {expected}")
+        if kind == "weight" and act is not None:
+            act_name = name[: match.start("kind")] + act.name
+            if act_name in weight_map:
+                wanted[act_name] = (bank_layer, int(match.group("expert")), spec.proj_to_role[proj] + "_input")
+    experts = (_num_moe_layers(config) if keep is None else len(keep)) * config.num_experts
+    acts = [name for name, (_, _, role) in wanted.items() if role.endswith("_input")]
+    # a partial set serves no W4A4 kernel; the pieces then carry none
+    if len(acts) != experts * 3:
+        for name in acts:
+            del wanted[name]
+        acts = []
+    per_expert = 12 if acts else 9
+    if len(wanted) != experts * per_expert:
+        raise ValueError(f"{spec.desc}: found {len(wanted)} expert tensors, expected {experts * per_expert}")
 
     def _serial():
         by_shard: dict[str, list[str]] = collections.defaultdict(list)
@@ -121,21 +156,16 @@ def _serial():
             drop(path)
             with safetensors.safe_open(path, framework="pt", device="cpu") as f:
                 for name in by_shard[shard]:
-                    tensor = f.get_tensor(name)
-                    if wanted[name][2].endswith("_global"):
-                        tensor = _ingest_global(spec, tensor)
-                    yield name, tensor
+                    yield name, _ingest(spec, act, wanted[name][2], f.get_tensor(name))
             drop(path)
 
     def _parallel():
         from freetoken.models.weight import iter_expert_tensors_parallel
 
         for name, tensor in iter_expert_tensors_parallel(folder, lambda n: n in wanted, workers=workers, chunk=chunk):
-            if wanted[name][2].endswith("_global"):
-                tensor = _ingest_global(spec, tensor)
-            yield name, tensor
+            yield name, _ingest(spec, act, wanted[name][2], tensor)
 
-    return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=9)
+    return per_expert_pieces(_parallel() if parallel else _serial(), wanted.get, tensors_per_expert=per_expert)
 
 
 __all__ = ["Nvfp4ExpertSourceSpec", "iter_nvfp4_expert_pieces"]

python/freetoken/moe/expert_pieces.py

@@ -151,7 +151,7 @@ def per_expert_pieces(
     """Group per-expert tensors into one piece per expert, in whatever order they arrive.
 
     ``locate(name)`` -> ``(bank_layer, expert, role)`` or None to skip; a role tensor is one
-    expert's ``[...]`` row and becomes the piece's ``[1, ...]`` entry. Globals may be scalars.
+    expert's ``[...]`` row and becomes the piece's ``[1, ...]`` entry. Globals and activation scales may be scalars.
     """
     pending: dict[tuple[int, int], dict[str, torch.Tensor]] = {}
     for name, tensor in tensors:
@@ -160,7 +160,7 @@ def per_expert_pieces(
             continue
         bank_layer, expert, role = where
         piece = pending.setdefault((bank_layer, expert), {})
-        piece[role] = tensor.reshape(1, -1) if role.endswith("_global") else tensor.unsqueeze(0)
+        piece[role] = tensor.reshape(1, -1) if role.endswith(("_global", "_input")) else tensor.unsqueeze(0)
         if len(piece) == tensors_per_expert:
             del pending[(bank_layer, expert)]
             yield bank_layer, expert, expert + 1, piece