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 442c1c43a3..988cd10de5 100644 --- a/mlx/fast.cpp +++ b/mlx/fast.cpp @@ -1187,7 +1187,8 @@ std::vector gated_delta_update( return std::vector{out, state}; }; - if (!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) 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( 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