feat(v4): KDA Path E via the SAME mlx-lm gated_delta kernel as GDN

Earlier I claimed KDA has no Path E because the upstream mlx-lm
gated_delta_update only takes a per-head scalar gate. That was wrong —
mlx_lm/models/gated_delta.py:_make_gated_delta_kernel has a `vectorized`
flag, and gated_delta_kernel auto-selects the vectorised-gate variant
when g.ndim == 4. mlx-lm itself drives this branch from
kimi_linear.py::KimiDeltaAttention, which is the KDA-style attention.

This commit adds the matching adapter + dispatch wiring:

  - cppmega_v4/nn/_external/mlx_lm_kda_update.py: thin adapter that
    maps our FLA-style (q, k, v, g_per_K, beta) onto the upstream
    (q_scaled, k, v, g_decay=exp(g), beta, state_transposed) call.
    State layout flips between FLA [B,HV,K,V] and upstream [B,Hv,Dv,Dk].
    Smaller dims (Dk%32!=0 or Dv%4!=0) drop to the pure-ops reference.
  - cppmega_v4/_tilelang/kda_paths.py: _path_e_status / _path_e_call /
    register in kda_path_statuses(); preference tuple updated to
    ("path_c", "path_b", "path_e", "path_d", "path_a"); the previous
    'KDA has no Path E' ValueError on env-forced path_e is gone.
  - cppmega_v4/_tilelang/benchmark_matrix.py: KDA_PATHS includes path_e.
  - tests:
    - new KDA Path E tests: status, kernel-path parity, ops-fallback
      parity, env-forced dispatch
    - test_kda_statuses_keys (both copies) updated to expect the 5-key set
    - test_kda_auto_mode_rejects_path_e replaced with
      test_kda_auto_mode_accepts_path_e
    - test_run_matrix_kda_covers_4_paths_no_path_e renamed +
      assertion bumped to 5 cells

Parity smoke: KDA path_e vs path_a maxdiff ~2e-6 (B=1,T=4,H=2,HV=4,K=V=32).

Suite: 371 passed / 0 skipped (was 367 passed / 0 skipped).

Author: apstenku123

cppmega_v4/_tilelang/benchmark_matrix.py

@@ -51,7 +51,7 @@
 }
 
 GDN_PATHS: tuple[Path_, ...] = ("path_a", "path_b", "path_c", "path_d", "path_e")
-KDA_PATHS: tuple[Path_, ...] = ("path_a", "path_b", "path_c", "path_d")
+KDA_PATHS: tuple[Path_, ...] = ("path_a", "path_b", "path_c", "path_d", "path_e")
 
 
 @dataclass

cppmega_v4/_tilelang/kda_paths.py

@@ -1,7 +1,12 @@
-"""KDA multi-path dispatch (Paths B/C/D + auto-mode).
+"""KDA multi-path dispatch (Paths B/C/D/E + auto-mode).
 
-Same shape as ``linear_attention_paths.py`` but for the KDA backend. KDA
-has no mlx-lm equivalent op (no Path E).
+Same shape as ``linear_attention_paths.py`` but for the KDA backend. Path E
+reuses the SAME mlx-lm gated_delta Metal kernel as GDN — the kernel
+auto-selects the vectorised-gate variant when ``g.ndim == 4``
+(``mlx_lm/models/gated_delta.py::gated_delta_kernel``), which is exactly
+KDA's per-K log-decay shape. Mirrors how
+``mlx_lm/models/kimi_linear.py::KimiDeltaAttention`` drives the same
+upstream kernel from KDA-shaped inputs.
 
 Backend status (May 2026):
     - Path A: pure-MLX naive recurrent KDA (golden) — always available.
@@ -16,6 +21,10 @@
     - Path D: ``poc.triton_frontend.from_triton_kernel`` over FLA KDA
       chunk kernels. cppmega owns the runtime adapter for FLA's forward
       multi-kernel launch, including packed varlen metadata.
+    - Path E: vendored mlx-lm gated_delta vectorised-gate Metal kernel
+      (the same kernel that powers GDN Path E; KDA just hits the
+      ``g.ndim == 4`` branch). Requires Dk % 32 == 0 and Dv % 4 == 0;
+      smaller dims drop to the pure-ops reference path.
 
 Env override: ``CPPMEGA_V4_KERNEL_PATH__KDA``.
 """
@@ -137,6 +146,43 @@ def _path_d_call(*args, **kwargs):
         return _path_a_call(*args, **kwargs)
 
 
+# ----- Path E (vendored mlx-lm gated_delta vectorised-gate kernel) -----
+
+
+def _path_e_status() -> PathStatus:
+    try:
+        importlib.import_module(
+            "cppmega_v4.nn._external.mlx_lm_kda_update"
+        )
+    except Exception as exc:
+        return PathStatus(
+            path="path_e", available=False,
+            reason=(
+                "mlx-lm gated_delta vectorised-gate adapter not importable: "
+                f"{exc}"
+            ),
+        )
+    return PathStatus(
+        path="path_e", available=True,
+        reason=(
+            "vendored mlx-lm gated_delta vectorised-gate Metal kernel "
+            "(Dk%32==0 & Dv%4==0; smaller dims fall back to the pure-ops "
+            "reference). The same kernel as GDN Path E — KDA hits the "
+            "g.ndim==4 branch in gated_delta_kernel."
+        ),
+    )
+
+
+def _path_e_call(*args, **kwargs):
+    if not _path_e_status().available:
+        return _path_a_call(*args, **kwargs)
+    try:
+        from cppmega_v4.nn._external.mlx_lm_kda_update import kda_update
+        return kda_update(*args, **kwargs)
+    except Exception:
+        return _path_a_call(*args, **kwargs)
+
+
 # ----- Public dispatch -----
 
 
@@ -146,20 +192,17 @@ def kda_path_statuses() -> dict[PathName, PathStatus]:
         "path_b": _path_b_status(),
         "path_c": _path_c_status(),
         "path_d": _path_d_status(),
+        "path_e": _path_e_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"),
+        preference=("path_c", "path_b", "path_e", "path_d", "path_a"),
     )
 
 
@@ -171,6 +214,7 @@ def kda_recurrent_dispatch(*args, **kwargs):
         "path_b": _path_b_call,
         "path_c": _path_c_call,
         "path_d": _path_d_call,
+        "path_e": _path_e_call,
     }[path]
     return fn(*args, **kwargs)
 

cppmega_v4/nn/_external/mlx_lm_kda_update.py

@@ -0,0 +1,134 @@
+"""KDA Path E adapter — wraps mlx-lm's gated_delta kernel (vectorised gate).
+
+The same upstream Metal kernel (``mlx_lm/models/gated_delta.py``,
+``_make_gated_delta_kernel(vectorized=True)``) that powers GDN Path E also
+handles the per-K vectorised gate used by KDA-style attention — see
+``mlx_lm/models/kimi_linear.py::KimiDeltaAttention`` which calls
+``gated_delta_update`` with ``a_logits`` shaped ``[B, T, num_heads, head_dim]``.
+
+The kernel selection happens automatically inside ``gated_delta_kernel``:
+when ``g.ndim == 4`` it picks ``_gated_delta_kernel_vec``; when ``g.ndim == 3``
+it picks the scalar variant used by GDN. Group expansion (Hv > Hk) is also
+handled inside the kernel (``hk_idx = hv_idx / (Hv / Hk)``).
+
+Our naive_recurrent_kda signature (FLA convention):
+    naive_recurrent_kda(q, k, v, g, beta, scale=None,
+                        initial_state=None, output_final_state=False)
+                        -> (o, S)
+    where:
+      - q, k: [B, T, H,  K]                 (un-repeated; HV/H expansion is kernel-side)
+      - v:    [B, T, HV, V]
+      - g:    [B, T, HV, K]  (per-K log-decay; state is multiplied by exp(g))
+      - beta: [B, T, HV]
+      - S:    [B, HV, K, V]  (FLA layout — K rows, V cols)
+    FLA pre-scales q by 1/sqrt(K).
+
+Upstream gated_delta_kernel signature:
+    gated_delta_kernel(q, k, v, g, beta, state, mask=None)
+      - q, k: [B, T, Hk, Dk]
+      - v:    [B, T, Hv, Dv]
+      - g:    [B, T, Hv, Dk]  (vectorized)  OR  [B, T, Hv] (scalar)
+      - beta: [B, T, Hv]      (kernel applies it as-is, no sigmoid)
+      - state: [B, Hv, Dv, Dk]  (upstream layout — V rows, K cols; transposed vs FLA)
+    Upstream does NOT pre-scale q. Decay is plain exp(g) where g is the value
+    we pass in directly (kernel does ``state[i] = state[i] * g_[s_idx]``),
+    so we pass ``g_decay = exp(our_g)`` directly here.
+
+Constraints:
+    - Upstream Metal kernel requires Dk % 32 == 0 and Dv % 4 == 0; smaller
+      dims fall back to the pure-ops reference path.
+"""
+
+from __future__ import annotations
+
+import math
+
+import mlx.core as mx
+
+from cppmega_v4.nn._external._mlx_lm_gated_delta_vendored import (
+    gated_delta_kernel as _upstream_kernel,
+    gated_delta_ops as _upstream_ops,
+)
+
+
+def kda_update(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    *,
+    scale: float | None = None,
+    initial_state: mx.array | None = None,
+    output_final_state: bool = False,
+):
+    """Path E entry — same signature as ``naive_recurrent_kda``.
+
+    Args:
+        q, k:          [B, T, H,  K]  (FLA convention; HV/H expansion is kernel-side)
+        v:             [B, T, HV, V]
+        g:             [B, T, HV, K]  per-K log-decay (state *= exp(g) per step)
+        beta:          [B, T, HV]
+        scale:         optional — defaults to 1/sqrt(K)
+        initial_state: optional [B, HV, K, V]
+        output_final_state: if True, return (o, S_final[B, HV, K, V])
+
+    Returns:
+        (o[B, T, HV, V], S_final or None)
+    """
+    if q.ndim != 4 or k.shape != q.shape:
+        raise ValueError(
+            f"q/k must be [B, T, H, K]; got q={q.shape}, k={k.shape}"
+        )
+    if v.ndim != 4 or v.shape[:2] != q.shape[:2]:
+        raise ValueError(f"v must be [B, T, HV, V]; got v={v.shape}")
+    if g.shape != (*v.shape[:3], k.shape[-1]):
+        raise ValueError(
+            f"g must be [B, T, HV, K]; got g={g.shape}, expected {(*v.shape[:3], k.shape[-1])}"
+        )
+    if beta.shape != v.shape[:3]:
+        raise ValueError(f"beta must be [B, T, HV]; got beta={beta.shape}")
+
+    b_size, t_size, h_size, kdim = q.shape
+    hv_size, vdim = v.shape[2], v.shape[-1]
+    if hv_size % h_size != 0:
+        raise ValueError(f"HV ({hv_size}) must be divisible by H ({h_size})")
+
+    # FLA pre-scales q by 1/sqrt(K); upstream does not.
+    fla_scale = scale if scale is not None else 1.0 / math.sqrt(kdim)
+    q_scaled = (q.astype(mx.float32) * fla_scale).astype(q.dtype)
+
+    # Upstream kernel multiplies state by `g` as-is each step. FLA convention:
+    # state is multiplied by `exp(g)`. Pre-exponentiate.
+    g_decay = mx.exp(g.astype(mx.float32))
+
+    # Upstream state layout: [B, Hv, Dv, Dk]. FLA: [B, HV, K, V]. Transpose
+    # last two axes when handing off.
+    if initial_state is None:
+        state = mx.zeros((b_size, hv_size, vdim, kdim), dtype=mx.float32)
+    else:
+        state = mx.transpose(initial_state.astype(mx.float32), (0, 1, 3, 2))
+
+    beta_f = beta.astype(mx.float32)
+
+    # Upstream Metal kernel needs Dk % 32 == 0 and Dv % 4 == 0; otherwise
+    # fall through to the pure-ops reference (which also handles vector-gate).
+    use_kernel = (kdim % 32 == 0) and (vdim % 4 == 0) and mx.metal.is_available()
+
+    if use_kernel:
+        y, new_state = _upstream_kernel(
+            q_scaled, k, v, g_decay, beta_f, state, mask=None,
+        )
+    else:
+        y, new_state = _upstream_ops(
+            q_scaled, k, v, g_decay, beta_f, state, mask=None,
+        )
+
+    final = None
+    if output_final_state:
+        # Convert back to FLA layout [B, HV, K, V].
+        final = mx.transpose(new_state, (0, 1, 3, 2))
+    return y, final
+
+
+__all__ = ["kda_update"]

tests/v4/test_benchmark_matrix.py

@@ -125,13 +125,13 @@ def test_run_matrix_gdn_covers_all_5_paths():
     assert receipt.promotion.env_var.endswith("LINEAR_ATTENTION")
 
 
-def test_run_matrix_kda_covers_4_paths_no_path_e():
+def test_run_matrix_kda_covers_all_5_paths():
     shape = CellShape(
         block="kda", path="path_a", batch=1, seq_len=4,
         num_heads=2, head_dim_k=4, head_dim_v=4, num_v_heads=2,
     )
     receipt = run_matrix(shape, warmup=1, iters=2)
-    assert len(receipt.cells) == 4
+    assert len(receipt.cells) == 5
     paths_seen = {c.path for c in receipt.cells}
     assert paths_seen == set(KDA_PATHS)
     assert receipt.promotion.env_var.endswith("KDA")

tests/v4/test_kda_paths.py

@@ -326,5 +326,63 @@ def test_kda_path_d_try_lower_returns_seam_message():
 # ----- Dispatch sanity -----
 
 
+# ----- Path E (vendored mlx-lm gated_delta vec-gate kernel) -----
+
+
+def test_kda_path_e_status_available():
+    from cppmega_v4._tilelang.kda_paths import _path_e_status
+    st = _path_e_status()
+    assert st.available, st.reason
+    assert "gated_delta" in st.reason
+
+
+def test_kda_path_e_parity_with_path_a_kernel_path():
+    """Dk%32==0 + Dv%4==0 hits the Metal kernel — must match Path A."""
+    from cppmega_v4.nn._external.mlx_lm_kda_update import kda_update
+    B, T, H, HV, K, V = 1, 4, 2, 4, 32, 32
+    rng = np.random.default_rng(123)
+    q = mx.array(rng.standard_normal((B, T, H, K)).astype(np.float32))
+    k = mx.array(rng.standard_normal((B, T, H, K)).astype(np.float32))
+    v = mx.array(rng.standard_normal((B, T, HV, V)).astype(np.float32))
+    g = mx.array(-rng.uniform(0.01, 0.2, (B, T, HV, K)).astype(np.float32))
+    beta = mx.array(rng.uniform(0.1, 0.9, (B, T, HV)).astype(np.float32))
+    o_e, sf_e = kda_update(q, k, v, g, beta, output_final_state=True)
+    o_a, sf_a = naive_recurrent_kda(q, k, v, g, beta, output_final_state=True)
+    mx.eval(o_e, o_a, sf_e, sf_a)
+    np.testing.assert_allclose(np.array(o_e), np.array(o_a), atol=1e-4, rtol=1e-4)
+    np.testing.assert_allclose(np.array(sf_e), np.array(sf_a), atol=1e-4, rtol=1e-4)
+
+
+def test_kda_path_e_parity_ops_fallback():
+    """Small dims (Dk%32!=0) fall through to the pure-ops reference."""
+    from cppmega_v4.nn._external.mlx_lm_kda_update import kda_update
+    B, T, H, HV, K, V = 1, 3, 2, 4, 8, 8
+    rng = np.random.default_rng(131)
+    q = mx.array(rng.standard_normal((B, T, H, K)).astype(np.float32))
+    k = mx.array(rng.standard_normal((B, T, H, K)).astype(np.float32))
+    v = mx.array(rng.standard_normal((B, T, HV, V)).astype(np.float32))
+    g = mx.array(-rng.uniform(0.01, 0.2, (B, T, HV, K)).astype(np.float32))
+    beta = mx.array(rng.uniform(0.1, 0.9, (B, T, HV)).astype(np.float32))
+    o_e, _ = kda_update(q, k, v, g, beta)
+    o_a, _ = naive_recurrent_kda(q, k, v, g, beta)
+    np.testing.assert_allclose(np.array(o_e), np.array(o_a), atol=1e-4, rtol=1e-4)
+
+
+def test_kda_path_e_forced_via_env(monkeypatch):
+    monkeypatch.setenv(KDA_ENV, "path_e")
+    B, T, H, HV, K, V = 1, 4, 2, 4, 32, 32
+    rng = np.random.default_rng(141)
+    q = mx.array(rng.standard_normal((B, T, H, K)).astype(np.float32))
+    k = mx.array(rng.standard_normal((B, T, H, K)).astype(np.float32))
+    v = mx.array(rng.standard_normal((B, T, HV, V)).astype(np.float32))
+    g = mx.array(-rng.uniform(0.01, 0.2, (B, T, HV, K)).astype(np.float32))
+    beta = mx.array(rng.uniform(0.1, 0.9, (B, T, HV)).astype(np.float32))
+    o, _ = kda_recurrent_dispatch(q, k, v, g, beta)
+    assert o.shape == (B, T, HV, V)
+    assert not bool(mx.any(mx.isnan(o)).item())
+
+
 def test_kda_statuses_keys_unchanged():
-    assert set(kda_path_statuses().keys()) == {"path_a", "path_b", "path_c", "path_d"}
+    assert set(kda_path_statuses().keys()) == {
+        "path_a", "path_b", "path_c", "path_d", "path_e",
+    }

tests/v4/test_path_dispatch.py

@@ -145,19 +145,19 @@ def test_gdn_dispatch_each_path_runs(monkeypatch, path):
 
 
 def test_kda_statuses_keys():
-    # KDA has no Path E.
+    # KDA now has Path E (same upstream Metal kernel as GDN; KDA hits the
+    # vectorised-gate branch via g.ndim==4).
     statuses = kda_path_statuses()
-    assert set(statuses.keys()) == {"path_a", "path_b", "path_c", "path_d"}
+    assert set(statuses.keys()) == {"path_a", "path_b", "path_c", "path_d", "path_e"}
 
 
 def test_kda_path_a_always_available():
     assert kda_path_statuses()["path_a"].available
 
 
-def test_kda_auto_mode_rejects_path_e(monkeypatch):
+def test_kda_auto_mode_accepts_path_e(monkeypatch):
     monkeypatch.setenv(KDA_ENV, "path_e")
-    with pytest.raises(ValueError, match="no Path E"):
-        kda_auto_mode_for_inputs()
+    assert kda_auto_mode_for_inputs() == "path_e"
 
 
 def test_kda_dispatch_returns_same_as_path_a(monkeypatch):