feat(v4): wire GDN/KDA blocks through path dispatch — env-override actually works now

LinearAttentionBlock.__call__ and KimiDeltaAttentionBlock.__call__
previously hardcoded naive_recurrent_* (Path A) regardless of the
CPPMEGA_V4_KERNEL_PATH__LINEAR_ATTENTION / __KDA env var. All the
Path B/C/E speedup work was effectively dead inside the
UnifiedSuperblockV4 stack.

Both blocks now route through gated_delta_recurrent_dispatch /
kda_recurrent_dispatch, which honour the env override and auto-pick
the fastest available backend (Path B/C/E when Metal/tilelang/vendored
mlx-lm op are reachable, falling back to Path A).

The doc_ids state-reset path also dispatches through the same routing
(per-segment).

Tests (9 new in test_block_dispatch_env_override.py):
  - GDN block honours env override for path_a / path_b
  - KDA block honours env override for path_a / path_b
  - path_a vs path_b outputs close within float32 simd_sum drift
    (atol=5e-4 rtol=5e-3)
  - default auto-mode produces valid output (no env override)
  - doc_ids path also dispatches through the path system

v4 dispatch tests: 9/9 passed.

Author: apstenku123

cppmega_v4/nn/kimi_delta_attention.py

@@ -150,7 +150,8 @@ def __call__(self, x: mx.array, *, doc_ids: mx.array | None = None) -> mx.array:
         g = g.reshape(batch, seq_len, cfg._num_v_heads, cfg.head_k_dim)
 
         if doc_ids is None:
-            o, _ = naive_recurrent_kda(q, k, v, g, beta)
+            from cppmega_v4._tilelang.kda_paths import kda_recurrent_dispatch
+            o, _ = kda_recurrent_dispatch(q, k, v, g, beta)
         else:
             o = self._recurrent_with_doc_reset(q, k, v, g, beta, doc_ids)
 
@@ -178,7 +179,8 @@ def _recurrent_with_doc_reset(q, k, v, g, beta, doc_ids: mx.array) -> mx.array:
             runs.append((start, seq_len))
             outs = []
             for s, e in runs:
-                ob, _ = naive_recurrent_kda(
+                from cppmega_v4._tilelang.kda_paths import kda_recurrent_dispatch
+                ob, _ = kda_recurrent_dispatch(
                     q[b:b + 1, s:e],
                     k[b:b + 1, s:e],
                     v[b:b + 1, s:e],

cppmega_v4/nn/linear_attention.py

@@ -181,10 +181,17 @@ def __call__(self, x: mx.array, *, doc_ids: mx.array | None = None) -> mx.array:
         g = mx.log(mx.sigmoid(self.a_proj(x)) + cfg.norm_eps)  # (B, T, num_v_heads)
 
         if doc_ids is None:
-            o, _ = naive_recurrent_gated_delta_rule(q, k, v, beta, g)
+            # Dispatch through the path system: env-overridable via
+            # CPPMEGA_V4_KERNEL_PATH__LINEAR_ATTENTION, auto-mode picks the
+            # fastest available backend (Path B/C/E if Metal/tilelang/mlx-lm
+            # vendored op are reachable; falls back to Path A).
+            from cppmega_v4._tilelang.linear_attention_paths import (
+                gated_delta_recurrent_dispatch,
+            )
+            o, _ = gated_delta_recurrent_dispatch(q, k, v, beta, g)
         else:
             # Document-boundary state reset: split into runs of contiguous
-            # doc_id, recur within each, concatenate outputs.
+            # doc_id, dispatch each run through the path system, concat.
             o = self._recurrent_with_doc_reset(q, k, v, beta, g, doc_ids)
 
         # o has shape [B, T, H_v, V_dim]; flatten head axis for o_proj.
@@ -232,7 +239,10 @@ def _recurrent_with_doc_reset(
                 vb = v[b:b + 1, s:e]
                 bb = beta[b:b + 1, s:e]
                 gb = g[b:b + 1, s:e]
-                ob, _ = naive_recurrent_gated_delta_rule(qb, kb, vb, bb, gb)
+                from cppmega_v4._tilelang.linear_attention_paths import (
+                    gated_delta_recurrent_dispatch,
+                )
+                ob, _ = gated_delta_recurrent_dispatch(qb, kb, vb, bb, gb)
                 run_outputs.append(ob)
             per_batch_outputs.append(mx.concatenate(run_outputs, axis=1))
         return mx.concatenate(per_batch_outputs, axis=0)

tests/v4/test_block_dispatch_env_override.py

@@ -0,0 +1,115 @@
+"""Tests that LinearAttentionBlock / KimiDeltaAttentionBlock honour the
+CPPMEGA_V4_KERNEL_PATH__* env override (i.e. the block dispatches through
+linear_attention_paths / kda_paths, not the naive recurrent kernel directly).
+"""
+
+import mlx.core as mx
+import numpy as np
+import pytest
+
+from cppmega_v4._tilelang.kda_paths import ENV_VAR as KDA_ENV
+from cppmega_v4._tilelang.linear_attention_paths import ENV_VAR as GDN_ENV
+from cppmega_v4.nn.kimi_delta_attention import (
+    KimiDeltaAttentionBlock,
+    KimiDeltaAttentionConfig,
+)
+from cppmega_v4.nn.linear_attention import (
+    LinearAttentionBlock,
+    LinearAttentionConfig,
+)
+
+
+def _gdn_block(d=64, h=4):
+    cfg = LinearAttentionConfig(
+        hidden_size=d, num_heads=h, head_dim=d // h, use_short_conv=False,
+    )
+    blk = LinearAttentionBlock(cfg)
+    # Perturb o_proj so block delta is non-zero (default is zero-init).
+    rng = np.random.default_rng(0)
+    blk.o_proj.weight = mx.array(
+        rng.standard_normal(blk.o_proj.weight.shape).astype(np.float32) * 0.05
+    )
+    return blk
+
+
+def _kda_block(d=64, h=4):
+    cfg = KimiDeltaAttentionConfig(
+        hidden_size=d, num_heads=h, head_dim=d // h, use_short_conv=False,
+    )
+    blk = KimiDeltaAttentionBlock(cfg)
+    rng = np.random.default_rng(0)
+    blk.o_proj.weight = mx.array(
+        rng.standard_normal(blk.o_proj.weight.shape).astype(np.float32) * 0.05
+    )
+    return blk
+
+
+@pytest.mark.parametrize("path", ["path_a", "path_b"])
+def test_gdn_block_honours_env_override(monkeypatch, path):
+    monkeypatch.setenv(GDN_ENV, path)
+    blk = _gdn_block()
+    x = mx.random.normal((1, 8, blk.config.hidden_size))
+    out = blk(x)
+    assert out.shape == x.shape
+    assert not bool(mx.any(mx.isnan(out)).item())
+
+
+@pytest.mark.parametrize("path", ["path_a", "path_b"])
+def test_kda_block_honours_env_override(monkeypatch, path):
+    monkeypatch.setenv(KDA_ENV, path)
+    blk = _kda_block()
+    x = mx.random.normal((1, 8, blk.config.hidden_size))
+    out = blk(x)
+    assert out.shape == x.shape
+    assert not bool(mx.any(mx.isnan(out)).item())
+
+
+def test_gdn_block_path_a_and_path_b_close(monkeypatch):
+    """Block output via Path A and Path B should be numerically close."""
+    blk = _gdn_block()
+    x = mx.random.normal((1, 8, blk.config.hidden_size))
+
+    monkeypatch.setenv(GDN_ENV, "path_a")
+    out_a = np.array(blk(x))
+    monkeypatch.setenv(GDN_ENV, "path_b")
+    out_b = np.array(blk(x))
+    np.testing.assert_allclose(out_a, out_b, atol=5e-4, rtol=5e-3)
+
+
+def test_kda_block_path_a_and_path_b_close(monkeypatch):
+    blk = _kda_block()
+    x = mx.random.normal((1, 8, blk.config.hidden_size))
+
+    monkeypatch.setenv(KDA_ENV, "path_a")
+    out_a = np.array(blk(x))
+    monkeypatch.setenv(KDA_ENV, "path_b")
+    out_b = np.array(blk(x))
+    np.testing.assert_allclose(out_a, out_b, atol=5e-4, rtol=5e-3)
+
+
+def test_gdn_block_default_auto_picks_available_path(monkeypatch):
+    """Without env override, auto-mode should still produce valid output."""
+    monkeypatch.delenv(GDN_ENV, raising=False)
+    blk = _gdn_block()
+    x = mx.random.normal((1, 4, blk.config.hidden_size))
+    out = blk(x)
+    assert not bool(mx.any(mx.isnan(out)).item())
+
+
+def test_kda_block_default_auto_picks_available_path(monkeypatch):
+    monkeypatch.delenv(KDA_ENV, raising=False)
+    blk = _kda_block()
+    x = mx.random.normal((1, 4, blk.config.hidden_size))
+    out = blk(x)
+    assert not bool(mx.any(mx.isnan(out)).item())
+
+
+def test_gdn_block_with_doc_ids_uses_dispatch(monkeypatch):
+    """doc_ids path also dispatches through linear_attention_paths."""
+    monkeypatch.setenv(GDN_ENV, "path_b")
+    blk = _gdn_block()
+    x = mx.random.normal((1, 8, blk.config.hidden_size))
+    doc_ids = mx.array([[0, 0, 0, 0, 1, 1, 1, 1]], dtype=mx.int32)
+    out = blk(x, doc_ids=doc_ids)
+    assert out.shape == x.shape
+    assert not bool(mx.any(mx.isnan(out)).item())