Route of per-channel gate in gated_delta_update fixed

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;
 }

mlx/backend/metal/gated_delta_update.cpp

@@ -71,13 +71,19 @@ 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;
   }
   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);
 }
 

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;
 }

mlx/fast.cpp

@@ -1187,7 +1187,8 @@ std::vector<array> gated_delta_update(
     return std::vector<array>{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},

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<array>& inputs, std::vector<array>& outputs)

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(

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