feat(v4): bwd for GDN/KDA Path B + wire Path E training=True (PR #1217 VJP)

Three landings, all enabling actual training (fwd+bwd):

1. cppmega_v4/_tilelang/linear_attention_path_b_bwd.py
   - gdn_apply_path_b(q, k, v, beta, g) wrapped as @mx.custom_function.
   - Forward = fast Path B Metal kernel (3-40× over A from earlier bench).
   - VJP = mx.grad through Path A reference (algebraically identical),
     guarantees correct grads. Pattern named "path_b_fwd_path_a_bwd"
     mirroring mamba3's mamba3_mimo_apply_with_state_path_c_fwd_path_b_bwd.
   - Future revision swaps the VJP for a hand-MSL backward kernel; this
     wrapper already unlocks fwd+bwd training today.

2. cppmega_v4/_tilelang/kda_path_b_bwd.py — same pattern for KDA.

3. Path E training=True wiring:
   - Vendored mlx-lm PR #1217's gated_delta_vjp.py (179 LoC, Python
     chunked-VJP reference) and gated_delta_vjp_metal.py (772 LoC, Metal
     backward kernel — 8-11× over Python).
   - Patched _mlx_lm_gated_delta_vendored.py imports to point at the
     vendored module names (was `from .gated_delta_vjp_metal`).
   - Added `training=True` flag to our Path E adapter
     (mlx_lm_gated_delta_update.py); when set, dispatches through the
     chunked Metal VJP path for finite-memory training at T≥2048.

Tests (10 new):
  - test_path_b_bwd.py (4): fwd parity vs Path A; bwd grad parity for
    (dq, dk, dv, dbeta, dg); fwd+bwd produces finite grads on larger
    shape; chain through tanh.sum() works.
  - test_kda_path_b_bwd.py (3): same for KDA.
  - test_path_e_training.py (3): training=True produces finite output;
    matches inference within chunk tolerance (atol 1e-2); grads
    propagate through mx.grad.

v4 suite: 273 passed / 2 skipped.

Author: apstenku123

cppmega_v4/_tilelang/kda_path_b_bwd.py

@@ -0,0 +1,30 @@
+"""KDA Path B forward + backward — same fwd-fast / bwd-correct pattern as GDN."""
+
+import mlx.core as mx
+
+from cppmega_v4._tilelang.kda_path_b import kda_forward_path_b
+from cppmega_v4.nn._external.fla_naive_kda import naive_recurrent_kda
+
+
+@mx.custom_function
+def kda_apply_path_b(
+    q: mx.array, k: mx.array, v: mx.array, g: mx.array, beta: mx.array,
+) -> mx.array:
+    """Forward via fast Path B Metal kernel; backward via Path A reference grad."""
+    y, _ = kda_forward_path_b(q, k, v, g, beta, output_final_state=False)
+    return y
+
+
+@kda_apply_path_b.vjp
+def _kda_apply_path_b_vjp(primals, cotangent, output):
+    del output
+    q, k, v, g, beta = primals
+
+    def _loss_proxy(q_, k_, v_, g_, beta_):
+        y, _ = naive_recurrent_kda(q_, k_, v_, g_, beta_)
+        return (y * cotangent).sum()
+
+    return mx.grad(_loss_proxy, argnums=(0, 1, 2, 3, 4))(q, k, v, g, beta)
+
+
+__all__ = ["kda_apply_path_b"]

cppmega_v4/_tilelang/linear_attention_path_b_bwd.py

@@ -0,0 +1,66 @@
+"""GDN Path B forward + backward — fwd via fast Metal kernel, bwd via
+mx.grad through Path A reference.
+
+This is the *fwd-fast / bwd-correct* path, named ``path_b_fwd_path_a_bwd`` in
+the spirit of mamba3's ``mamba3_mimo_apply_with_state_path_c_fwd_path_b_bwd``.
+The fast hand-MSL kernel handles forward (3-10× over Path A), while
+backward delegates to ``mx.grad`` through the FLA-naive reference so
+training is at least *correct* even if not yet fully accelerated.
+
+A future revision (tracked: cppmega_v4._tilelang.linear_attention_path_b_bwd_metal)
+replaces the backward with a fused Metal kernel; until then this wrapper
+unlocks full V4-stack training with measurable forward speedup.
+
+API:
+    forward_with_grad(q, k, v, beta, g) -> (o, h_last)
+        Implemented as ``mx.custom_function`` so ``mx.grad`` traces through.
+"""
+
+from typing import Optional
+
+import mlx.core as mx
+
+from cppmega_v4._tilelang.linear_attention_path_b import gdn_forward_path_b
+from cppmega_v4.nn._external.fla_naive_gated_delta_rule import (
+    naive_recurrent_gated_delta_rule,
+)
+
+
+@mx.custom_function
+def gdn_apply_path_b(
+    q: mx.array, k: mx.array, v: mx.array, beta: mx.array, g: mx.array,
+) -> mx.array:
+    """Forward-only call returning y (no state). Differentiable via custom VJP.
+
+    The fast Path B Metal kernel produces y; the VJP differentiates the
+    Path A reference (algebraically identical), so gradients are correct.
+    """
+    y, _ = gdn_forward_path_b(q, k, v, beta, g, output_final_state=False)
+    return y
+
+
+@gdn_apply_path_b.vjp
+def _gdn_apply_path_b_vjp(
+    primals: tuple,
+    cotangent: mx.array,
+    output: mx.array,
+) -> tuple:
+    """Backward via mx.grad through the FLA naive reference.
+
+    Path A reference is algebraically identical to Path B (Path B is just
+    a faster MSL implementation of the same recurrence), so grads from
+    Path A are correct grads for Path B. This is the "path_b_fwd_path_a_bwd"
+    pattern — see mamba3_path_c's analogous ``path_c_fwd_path_b_bwd`` mode.
+    """
+    del output  # unused — VJP re-runs forward through Path A
+    q, k, v, beta, g = primals
+
+    def _loss_proxy(q_, k_, v_, beta_, g_):
+        y, _ = naive_recurrent_gated_delta_rule(q_, k_, v_, beta_, g_)
+        return (y * cotangent).sum()
+
+    grad_fn = mx.grad(_loss_proxy, argnums=(0, 1, 2, 3, 4))
+    return grad_fn(q, k, v, beta, g)
+
+
+__all__ = ["gdn_apply_path_b"]

cppmega_v4/nn/_external/_mlx_lm_gated_delta_vendored.py

@@ -308,15 +308,15 @@ def gated_delta_update(
         )
         if can_use_metal:
             try:
-                from .gated_delta_vjp_metal import gated_delta_update_vjp_metal
-
+                from ._mlx_lm_gated_delta_vjp_metal_vendored import (
+                    gated_delta_update_vjp_metal,
+                )
                 return gated_delta_update_vjp_metal(
                     q, k, v, a, b, A_log, dt_bias, state, mask
                 )
             except ImportError:
                 pass
-        from .gated_delta_vjp import gated_delta_update_vjp
-
+        from ._mlx_lm_gated_delta_vjp_vendored import gated_delta_update_vjp
         return gated_delta_update_vjp(q, k, v, a, b, A_log, dt_bias, state, mask)
 
     beta = mx.sigmoid(b)