feat(v4): vendor mlx-lm PR #1224 FP8 block-dequant as generic utility

Pulled the FP8 dequant snippet out of qwen3_5_moe.py's Model.sanitize
(which mirrors the same pattern in models/{deepseek_v3,deepseek_v32,
minimax,mimo_v2_flash,ministral3}.py) and exposed it as a standalone
utility under cppmega_v4/nn/_external/_mlx_lm_fp8_dequant_vendored.py:

  - dequant_block_fp8(weight, scale_inv) -> bfloat16
    Block size 128; uses mx.from_fp8 + per-block scale_inv broadcast.
    Handles non-block-aligned shapes via pad → reshape → unpad.

  - sanitize_fp8_weights(weights) -> dict
    Iterates a state-dict, dequants every (<prefix>.weight,
    <prefix>.weight_scale_inv) pair, drops activation_scale PTQ
    artefacts, leaves non-fp8 tensors untouched. No-op if no
    weight_scale_inv keys present.

Rationale for vendoring as utility rather than model-specific copy:
ROI 7 (Lightning Indexer FP8) needs the same dequant path, and we
don't want to re-implement it in each loader.

PRs #1199 (EngGPT MoE), #991 (sarvam), #560 (LLada2): inspected and
skipped — those are vanilla Mixtral-style (softmax routing +
SwitchGLU). Our moe_v4.py is already more advanced (aux-loss-free +
sqrtsoftplus + hash routing) so vendoring them would be a regression.

Tests (5 new): shape round-trip, signal recovery within fp8 precision,
non-block-aligned shapes, sanitize replaces paired keys + drops
activation_scale, sanitize no-op on non-fp8 weights.

v4 suite: 195 passed / 2 skipped.

Author: apstenku123

cppmega_v4/nn/_external/_mlx_lm_fp8_dequant_vendored.py

@@ -0,0 +1,73 @@
+# Vendored from ml-explore/mlx-lm PR #1224 (mlx_lm/models/qwen3_5_moe.py
+# `Model.sanitize` FP8 dequant block), generalized to a standalone utility.
+# Upstream license: MIT, © 2025-2026 Apple Inc.
+#
+# Mirrors the inline FP8 dequant pattern that lives in models/{deepseek_v3,
+# deepseek_v32, minimax, mimo_v2_flash, ministral3, qwen3_5_moe}.py — pulled
+# out so cppmega_v4 (MoE + Lightning Indexer ROI 7) can dequant FP8
+# checkpoints without re-copying the snippet in every loader.
+
+"""FP8 block-scaled dequant utility (block size 128, paired weight_scale_inv).
+
+Triggered by HF checkpoints with quant_method=fp8 (e.g. Qwen/Qwen3.6-27B-FP8,
+DeepSeek-V3-FP8). Pairs each weight tensor with its `weight_scale_inv` and
+multiplies them block-wise to recover bfloat16 weights.
+"""
+
+from typing import Mapping
+
+import mlx.core as mx
+
+_BLOCK_SIZE = 128
+
+
+def dequant_block_fp8(weight: mx.array, scale_inv: mx.array) -> mx.array:
+    """Block-dequantize a single FP8 weight tensor.
+
+    Args:
+        weight: [m, n] in fp8.
+        scale_inv: [ceil(m/128), ceil(n/128)] in fp32; broadcast across
+            the [128, 128] block.
+
+    Returns:
+        [m, n] in bfloat16.
+    """
+    weight = mx.from_fp8(weight, dtype=mx.bfloat16)
+    bs = _BLOCK_SIZE
+    m, n = weight.shape
+    pad_bottom = (-m) % bs
+    pad_side = (-n) % bs
+    weight = mx.pad(weight, ((0, pad_bottom), (0, pad_side)))
+    weight = weight.reshape(
+        ((m + pad_bottom) // bs, bs, (n + pad_side) // bs, bs)
+    )
+    weight = (weight * scale_inv[:, None, :, None]).reshape(
+        m + pad_bottom, n + pad_side
+    )
+    return weight[:m, :n].astype(mx.bfloat16)
+
+
+def sanitize_fp8_weights(weights: Mapping[str, mx.array]) -> dict[str, mx.array]:
+    """Dequantize all FP8-quantized tensors in a state-dict.
+
+    Iterates over keys and for every ``<prefix>.weight_scale_inv`` pair with a
+    matching ``<prefix>.weight`` in fp8, replaces the fp8 weight with its
+    bfloat16 dequant. ``activation_scale`` keys are dropped (PTQ artefact).
+    Returns the rewritten state-dict (or the input unchanged if no fp8 keys).
+    """
+    if not any("weight_scale_inv" in k for k in weights):
+        return dict(weights)
+
+    new_weights: dict[str, mx.array] = {}
+    for k, v in weights.items():
+        if "weight_scale_inv" in k:
+            wk = k.replace("_scale_inv", "")
+            new_weights[wk] = dequant_block_fp8(weights[wk], v)
+        elif "activation_scale" in k:
+            continue
+        elif k not in new_weights:
+            new_weights[k] = v
+    return new_weights
+
+
+__all__ = ["dequant_block_fp8", "sanitize_fp8_weights"]

tests/v4/test_fp8_dequant.py

@@ -0,0 +1,98 @@
+"""Tests for vendored mlx-lm PR #1224 FP8 block-dequant utility."""
+
+import mlx.core as mx
+import numpy as np
+import pytest
+
+from cppmega_v4.nn._external._mlx_lm_fp8_dequant_vendored import (
+    dequant_block_fp8,
+    sanitize_fp8_weights,
+)
+
+
+def _make_fp8_weight(m: int, n: int, seed: int = 0) -> tuple[mx.array, mx.array]:
+    """Build a (weight_fp8, scale_inv) pair from a bf16 reference tensor."""
+    rng = np.random.default_rng(seed)
+    bs = 128
+    blocks_m = (m + bs - 1) // bs
+    blocks_n = (n + bs - 1) // bs
+    scale_inv = mx.array(rng.uniform(0.5, 1.5, (blocks_m, blocks_n)).astype(np.float32))
+    bf16_ref = mx.array(rng.standard_normal((m, n)).astype(np.float32) * 0.1).astype(
+        mx.bfloat16
+    )
+    # Simulate fp8 quantization: divide by per-block scale, cast to fp8.
+    pad_bottom = (-m) % bs
+    pad_side = (-n) % bs
+    padded = mx.pad(bf16_ref.astype(mx.float32), ((0, pad_bottom), (0, pad_side)))
+    blocks = padded.reshape(blocks_m, bs, blocks_n, bs)
+    scaled = (blocks / scale_inv[:, None, :, None]).reshape(
+        m + pad_bottom, n + pad_side
+    )[:m, :n]
+    fp8 = mx.to_fp8(scaled)  # uint8 storage
+    return fp8, scale_inv
+
+
+def test_dequant_block_fp8_round_trip_shape():
+    m, n = 256, 384
+    fp8, scale_inv = _make_fp8_weight(m, n)
+    out = dequant_block_fp8(fp8, scale_inv)
+    assert out.shape == (m, n)
+    assert out.dtype == mx.bfloat16
+
+
+def test_dequant_block_fp8_recovers_signal():
+    """Quant→dequant must preserve signal within fp8 precision (~5% rel error)."""
+    m, n = 128, 128
+    fp8, scale_inv = _make_fp8_weight(m, n, seed=7)
+    out = dequant_block_fp8(fp8, scale_inv)
+    # Build the expected value: scale_inv * fp8_as_fp32
+    expected = (
+        mx.from_fp8(fp8, dtype=mx.bfloat16).astype(mx.float32)
+        * scale_inv[0, 0, None, None]
+    ).astype(mx.bfloat16)
+    np.testing.assert_allclose(
+        np.array(out.astype(mx.float32)),
+        np.array(expected.astype(mx.float32)),
+        atol=1e-3,
+    )
+
+
+def test_dequant_block_fp8_handles_non_block_aligned_shapes():
+    """Shapes not divisible by 128 must still produce correct (m, n) output."""
+    m, n = 100, 200  # both < 128 multiples, edge of pad path
+    fp8, scale_inv = _make_fp8_weight(m, n)
+    out = dequant_block_fp8(fp8, scale_inv)
+    assert out.shape == (m, n)
+
+
+def test_sanitize_fp8_weights_replaces_paired_keys():
+    fp8_a, scale_a = _make_fp8_weight(64, 64, seed=1)
+    fp8_b, scale_b = _make_fp8_weight(128, 128, seed=2)
+    other = mx.array(np.random.randn(8, 8).astype(np.float32))
+    weights = {
+        "model.layers.0.mlp.gate_proj.weight": fp8_a,
+        "model.layers.0.mlp.gate_proj.weight_scale_inv": scale_a,
+        "model.layers.1.mlp.gate_proj.weight": fp8_b,
+        "model.layers.1.mlp.gate_proj.weight_scale_inv": scale_b,
+        "model.norm.weight": other,
+        # PTQ artefact that must be dropped:
+        "model.layers.0.mlp.gate_proj.activation_scale": mx.array([1.0]),
+    }
+    out = sanitize_fp8_weights(weights)
+    assert "model.layers.0.mlp.gate_proj.weight" in out
+    assert "model.layers.0.mlp.gate_proj.weight_scale_inv" not in out
+    assert "model.layers.0.mlp.gate_proj.activation_scale" not in out
+    assert out["model.layers.0.mlp.gate_proj.weight"].dtype == mx.bfloat16
+    assert out["model.norm.weight"] is other  # untouched
+
+
+def test_sanitize_fp8_weights_no_op_when_no_fp8():
+    """Without any weight_scale_inv key, return a dict copy unchanged."""
+    weights = {
+        "model.norm.weight": mx.array([1.0, 2.0]),
+        "model.embed_tokens.weight": mx.array([[0.0, 1.0]]),
+    }
+    out = sanitize_fp8_weights(weights)
+    assert out == dict(weights)
+    # Different dict object (defensive copy), same values
+    assert out is not weights