feat(v4): fused FP8 GEMM Metal kernel + StreamingPoolCache for CSA/HCA

Two next-level optimizations:

1. cppmega_v4/_tilelang/fused_fp8_gemm.py — fused dequant + GEMM
   Metal kernel. Replaces the "dequant to bf16, then matmul" round-trip
   in LightningIndexerFP8._wq_b_apply with one MSL pass per output
   column. Per-block (128x128) scale_inv lookup is amortized across
   128 K-elements; per-thread accumulator stays in fp32 registers.
   FP8 e4m3 decoder is inlined into MSL (sign/expt/mant bit-fields,
   subnormal & NaN handling) so no MLX op call inside the kernel.

   Tests (5): shape, fused-vs-unfused parity within fp8 precision
   (atol 2e-2 rtol 5e-2), dtype/shape rejection, non-128-aligned shapes.

2. cppmega_v4/nn/streaming_kv_cache.py — StreamingPoolCache for
   CSA/HCA/NSA-Compress branches. Avoids the per-forward recompute of
   mean-pool over historical tokens during decode. Maintains:
     - completed [B, n_super, H, D] of finalized super-tokens
     - partial [B, H, D] sum + count of the in-progress block
   `append(k, v)` walks tokens, flushing when a block fills.
   `snapshot(include_partial=True/False)` returns the current view.

   Tests (7): construct/append validation, streaming-vs-oneshot mean-pool
   parity, partial-block finalization, total/n_super counters, reset
   clears state, single-token append (decode-style streaming).

Plus: two background agents spawned to extend triton_frontend op_mapping
(tt.dot/tt.exp emitters for FLA chunk kernel) and to build the C++
PtrAnalysis shim — both work in our /Volumes/external/sources/tilelang
fork.

v4 suite: 285 passed / 2 skipped.

Author: apstenku123

cppmega_v4/_tilelang/fused_fp8_gemm.py

@@ -0,0 +1,149 @@
+"""Fused FP8 GEMM Metal kernel — block-fp8 weight × bf16/fp32 activation.
+
+Used by ROI 7 (Lightning Indexer FP8) to fuse dequant + GEMM instead of
+materializing the full bf16 weight tile on every forward.
+
+Layout:
+  - W_fp8: [M, K] uint8 (e4m3 storage); each [128, 128] block has a single
+    fp32 scale_inv stored in W_scale_inv: [ceil(M/128), ceil(K/128)].
+  - A:     [..., K] fp32/bf16 (caller's choice — internally promoted to fp32).
+  - Out:   [..., M] same dtype as A.
+
+The kernel computes (per output column m):
+
+    acc = 0
+    for k in [0, K):
+        block_m = m // 128
+        block_k = k // 128
+        acc += dequant(W_fp8[m, k], W_scale_inv[block_m, block_k]) * A[..., k]
+    Out[..., m] = acc
+
+One thread per output column m × batch row. Block_size=128 amortizes the
+scale_inv lookup across 128 K-elements per (m, batch_row).
+"""
+
+from typing import Tuple
+
+import mlx.core as mx
+
+from cppmega_v4._tilelang._kernel_cache import get_or_build_kernel
+
+_BLOCK = 128
+
+
+def fused_fp8_gemm(
+    w_fp8: mx.array,
+    w_scale_inv: mx.array,
+    a: mx.array,
+) -> mx.array:
+    """Fused dequant + GEMM: out = a @ W.T where W = dequant(w_fp8, w_scale_inv).
+
+    Args:
+        w_fp8: [M, K] uint8 fp8 weight.
+        w_scale_inv: [ceil(M/128), ceil(K/128)] fp32 per-block inverse-scale.
+        a: [..., K] bf16/fp32 activation.
+
+    Returns:
+        out: [..., M] same dtype as a.
+    """
+    if w_fp8.dtype != mx.uint8:
+        raise TypeError(f"w_fp8 must be uint8 (fp8 storage); got {w_fp8.dtype}")
+    if w_fp8.ndim != 2:
+        raise ValueError(f"w_fp8 must be 2D [M, K]; got {w_fp8.shape}")
+    M, K = w_fp8.shape
+    if a.shape[-1] != K:
+        raise ValueError(f"a.shape[-1] ({a.shape[-1]}) must equal w_fp8.shape[1] ({K})")
+    bs = _BLOCK
+    blocks_m_expected = (M + bs - 1) // bs
+    blocks_k_expected = (K + bs - 1) // bs
+    if w_scale_inv.shape != (blocks_m_expected, blocks_k_expected):
+        raise ValueError(
+            f"w_scale_inv shape {w_scale_inv.shape} != expected "
+            f"({blocks_m_expected}, {blocks_k_expected}) for W={M, K} with block={bs}"
+        )
+
+    out_dtype = a.dtype
+    a_fp32 = a.astype(mx.float32)
+    # Flatten activations to [N, K] for the kernel.
+    leading = a.shape[:-1]
+    n = 1
+    for d in leading:
+        n *= d
+    a_flat = a_fp32.reshape(n, K)
+
+    blocks_m = blocks_m_expected
+    blocks_k = blocks_k_expected
+    w_flat = w_fp8.reshape(-1)             # [M*K]
+    s_flat = w_scale_inv.reshape(-1)        # [blocks_m * blocks_k]
+    a_flat_1d = a_flat.reshape(-1)         # [N*K]
+
+    # Convert per-row scale_inv lookup to a flat index inside the kernel.
+    # Each thread = (m, row); inner loop over k, reads s_flat[(m/128)*blocks_k + (k/128)]
+    # and the corresponding fp8 byte from w_flat[m*K + k].
+    source = f"""
+        uint m   = thread_position_in_grid.x;
+        uint row = thread_position_in_grid.y;
+        if (m >= {M}u || row >= {n}u) return;
+
+        uint block_m = m / {bs}u;
+        float acc = 0.0f;
+
+        // Inner K-loop: unrolled per-block to amortize scale_inv lookup.
+        for (uint kb = 0; kb < {blocks_k}u; ++kb) {{
+            float scale = s_flat[block_m * {blocks_k}u + kb];
+            uint k_start = kb * {bs}u;
+            uint k_end_full = k_start + {bs}u;
+            uint k_end = k_end_full < {K}u ? k_end_full : {K}u;
+            for (uint k = k_start; k < k_end; ++k) {{
+                // mx.from_fp8 conversion: e4m3 fp8 byte → fp32.
+                // Reproduce the bit-decoder inline so this kernel stays
+                // self-contained (no MLX op call inside MSL).
+                uint  byte = (uint)w_flat[m * {K}u + k];
+                int   sign = (byte >> 7) & 0x1;
+                int   expt = (byte >> 3) & 0xF;
+                int   mant = byte & 0x7;
+                float val;
+                if (expt == 0) {{
+                    // subnormal: val = (-1)^s * 2^-6 * (mant/8)
+                    val = (float)mant / 8.0f * 0.015625f;  // 2^-6 = 1/64
+                }} else if (expt == 0xF && mant == 0x7) {{
+                    // NaN
+                    val = 0.0f;  // treat as zero in GEMM accumulation
+                }} else {{
+                    // normal: val = (-1)^s * 2^(expt-7) * (1 + mant/8)
+                    float mantissa = 1.0f + (float)mant / 8.0f;
+                    int   bias_exp = expt - 7;
+                    // 2^bias_exp via metal::ldexp.
+                    val = metal::ldexp(mantissa, bias_exp);
+                }}
+                if (sign) val = -val;
+                acc += val * scale * a_flat[row * {K}u + k];
+            }}
+        }}
+        out[row * {M}u + m] = acc;
+    """
+
+    name = f"v4_fused_fp8_gemm_{M}_{K}_{n}"
+    kernel = get_or_build_kernel(
+        name=name,
+        input_names=["w_flat", "s_flat", "a_flat"],
+        output_names=["out"],
+        source=source,
+    )
+
+    grid = (M, n, 1)
+    tg_x = min(M, 32)
+    threadgroup = (tg_x, 1, 1)
+
+    (out_flat,) = kernel(
+        inputs=[w_flat, s_flat, a_flat_1d],
+        output_shapes=[(n * M,)],
+        output_dtypes=[mx.float32],
+        grid=grid,
+        threadgroup=threadgroup,
+    )
+    out = out_flat.reshape(*leading, M).astype(out_dtype)
+    return out
+
+
+__all__ = ["fused_fp8_gemm"]

cppmega_v4/nn/streaming_kv_cache.py

@@ -0,0 +1,157 @@
+"""Streaming KV cache for CSA / HCA / NSA-Compress branches.
+
+Avoids recomputing the mean-pool over historical tokens on every forward.
+Maintains a running compressed K/V buffer of super-tokens (size ⌈S_seen/m⌉)
+plus a partial super-token-in-progress with its accumulated mean and count.
+
+Three operations:
+    cache = StreamingPoolCache(m=m, n_heads=H, head_dim=D, dtype=mx.float32)
+    cache.append(new_k, new_v)        # append [B, T_new, H, D]
+    k_super, v_super = cache.snapshot()   # [B, n_super, H, D] — current view
+
+Pure-MLX implementation (no Metal kernel needed). A future Metal/TileLang
+fused kernel could maintain the running mean in shared memory; current
+impl is allocation-friendly enough for inference on Apple Silicon.
+"""
+
+from typing import Optional, Tuple
+
+import mlx.core as mx
+
+
+class StreamingPoolCache:
+    """Per-(B, H) running mean-pool over m tokens for compressed-attention.
+
+    State:
+        - completed: [B, n_super, H, D] fp32, the per-block means already finalized.
+        - partial_sum: [B, H, D] fp32, sum-so-far of the in-progress super-token.
+        - partial_count: int — number of tokens already added to partial_sum.
+
+    Calling `snapshot()` returns the completed block + the in-progress block
+    finalized as `partial_sum / partial_count` so the caller sees a coherent
+    [B, n_super_total, H, D] view.
+    """
+
+    def __init__(
+        self,
+        *,
+        m: int,
+        batch: int,
+        n_heads: int,
+        head_dim: int,
+        dtype: mx.Dtype = mx.float32,
+    ):
+        if m <= 0:
+            raise ValueError(f"m must be positive, got {m}")
+        if batch <= 0 or n_heads <= 0 or head_dim <= 0:
+            raise ValueError("batch / n_heads / head_dim must all be positive")
+        self.m = m
+        self.batch = batch
+        self.n_heads = n_heads
+        self.head_dim = head_dim
+        self.dtype = dtype
+        self._completed_k: Optional[mx.array] = None   # [B, n_super, H, D]
+        self._completed_v: Optional[mx.array] = None
+        self._partial_k = mx.zeros((batch, n_heads, head_dim), dtype=dtype)
+        self._partial_v = mx.zeros((batch, n_heads, head_dim), dtype=dtype)
+        self._partial_count = 0
+        self._n_super = 0
+        self._total_tokens = 0
+
+    @property
+    def total_tokens(self) -> int:
+        return self._total_tokens
+
+    @property
+    def n_super_completed(self) -> int:
+        return self._n_super
+
+    def append(self, k: mx.array, v: mx.array) -> None:
+        """Append [B, T_new, H, D] new K/V tokens to the running cache."""
+        if k.shape != v.shape:
+            raise ValueError(f"k.shape ({k.shape}) != v.shape ({v.shape})")
+        if k.ndim != 4:
+            raise ValueError(f"expected [B, T, H, D]; got {k.shape}")
+        b, t, h, d = k.shape
+        if b != self.batch or h != self.n_heads or d != self.head_dim:
+            raise ValueError(
+                f"k shape {k.shape} incompatible with cache "
+                f"(B={self.batch}, H={self.n_heads}, D={self.head_dim})"
+            )
+        k = k.astype(self.dtype)
+        v = v.astype(self.dtype)
+
+        # Walk tokens one at a time: simpler than partial-block arithmetic
+        # because we need to flush completed super-tokens mid-batch.
+        token_idx = 0
+        while token_idx < t:
+            tokens_to_fill = self.m - self._partial_count
+            tokens_available = t - token_idx
+            take = min(tokens_to_fill, tokens_available)
+            # Accumulate take tokens into partial.
+            chunk_k = k[:, token_idx : token_idx + take]   # [B, take, H, D]
+            chunk_v = v[:, token_idx : token_idx + take]
+            self._partial_k = self._partial_k + chunk_k.sum(axis=1)
+            self._partial_v = self._partial_v + chunk_v.sum(axis=1)
+            self._partial_count += take
+            token_idx += take
+            self._total_tokens += take
+            if self._partial_count == self.m:
+                # Flush completed super-token.
+                mean_k = (self._partial_k / float(self.m))[:, None, :, :]
+                mean_v = (self._partial_v / float(self.m))[:, None, :, :]
+                if self._completed_k is None:
+                    self._completed_k = mean_k
+                    self._completed_v = mean_v
+                else:
+                    self._completed_k = mx.concatenate(
+                        [self._completed_k, mean_k], axis=1
+                    )
+                    self._completed_v = mx.concatenate(
+                        [self._completed_v, mean_v], axis=1
+                    )
+                self._n_super += 1
+                self._partial_k = mx.zeros_like(self._partial_k)
+                self._partial_v = mx.zeros_like(self._partial_v)
+                self._partial_count = 0
+
+    def snapshot(
+        self, *, include_partial: bool = True,
+    ) -> Tuple[mx.array, mx.array]:
+        """Return current compressed K/V super-tokens [B, n_super_total, H, D].
+
+        With ``include_partial=True`` (default), the in-progress block is
+        finalized as ``partial_sum / partial_count``. With False, only the
+        completed super-tokens are returned.
+        """
+        completed_k = (
+            self._completed_k if self._completed_k is not None
+            else mx.zeros((self.batch, 0, self.n_heads, self.head_dim),
+                          dtype=self.dtype)
+        )
+        completed_v = (
+            self._completed_v if self._completed_v is not None
+            else mx.zeros((self.batch, 0, self.n_heads, self.head_dim),
+                          dtype=self.dtype)
+        )
+        if include_partial and self._partial_count > 0:
+            mean_k = (self._partial_k / float(self._partial_count))[:, None, :, :]
+            mean_v = (self._partial_v / float(self._partial_count))[:, None, :, :]
+            completed_k = mx.concatenate([completed_k, mean_k], axis=1)
+            completed_v = mx.concatenate([completed_v, mean_v], axis=1)
+        return completed_k, completed_v
+
+    def reset(self) -> None:
+        """Drop all state (start a new sequence)."""
+        self._completed_k = None
+        self._completed_v = None
+        self._partial_k = mx.zeros((self.batch, self.n_heads, self.head_dim),
+                                    dtype=self.dtype)
+        self._partial_v = mx.zeros((self.batch, self.n_heads, self.head_dim),
+                                    dtype=self.dtype)
+        self._partial_count = 0
+        self._n_super = 0
+        self._total_tokens = 0
+
+
+__all__ = ["StreamingPoolCache"]

tests/v4/test_fused_fp8_gemm.py

@@ -0,0 +1,78 @@
+"""Tests for fused FP8 GEMM Metal kernel."""
+
+import mlx.core as mx
+import numpy as np
+import pytest
+
+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
+
+
+def _make_fp8(m, n, seed=0):
+    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))
+    bf = mx.array(rng.standard_normal((m, n)).astype(np.float32) * 0.1).astype(mx.bfloat16)
+    pad_b = (-m) % bs
+    pad_s = (-n) % bs
+    padded = mx.pad(bf.astype(mx.float32), ((0, pad_b), (0, pad_s)))
+    blocks = padded.reshape(blocks_m, bs, blocks_n, bs)
+    scaled = (blocks / scale_inv[:, None, :, None]).reshape(m + pad_b, n + pad_s)[:m, :n]
+    fp8 = mx.to_fp8(scaled)
+    return fp8, scale_inv
+
+
+def test_fused_fp8_gemm_shape():
+    M, K, B = 128, 128, 4
+    w, s = _make_fp8(M, K)
+    a = mx.random.normal((B, K))
+    out = fused_fp8_gemm(w, s, a)
+    assert out.shape == (B, M)
+
+
+def test_fused_fp8_gemm_matches_dequant_then_matmul():
+    """Fused kernel must match the unfused dequant→matmul path within fp8 precision."""
+    M, K, B = 128, 128, 8
+    w, s = _make_fp8(M, K, seed=7)
+    rng = np.random.default_rng(99)
+    a = mx.array(rng.standard_normal((B, K)).astype(np.float32))
+
+    out_fused = fused_fp8_gemm(w, s, a)
+    # Reference path: dequant then matmul.
+    w_bf16 = dequant_block_fp8(w, s)
+    out_ref = (a.astype(mx.float32) @ w_bf16.T.astype(mx.float32))
+
+    np.testing.assert_allclose(
+        np.array(out_fused.astype(mx.float32)),
+        np.array(out_ref),
+        atol=2e-2, rtol=5e-2,  # fp8 rounding inside the kernel
+    )
+
+
+def test_fused_fp8_gemm_rejects_dtype_mismatch():
+    M, K = 128, 128
+    w = mx.zeros((M, K), dtype=mx.float32)
+    s = mx.ones((1, 1), dtype=mx.float32)
+    a = mx.zeros((1, K))
+    with pytest.raises(TypeError, match="uint8"):
+        fused_fp8_gemm(w, s, a)
+
+
+def test_fused_fp8_gemm_rejects_shape_mismatch():
+    M, K = 128, 128
+    w, s = _make_fp8(M, K)
+    a = mx.zeros((1, K + 1))
+    with pytest.raises(ValueError, match="a.shape"):
+        fused_fp8_gemm(w, s, a)
+
+
+def test_fused_fp8_gemm_handles_non_block_aligned():
+    """Shapes not divisible by 128 must still produce correct (n, M) output."""
+    M, K, B = 100, 200, 2
+    w, s = _make_fp8(M, K)
+    a = mx.random.normal((B, K))
+    out = fused_fp8_gemm(w, s, a)
+    assert out.shape == (B, M)
+    assert np.all(np.isfinite(np.array(out)))