From 5c800eb1ed9420b7d0de13bf62ee210fc9e7a23b Mon Sep 17 00:00:00 2001 From: Chandrahas Aroori Date: Fri, 16 Oct 2026 07:20:42 -0700 Subject: [PATCH 1/3] Route per-channel gate in gated_delta_update to the fallback The fused Metal kernels (forward and VJP) read the gate as [B, T, Hv]. A per-channel gate [B, T, Hv, Dk] that passes the head and dim checks was dispatched to them anyway and read with the wrong layout, giving silently wrong outputs and gradients. The ops fallback already handles both gate ranks, so route a 4-D gate there. Adds a test that compares the default stream against the CPU stream for a 4-D gate at a shape the fused kernels accept; it fails before this change and passes after it. Co-Authored-By: Claude Fable 5.1 --- mlx/fast.cpp | 6 ++++- python/tests/test_fast_gated_delta.py | 34 +++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 442c1c43a3..31863eadd9 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1187,7 +1187,11 @@ std::vector gated_delta_update( return std::vector{out, state}; }; - if (!GatedDeltaUpdate::use_fallback(Hk, Dk, Hv, Dv, has_mask, s)) { + // The fused Metal kernels read the gate as [B, T, Hv]. A per-channel gate + // [B, T, Hv, Dk] must take the fallback, which handles both ranks. + const bool vector_gate = g.ndim() == 4; + if (!vector_gate && + !GatedDeltaUpdate::use_fallback(Hk, Dk, Hv, Dv, has_mask, s)) { auto result = array::make_arrays( /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dv, Dk}}, /* dtypes */ {out_dtype, float32}, diff --git a/python/tests/test_fast_gated_delta.py b/python/tests/test_fast_gated_delta.py index f0a5255ff6..b3a84caeb7 100644 --- a/python/tests/test_fast_gated_delta.py +++ b/python/tests/test_fast_gated_delta.py @@ -173,6 +173,40 @@ def test_gated_delta_fallback(self): mx.allclose(hf_ref, hf, atol=1e-4, rtol=1e-4), msg="State " + msg ) + def test_gated_delta_vector_gate(self): + # A per-channel gate [B, T, Hv, Dk] must not reach the fused Metal + # kernels, which read the gate as [B, T, Hv]. Compare the default + # stream against the CPU stream, which always takes the fallback. + B, Hk, Hv, T, Dk, Dv = 1, 32, 32, 64, 128, 128 + mx.random.seed(0) + q = mx.random.normal(shape=(B, T, Hk, Dk)) + k = mx.random.normal(shape=(B, T, Hk, Dk)) + k = k / (mx.linalg.norm(k, axis=-1, keepdims=True) + 1e-6) + v = mx.random.normal(shape=(B, T, Hv, Dv)) + g = mx.random.uniform(low=0.85, high=1.0, shape=(B, T, Hv, Dk)) + b = mx.sigmoid(mx.random.normal(shape=(B, T, Hv))) + h0 = mx.random.normal((B, Hv, Dv, Dk), dtype=mx.float32) + primals = [q, k, v, g, b, h0] + cotans = [ + mx.random.normal(shape=(B, T, Hv, Dv)) * 1e-2, + mx.random.normal(shape=(B, Hv, Dv, Dk)) * 1e-2, + ] + + def f(q, k, v, g, b, h0): + return mx.fast.gated_delta_update(q, k, v, g, b, initial_state=h0) + + outs, vjps = mx.vjp(f, primals, cotans) + with mx.stream(mx.cpu): + outs_ref, vjps_ref = mx.vjp(f, primals, cotans) + for i, (o, o_ref) in enumerate(zip(outs, outs_ref)): + self.assertTrue( + mx.allclose(o_ref, o, atol=1e-4, rtol=1e-4), msg=f"Output {i}" + ) + for i, (vj, vj_ref) in enumerate(zip(vjps, vjps_ref)): + self.assertTrue( + mx.allclose(vj_ref, vj, atol=1e-4, rtol=1e-4), msg=f"VJP {i}" + ) + def test_gated_delta_fallback_masked(self): for dims in self.fallback_dims: B, Hk, Hv, T, Dk, Dv = dims From 5d0a44faf69e0556435ef49d81909458ad536c2d Mon Sep 17 00:00:00 2001 From: Chandrahas Aroori Date: Fri, 16 Oct 2026 07:20:51 -0700 Subject: [PATCH 2/3] Fix gated_delta_update docstring shapes initial_state and final_state are [B, Hv, Dv, Dk], as the C++ op and the tests use, not [B, T, Hk, Dv]. Document the per-channel gate rank. Co-Authored-By: Claude Fable 5.1 --- python/src/fast.cpp | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/python/src/fast.cpp b/python/src/fast.cpp index 31d0f27323..0cca0ac239 100644 --- a/python/src/fast.cpp +++ b/python/src/fast.cpp @@ -420,12 +420,13 @@ void init_fast(nb::module_& parent_module) { q: Queries [B, T, Hk, Dk] k: Keys [B, T, Hk, Dk] v: Values [B, T, Hv, Dv] - gamma: Decay rate in linear space [B, T, Hv] + gamma: Decay rate in linear space, [B, T, Hv] for a scalar + gate per head or [B, T, Hv, Dk] for a per-channel gate beta: Delta update rates [B, T, Hv] - initial_state: Optional initial hidden state [B, T, Hk, Dv] + initial_state: Optional initial hidden state [B, Hv, Dv, Dk] mask: Optional Returns: - Tuple of (output [B, T, Hv, Dv], final_state [B, T, Hk, Dv]) + Tuple of (output [B, T, Hv, Dv], final_state [B, Hv, Dv, Dk]) )"); m.def( From 5359b2ac5e38e014335ea23ef4d4c3b2eb966c60 Mon Sep 17 00:00:00 2001 From: Chandrahas Aroori Date: Tue, 27 Oct 2026 19:01:31 -0700 Subject: [PATCH 3/3] Move the per-channel gate check into GatedDeltaUpdate::use_fallback Pass whether the gate is per-channel to use_fallback, as requested in review, instead of checking it at the call site. Co-Authored-By: Claude Opus 5.5 --- mlx/backend/cuda/primitives.cpp | 1 + mlx/backend/metal/gated_delta_update.cpp | 6 ++++++ mlx/backend/no_gpu/primitives.cpp | 1 + mlx/fast.cpp | 7 ++----- mlx/fast_primitives.h | 1 + 5 files changed, 11 insertions(+), 5 deletions(-) diff --git a/mlx/backend/cuda/primitives.cpp b/mlx/backend/cuda/primitives.cpp index 104c7fd3fc..6c18a60a46 100644 --- a/mlx/backend/cuda/primitives.cpp +++ b/mlx/backend/cuda/primitives.cpp @@ -30,6 +30,7 @@ bool fast::GatedDeltaUpdate::use_fallback( const int Hv, const int Dv, const bool has_mask, + const bool vector_gate, Stream s) { return true; } diff --git a/mlx/backend/metal/gated_delta_update.cpp b/mlx/backend/metal/gated_delta_update.cpp index 2a3c3b1f35..8f41b5db11 100644 --- a/mlx/backend/metal/gated_delta_update.cpp +++ b/mlx/backend/metal/gated_delta_update.cpp @@ -71,6 +71,7 @@ bool GatedDeltaUpdate::use_fallback( const int Hv, const int Dv, const bool has_mask, + const bool vector_gate, Stream s) { if (s.device == Device::cpu) { return true; @@ -78,6 +79,11 @@ bool GatedDeltaUpdate::use_fallback( if (has_mask) { return true; } + // The fused kernels read the gate as [B, T, Hv]. A per-channel gate + // [B, T, Hv, Dk] must take the fallback, which handles both ranks. + if (vector_gate) { + return true; + } return !supported_gated_delta_shape(Hk, Dk, Hv, Dv); } diff --git a/mlx/backend/no_gpu/primitives.cpp b/mlx/backend/no_gpu/primitives.cpp index 4cf04348ae..b4fb081e28 100644 --- a/mlx/backend/no_gpu/primitives.cpp +++ b/mlx/backend/no_gpu/primitives.cpp @@ -60,6 +60,7 @@ bool fast::GatedDeltaUpdate::use_fallback( const int Hv, const int Dv, const bool has_mask, + const bool vector_gate, Stream s) { return true; } diff --git a/mlx/fast.cpp b/mlx/fast.cpp index 31863eadd9..988cd10de5 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1187,11 +1187,8 @@ std::vector gated_delta_update( return std::vector{out, state}; }; - // The fused Metal kernels read the gate as [B, T, Hv]. A per-channel gate - // [B, T, Hv, Dk] must take the fallback, which handles both ranks. - const bool vector_gate = g.ndim() == 4; - if (!vector_gate && - !GatedDeltaUpdate::use_fallback(Hk, Dk, Hv, Dv, has_mask, s)) { + if (!GatedDeltaUpdate::use_fallback( + Hk, Dk, Hv, Dv, has_mask, g.ndim() == 4, s)) { auto result = array::make_arrays( /* output shapes */ {{B, T, Hv, Dv}, {B, Hv, Dv, Dk}}, /* dtypes */ {out_dtype, float32}, diff --git a/mlx/fast_primitives.h b/mlx/fast_primitives.h index 56daea5a68..1c7d169610 100644 --- a/mlx/fast_primitives.h +++ b/mlx/fast_primitives.h @@ -414,6 +414,7 @@ class GatedDeltaUpdate : public Custom { const int Hv, const int Dv, const bool has_mask, + const bool vector_gate, Stream s); void eval_cpu(const std::vector& inputs, std::vector& outputs)