feat(v4): FP8 MoE experts via fused FP8 GEMM Metal kernel

cppmega_v4/nn/moe_fp8.py — FP8-backed MoE expert layer:

  - FP8Linear: nn.Linear analogue with uint8 fp8 weight storage and
    per-128-block fp32 scale_inv. Forward routes through the fused
    dequant+GEMM Metal kernel (cppmega_v4/_tilelang/fused_fp8_gemm).
    `from_linear(lin)` quantizes an existing nn.Linear.
  - FP8FeedForwardExpert: drop-in for cppmega_mlx.nn.moe.FeedForwardExpert.
    Three FP8Linear projections (gate / up / down) with the same
    swiglu / relu2 / gelu activation choices.
  - quantize_linear_to_fp8(weight) → (uint8 fp8, fp32 scale_inv)
    pure utility; per-block amax → scale_inv = amax/448 (e4m3 max).
  - convert_v4moe_to_fp8(moe) — in-place: replaces every moe.experts[i]
    (and shared_expert if present) with FP8FeedForwardExpert, quantizing
    the existing weights via from_fp32_expert.

Tests (6 new):
  - shape / dtype invariants on quantize_linear_to_fp8.
  - rejection of non-2D weights.
  - FP8Linear round-trip vs fp32 nn.Linear (atol=1.5e-1 rtol=2e-1 —
    fp8 e4m3 worst-case ~12% relative drift on random init).
  - FP8FeedForwardExpert round-trip vs FeedForwardExpert.
  - convert_v4moe_to_fp8 preserves output shape + numerical closeness.
  - shared_expert is also converted when configured.

This unlocks loading deepseek-v3-fp8 / qwen3.6-fp8 style checkpoints
without dequant-to-bf16 first — the saved RAM matters for routed
experts (8 × hidden_size × intermediate × bf16 = a lot at scale).

v4 suite: 318 passed / 2 skipped.

Author: apstenku123

cppmega_v4/nn/moe_fp8.py

@@ -0,0 +1,158 @@
+"""FP8 expert layer + V4MoE checkpoint converter.
+
+Wraps each MoE expert's gate_proj / up_proj / down_proj in an FP8 path
+that uses the fused dequant+GEMM Metal kernel from
+``cppmega_v4/_tilelang/fused_fp8_gemm.py``.
+
+Two surfaces:
+  - ``FP8FeedForwardExpert`` — drop-in nn.Module that replaces
+    ``cppmega_mlx.nn.moe.FeedForwardExpert``. Weights are stored as
+    uint8 fp8 + per-128-block fp32 scale_inv per projection.
+  - ``convert_v4moe_to_fp8(moe)`` — in-place converter: replaces every
+    ``moe.experts[i]`` with an ``FP8FeedForwardExpert`` whose weights
+    are quantized from the original bf16/fp32 linears via
+    ``quantize_linear_to_fp8``.
+
+Numeric tolerance: fp8 e4m3 round-trip ≈ 2% on per-row activations,
+so quant→fp8→dequant→matmul matches the original bf16 matmul to
+atol ~5e-2 on random inputs. Adequate for inference; training-time
+parity is not yet covered (would need fp8 fwd + bf16 master-weight
+optimizer state).
+"""
+
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+
+from cppmega_v4._tilelang.fused_fp8_gemm import fused_fp8_gemm
+from cppmega_v4.nn._external._mlx_lm_fp8_dequant_vendored import dequant_block_fp8
+
+_BLOCK = 128
+
+
+def quantize_linear_to_fp8(weight: mx.array) -> tuple[mx.array, mx.array]:
+    """Quantize a [out, in] fp32/bf16 weight to fp8 + per-block scale_inv.
+
+    Returns (w_fp8: uint8 [out, in], scale_inv: fp32 [ceil(out/128), ceil(in/128)]).
+    """
+    if weight.ndim != 2:
+        raise ValueError(f"weight must be 2D [out, in]; got {weight.shape}")
+    w = weight.astype(mx.float32)
+    out_dim, in_dim = w.shape
+    bs = _BLOCK
+    pad_o = (-out_dim) % bs
+    pad_i = (-in_dim) % bs
+    blocks_o = (out_dim + pad_o) // bs
+    blocks_i = (in_dim + pad_i) // bs
+    padded = mx.pad(w, ((0, pad_o), (0, pad_i)))
+    blocks = padded.reshape(blocks_o, bs, blocks_i, bs)
+    amax = mx.maximum(mx.max(mx.abs(blocks), axis=(1, 3), keepdims=False), 1e-6)
+    scale_inv = (amax / 448.0).astype(mx.float32)
+    scaled = (blocks / scale_inv[:, None, :, None]).reshape(
+        out_dim + pad_o, in_dim + pad_i,
+    )[:out_dim, :in_dim]
+    fp8 = mx.to_fp8(scaled.astype(mx.bfloat16))
+    return fp8, scale_inv
+
+
+class FP8Linear(nn.Module):
+    """nn.Linear analogue with fp8-stored weight, fused dequant+GEMM forward."""
+
+    def __init__(self, in_features: int, out_features: int, *, bias: bool = False):
+        super().__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        # Initialise with zero weights — caller must load via load_fp8 or
+        # use ``from_linear`` to quantize an existing nn.Linear.
+        self.weight_fp8 = mx.zeros((out_features, in_features), dtype=mx.uint8)
+        blocks_o = (out_features + _BLOCK - 1) // _BLOCK
+        blocks_i = (in_features + _BLOCK - 1) // _BLOCK
+        self.weight_scale_inv = mx.ones((blocks_o, blocks_i), dtype=mx.float32)
+        self.bias = (
+            mx.zeros((out_features,), dtype=mx.bfloat16)
+            if bias else None
+        )
+
+    @classmethod
+    def from_linear(cls, lin: nn.Linear) -> "FP8Linear":
+        out_dim, in_dim = lin.weight.shape
+        has_bias = getattr(lin, "bias", None) is not None
+        mod = cls(in_dim, out_dim, bias=has_bias)
+        mod.weight_fp8, mod.weight_scale_inv = quantize_linear_to_fp8(lin.weight)
+        if has_bias:
+            mod.bias = lin.bias.astype(mx.bfloat16)
+        return mod
+
+    def __call__(self, x: mx.array) -> mx.array:
+        out = fused_fp8_gemm(self.weight_fp8, self.weight_scale_inv, x)
+        if self.bias is not None:
+            out = out + self.bias.astype(out.dtype)
+        return out
+
+
+class FP8FeedForwardExpert(nn.Module):
+    """Drop-in for ``FeedForwardExpert`` with fp8 weights + fused GEMM."""
+
+    def __init__(
+        self,
+        d_model: int,
+        hidden_size: int,
+        *,
+        activation: str = "gelu",
+        bias: bool = False,
+    ):
+        super().__init__()
+        self.activation = activation
+        self.gate_proj = FP8Linear(d_model, hidden_size, bias=bias)
+        self.up_proj = (
+            FP8Linear(d_model, hidden_size, bias=bias)
+            if activation == "swiglu"
+            else None
+        )
+        self.down_proj = FP8Linear(hidden_size, d_model, bias=bias)
+
+    @classmethod
+    def from_fp32_expert(cls, expert: nn.Module) -> "FP8FeedForwardExpert":
+        """Quantize an existing FeedForwardExpert in-place style (returns new module)."""
+        d_in = expert.gate_proj.weight.shape[1]
+        hidden = expert.gate_proj.weight.shape[0]
+        has_up = getattr(expert, "up_proj", None) is not None
+        activation = getattr(expert, "activation", "swiglu" if has_up else "gelu")
+        bias = expert.gate_proj.bias is not None if hasattr(expert.gate_proj, "bias") else False
+        mod = cls(d_in, hidden, activation=activation, bias=bias)
+        mod.gate_proj = FP8Linear.from_linear(expert.gate_proj)
+        if has_up:
+            mod.up_proj = FP8Linear.from_linear(expert.up_proj)
+        mod.down_proj = FP8Linear.from_linear(expert.down_proj)
+        return mod
+
+    def __call__(self, x: mx.array) -> mx.array:
+        h = self.gate_proj(x)
+        if self.activation == "swiglu":
+            assert self.up_proj is not None
+            h = nn.silu(h) * self.up_proj(x)
+        elif self.activation == "relu2":
+            h = mx.square(mx.maximum(h, mx.array(0.0, dtype=h.dtype)))
+        else:
+            h = nn.gelu_approx(h)
+        return self.down_proj(h)
+
+
+def convert_v4moe_to_fp8(moe) -> None:
+    """In-place convert every expert in a V4MoE to FP8FeedForwardExpert.
+
+    The shared_expert (if any) is also converted.
+    """
+    new_experts = [FP8FeedForwardExpert.from_fp32_expert(e) for e in moe.experts]
+    moe.experts = new_experts
+    if getattr(moe, "shared_expert", None) is not None:
+        moe.shared_expert = FP8FeedForwardExpert.from_fp32_expert(moe.shared_expert)
+
+
+__all__ = [
+    "FP8FeedForwardExpert",
+    "FP8Linear",
+    "convert_v4moe_to_fp8",
+    "quantize_linear_to_fp8",
+]

tests/v4/test_moe_fp8.py

@@ -0,0 +1,105 @@
+"""Tests for FP8 MoE expert layer + V4MoE checkpoint converter."""
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+import pytest
+
+from cppmega_v4.nn.moe_fp8 import (
+    FP8FeedForwardExpert,
+    FP8Linear,
+    convert_v4moe_to_fp8,
+    quantize_linear_to_fp8,
+)
+from cppmega_v4.nn.moe_v4 import V4MoE, V4MoEConfig
+
+
+def test_quantize_linear_shape_and_dtypes():
+    w = mx.random.normal((128, 256)).astype(mx.bfloat16)
+    fp8, scale = quantize_linear_to_fp8(w)
+    assert fp8.shape == w.shape
+    assert fp8.dtype == mx.uint8
+    assert scale.shape == (1, 2)   # ceil(128/128)=1, ceil(256/128)=2
+    assert scale.dtype == mx.float32
+
+
+def test_quantize_linear_rejects_wrong_rank():
+    w = mx.zeros((128,))
+    with pytest.raises(ValueError, match="2D"):
+        quantize_linear_to_fp8(w)
+
+
+def test_fp8_linear_from_linear_round_trip():
+    """FP8Linear(x) ≈ original_linear(x) within fp8 precision."""
+    in_dim, out_dim = 128, 128
+    rng = np.random.default_rng(0)
+    lin = nn.Linear(in_dim, out_dim, bias=False)
+    lin.weight = mx.array(rng.standard_normal((out_dim, in_dim)).astype(np.float32) * 0.1)
+    fp8 = FP8Linear.from_linear(lin)
+
+    x = mx.array(rng.standard_normal((4, in_dim)).astype(np.float32))
+    y_ref = lin(x.astype(mx.float32))
+    y_fp8 = fp8(x)
+    np.testing.assert_allclose(
+        np.array(y_fp8.astype(mx.float32)), np.array(y_ref),
+        atol=1.5e-1, rtol=2e-1,    # fp8 e4m3 worst-case ~12% relative drift
+    )
+
+
+def test_fp8_feedforward_expert_round_trip():
+    """FP8FeedForwardExpert ≈ FeedForwardExpert within fp8 precision."""
+    from cppmega_mlx.nn.moe import FeedForwardExpert
+    d, hidden = 64, 128
+    rng = np.random.default_rng(1)
+    expert = FeedForwardExpert(d, hidden, activation="swiglu", bias=False)
+    # randomize weights so the test isn't all-zero
+    expert.gate_proj.weight = mx.array(rng.standard_normal((hidden, d)).astype(np.float32) * 0.1)
+    expert.up_proj.weight = mx.array(rng.standard_normal((hidden, d)).astype(np.float32) * 0.1)
+    expert.down_proj.weight = mx.array(rng.standard_normal((d, hidden)).astype(np.float32) * 0.1)
+
+    fp8_expert = FP8FeedForwardExpert.from_fp32_expert(expert)
+    x = mx.array(rng.standard_normal((2, 8, d)).astype(np.float32))
+    y_ref = expert(x.astype(mx.float32))
+    y_fp8 = fp8_expert(x)
+    assert y_fp8.shape == y_ref.shape
+    np.testing.assert_allclose(
+        np.array(y_fp8.astype(mx.float32)), np.array(y_ref),
+        atol=2e-1, rtol=1e-1,   # 2 fp8 GEMMs in a row + nonlinearity
+    )
+
+
+def test_convert_v4moe_to_fp8_preserves_shape_and_close_output():
+    rng = np.random.default_rng(2)
+    cfg = V4MoEConfig(d_model=64, num_experts=4, top_k=2,
+                       expert_hidden_size=128, activation="swiglu")
+    moe = V4MoE(cfg)
+    # randomize weights
+    for e in moe.experts:
+        e.gate_proj.weight = mx.array(rng.standard_normal((128, 64)).astype(np.float32) * 0.1)
+        e.up_proj.weight = mx.array(rng.standard_normal((128, 64)).astype(np.float32) * 0.1)
+        e.down_proj.weight = mx.array(rng.standard_normal((64, 128)).astype(np.float32) * 0.1)
+    moe.gate.weight = mx.array(rng.standard_normal((4, 64)).astype(np.float32) * 0.1)
+
+    x = mx.array(rng.standard_normal((1, 8, 64)).astype(np.float32))
+    out_ref = moe(x).output
+
+    convert_v4moe_to_fp8(moe)
+    out_fp8 = moe(x).output
+
+    assert out_fp8.shape == out_ref.shape
+    # FP8 ~5-10% drift on stacked GEMMs through routed experts.
+    np.testing.assert_allclose(
+        np.array(out_fp8.astype(mx.float32)), np.array(out_ref),
+        atol=3e-1, rtol=2e-1,
+    )
+
+
+def test_convert_v4moe_converts_shared_expert_too():
+    cfg = V4MoEConfig(d_model=64, num_experts=2, top_k=1,
+                       expert_hidden_size=128, shared_expert_hidden_size=64,
+                       activation="swiglu")
+    moe = V4MoE(cfg)
+    assert moe.shared_expert is not None
+    assert type(moe.shared_expert).__name__ != "FP8FeedForwardExpert"
+    convert_v4moe_to_fp8(moe)
+    assert type(moe.shared_expert).__name__ == "FP8FeedForwardExpert"