Fix indexing bug in slice update with op - #3483

Fix indexing bug in slice update with op (ml-explore#3483)

mx .at[:, t].add scatters read the update tensor with a wrong batch
stride on the Metal stream, silently corrupting d_q/d_k/d_v/d_g/d_beta
for every batch element except index 0. Collect per-step grads in lists
and stack instead; verified against plain autodiff at B in {2,3,4}.

The underlying Metal scatter indexing bug is fixed upstream in
ml-explore/mlx#3483 but is not in any release yet (latest: 0.31.2).
The list+stack form could revert to the scatter once a fixed mlx is
the minimum supported version, though it is also the cheaper op here.

mx .at[:, t].add scatters read the update tensor with a wrong batch
stride on the Metal stream, silently corrupting d_q/d_k/d_v/d_g/d_beta
for every batch element except index 0. Collect per-step grads in lists
and stack instead; verified against plain autodiff at B in {2,3,4}.

The underlying Metal scatter indexing bug is fixed upstream in
ml-explore/mlx#3483 but is not in any release yet (latest: 0.31.2).
The list+stack form could revert to the scatter once a fixed mlx is
the minimum supported version, though it is also the cheaper op here.

MLX 0.31.2 scatter-add can read the update tensor with a wrong batch stride for mx .at[:, t].add, silently corrupting d_q/d_k/d_v/d_g/d_beta for every batch element except index 0. Collect per-step grads in lists and stack instead.

Verified against a plain autodiff recurrence: the old scatter path matches on mlx 0.31.1, diverges on mlx 0.31.2 for B>1, and the list+stack path matches on both. The underlying MLX regression is fixed upstream in ml-explore/mlx#3483, but our release venv is still on mlx 0.31.2.

MLX 0.31.2 scatter-add can read the update tensor with a wrong batch stride for mx .at[:, t].add, silently corrupting d_q/d_k/d_v/d_g/d_beta for every batch element except index 0. Collect per-step grads in lists and stack instead.

Verified against a plain autodiff recurrence: the old scatter path matches on mlx 0.31.1, diverges on mlx 0.31.2 for B>1, and the list+stack path matches on both. The underlying MLX regression is fixed upstream in ml-explore/mlx#3483, but our release venv is still on mlx 0.31.2.

#754)

* fix(mlx): rebind stale gated_delta_update consumer imports

patch_gated_delta only rebound mlx_lm.models.gated_delta, but every
consumer (mlx_lm qwen3_5 / qwen3_next / kimi_linear, mlx_vlm qwen3_5)
binds gated_delta_update via from-import at module load, so the
memory-efficient VJP never reached any call site and training ran the
raw sequential ops loop. Sweep already-imported consumer modules and
rebind by identity; warn once about foreign implementations.

* fix(mlx): wrong gated-delta grads for batch rows past the first

mx .at[:, t].add scatters read the update tensor with a wrong batch
stride on the Metal stream, silently corrupting d_q/d_k/d_v/d_g/d_beta
for every batch element except index 0. Collect per-step grads in lists
and stack instead; verified against plain autodiff at B in {2,3,4}.

The underlying Metal scatter indexing bug is fixed upstream in
ml-explore/mlx#3483 but is not in any release yet (latest: 0.31.2).
The list+stack form could revert to the scatter once a fixed mlx is
the minimum supported version, though it is also the cheaper op here.

* fix(mlx): install the gated-delta VJP for qwen3_next and kimi_linear too

The patch trigger was gated on 'qwen3_5' in model_type, but qwen3_next
and kimi_linear share the same gated_delta_update call sites and hit the
identical slow path. Detect gated-delta layers structurally (class name
contains 'delta' + the A_log/dt_bias pair; Mamba/SSM mixers carry the
parameters but never the name) and trigger on that instead.

* feat(mlx): fused Metal backward kernels for gated-delta training

Forward reuses mlx-lm's fused gated_delta_kernel per 64-step chunk
(boundary states fall out as chunk outputs); backward replays chunk
states into scratch (K1) then reverse-scans with atomic gradient
accumulation (K2), translating the validated ops BPTT into Metal.
Graph per layer drops from ~25k unrolled ops to ~100 kernel launches:
Qwen3.5-0.8B LoRA seq-2048 goes 151 -> ~900 tok/s at 4.2 GB peak (was
12.8 GB). Scalar gating, no mask, Dk%32==0; anything else falls back to
the ops VJP automatically. Gradient parity vs plain autodiff: <=1e-7
fp32 across GQA, multi-chunk, uneven tails, and B in {2,3}.

* feat(mlx): vectorized-gating support in the fused gated-delta backward

Extends K1/K2 to kimi_linear's per-column gating (g: [B, T, Hv, Dk])
via the same template-flag pattern mlx-lm uses for its forward kernel.
The vectorized d_g is per-column atomic accumulation (simpler than the
scalar simdgroup reduction); compute_g stays outside the kernel so
d_A_log/d_dt_bias flow through autodiff unchanged. With this, every
gated-delta architecture in mlx-lm trains on the kernel path; the ops
VJP remains only as a guard fallback (no Metal, masked calls, unaligned
head dims, or UNSLOTH_DISABLE_GD_KERNEL_VJP=1 — pass compile=False with
it, mx.compile cannot fuse the unrolled ops graph).
Gradient parity vs plain autodiff: <=2.4e-7 fp32 incl. vectorized+GQA;
scalar path regression-checked bit-identical on Qwen3.5-0.8B.

* test(mlx): consolidate gated-delta tests into one file

Merges the binding-sweep/detection suite and the gradient-parity suite
into tests/test_mlx_gated_delta_vjp.py. The torch shim now installs only
when real mlx is absent (CI), parity tests skip without Metal, and the
module-restore fixture keeps the file order-independent against
shim-based sibling files. Folds source-patched/foreign-untouched asserts
into the sweep test, merges detection cases, and parametrizes parity
over both VJP implementations.

* docs(mlx): pin the kernel-section placement constraint in a comment

* fix(mlx): unify gated-delta patching with #738's VLM patch

After rebasing onto #738: drop its legacy-layout branch (the binding
sweep already rebinds mlx-vlm 0.4-0.5 from-imports), route its
mlx-vlm >= 0.6 training branch through the fused-kernel dispatch (ops
VJP under whole-step mx.compile would reintroduce the compile_fuse
wedge the kernels eliminated), and teach the sweep to recognize the
sibling patch instead of warning about a foreign implementation.

* fix(mlx): avoid stale gated-delta kernel pointers

* perf(mlx): reduce gated-delta backward atomics

* fix(mlx): guard kernel dispatch when Dv is not a threadgroup-row multiple

The backward kernel's shared-memory pre-reduction has row 0 read every
threadgroup row slot; a partial trailing threadgroup (Dv % 4 != 0)
would read uninitialized slots. No current architecture hits this
(Dv is 64/128 everywhere), so route such shapes to the ops VJP.

* fix(mlx): harden gated-delta patch entry points

Review feedback: tolerate a pre-existing patch flag without a recorded
original (skip the identity sweep instead of AttributeError), raise a
clear ValueError when gated_delta_kernel_efficient is called outside
kernel support instead of failing on a None kernel, and parameterize
the warned-foreign set annotation.

* fix(mlx): preserve legacy vlm gated-delta binding

* Tighten MLX gated-delta comments

* Add AGPL-3.0 license header to new MLX gated-delta test

---------

Co-authored-by: danielhanchen <[email protected]>
Co-authored-by: Daniel Han <[email protected]>