feat(v4): vendor mlx-lm PR #1066 + fix Path C PEP-563 closure-locals

Three landings in one commit, all enabled by switching the test runner
to the nanochat venv (Python 3.13 + tilelang already installed there):

1. mlx-lm PR #1066 — Kahan-compensated kv_mem accumulation + 4-way
   time-loop unroll. Vendored into both:
   - cppmega_v4/nn/_external/_mlx_lm_gated_delta_vendored.py (Path E)
   - cppmega_v4/_tilelang/linear_attention_path_b.py (our hand-MSL)
   Without this, long-T runs accumulate bf16 rounding into kv_mem and
   the delta-corrected state drifts. Cost: 3 FLOPs/iter on the inner
   reduction; the unroll amortizes the loop overhead.

2. Removed `from __future__ import annotations` from
   linear_attention_path_c.py and kda_path_c.py. tilelang's
   @T.prim_func builder calls get_type_hints which evaluates
   T.Tensor((BATCH, SEQ, ...)) annotations against globals — under
   PEP 563 the closure-local BATCH/SEQ/... become unresolvable
   strings and the kernel build fails with NameError. mamba3_path_c.py
   avoids future-annotations for exactly this reason; we now match.

3. Dispatch tests updated: Path C is real-or-fallback (real when
   tilelang importable, fallback otherwise); auto-mode preference
   accepts path_c when available. Renamed
   test_gdn_paths_c_d_currently_deferred → test_gdn_path_d_still_
   deferred_without_triton (Path C no longer in that bucket).

Real Path C prim_func parity test (was skipped) now PASSES against
Path A on the nanochat venv. v4 suite: 178 passed / 2 skipped (was
175 passed / 3 failed).

Author: apstenku123

cppmega_v4/_tilelang/kda_path_c.py

@@ -16,13 +16,15 @@
 host ``cppmega_mlx`` package read-only — never modifies them.
 """
 
-from __future__ import annotations
-
 from functools import lru_cache
-from typing import Any
+from typing import Any, Optional
 
 import mlx.core as mx
 
+# Deliberately NOT using `from __future__ import annotations`: tilelang's
+# @T.prim_func builder evaluates T.Tensor((BATCH, SEQ, ...)) annotations
+# via get_type_hints; PEP 563 string-annotations break closure-local lookup.
+
 
 def _tilelang_importable() -> tuple[bool, str]:
     try:
@@ -139,8 +141,8 @@ def _kda_fwd_path_c_call(
     g: mx.array,
     beta: mx.array,
     *,
-    scale: float | None = None,
-    initial_state: mx.array | None = None,
+    scale: Optional[float] = None,
+    initial_state: Optional[mx.array] = None,
     output_final_state: bool = False,
 ):
     """KDA Path C entry — same signature as ``naive_recurrent_kda``."""

cppmega_v4/_tilelang/linear_attention_path_b.py

@@ -97,10 +97,18 @@ def _gdn_forward_kernel(
             for (int i = 0; i < {dh}; i++) state[i] *= alpha_t;
 
             // Phase 2: kth_S_j = sum_i k[i] * state[i, j]  (own column)
+            // mlx-lm PR #1066: Kahan-compensated summation. Without this,
+            // long-T runs accumulate bf16 rounding into kv_mem and the
+            // delta-corrected state drifts. Compensation costs 3 FLOPs/iter.
             float kth_S_j = 0.0f;
+            float kth_c   = 0.0f;
             for (int i = 0; i < {dh}; i++) {{
                 float k_i = k[kv_base + i];
-                kth_S_j += k_i * state[i];
+                float p   = k_i * state[i];
+                float a   = p - kth_c;
+                float b   = kth_S_j + a;
+                kth_c     = (b - kth_S_j) - a;
+                kth_S_j   = b;
             }}
 
             // Phase 3: v_eff = beta * (v - kth_S)

cppmega_v4/_tilelang/linear_attention_path_c.py

@@ -35,13 +35,16 @@
     y[b, t, h, v_idx] = out
 """
 
-from __future__ import annotations
-
 from functools import lru_cache
-from typing import Any
+from typing import Any, Optional
 
 import mlx.core as mx
 
+# Deliberately NOT using `from __future__ import annotations` here:
+# tilelang's @T.prim_func builder calls get_type_hints which evaluates
+# the T.Tensor((BATCH, SEQ, ...)) annotations against globals — under
+# PEP 563 the closure locals (BATCH/SEQ/...) become unresolvable strings.
+
 
 def _tilelang_importable() -> tuple[bool, str]:
     """Probe whether the full tilelang stack is loadable.
@@ -171,8 +174,8 @@ def _gdn_fwd_path_c_call(
     beta: mx.array,
     g: mx.array,
     *,
-    scale: float | None = None,
-    initial_state: mx.array | None = None,
+    scale: Optional[float] = None,
+    initial_state: Optional[mx.array] = None,
     output_final_state: bool = False,
 ):
     """Path C entry — same signature as ``naive_recurrent_gated_delta_rule``.

cppmega_v4/nn/_external/_mlx_lm_gated_delta_vendored.py

@@ -71,39 +71,47 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False):
         {g_setup}
         auto beta_ = beta + b_idx * T * Hv;
 
-        for (int t = 0; t < T; ++t) {{
-          if ({mask_source}) {{
-            float kv_mem = 0.0f;
-            for (int i = 0; i < n_per_t; ++i) {{
-              auto s_idx = n_per_t * dk_idx + i;
-              state[i] = state[i] * {g_access};
-              kv_mem += state[i] * k_[s_idx];
-            }}
-            kv_mem = simd_sum(kv_mem);
-
-            auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
-
-            float out = 0.0f;
-            for (int i = 0; i < n_per_t; ++i) {{
-              auto s_idx = n_per_t * dk_idx + i;
-              state[i] = state[i] + k_[s_idx] * delta;
-              out += state[i] * q_[s_idx];
-            }}
-            out = simd_sum(out);
-            if (thread_index_in_simdgroup == 0) {{
-              y[dv_idx] = static_cast<InT>(out);
-            }}
-          }} else {{
-            y[dv_idx] = static_cast<InT>(0);
-          }}
-          // Increment data pointers to next time step
-          q_ += Hk * Dk;
-          k_ += Hk * Dk;
-          v_ += Hv * Dv;
-          y += Hv * Dv;
-          {g_advance}
-          beta_ += Hv;
+        // mlx-lm PR #1066: Kahan-compensated kv_mem accumulation +
+        // 4-way time-loop unroll. Fixes loss-of-precision on long T where
+        // bf16 state values cause kv_mem to drift away from the true sum.
+        #define BODY() {{ \
+            float kv_mem = 0.0f, kv_c = 0.0f; \
+            for (int i = 0; i < n_per_t; ++i) {{ \
+              auto s_idx = n_per_t * dk_idx + i; \
+              state[i] *= {g_access}; \
+              auto p = state[i] * k_[s_idx]; \
+              auto a = p - kv_c; \
+              auto b = kv_mem + a; \
+              kv_c = (b - kv_mem) - a; \
+              kv_mem = b; \
+            }} \
+            kv_mem = simd_sum(kv_mem); \
+            auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx]; \
+            float out = 0.0f; \
+            for (int i = 0; i < n_per_t; ++i) {{ \
+              auto s_idx = n_per_t * dk_idx + i; \
+              state[i] += k_[s_idx] * delta; \
+              out += state[i] * q_[s_idx]; \
+            }} \
+            out = simd_sum(out); \
+            if (thread_index_in_simdgroup == 0) \
+              y[dv_idx] = static_cast<InT>(out); \
         }}
+        #define ADV() q_ += Hk * Dk; k_ += Hk * Dk; v_ += Hv * Dv; y += Hv * Dv; {g_advance} beta_ += Hv;
+
+        int t = 0;
+        for (; t + 3 < T; t += 4) {{
+          if ({mask_source}) BODY() ADV()
+          if ({"mask[b_idx * T + t + 1]" if has_mask else "true"}) BODY() ADV()
+          if ({"mask[b_idx * T + t + 2]" if has_mask else "true"}) BODY() ADV()
+          if ({"mask[b_idx * T + t + 3]" if has_mask else "true"}) BODY() ADV()
+        }}
+        for (; t < T; ++t) {{
+          if ({mask_source}) BODY()
+          ADV()
+        }}
+        #undef BODY
+        #undef ADV
         for (int i = 0; i < n_per_t; ++i) {{
           auto s_idx = n_per_t * dk_idx + i;
           o_state[s_idx] = static_cast<StT>(state[i]);

tests/v4/test_path_dispatch.py

@@ -82,22 +82,26 @@ def test_gdn_path_a_always_available():
     assert statuses["path_a"].available
 
 
-def test_gdn_paths_c_d_currently_deferred():
+def test_gdn_path_d_still_deferred_without_triton():
+    """Path D needs triton — unavailable on Apple Silicon. Path C is now
+    real-or-fallback depending on whether tilelang is importable."""
     statuses = linear_attention_path_statuses()
-    # Path C/D still deferred (need TileLang/Triton wiring); Path B is now
-    # real Metal kernel as of the path_b ROI commit.
-    for p in ("path_c", "path_d"):
-        st = statuses[p]
-        assert not st.available
-        assert len(st.reason) > 20  # non-empty rationale
+    st_d = statuses["path_d"]
+    assert not st_d.available
+    assert len(st_d.reason) > 20
+    # Path C reason must always name the lowering pipeline, regardless of
+    # whether tilelang is reachable.
+    st_c = statuses["path_c"]
+    assert "tvm_ffi" in st_c.reason and "metal" in st_c.reason.lower()
 
 
 def test_gdn_auto_mode_default_picks_first_available(monkeypatch):
     monkeypatch.delenv(GDN_ENV, raising=False)
     chosen = linear_attention_auto_mode_for_inputs()
-    # Path B is available (real Metal kernel) — auto picks it before falling
-    # back to A in the preference order (c > b > e > d > a).
-    assert chosen in ("path_a", "path_b")
+    # Preference order: c > b > e > d > a. Path C wins when tilelang is
+    # available; Path E wins next; Path B (real Metal) wins on Apple Silicon
+    # without tilelang; Path A is the universal floor.
+    assert chosen in ("path_a", "path_b", "path_c", "path_e")
 
 
 def test_gdn_auto_mode_respects_env(monkeypatch):