perf(gated_delta): chunk-parallel differentiable training path for Ga…

.gitignore

@@ -140,3 +140,4 @@ dmypy.json
 
 # .DS_Store files
 .DS_Store
+mlx-recon/

mlx_lm/models/gated_delta.py

@@ -1,5 +1,4 @@
 # Copyright © 2025 Apple Inc.
-
 import os
 from functools import partial
 from typing import Optional, Tuple
@@ -602,6 +601,123 @@ def gated_delta_ops(
     return y, state
 
 
+def _tri_inv(A: mx.array) -> mx.array:
+    """Inverse of ``A = I + N`` where ``N`` is strictly lower triangular.
+
+    Block-Schur recursion: for a unit lower-triangular
+    ``A = [[A11, 0], [A21, A22]]`` the exact inverse is
+    ``[[A11^-1, 0], [-A22^-1 @ A21 @ A11^-1, A22^-1]]``, recursing to 1x1
+    blocks (which are ``[[1]]`` for a unit diagonal), then two Newton polish
+    steps for the matmul rounding. Pure matmuls, so it runs on the GPU and is
+    differentiable by autograd.
+
+    A Newton-Schulz iteration from the identity instead diverges when A is
+    near-singular even at modest ``||N||`` -- e.g. near-duplicate keys at
+    chunk_size 64 drive its error to ~1e9 while this recursion stays at
+    machine precision. (Both lose accuracy only once the true inverse itself
+    overflows fp32, which the normalize_qk contract keeps real inputs clear
+    of.)
+    """
+
+    def _rec(M: mx.array) -> mx.array:
+        n = M.shape[-1]
+        if n == 1:
+            return mx.ones_like(M)
+        h = n // 2
+        A11 = _rec(M[..., :h, :h])
+        A22 = _rec(M[..., h:, h:])
+        A21 = -A22 @ (M[..., h:, :h] @ A11)
+        top = mx.concatenate([A11, mx.zeros_like(A21).swapaxes(-1, -2)], axis=-1)
+        bot = mx.concatenate([A21, A22], axis=-1)
+        return mx.concatenate([top, bot], axis=-2)
+
+    C = A.shape[-1]
+    X = _rec(A)
+    eye = mx.eye(C, dtype=A.dtype)
+    for _ in range(2):  # Newton polish for matmul rounding
+        X = X @ (2.0 * eye - A @ X)
+    return X
+
+
+##     Numerical stability requires normalized keys (the normalize_qk contract); unnormalized keys can make the chunk inverse diverge.
+def gated_delta_chunked(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: Optional[mx.array] = None,
+    chunk_size: int = 64,
+) -> Tuple[mx.array, mx.array]:
+    """Chunk-parallel, fully-differentiable gated delta rule for training.
+
+    Mathematically equivalent to ``gated_delta_ops`` but replaces the per-token
+    Python loop with a chunkwise UT (WY) transform, so the backward pass has
+    ``T / chunk_size`` sequential steps instead of ``T``. Scalar gate and no
+    padding mask only; other cases use ``gated_delta_ops``.
+
+    Shapes:
+      - q, k: [B, T, Hk, Dk]   (normalized per ``normalize_qk``)
+      - v: [B, T, Hv, Dv]
+      - g: [B, T, Hv]          (scalar decay in (0, 1])
+      - beta: [B, T, Hv]
+      - state: [B, Hv, Dv, Dk]
+    """
+    B, T, Hk, Dk = q.shape
+    Hv, Dv = v.shape[-2:]
+    if state is None:
+        state = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
+    if (repeat_factor := Hv // Hk) > 1:
+        q = mx.repeat(q, repeat_factor, -2)
+        k = mx.repeat(k, repeat_factor, -2)
+
+    ys = []
+    for start in range(0, T, chunk_size):
+        end = min(start + chunk_size, T)
+        qc = q[:, start:end].swapaxes(1, 2).astype(mx.float32)  # [B, H, C, Dk]
+        kc = k[:, start:end].swapaxes(1, 2).astype(mx.float32)
+        vc = v[:, start:end].swapaxes(1, 2).astype(mx.float32)  # [B, H, C, Dv]
+        gc = mx.maximum(
+            g[:, start:end].swapaxes(1, 2).astype(mx.float32), 1e-8
+        )  # [B, H, C]
+        bc = beta[:, start:end].swapaxes(1, 2).astype(mx.float32)
+        C = qc.shape[2]
+
+        idx = mx.arange(C)
+        strict_lower = idx[:, None] > idx[None, :]
+        incl_lower = idx[:, None] >= idx[None, :]
+
+        log_gamma = mx.cumsum(mx.log(gc), axis=-1)  # [B, H, C]
+        gamma = mx.exp(log_gamma)
+        # Upper triangle (i<j) is masked away below, but exp(+sum log g)
+        # overflows to inf and its backward (0*inf) poisons every gradient.
+        # Clamp to <=0: a no-op on the used lower triangle (i>=j, exponent <=0).
+        decay_ratio = mx.exp(
+            mx.minimum(log_gamma[..., :, None] - log_gamma[..., None, :], 0.0)
+        )
+
+        kkt = kc @ kc.swapaxes(-1, -2)
+        A = mx.eye(C) + bc[..., :, None] * mx.where(
+            strict_lower, decay_ratio * kkt, 0.0
+        )
+        A_inv = _tri_inv(A)
+
+        v_prime = vc - gamma[..., :, None] * (kc @ state.swapaxes(-1, -2))
+        u = A_inv @ (bc[..., :, None] * v_prime)  # [B, H, C, Dv]
+
+        qkt = qc @ kc.swapaxes(-1, -2)
+        p = mx.where(incl_lower, decay_ratio * qkt, 0.0)
+        yc = gamma[..., :, None] * (qc @ state.swapaxes(-1, -2)) + p @ u
+
+        gamma_to_end = mx.exp(log_gamma[..., -1:] - log_gamma)
+        w = gamma_to_end[..., :, None] * u
+        state = gamma[..., -1][..., None, None] * state + w.swapaxes(-1, -2) @ kc
+
+        ys.append(yc.swapaxes(1, 2))  # [B, C, H, Dv]
+
+    return mx.concatenate(ys, axis=1).astype(q.dtype), state
+
+
 def gated_delta_update(
     q: mx.array,
     k: mx.array,
@@ -636,6 +752,13 @@ def gated_delta_update(
         Hv, Dv = v.shape[-2:]
         state = mx.zeros((B, Hv, Dv, Dk), dtype=mx.float32)
 
+    # Training (use_kernel=False): the Metal kernel has no VJP, so the only
+    # differentiable fallback was the per-token gated_delta_ops loop, whose
+    # backward scales with T. Use the chunk-parallel path when it applies
+    # (scalar gate, no padding mask, more than one step).
+    if not use_kernel and mask is None and g.ndim == 3 and q.shape[1] > 1:
+        return gated_delta_chunked(q, k, v, g, beta, state)
+
     if (
         not use_kernel
         or mx.default_device() != mx.gpu

tests/test_gated_delta.py

@@ -6,10 +6,13 @@
 
 import mlx_lm.models.gated_delta as gated_delta
 from mlx_lm.models.gated_delta import (
+    gated_delta_chunked,
     gated_delta_kernel,
     gated_delta_kernel_unpacked,
     gated_delta_kernel_xtree,
     gated_delta_ops,
+    gated_delta_update,
+    normalize_qk,
 )
 
 
@@ -169,6 +172,231 @@ def test_unsupported_shapes_fall_back(self):
         self.assertTrue(mx.array_equal(y_p, y_u))
         self.assertTrue(mx.array_equal(s_p, s_u))
 
+    def test_chunked_matches_ops_forward(self):
+        if mx.default_device() != mx.gpu:
+            raise unittest.SkipTest("gated delta kernels are GPU only")
+        cases = [
+            (1, 16, 32, 128, 128),  # Qwen3.5/3.6 shape
+            (1, 4, 8, 128, 128),  # fewer heads
+            (2, 16, 32, 128, 128),  # batched
+            (1, 8, 16, 128, 64),  # Dv != Dk
+        ]
+        for B, Hk, Hv, Dk, Dv in cases:
+            for T in (7, 64, 128, 130):
+                for cs in (32, 64):
+                    with self.subTest(B=B, Hv=Hv, T=T, cs=cs):
+                        q, k, v, g, beta, state = self._inputs(
+                            B, T, Hk, Hv, Dk, Dv, mx.float32
+                        )
+                        y_c, s_c = gated_delta_chunked(
+                            q, k, v, g, beta, state, chunk_size=cs
+                        )
+                        y_r, s_r = gated_delta_ops(q, k, v, g, beta, state)
+                        mx.eval(y_c, s_c, y_r, s_r)
+                        self.assertLess(_rel_l2(y_c, y_r), 1e-4)
+                        self.assertLess(_rel_l2(s_c, s_r), 1e-4)
+
+    def test_chunked_matches_ops_backward(self):
+        if mx.default_device() != mx.gpu:
+            raise unittest.SkipTest("gated delta kernels are GPU only")
+        q, k, v, g, beta, state = self._inputs(1, 64, 16, 32, 128, 128, mx.float32)
+
+        def loss_ops(q, k, v, g, beta):
+            return gated_delta_ops(q, k, v, g, beta, state)[0].sum()
+
+        def loss_chunked(q, k, v, g, beta):
+            return gated_delta_chunked(q, k, v, g, beta, state, chunk_size=32)[0].sum()
+
+        (_,), gr = mx.vjp(loss_ops, [q, k, v, g, beta], [mx.array(1.0)])
+        (_,), gc = mx.vjp(loss_chunked, [q, k, v, g, beta], [mx.array(1.0)])
+        mx.eval(*gr, *gc)
+        for name, a, b in zip(("q", "k", "v", "g", "beta"), gc, gr):
+            with self.subTest(grad=name):
+                self.assertLess(_rel_l2(a, b), 1e-4)
+
+    def test_chunked_allows_neg_eigval_beta(self):
+        # allow_neg_eigval puts beta in [0, 2]; the UT transform must stay stable.
+        if mx.default_device() != mx.gpu:
+            raise unittest.SkipTest("gated delta kernels are GPU only")
+        q, k, v, g, beta, state = self._inputs(1, 128, 16, 32, 128, 128, mx.float32)
+        beta = beta * 2.0
+        y_c, s_c = gated_delta_chunked(q, k, v, g, beta, state, chunk_size=64)
+        y_r, s_r = gated_delta_ops(q, k, v, g, beta, state)
+        mx.eval(y_c, s_c, y_r, s_r)
+        self.assertLess(_rel_l2(y_c, y_r), 1e-4)
+        self.assertLess(_rel_l2(s_c, s_r), 1e-4)
+
+    def test_update_training_routes_to_chunked_and_differentiable(self):
+        # gated_delta_update(use_kernel=False) must match the kernel path and,
+        # unlike the kernel, be differentiable (the point of the chunked path).
+        if mx.default_device() != mx.gpu:
+            raise unittest.SkipTest("gated delta kernels are GPU only")
+        mx.random.seed(0)
+        B, T, Hk, Hv, Dk, Dv = 1, 128, 16, 32, 128, 128
+        q = mx.random.normal((B, T, Hk, Dk)).astype(mx.bfloat16)
+        k = mx.random.normal((B, T, Hk, Dk)).astype(mx.bfloat16)
+        q, k = normalize_qk(q, k, inv_scale=Dk**-0.5, eps=1e-6)
+        v = mx.random.normal((B, T, Hv, Dv)).astype(mx.bfloat16)
+        a = mx.random.normal((B, T, Hv))
+        b = mx.random.normal((B, T, Hv))
+        A_log = mx.log(mx.random.uniform(low=0.01, high=16, shape=(Hv,)))
+        dt_bias = mx.ones((Hv,))
+
+        y_k, s_k = gated_delta_update(q, k, v, a, b, A_log, dt_bias, use_kernel=True)
+        y_c, s_c = gated_delta_update(q, k, v, a, b, A_log, dt_bias, use_kernel=False)
+        mx.eval(y_k, s_k, y_c, s_c)
+        self.assertLess(_rel_l2(y_c, y_k), 5e-3)
+        self.assertLess(_rel_l2(s_c, s_k), 5e-3)
+
+        def loss(q, k, v, a, b):
+            return gated_delta_update(q, k, v, a, b, A_log, dt_bias, use_kernel=False)[
+                0
+            ].sum()
+
+        (_,), grads = mx.vjp(loss, [q, k, v, a, b], [mx.array(1.0)])
+        mx.eval(*grads)
+        self.assertTrue(all((gr != 0).sum().item() > 0 for gr in grads))
+
+    def test_chunked_numerical_extremes(self):
+        # Chunked must track ops across decay/beta/value extremes and a large
+        # single chunk -- the operating envelope a real training run can hit.
+        if mx.default_device() != mx.gpu:
+            raise unittest.SkipTest("gated delta kernels are GPU only")
+        q, k, v, g, beta, state = self._inputs(1, 128, 16, 32, 128, 128, mx.float32)
+
+        variants = {
+            "strong_decay": (
+                q,
+                k,
+                v,
+                mx.exp(-mx.random.uniform(shape=g.shape) * 5),
+                beta,
+            ),
+            "extreme_decay": (
+                q,
+                k,
+                v,
+                mx.exp(-mx.random.uniform(shape=g.shape) * 20),
+                beta,
+            ),
+            "beta_zero": (q, k, v, g, mx.zeros_like(beta)),
+            "beta_double": (q, k, v, g, beta * 2.0),  # allow_neg_eigval max
+            "tiny_v": (q, k, v * 1e-4, g, beta),
+            "huge_v": (q, k, v * 1e4, g, beta),
+        }
+        for name, (qq, kk, vv, gg, bb) in variants.items():
+            with self.subTest(case=name):
+                y_c, s_c = gated_delta_chunked(qq, kk, vv, gg, bb, state, chunk_size=64)
+                y_r, s_r = gated_delta_ops(qq, kk, vv, gg, bb, state)
+                mx.eval(y_c, s_c, y_r, s_r)
+                self.assertLess(_rel_l2(y_c, y_r), 1e-3)
+                self.assertLess(_rel_l2(s_c, s_r), 1e-3)
+
+        # a single chunk spanning the whole sequence (Newton-Schulz at C=256)
+        q2, k2, v2, g2, beta2, st2 = self._inputs(1, 256, 16, 32, 128, 128, mx.float32)
+        y_c, s_c = gated_delta_chunked(q2, k2, v2, g2, beta2, st2, chunk_size=256)
+        y_r, s_r = gated_delta_ops(q2, k2, v2, g2, beta2, st2)
+        mx.eval(y_c, s_c, y_r, s_r)
+        self.assertLess(_rel_l2(y_c, y_r), 1e-3)
+        self.assertLess(_rel_l2(s_c, s_r), 1e-3)
+
+    def test_update_decode_and_mask_fall_back_to_ops(self):
+        # T==1 (decode) and masked inputs must NOT take the chunked path; they
+        # fall back and stay consistent with the kernel path.
+        if mx.default_device() != mx.gpu:
+            raise unittest.SkipTest("gated delta kernels are GPU only")
+
+        def make(T):
+            mx.random.seed(2)
+            q = mx.random.normal((1, T, 16, 128)).astype(mx.bfloat16)
+            k = mx.random.normal((1, T, 16, 128)).astype(mx.bfloat16)
+            q, k = normalize_qk(q, k, inv_scale=128**-0.5, eps=1e-6)
+            v = mx.random.normal((1, T, 32, 128)).astype(mx.bfloat16)
+            a = mx.random.normal((1, T, 32))
+            b = mx.random.normal((1, T, 32))
+            A_log = mx.log(mx.random.uniform(low=0.01, high=16, shape=(32,)))
+            dt_bias = mx.ones((32,))
+            return q, k, v, a, b, A_log, dt_bias
+
+        # decode: T == 1 falls back (guard q.shape[1] > 1)
+        args = make(1)
+        yf, _ = gated_delta_update(*args, use_kernel=False)
+        yt, _ = gated_delta_update(*args, use_kernel=True)
+        mx.eval(yf, yt)
+        self.assertLess(_rel_l2(yf, yt), 5e-3)
+
+        # mask present falls back (guard mask is None)
+        q, k, v, a, b, A_log, dt_bias = make(64)
+        mask = mx.arange(64)[None] < mx.array([[50]])
+        yf, _ = gated_delta_update(
+            q, k, v, a, b, A_log, dt_bias, mask=mask, use_kernel=False
+        )
+        yt, _ = gated_delta_update(
+            q, k, v, a, b, A_log, dt_bias, mask=mask, use_kernel=True
+        )
+        valid = mask[..., None, None]
+        yf = mx.where(valid, yf, 0)
+        yt = mx.where(valid, yt, 0)
+        mx.eval(yf, yt)
+        self.assertLess(_rel_l2(yf, yt), 5e-3)
+
+    def test_chunked_decay_nan_free_hard_gates(self):
+        # Regression: hard-decaying gates make decay_ratio's masked upper
+        # triangle overflow to inf; its backward (0*inf) produced NaN grads in
+        # real training. The exp(min(., 0)) clamp fixes it without changing the
+        # used lower triangle. Fails on the pre-clamp code, passes now.
+        mx.random.seed(0)
+        B, Hk, Hv, T, Dk, Dv = 1, 2, 4, 64, 32, 32
+        base = mx.random.normal((B, T, Hk, Dk))
+        k = 0.9 * base + 0.1 * mx.random.normal((B, T, Hk, Dk))
+        q = mx.random.normal((B, T, Hk, Dk))
+        v = mx.random.normal((B, T, Hv, Dv))
+        k = k / mx.linalg.norm(k, axis=-1, keepdims=True)
+        q = q / mx.linalg.norm(q, axis=-1, keepdims=True)
+        # Sum(log g) ~ -170 over the 64-token chunk.
+        g = mx.exp(mx.random.uniform(-5.4, 0.0, (B, T, Hv)))
+        beta = mx.random.uniform(0.85, 1.0, (B, T, Hv))
+
+        def loss(q, k, v, g, beta):
+            y, _ = gated_delta_chunked(q, k, v, g, beta)
+            return (y**2).sum()
+
+        y, _ = gated_delta_chunked(q, k, v, g, beta)
+        self.assertFalse(bool(mx.isnan(y).any()))
+        grads = mx.grad(loss, argnums=(0, 1, 2, 3, 4))(q, k, v, g, beta)
+        for gr in grads:
+            self.assertFalse(bool(mx.isnan(gr).any()))
+            self.assertFalse(bool(mx.isinf(gr).any()))
+
+    def test_tri_inv_near_singular_identical_keys(self):
+        # Regression: a Newton-Schulz inverse from the identity diverges when
+        # A = I + N is near-singular even at modest ||N||. The deterministic
+        # worst case is identical keys (KKt = ones) with chunk_size 64: NS
+        # reached ~1e9 error there, the block-Schur recursion is exact. No
+        # randomness, so this pins the failure the recursion fixes.
+        from mlx_lm.models.gated_delta import _tri_inv
+
+        C = 64
+        g = 0.95
+        idx = mx.arange(C)
+        log_gamma = (idx + 1) * mx.log(mx.array(g))
+        decay = mx.exp(log_gamma[:, None] - log_gamma[None, :])
+        strict_lower = idx[:, None] > idx[None, :]
+        A = mx.eye(C) + mx.where(strict_lower, decay * mx.ones((C, C)), 0.0)
+
+        X = _tri_inv(A)
+        err = float(mx.abs(A @ X - mx.eye(C)).max())
+        self.assertFalse(bool(mx.isnan(mx.array(err))))
+        self.assertLess(err, 1e-6)
+
+        # Differentiable with no custom VJP: finite gradients.
+        def f(A):
+            return (_tri_inv(A) ** 2).sum()
+
+        gA = mx.grad(f)(A)
+        self.assertFalse(bool(mx.isnan(gA).any()))
+        self.assertFalse(bool(mx.isinf(gA).any()))
+
 
 if __name__ == "__main__":
     unittest.main()