feat(v4-path-e): remove Dk%32 forward shape limit via in-MSL remainder-mask

The fast gated_delta Metal kernel silently truncated the Dk tail (n_per_t=Dk/32 floor dropped
the trailing Dk%32 keys -> wrong answer for non-mult-32 Dk). Fixed in the actual MSL: ceil
tiling n_per_t=(Dk+31)/32 + per-i s_idx<Dk guards on every k/q/state load (additive identity 0
past Dk), and the hard part -- the VECTORIZED KDA gate tail-masked to decay exp(0)=1 (never reads
g_[s_idx] OOB) so the pre-slice simd_sum(kv_mem) is not poisoned. Same mask in the VJP fwd-save +
both backward kernels (incl. guarded [4*Dk] shared-tile writes). Forward now runs for ANY Dk.

Eligibility: shape_uses_fast_kernel True for all shapes (forward); new shape_uses_fast_kernel_backward
keeps Dv%4 for the VJP cooperative reduction. Training still fail-closes to the Python-ops VJP for
Dv%4!=0 (deferred). Gate-sign behavior UNCHANGED (inherent GDN model-math: the kernel has no clamp;
KDA feeds exp(g)>1 into the same kernel and amplifies; the (0,1] bound is the GDN parameterization
a=softplus_inverse(-g)); only the reason string reworded. Provenance: the two kernel files
reclassified verbatim->derived with upstream_base_sha256 + delta note.

Metal parity vs gated_delta_ops (fp32): scalar fwd 65/65 rel<1e-6 (incl Dk 40/96/100/192); KDA
vectorized-gate Dk%32!=0 rel ~1e-7..1e-6; asymmetric, group-expansion, state round-trip bit-exact;
VJP odd-Dk (Dv%4==0) rel ~1e-6. 133 tests pass (65 new odd-shape); 6 unrelated pre-existing; v3 untouched.

Author: jeweldave

cppmega_v4/_tilelang/linear_attention_paths.py

@@ -219,9 +219,10 @@ def _path_e_status() -> PathStatus:
         return PathStatus(
             path="path_e", available=True,
             reason=(
-                "vendored mlx-lm PR #1217 gated_delta_update (fast Metal kernel "
-                "for Dk%32==0 & Dv%4==0 and gate g<=0; fails closed otherwise "
-                "so the dispatcher falls back to Path B/A)"
+                "vendored mlx-lm PR #1217 gated_delta_update (fast Metal kernel; "
+                "forward runs for ANY Dk via the in-MSL remainder-mask and any "
+                "Dv; gate must be g<=0 for GDN, otherwise fails closed so the "
+                "dispatcher falls back to Path B/A)"
             ),
         )
     except Exception:

cppmega_v4/nn/_external/VENDORED_MANIFEST.json

@@ -19,10 +19,11 @@
       "pr": 1217,
       "pr_url": "https://github.com/ml-explore/mlx-lm/pull/1217",
       "upstream_path": "mlx_lm/models/gated_delta.py",
-      "kind": "verbatim",
+      "kind": "derived",
       "used_by": "GDN/KDA Path E (gated_delta_update, gated_delta_kernel, gated_delta_ops, compute_g)",
-      "sha256": "18b3a77250b246d46e63d8fa9e8208936ea09bd46b97b036c2ec34edfc0a176b",
-      "notes": "Snapshot includes the PR #1066 Kahan-compensated kv_mem accumulation + 4-way time-loop unroll; this is NOT identical to an older mlx-lm checkout that predates PR #1066."
+      "upstream_base_sha256": "18b3a77250b246d46e63d8fa9e8208936ea09bd46b97b036c2ec34edfc0a176b",
+      "sha256": "d2f3c76ca2870c98dfe95390287f01a4ee75d048227f0784076826938514ea08",
+      "notes": "Was 'verbatim' (PR #1217 snapshot, incl. PR #1066 Kahan kv_mem + 4-way unroll). Reclassified 'derived': cppmega added an in-MSL Dk remainder-mask (ceil tiling n_per_t=(Dk+31)/32 + per-i s_idx<Dk guards; vectorized gate uses decay=1.0 on the masked tail) so the fast Metal forward kernel is correct for ANY Dk, not just Dk%32==0 — upstream silently truncated the trailing Dk%32 keys. upstream_base_sha256 records the pre-delta PR #1217 snapshot; the only delta is this remainder-mask."
     },
     {
       "file": "_mlx_lm_fp8_dequant_vendored.py",
@@ -41,8 +42,9 @@
       "upstream_path": "mlx_lm/models/gated_delta.py (kernel reused; VJP is cppmega-authored)",
       "kind": "derived",
       "used_by": "GDN Path E training backward (fused Metal VJP)",
-      "sha256": "dfe959d0ed50d60d75d93532a80b0b6f769aad74ce39bc00a0770128c9903f4e",
-      "notes": "cppmega-authored fused Metal VJP built on the vendored PR #1217 gated_delta_kernel; upstream has no registered VJP."
+      "upstream_base_sha256": "dfe959d0ed50d60d75d93532a80b0b6f769aad74ce39bc00a0770128c9903f4e",
+      "sha256": "e4448b96fef0830c02c88bae918cd02e7c053892bde7f9021230d18384ea9386",
+      "notes": "cppmega-authored fused Metal VJP built on the vendored PR #1217 gated_delta_kernel; upstream has no registered VJP. DELTA over upstream_base_sha256: same in-MSL Dk remainder-mask as the forward (ceil tiling + per-i s_idx<Dk guards in fwd-save and both bwd kernels; vectorized gate tail decay=1.0) so the VJP handles any Dk. Still requires Dv%4==0 (four-SIMD-group cooperative reduction over a fixed [4*Dk] tile) — ragged-Dv backward deferred."
     },
     {
       "file": "_mlx_lm_gated_delta_vjp_vendored.py",

cppmega_v4/nn/_external/_mlx_lm_gated_delta_vendored.py

@@ -1,14 +1,26 @@
-# Vendored verbatim from ml-explore/mlx-lm PR #1217
-#   (mlx_lm/models/gated_delta.py)
+# Derived from ml-explore/mlx-lm PR #1217 (mlx_lm/models/gated_delta.py).
 # Upstream license: MIT, Copyright © 2023 Apple Inc.
 #
 # Public surface used by cppmega_v4 Path E:
 #   - gated_delta_update(...) -> (y, state)
 #   - gated_delta_kernel(...) — low-level Metal kernel entry
 #   - gated_delta_ops(...) — pure-MLX reference for prefill
 #   - compute_g(A_log, a, dt_bias) — gate transform
-# No edits — pure copy. cppmega_v4/nn/path_e_adapter.py wraps our (g, beta)
-# API to this module's (a, b, A_log, dt_bias) API.
+# cppmega_v4/nn/path_e_adapter.py wraps our (g, beta) API to this module's
+# (a, b, A_log, dt_bias) API.
+#
+# CPPMEGA DELTA over the PR #1217 snapshot (see VENDORED_MANIFEST.json):
+#   In-MSL Dk remainder-mask so the fast Metal kernel is CORRECT for ANY Dk,
+#   not just Dk % 32 == 0. Upstream uses `n_per_t = Dk / 32` (integer floor),
+#   so the trailing Dk % 32 keys owned past lane 31's slice were NEVER loaded
+#   -> silent truncation / wrong answer for non-multiple-of-32 Dk. We switch
+#   to a ceil tiling `n_per_t = (Dk + 31) / 32` and guard every per-`i` load
+#   with `s_idx < Dk`, substituting additive/multiplicative identities for the
+#   out-of-range tail (0 for k/q/state contributions; the VECTORIZED gate uses
+#   decay = exp(0) = 1, i.e. g_access expands to 0.0 on the tail so that the
+#   pre-reduction simd_sum is never poisoned). The scalar GDN gate is per-(b,t,
+#   hv) broadcast and so unaffected. Everything else is the verbatim PR #1217
+#   snapshot (incl. the PR #1066 Kahan-compensated kv_mem + 4-way unroll).
 
 from functools import partial
 from typing import Optional, Tuple
@@ -27,11 +39,17 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False):
         return None
     mask_source = "mask[b_idx * T + t]" if has_mask else "true"
 
-    # Configure g indexing based on whether gating is vectorized
+    # Configure g indexing based on whether gating is vectorized.
+    # Vectorized g is [B,T,Hv,Dk] and is read per-key (s_idx). On the Dk
+    # remainder tail (s_idx >= Dk) the load would be OOB, so we mask it and
+    # substitute decay = 1.0 (== exp(0)) for that lane's tail entry. Using 1.0
+    # leaves state[i] unchanged; since the tail state[i] is itself 0 (masked
+    # load) and the tail k_/q_ contributions are 0, the tail never poisons the
+    # pre-reduction simd_sum(kv_mem)/simd_sum(out).
     if vectorized:
         g_comment = "// g: [B, T, Hv, Dk]"
         g_setup = "auto g_ = g + (b_idx * T * Hv + hv_idx) * Dk;"
-        g_access = "g_[s_idx]"
+        g_access = "(s_idx < Dk ? static_cast<float>(g_[s_idx]) : 1.0f)"
         g_advance = "g_ += Hv * Dk;"
     else:
         g_comment = "// g: [B, T, Hv]"
@@ -44,7 +62,9 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False):
         auto b_idx = n / Hv;
         auto hv_idx = n % Hv;
         auto hk_idx = hv_idx / (Hv / Hk);
-        constexpr int n_per_t = Dk / 32;
+        // cppmega delta: ceil tiling so the trailing Dk % 32 keys are owned by
+        // a lane and loaded (upstream used floor `Dk / 32`, dropping the tail).
+        constexpr int n_per_t = (Dk + 31) / 32;
 
         // q, k: [B, T, Hk, Dk]
         auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
@@ -64,7 +84,9 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False):
         float state[n_per_t];
         for (int i = 0; i < n_per_t; ++i) {{
           auto s_idx = n_per_t * dk_idx + i;
-          state[i] = static_cast<float>(i_state[s_idx]);
+          // Dk remainder-mask: tail lanes hold 0 (additive identity) so they
+          // contribute nothing to any simd_sum reduction.
+          state[i] = (s_idx < Dk) ? static_cast<float>(i_state[s_idx]) : 0.0f;
         }}
 
         {g_comment}
@@ -78,6 +100,7 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False):
             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; \
+              if (s_idx >= Dk) continue; \
               state[i] *= {g_access}; \
               auto p = state[i] * k_[s_idx]; \
               auto a = p - kv_c; \
@@ -90,6 +113,7 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False):
             float out = 0.0f; \
             for (int i = 0; i < n_per_t; ++i) {{ \
               auto s_idx = n_per_t * dk_idx + i; \
+              if (s_idx >= Dk) continue; \
               state[i] += k_[s_idx] * delta; \
               out += state[i] * q_[s_idx]; \
             }} \
@@ -114,7 +138,7 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False):
         #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]);
+          if (s_idx < Dk) o_state[s_idx] = static_cast<StT>(state[i]);
         }}
     """
     inputs = ["q", "k", "v", "g", "beta", "state_in", "T"]
@@ -295,15 +319,15 @@ def gated_delta_update(
     if training:
         # Chunked VJP path with O(T/chunk) autodiff graph — fits T≥2048
         # on 36 GB Apple Silicon where the Python-ops path OOMs.
-        # Metal backward kernel is 8–11× faster than the Python reference,
-        # but only handles the unmasked GPU path with Dk%32==0 and Dv%4==0.
-        Dk = q.shape[-1]
+        # Metal backward kernel is 8–11× faster than the Python reference. It
+        # now handles ANY Dk (in-MSL remainder-mask), but still needs Dv%4==0
+        # for the four-SIMD-group cooperative reduction (ragged Dv backward
+        # deferred). The unmasked GPU path is also required.
         Dv = v.shape[-1]
         can_use_metal = (
             mx.metal.is_available()
             and mx.default_device() == mx.gpu
             and mask is None
-            and Dk % 32 == 0
             and Dv % 4 == 0
         )
         if can_use_metal: