Gated-Delta Fused Kernel (Qwen3Next) (#454)

ivanfioravanti · Goekdeniz-Guelmez · awni · web-flow · commit a7f534c3f5cc · 2025-09-17T14:23:40.000-07:00
* apply gating in recurrent_gated_delta_rule

* update cache with new state

* prealocate outputs in recurrent_gated_delta_rule

* feat(kernel): gated-delta kernel scaffolding with CPU fallbacks and tests; integrate in Qwen3Next behind flag

* feat(kernel): implement Metal kernel for gated delta prefill with time iteration to optimize performance

* faster single time step kernel

* use kernel for prefill

* version bump

---------

Co-authored-by: Goekdeniz-Guelmez &lt;gulmezdeniz1999@gmail.com&gt;
Co-authored-by: Awni Hannun &lt;awni@apple.com&gt;
diff --git a/mlx_lm/_version.py b/mlx_lm/_version.py
@@ -1,3 +1,3 @@
 # Copyright © 2023-2025 Apple Inc.
 
-__version__ = "0.27.1"
+__version__ = "0.28.0"
diff --git a/mlx_lm/models/gated_delta.py b/mlx_lm/models/gated_delta.py
@@ -0,0 +1,219 @@
+from functools import partial
+from typing import Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+@partial(mx.compile, shapeless=True)
+def compute_g(A_log, a, dt_bias):
+    return mx.exp(
+        -mx.exp(A_log.astype(mx.float32)) * nn.softplus(a + dt_bias).astype(A_log.dtype)
+    )
+
+
+def _make_gated_delta_kernel():
+    if not mx.metal.is_available():
+        return None
+    source = """
+        auto n = thread_position_in_grid.z;
+        auto b_idx = n / Hv;
+        auto hv_idx = n % Hv;
+        auto hk_idx = hv_idx / (Hv / Hk);
+        constexpr int n_per_t = Dk / 32;
+
+        // q, k: [B, T, Hk, Dk]
+        auto q_ = q + b_idx * T * Hk * Dk + hk_idx * Dk;
+        auto k_ = k + b_idx * T * Hk * Dk + hk_idx * Dk;
+
+        // v, y: [B, T, Hv, Dv]
+        auto v_ = v + b_idx * T * Hv * Dv + hv_idx * Dv;
+        y += b_idx * T * Hv * Dv + hv_idx * Dv;
+
+        auto dk_idx = thread_position_in_threadgroup.x;
+        auto dv_idx = thread_position_in_grid.y;
+
+        // state_in, state_out: [B, Hv, Dv, Dk]
+        auto i_state = state_in + (n * Dv + dv_idx) * Dk;
+        auto o_state = state_out + (n * Dv + dv_idx) * Dk;
+
+        float state[n_per_t];
+        for (int i = 0; i < n_per_t; ++i) {
+           auto s_idx = n_per_t * dk_idx + i;
+           state[i] = static_cast<float>(i_state[s_idx]);
+        }
+
+        // beta, g: [B, T, Hv]
+        auto g_ = g + b_idx * T * Hv;
+        auto beta_ = beta + b_idx * T * Hv;
+
+        for (int t = 0; t < T; ++t) {
+          float kv_mem = 0.0f;
+          for (int i = 0; i < n_per_t; ++i) {
+            auto s_idx = n_per_t * dk_idx + i;
+            state[i] = state[i] * g_[hv_idx];
+            kv_mem += state[i] * k_[s_idx];
+          }
+          kv_mem = simd_sum(kv_mem);
+
+          auto delta = (v_[dv_idx] - kv_mem) * beta_[hv_idx];
+
+          float out = 0.0f;
+          for (int i = 0; i < n_per_t; ++i) {
+            auto s_idx = n_per_t * dk_idx + i;
+            state[i] = state[i] + k_[s_idx] * delta;
+            out += state[i] * q_[s_idx];
+          }
+          out = simd_sum(out);
+          if (thread_index_in_simdgroup == 0) {
+            y[dv_idx] = static_cast<InT>(out);
+          }
+          // Increment data pointers to next time step
+          q_ += Hk * Dk;
+          k_ += Hk * Dk;
+          v_ += Hv * Dv;
+          y += Hv * Dv;
+          g_ += Hv;
+          beta_ += Hv;
+        }
+        for (int i = 0; i < n_per_t; ++i) {
+          auto s_idx = n_per_t * dk_idx + i;
+          o_state[s_idx] = static_cast<InT>(state[i]);
+        }
+    """
+    return mx.fast.metal_kernel(
+        name="gated_delta_step",
+        input_names=["q", "k", "v", "g", "beta", "state_in", "T"],
+        output_names=["y", "state_out"],
+        source=source,
+    )
+
+
+_gated_delta_kernel = _make_gated_delta_kernel()
+
+
+def _gated_delta_step_ops(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+) -> Tuple[mx.array, mx.array]:
+    """
+    Ops-based reference implementation for a single recurrent step.
+
+    Shapes:
+      - q, k: [B, H, Dk]
+      - v: [B, H, Dv]
+      - g, beta: [B, H]
+      - state: [B, H, Dv, Dk]
+    Returns:
+      - y: [B, H, Dv]
+      - new_state: [B, H, Dv, Dk]
+    """
+
+    # Decay
+    state = state * g[..., None, None]
+    kv_mem = (state * k[..., None, :]).sum(axis=-1)  # [B, H, Dv]
+    delta = (v - kv_mem) * beta[..., None]  # [B, H, Dv]
+    state = state + k[..., None, :] * delta[..., None]
+    # Output projection along key dim with q
+    y = (state * q[..., None, :]).sum(axis=-1)  # [B, H, Dv]
+    return y, state
+
+
+def gated_delta_kernel(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: mx.array,
+) -> Tuple[mx.array, mx.array]:
+    B, T, Hk, Dk = k.shape
+    Hv, Dv = v.shape[2:]
+    input_type = q.dtype
+    return _gated_delta_kernel(
+        inputs=[q, k, v, g, beta, state, T],
+        template=[
+            ("InT", input_type),
+            ("Dk", Dk),
+            ("Dv", Dv),
+            ("Hk", Hk),
+            ("Hv", Hv),
+        ],
+        grid=(32, Dv, B * Hv),
+        threadgroup=(32, 4, 1),
+        output_shapes=[(B, T, Hv, Dv), state.shape],
+        output_dtypes=[input_type, input_type],
+    )
+
+
+def gated_delta_ops(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    g: mx.array,
+    beta: mx.array,
+    state: Optional[mx.array] = None,
+) -> Tuple[mx.array, mx.array]:
+    """
+    Ops-based reference implementation for prompt prefill (sequential loop).
+
+    Shapes:
+      - q, k: [B, T, Hk, Dk]
+      - v: [B, T, Hv, Dv]
+      - g, beta: [B, T, Hv]
+      - state: [B, Hv, Dk, Dv]
+    Returns:
+      - y: [B, T, Hv, Dv]
+      - state: [B, Hv, Dk, Dv]
+    """
+    B, T, Hk, Dk = q.shape
+    Hv, Dv = v.shape[-2:]
+    if state is None:
+        state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
+
+    if (repeat_factor := Hv // Hk) > 1:
+        q = mx.repeat(q, repeat_factor, -2)
+        k = mx.repeat(k, repeat_factor, -2)
+
+    ys = []
+    for t in range(T):
+        y, state = _gated_delta_step_ops(
+            q[:, t],
+            k[:, t],
+            v[:, t],
+            g[:, t],
+            beta[:, t],
+            state,
+        )
+        ys.append(y)
+    y = mx.stack(ys, axis=1)
+    return y, state
+
+
+def gated_delta_update(
+    q: mx.array,
+    k: mx.array,
+    v: mx.array,
+    a: mx.array,
+    b: mx.array,
+    A_log: mx.array,
+    dt_bias: mx.array,
+    state: Optional[mx.array] = None,
+) -> Tuple[mx.array, mx.array]:
+
+    beta = mx.sigmoid(b)
+    g = compute_g(A_log, a, dt_bias)
+    if state is None:
+        B, _, Hk, Dk = q.shape
+        Hv, Dv = v.shape[-2:]
+        if state is None:
+            state = mx.zeros((B, Hv, Dv, Dk), dtype=q.dtype)
+
+    if mx.default_device() != mx.gpu or not mx.metal.is_available():
+        return gated_delta_ops(q, k, v, g, beta, state)
+    else:
+        return gated_delta_kernel(q, k, v, g, beta, state)
diff --git a/mlx_lm/models/qwen3_next.py b/mlx_lm/models/qwen3_next.py
@@ -1,14 +1,14 @@
 # Copyright © 2025 Apple Inc.
 
 from dataclasses import dataclass
-from functools import partial
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import mlx.core as mx
 import mlx.nn as nn
 
 from .base import BaseModelArgs, create_attention_mask, scaled_dot_product_attention
 from .cache import KVCache, MambaCache
+from .gated_delta import gated_delta_update
 from .rope_utils import initialize_rope
 from .switch_layers import SwitchGLU
 
@@ -45,52 +45,6 @@ class ModelArgs(BaseModelArgs):
     full_attention_interval: int = 4
 
 
-@partial(mx.compile, shapeless=True)
-def compute_g(A_log, a, dt_bias):
-    return mx.exp(-mx.exp(A_log.astype(mx.float32)) * nn.softplus(a + dt_bias)).astype(
-        A_log.dtype
-    )
-
-
-def recurrent_gated_delta_rule(
-    query: mx.array,
-    key: mx.array,
-    value: mx.array,
-    a: mx.array,
-    b: mx.array,
-    A_log: mx.array,
-    dt_bias: mx.array,
-    state: mx.array,
-    use_qk_l2norm_in_kernel: bool = False,
-) -> Tuple[mx.array, mx.array]:
-    B, S, Hk, Dk = key.shape
-    Hv, Dv = value.shape[2:]
-    inv_scale = Dk**-0.5
-
-    if use_qk_l2norm_in_kernel:
-        query = (inv_scale**2) * mx.fast.rms_norm(query, None, 1e-6)
-        key = inv_scale * mx.fast.rms_norm(key, None, 1e-6)
-    else:
-        query = inv_scale * query
-
-    input_type = query.dtype
-    if (repeat_factor := Hv // Hk) > 1:
-        query = mx.repeat(query, repeat_factor, 2)
-        key = mx.repeat(key, repeat_factor, 2)
-
-    beta = mx.sigmoid(b)
-    g = compute_g(A_log, a, dt_bias)
-
-    out = mx.zeros((B, S, Hv, Dv), dtype=input_type)
-    for i in range(S):
-        state *= g[:, i, :, None, None]
-        kv_mem = (state * key[:, i, :, :, None]).sum(axis=-2)
-        delta = (value[:, i] - kv_mem) * beta[:, i, :, None]
-        state += key[:, i, :, :, None] * delta[..., None, :]
-        out[:, i] = (state * query[:, i, :, :, None]).sum(axis=-2)
-    return out, state
-
-
 class Qwen3NextRMSNormGated(nn.Module):
     def __init__(self, hidden_size: int, eps: float = 1e-6):
         super().__init__()
@@ -297,25 +251,14 @@ def __call__(
             )
         ]
 
-        if cache is not None and cache[1] is not None:
+        if cache is not None:
             state = cache[1]
-        else:
-            state = mx.zeros(
-                (B, self.num_v_heads, self.head_k_dim, self.head_v_dim),
-                dtype=inputs.dtype,
-            )
 
-        out, state = recurrent_gated_delta_rule(
-            q,
-            k,
-            v,
-            a,
-            b,
-            self.A_log,
-            self.dt_bias,
-            state,
-            use_qk_l2norm_in_kernel=True,
-        )
+        inv_scale = k.shape[-1] ** -0.5
+        q = (inv_scale**2) * mx.fast.rms_norm(q, None, 1e-6)
+        k = inv_scale * mx.fast.rms_norm(k, None, 1e-6)
+
+        out, state = gated_delta_update(q, k, v, a, b, self.A_log, self.dt_bias, state)
 
         if cache is not None:
             cache[1] = state
diff --git a/mlx_lm/models/ssm.py b/mlx_lm/models/ssm.py
@@ -39,18 +39,18 @@ def make_ssm_kernel():
         float acc = 0.0;
         auto x_ = static_cast<float>(x[d_idx]);
 
-        for (int i = 0; i < n_per_t; ++i) {{
+        for (int i = 0; i < n_per_t; ++i) {
             auto s_idx = n_per_t * ds_idx + i;
             auto idx = d_idx * Ds + s_idx;
             auto dB_by_x = x_ * dt_ * static_cast<float>(B_[s_idx]);
             auto state = dA * i_state[idx] + dB_by_x;
             o_state[idx] = static_cast<T>(state);
             acc += state * C_[s_idx];
-        }}
+        }
         acc = simd_sum(acc);
-        if (thread_index_in_simdgroup == 0) {{
+        if (thread_index_in_simdgroup == 0) {
             out[d_idx] = static_cast<T>(acc + x_ * D[h_idx]);
-        }}
+        }
     """
     return mx.fast.metal_kernel(
         name="ssm_kernel",
diff --git a/tests/test_models.py b/tests/test_models.py
@@ -10,6 +10,7 @@
 from mlx_lm.models import rope_utils
 from mlx_lm.models.base import create_causal_mask, scaled_dot_product_attention
 from mlx_lm.models.cache import KVCache, RotatingKVCache, make_prompt_cache
+from mlx_lm.models.gated_delta import gated_delta_kernel, gated_delta_ops
 from mlx_lm.models.ssm import ssm_attn, ssm_update
 
 
@@ -1847,6 +1848,27 @@ def test_ssm_masked(self):
         self.assertTrue(mx.allclose(out, out_m, atol=1e-4, rtol=1e-4))
         self.assertTrue(mx.allclose(out_state, out_state_m, atol=1e-4, rtol=1e-4))
 
+    def test_gated_delta(self):
+        for B in [1, 2]:
+            for T in [1, 2]:
+                B = 1
+                Hk = 16
+                Hv = 32
+                Dk = 128
+                Dv = 128
+
+                q = mx.random.normal(shape=(B, T, Hk, Dk))
+                k = mx.random.normal(shape=(B, T, Hk, Dk))
+                v = mx.random.normal(shape=(B, T, Hv, Dv))
+                g = mx.random.normal(shape=(B, T, Hv))
+                beta = mx.random.normal(shape=(B, T, Hv))
+                state = mx.random.normal(shape=(B, Hv, Dk, Dv))
+
+                y_op, st_op = gated_delta_ops(q, k, v, g, beta, state)
+                y_c, st_c = gated_delta_kernel(q, k, v, g, beta, state)
+                self.assertTrue(mx.allclose(y_op, y_c, rtol=1e-4, atol=1e-4))
+                self.assertTrue(mx.allclose(st_op, st_c, rtol=1e-4, atol=1e-3))
+
 
 if __name__ == "__main__":
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,3 @@`
`1`	`1`	`# Copyright © 2023-2025 Apple Inc.`
`2`	`2`
`3`		`-__version__ = "0.27.1"`
	`3`	`+__version__ = "0.28.0"`