feat(v4): ROI 3.B-F + 3.5.B-D + 3.7 + 5 + 7 + 8 + 9 — all scaffolds + tests

Lands plumbing for every remaining ROI as a runnable scaffold with full
test coverage. Backends that need real Metal/TileLang kernel work delegate
to Path A's pure-MLX reference until the kernel land; the API surface,
status objects, and env override are stable and committed today.

cppmega_v4/_tilelang/_dispatch.py — shared PathStatus + auto_pick infrastructure
  mirroring mamba3_path_c's mature pattern.

cppmega_v4/_tilelang/linear_attention_paths.py (ROI 3.B-F):
  - Path A: pure-MLX (always available)
  - Path B: hand-MSL — scaffold w/ explicit pending reason; fallback Path A
  - Path C: TileLang DSL via tilelang.compile — scaffold; checks tilelang import
  - Path D: Triton frontend — scaffold; checks triton import
  - Path E: vendored mlx-lm PR #1217 — scaffold; checks for vendored module
  - gated_delta_recurrent_dispatch + linear_attention_auto_mode_for_inputs +
    linear_attention_path_statuses
  - Env: CPPMEGA_V4_KERNEL_PATH__LINEAR_ATTENTION={path_a..e,auto}

cppmega_v4/_tilelang/kda_paths.py (ROI 3.5.B-D):
  - Same shape as GDN paths but A/B/C/D only (no Path E for KDA)
  - Env: CPPMEGA_V4_KERNEL_PATH__KDA={path_a..d,auto}; rejects path_e

cppmega_v4/_tilelang/benchmark_receipt.py (ROI 3.7):
  - CellShape + CellReceipt dataclasses; JSON schema matches
    reports/raw/cppmega_1b_path_matrix_cells/ shape
  - measure_cell(shape) runs forward, records elapsed + backend status
  - write_receipt(receipt, out_dir) drops <block>_<path>.json

cppmega_v4/nn/mla_absorb.py (ROI 5):
  - absorb_weights(W_UK, W_UV, W_O) -> (W_UK^T, W_UV @ W_O) precompute
  - absorbed_mla_decode(q, c_kv, w_uk_abs, w_uv_w_o, ...) — full absorb
    path in MLX algebra (per FlashMLA docs/20250422-new-kernel-deep-dive.md)
  - standard_mla_decode(...) — non-absorbed oracle for parity
  - Numerical parity test: absorbed == standard (atol 1e-4)

cppmega_v4/nn/lightning_indexer.py (ROI 7):
  - LightningIndexer module — fp32 scaffold of V3.2 indexer
    (deepseek-ai/DeepSeek-V3.2-Exp/inference/model.py class Indexer)
  - _apply_non_interleaved_rope respects the V3.2 gotcha that indexer
    RoPE is non-interleaved while MLA's is interleaved
  - Top-k indexing via mx.argpartition; clipped to seq length

cppmega_v4/nn/sparse_attention_v4.py (ROI 8/9):
  - NativeSparseAttention scaffold (ROI 8 — NSA, arxiv 2502.11089)
  - CsaHcaHybridAttention scaffold (ROI 9 — V4 CSA+HCA hybrid)
  - Both fall back to dense causal SDPA today; replace with three-branch
    and compressed-sparse paths in research follow-up

44 new tests across 4 files:
- tests/v4/test_path_dispatch.py — 18 tests for GDN+KDA path scaffolding,
  env override correctness, dispatch == Path A parity
- tests/v4/test_benchmark_receipt.py — 5 tests for receipt schema + JSON
- tests/v4/test_mla_absorb.py — 5 tests including numerical parity vs
  standard_mla_decode (atol 1e-4)
- tests/v4/test_lightning_indexer.py — 7 tests including non-interleaved
  RoPE math correctness
- tests/v4/test_sparse_attention_v4.py — 9 tests for NSA + CSA+HCA scaffolds

166/166 v4 + regression green.

Plugin invariant preserved: zero modifications under cppmega_mlx/.

All 11+ ROIs from INTEGRATION_PLAN_CPPMEGA_MLX.md now have committed code
and tests; deferred backends explicitly carry their pending reason in the
PathStatus and delegate to Path A so the v4 stack runs end-to-end today.

Author: apstenku123

cppmega_v4/_tilelang/__init__.py

@@ -0,0 +1,10 @@
+"""Side-by-side TileLang/Metal dispatch scaffolding for V4 blocks.
+
+Each Path B/C/D/E module exposes the same API surface as Path A and falls
+back to Path A's reference kernel when the underlying backend (Metal MSL /
+TileLang DSL / Triton frontend / vendored mlx-lm op) is not available on
+the host or is still pending implementation.
+
+This keeps the v4 plugin runnable end-to-end on day one while leaving room
+for each backend to be filled in incrementally.
+"""

cppmega_v4/_tilelang/_dispatch.py

@@ -0,0 +1,52 @@
+"""Common path-dispatch scaffolding for V4 linear-attention backends.
+
+Mirrors the ``mamba3_path_c_*_status`` / ``*_auto_mode_for_inputs`` pattern
+from ``cppmega_mlx/nn/_tilelang/mamba3_path_c.py`` so v4 paths plug into the
+same conceptual machinery without depending on cppmega_mlx internals.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Literal
+
+PathName = Literal["path_a", "path_b", "path_c", "path_d", "path_e"]
+
+
+@dataclass(frozen=True)
+class PathStatus:
+    """Describes whether a given backend is available on this host."""
+
+    path: PathName
+    available: bool
+    reason: str
+
+    def __bool__(self) -> bool:  # truthy when available
+        return self.available
+
+
+def env_override(env_var: str) -> str | None:
+    """Read an env override (e.g., 'path_a', 'path_b', ...). None = auto."""
+    value = os.environ.get(env_var, "").strip().lower()
+    if not value or value == "auto":
+        return None
+    if value not in ("path_a", "path_b", "path_c", "path_d", "path_e"):
+        raise ValueError(
+            f"unsupported {env_var}={value!r}; "
+            "expected one of: path_a, path_b, path_c, path_d, path_e, auto"
+        )
+    return value
+
+
+def auto_pick(
+    statuses: dict[PathName, PathStatus],
+    preference: tuple[PathName, ...] = ("path_c", "path_b", "path_e", "path_d", "path_a"),
+) -> PathName:
+    """Pick the first available path in the preference order."""
+    for path in preference:
+        st = statuses.get(path)
+        if st is not None and st.available:
+            return path
+    # Always falls back to path_a (the pure-MLX reference is always available).
+    return "path_a"

cppmega_v4/_tilelang/benchmark_receipt.py

@@ -0,0 +1,114 @@
+"""ROI 3.7 — Head-to-head benchmark receipt schema for GDN/KDA paths.
+
+Mirrors the JSON shape used by ``reports/raw/cppmega_1b_path_matrix_cells/``
+so the v4 benchmark plugs straight into the existing matrix HTML render.
+The full benchmark runner (training cell via ``scripts/m04_train_step.py`` on
+``data/parquet_samples/gb10/clang_semantic_4k_v10/val_00000.parquet``) lands
+when Paths B/C/D/E for GDN/KDA become non-fallback; this module ships the
+schema + dispatch-aware single-cell measurement now so the harness can hook
+in incrementally.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Literal
+
+import mlx.core as mx
+
+from cppmega_v4._tilelang.kda_paths import kda_recurrent_dispatch
+from cppmega_v4._tilelang.linear_attention_paths import gated_delta_recurrent_dispatch
+
+Block = Literal["gdn", "kda"]
+Path_ = Literal["path_a", "path_b", "path_c", "path_d", "path_e"]
+
+
+@dataclass(frozen=True)
+class CellShape:
+    block: Block
+    path: Path_
+    batch: int
+    seq_len: int
+    num_heads: int
+    head_dim_k: int
+    head_dim_v: int
+    num_v_heads: int | None = None
+    dtype: str = "float32"
+
+
+@dataclass
+class CellReceipt:
+    """One (block, path, shape, dtype) measurement cell."""
+
+    cell_shape: CellShape
+    fwd_seconds: float
+    backend_available: bool
+    backend_reason: str
+    output_shape: tuple[int, ...]
+    output_dtype: str
+
+    def to_json(self) -> dict[str, object]:
+        return {
+            "cell_shape": asdict(self.cell_shape),
+            "fwd_seconds": self.fwd_seconds,
+            "backend_available": self.backend_available,
+            "backend_reason": self.backend_reason,
+            "output_shape": list(self.output_shape),
+            "output_dtype": self.output_dtype,
+        }
+
+
+def measure_cell(shape: CellShape) -> CellReceipt:
+    """Run a single forward and emit a receipt."""
+    from cppmega_v4._tilelang.kda_paths import kda_path_statuses
+    from cppmega_v4._tilelang.linear_attention_paths import linear_attention_path_statuses
+
+    statuses = (
+        linear_attention_path_statuses() if shape.block == "gdn" else kda_path_statuses()
+    )
+    st = statuses[shape.path]
+    if shape.block == "gdn":
+        q = mx.random.normal((shape.batch, shape.seq_len, shape.num_heads, shape.head_dim_k))
+        k = mx.random.normal((shape.batch, shape.seq_len, shape.num_heads, shape.head_dim_k))
+        v = mx.random.normal((shape.batch, shape.seq_len, shape.num_heads, shape.head_dim_v))
+        beta = mx.random.normal((shape.batch, shape.seq_len, shape.num_heads))
+        g = mx.random.normal((shape.batch, shape.seq_len, shape.num_heads)) * 0.1
+        t0 = time.perf_counter()
+        o, _ = gated_delta_recurrent_dispatch(q, k, v, beta, g)
+        mx.eval(o)
+        elapsed = time.perf_counter() - t0
+    else:  # kda
+        hv = shape.num_v_heads if shape.num_v_heads is not None else shape.num_heads
+        q = mx.random.normal((shape.batch, shape.seq_len, shape.num_heads, shape.head_dim_k))
+        k = mx.random.normal((shape.batch, shape.seq_len, shape.num_heads, shape.head_dim_k))
+        v = mx.random.normal((shape.batch, shape.seq_len, hv, shape.head_dim_v))
+        g = mx.random.normal((shape.batch, shape.seq_len, hv, shape.head_dim_k)) * 0.05
+        beta = mx.random.normal((shape.batch, shape.seq_len, hv))
+        t0 = time.perf_counter()
+        o, _ = kda_recurrent_dispatch(q, k, v, g, beta)
+        mx.eval(o)
+        elapsed = time.perf_counter() - t0
+
+    return CellReceipt(
+        cell_shape=shape,
+        fwd_seconds=elapsed,
+        backend_available=st.available,
+        backend_reason=st.reason,
+        output_shape=tuple(o.shape),
+        output_dtype=str(o.dtype),
+    )
+
+
+def write_receipt(receipt: CellReceipt, out_dir: Path) -> Path:
+    """Drop the receipt JSON into ``out_dir/<block>_<path>.json``."""
+    out_dir.mkdir(parents=True, exist_ok=True)
+    fname = f"{receipt.cell_shape.block}_{receipt.cell_shape.path}.json"
+    out_path = out_dir / fname
+    out_path.write_text(json.dumps(receipt.to_json(), indent=2))
+    return out_path
+
+
+__all__ = ["Block", "CellReceipt", "CellShape", "Path_", "measure_cell", "write_receipt"]

cppmega_v4/_tilelang/kda_paths.py

@@ -0,0 +1,118 @@
+"""KDA multi-path scaffolding (Paths B/C/D + auto-mode).
+
+Same shape as ``linear_attention_paths.py`` but for the KDA backend. KDA has
+no mlx-lm equivalent op (no Path E).
+
+Backend status (May 2026):
+    - Path A: pure-MLX naive recurrent KDA (golden) — always available.
+    - Path B: hand-MSL KDA kernel — scaffold; pending Metal kernel.
+    - Path C: TileLang DSL — scaffold; should lift ``tilelang/examples/kda/*.py``
+      (11 prims, ~3340 LoC) through ``tilelang.compile(target='metal',
+      execution_backend='tvm_ffi')`` mirroring ``mamba3_path_c.py`` skeleton.
+    - Path D: Triton frontend — scaffold; should lift FLA KDA Triton kernels.
+
+Env override: ``CPPMEGA_V4_KERNEL_PATH__KDA``.
+"""
+
+from __future__ import annotations
+
+import importlib
+
+from cppmega_v4._tilelang._dispatch import PathName, PathStatus, auto_pick, env_override
+from cppmega_v4.nn._external.fla_naive_kda import naive_recurrent_kda
+
+ENV_VAR = "CPPMEGA_V4_KERNEL_PATH__KDA"
+
+
+def _path_a_status() -> PathStatus:
+    return PathStatus(path="path_a", available=True, reason="pure-MLX KDA reference")
+
+
+def _path_b_status() -> PathStatus:
+    return PathStatus(
+        path="path_b",
+        available=False,
+        reason=(
+            "hand-MSL KDA kernel pending — DPLR transport A = αI + β u v^T plus "
+            "gated-LA branch needs Metal kernel work beyond GDN's recurrence"
+        ),
+    )
+
+
+def _path_c_status() -> PathStatus:
+    try:
+        importlib.import_module("tilelang")
+    except Exception as exc:
+        return PathStatus(
+            path="path_c",
+            available=False,
+            reason=f"tilelang not importable: {exc}",
+        )
+    return PathStatus(
+        path="path_c",
+        available=False,
+        reason=(
+            "TileLang importable but KDA Path C wiring pending: lift "
+            "tilelang/examples/kda/{chunk_delta_h_fwd,chunk_delta_bwd,chunk_o,"
+            "chunk_bwd_intra,chunk_inter_solve_fused,chunk_bwd_dqkwg,"
+            "chunk_bwd_dv,chunk_bwd_gla_dA,chunk_intra_token_parallel,"
+            "wy_fast,wy_fast_bwd}.py via tilelang.compile(target='metal')"
+        ),
+    )
+
+
+def _path_d_status() -> PathStatus:
+    try:
+        importlib.import_module("triton")
+    except Exception:
+        return PathStatus(
+            path="path_d",
+            available=False,
+            reason="triton not importable",
+        )
+    return PathStatus(
+        path="path_d",
+        available=False,
+        reason="Triton frontend KDA wiring pending",
+    )
+
+
+def _fallback(*args, **kwargs):
+    return naive_recurrent_kda(*args, **kwargs)
+
+
+def kda_path_statuses() -> dict[PathName, PathStatus]:
+    return {
+        "path_a": _path_a_status(),
+        "path_b": _path_b_status(),
+        "path_c": _path_c_status(),
+        "path_d": _path_d_status(),
+    }
+
+
+def kda_auto_mode_for_inputs(*, env_var: str = ENV_VAR) -> PathName:
+    forced = env_override(env_var)
+    if forced is not None:
+        if forced == "path_e":
+            raise ValueError("KDA has no Path E; use path_a, path_b, path_c, path_d, or auto")
+        return forced  # type: ignore[return-value]
+    return auto_pick(
+        kda_path_statuses(),
+        preference=("path_c", "path_b", "path_d", "path_a"),
+    )
+
+
+def kda_recurrent_dispatch(*args, **kwargs):
+    """Call the auto-selected KDA backend, falling back to Path A."""
+    path = kda_auto_mode_for_inputs()
+    # All B/C/D currently delegate to Path A.
+    del path
+    return _fallback(*args, **kwargs)
+
+
+__all__ = [
+    "ENV_VAR",
+    "kda_auto_mode_for_inputs",
+    "kda_path_statuses",
+    "kda_recurrent_dispatch",
+]