diff --git a/mlx_lm/models/gated_delta.py b/mlx_lm/models/gated_delta.py index fa6a2ed3f..b1cd782ea 100644 --- a/mlx_lm/models/gated_delta.py +++ b/mlx_lm/models/gated_delta.py @@ -1,9 +1,17 @@ +import os from functools import partial from typing import Optional, Tuple import mlx.core as mx import mlx.nn as nn +# For the shapes it supports, the packed kernel is bitwise-identical by +# construction to an explicit-tree comparator kernel that the tests pin it +# against (see _make_gated_delta_packed_kernel). Every other shape (masks, +# vector gates, Dk != 128) and the MLX_GDN_PACKED=0 kill-switch use the +# original simd_sum kernels, unchanged. +_ENABLE_GDN_PACKED = os.environ.get("MLX_GDN_PACKED", "1") != "0" + @partial(mx.compile, shapeless=True) def compute_g(A_log, a, dt_bias): @@ -115,12 +123,236 @@ def _make_gated_delta_kernel(has_mask=False, vectorized=False): ) +def _make_gated_delta_kernel_xtree(): + """Scalar-gate, unmasked kernel with an explicitly-written reduction. + + This is the unpacked comparator for the packed kernel: it replaces the + two simd_sum calls with the ascending butterfly (shuffle_xor 1,2,4,8,16) + written out in source, so the reduction order is a contract of this file + rather than of the simd_sum lowering. On current Apple GPUs this is the + same tree simd_sum lowers to, so it is bit-identical to the generic + kernel there. Only shapes eligible for the packed kernel ever run it; + masked, vector-gate and Dk != 128 paths keep the original kernels. + """ + 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(i_state[s_idx]); + } + + // g, beta: [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_shuffle_xor(kv_mem, 1); + kv_mem += simd_shuffle_xor(kv_mem, 2); + kv_mem += simd_shuffle_xor(kv_mem, 4); + kv_mem += simd_shuffle_xor(kv_mem, 8); + kv_mem += simd_shuffle_xor(kv_mem, 16); + + 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_shuffle_xor(out, 1); + out += simd_shuffle_xor(out, 2); + out += simd_shuffle_xor(out, 4); + out += simd_shuffle_xor(out, 8); + out += simd_shuffle_xor(out, 16); + if (thread_index_in_simdgroup == 0) { + y[dv_idx] = static_cast(out); + } + 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(state[i]); + } + """ + return mx.fast.metal_kernel( + name="gated_delta_step_xtree", + input_names=["q", "k", "v", "g", "beta", "state_in", "T"], + output_names=["y", "state_out"], + source=source, + ) + + +def _make_gated_delta_packed_kernel(): + """Make the scalar-gate Dk=128 prefill specialization. + + The generic kernel assigns one 32-lane SIMD-group to each value row. For + Dk=128 that leaves every lane with only four state elements and performs + two full-SIMD reductions per row and token. This kernel instead packs + eight independent value rows into a SIMD-group: four lanes own each row + and each lane keeps 32 contiguous state elements in registers. + + The reduction reproduces the unpacked comparator + (_make_gated_delta_kernel_xtree) bitwise BY CONSTRUCTION: both use the + same explicitly-written ascending butterfly (shuffle_xor 1,2,4,8,16) + rather than relying on how simd_sum lowers. The butterfly's first three + levels combine partials that live in a single packed lane (IEEE addition + is commutative, so the local pairwise tree is bit-identical), and the + last two levels map onto shuffle_xor(1) and shuffle_xor(2) within the + four-lane row group. Each 4-element partial keeps the comparator's + sequential order, so y and the state are bit-identical to it on any + device. On current Apple GPUs the explicit tree is also bit-identical + to the simd_sum-based generic kernel. + """ + if not mx.metal.is_available(): + return None + + source = r""" + constexpr int lanes_per_row = 4; + constexpr int rows_per_simdgroup = 32 / lanes_per_row; + constexpr int values_per_lane = Dk / lanes_per_row; + constexpr int partials_per_lane = values_per_lane / 4; + + auto n = thread_position_in_grid.z; + auto b_idx = n / Hv; + auto hv_idx = n % Hv; + auto hk_idx = hv_idx / (Hv / Hk); + + auto lane = thread_index_in_simdgroup; + auto row_in_simdgroup = lane / lanes_per_row; + auto lane_in_row = lane & (lanes_per_row - 1); + auto row_group = thread_position_in_grid.y; + auto dv_idx = row_group * rows_per_simdgroup + row_in_simdgroup; + + // q, k: [B, T, Hk, Dk] + auto q_ = q + (b_idx * T * Hk + hk_idx) * Dk + lane_in_row * values_per_lane; + auto k_ = k + (b_idx * T * Hk + hk_idx) * Dk + lane_in_row * values_per_lane; + + // v, y: [B, T, Hv, Dv] + auto v_ = v + (b_idx * T * Hv + hv_idx) * Dv; + y += (b_idx * T * Hv + hv_idx) * Dv; + + // state_in, state_out: [B, Hv, Dv, Dk] + auto i_state = state_in + (n * Dv + dv_idx) * Dk + lane_in_row * values_per_lane; + auto o_state = state_out + (n * Dv + dv_idx) * Dk + lane_in_row * values_per_lane; + + float state[values_per_lane]; + for (int i = 0; i < values_per_lane; ++i) { + state[i] = static_cast(i_state[i]); + } + + // g, beta: [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 gt = static_cast(g_[hv_idx]); + + // Partials mirror the generic kernel: each 4-element chain is one + // original lane's sequential accumulation. + float part[partials_per_lane]; + for (int pb = 0; pb < partials_per_lane; ++pb) { + float acc = 0.0f; + for (int i = 0; i < 4; ++i) { + int e = pb * 4 + i; + state[e] = state[e] * gt; + acc += state[e] * static_cast(k_[e]); + } + part[pb] = acc; + } + // Butterfly levels xor 1,2,4 stay inside this lane (commutative + // pairwise tree); levels xor 8,16 become the row-group shuffles. + float kv_mem = + ((part[0] + part[1]) + (part[2] + part[3])) + + ((part[4] + part[5]) + (part[6] + part[7])); + kv_mem += simd_shuffle_xor(kv_mem, 1); + kv_mem += simd_shuffle_xor(kv_mem, 2); + + auto delta = + (static_cast(v_[dv_idx]) - kv_mem) * + static_cast(beta_[hv_idx]); + + for (int pb = 0; pb < partials_per_lane; ++pb) { + float acc = 0.0f; + for (int i = 0; i < 4; ++i) { + int e = pb * 4 + i; + state[e] = state[e] + static_cast(k_[e]) * delta; + acc += state[e] * static_cast(q_[e]); + } + part[pb] = acc; + } + float out = + ((part[0] + part[1]) + (part[2] + part[3])) + + ((part[4] + part[5]) + (part[6] + part[7])); + out += simd_shuffle_xor(out, 1); + out += simd_shuffle_xor(out, 2); + if (lane_in_row == 0) { + y[dv_idx] = static_cast(out); + } + + q_ += Hk * Dk; + k_ += Hk * Dk; + v_ += Hv * Dv; + y += Hv * Dv; + g_ += Hv; + beta_ += Hv; + } + + for (int i = 0; i < values_per_lane; ++i) { + o_state[i] = static_cast(state[i]); + } + """ + return mx.fast.metal_kernel( + name="gated_delta_step_packed_btree", + input_names=["q", "k", "v", "g", "beta", "state_in", "T"], + output_names=["y", "state_out"], + source=source, + ) + + _gated_delta_kernel = _make_gated_delta_kernel(has_mask=False, vectorized=False) _gated_delta_kernel_masked = _make_gated_delta_kernel(has_mask=True, vectorized=False) _gated_delta_kernel_vec = _make_gated_delta_kernel(has_mask=False, vectorized=True) _gated_delta_kernel_vec_masked = _make_gated_delta_kernel( has_mask=True, vectorized=True ) +_gated_delta_kernel_xtree = _make_gated_delta_kernel_xtree() +_gated_delta_kernel_packed = _make_gated_delta_packed_kernel() @mx.compile @@ -168,7 +400,7 @@ def _gated_delta_step_ops( return y.astype(q.dtype), state -def gated_delta_kernel( +def _gated_delta_kernel_impl( q: mx.array, k: mx.array, v: mx.array, @@ -176,23 +408,48 @@ def gated_delta_kernel( beta: mx.array, state: mx.array, mask: Optional[mx.array] = None, + *, + allow_packed: bool, ) -> Tuple[mx.array, mx.array]: B, T, Hk, Dk = k.shape Hv, Dv = v.shape[2:] input_type = q.dtype state_type = state.dtype - if g.ndim == 4: + + # The packed kernel gives each lane Dk/4 state elements and packs 32/4 + # value rows into a SIMD-group, so it needs Dk == 128 and Dv divisible by + # 8. It is otherwise generic in B, Hk, Hv and the input element type. + # Vector gating and padding masks keep the original kernels untouched. + packed_eligible = ( + mask is None + and g.ndim == 3 + and Dk == 128 + and Dv % 8 == 0 + and g.dtype == mx.float32 + and state.dtype == mx.float32 + ) + + if packed_eligible and allow_packed and _ENABLE_GDN_PACKED: + kernel = _gated_delta_kernel_packed + inputs = [q, k, v, g, beta, state, T] + grid = (32, Dv // 8, B * Hv) + threadgroup = (32, 2, 1) + elif g.ndim == 4: kernel = _gated_delta_kernel_vec inputs = [q, k, v, g, beta, state, T] if mask is not None: kernel = _gated_delta_kernel_vec_masked inputs.append(mask) + grid = (32, Dv, B * Hv) + threadgroup = (32, 4, 1) else: kernel = _gated_delta_kernel inputs = [q, k, v, g, beta, state, T] if mask is not None: kernel = _gated_delta_kernel_masked inputs.append(mask) + grid = (32, Dv, B * Hv) + threadgroup = (32, 4, 1) return kernel( inputs=inputs, @@ -204,13 +461,73 @@ def gated_delta_kernel( ("Hk", Hk), ("Hv", Hv), ], + grid=grid, + threadgroup=threadgroup, + output_shapes=[(B, T, Hv, Dv), state.shape], + output_dtypes=[input_type, state_type], + ) + + +def gated_delta_kernel_xtree( + q: mx.array, + k: mx.array, + v: mx.array, + g: mx.array, + beta: mx.array, + state: mx.array, + mask: Optional[mx.array] = None, +) -> Tuple[mx.array, mx.array]: + """Explicit-tree comparator for the packed kernel (test use). + + Runs the unpacked layout with the same explicitly-written reduction tree + the packed kernel uses, defining its bitwise contract independently of + the simd_sum lowering. + """ + assert mask is None + B, T, Hk, Dk = k.shape + Hv, Dv = v.shape[2:] + return _gated_delta_kernel_xtree( + inputs=[q, k, v, g, beta, state, T], + template=[ + ("InT", q.dtype), + ("StT", state.dtype), + ("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, state_type], + output_dtypes=[q.dtype, state.dtype], ) +def gated_delta_kernel_unpacked( + q: mx.array, + k: mx.array, + v: mx.array, + g: mx.array, + beta: mx.array, + state: mx.array, + mask: Optional[mx.array] = None, +) -> Tuple[mx.array, mx.array]: + """Run the original one-value-row-per-SIMD-group kernel.""" + return _gated_delta_kernel_impl(q, k, v, g, beta, state, mask, allow_packed=False) + + +def gated_delta_kernel( + q: mx.array, + k: mx.array, + v: mx.array, + g: mx.array, + beta: mx.array, + state: mx.array, + mask: Optional[mx.array] = None, +) -> Tuple[mx.array, mx.array]: + return _gated_delta_kernel_impl(q, k, v, g, beta, state, mask, allow_packed=True) + + def gated_delta_ops( q: mx.array, k: mx.array, diff --git a/tests/test_gated_delta.py b/tests/test_gated_delta.py new file mode 100644 index 000000000..da7565d0c --- /dev/null +++ b/tests/test_gated_delta.py @@ -0,0 +1,174 @@ +# Copyright © 2026 Apple Inc. + +import unittest + +import mlx.core as mx + +import mlx_lm.models.gated_delta as gated_delta +from mlx_lm.models.gated_delta import ( + gated_delta_kernel, + gated_delta_kernel_unpacked, + gated_delta_kernel_xtree, + gated_delta_ops, +) + + +def _normed(shape, D, dtype): + x = mx.random.normal(shape) + return (mx.fast.rms_norm(x, None, 1e-6) * D**-0.5).astype(dtype) + + +def _rel_l2(a, b): + a = a.astype(mx.float32) + b = b.astype(mx.float32) + return (mx.linalg.norm(a - b) / mx.maximum(mx.linalg.norm(b), 1e-9)).item() + + +class TestGatedDelta(unittest.TestCase): + def setUp(self): + self._packed = gated_delta._ENABLE_GDN_PACKED + gated_delta._ENABLE_GDN_PACKED = True + + def tearDown(self): + gated_delta._ENABLE_GDN_PACKED = self._packed + + def test_kill_switch_restores_original_kernel(self): + # MLX_GDN_PACKED=0 routes back to the pre-existing simd_sum kernel. + if mx.default_device() != mx.gpu: + raise unittest.SkipTest("gated delta kernels are GPU only") + gated_delta._ENABLE_GDN_PACKED = False + args = self._inputs(1, 64, 16, 32, 128, 128, mx.bfloat16) + y1, s1 = gated_delta_kernel(*args, None) + y2, s2 = gated_delta_kernel_unpacked(*args, None) + mx.eval(y1, s1, y2, s2) + self.assertTrue(mx.array_equal(y1, y2)) + self.assertTrue(mx.array_equal(s1, s2)) + + def _inputs(self, B, T, Hk, Hv, Dk, Dv, dtype): + mx.random.seed(3) + q = _normed((B, T, Hk, Dk), Dk, dtype) + k = _normed((B, T, Hk, Dk), Dk, dtype) + v = mx.random.normal((B, T, Hv, Dv)).astype(dtype) + # decays in (0, 1], as produced by compute_g + g = mx.exp(-mx.random.uniform(shape=(B, T, Hv)) * 0.2).astype(mx.float32) + beta = mx.random.uniform(shape=(B, T, Hv)).astype(dtype) + state = (mx.random.normal((B, Hv, Dv, Dk)) * 0.3).astype(mx.float32) + mx.eval(q, k, v, g, beta, state) + return q, k, v, g, beta, state + + def test_packed_matches_unpacked(self): + if mx.default_device() != mx.gpu: + raise unittest.SkipTest("gated delta kernels are GPU only") + + cases = [ + # B, Hk, Hv, Dk, Dv, dtype + (1, 16, 32, 128, 128, mx.bfloat16), # Qwen3.5/3.6 shape + (2, 16, 32, 128, 128, mx.bfloat16), # batched + (1, 4, 8, 128, 128, mx.bfloat16), # fewer heads + (1, 8, 8, 128, 128, mx.bfloat16), # Hv == Hk + (1, 16, 32, 128, 128, mx.float16), + (1, 16, 32, 128, 128, mx.float32), + (1, 8, 16, 128, 64, mx.bfloat16), # Dv != Dk + (3, 2, 8, 128, 256, mx.bfloat16), # larger Dv + ] + for B, Hk, Hv, Dk, Dv, dtype in cases: + for T in (1, 7, 64, 257): # decode, ragged, aligned, spilling + with self.subTest(B=B, Hk=Hk, Hv=Hv, Dv=Dv, dtype=dtype, T=T): + args = self._inputs(B, T, Hk, Hv, Dk, Dv, dtype) + y_p, s_p = gated_delta_kernel(*args, None) + y_x, s_x = gated_delta_kernel_xtree(*args, None) + mx.eval(y_p, s_p, y_x, s_x) + # The packed kernel reproduces the comparator's explicit + # reduction tree, so this holds on any device + # independently of how simd_sum lowers. + self.assertTrue(mx.array_equal(y_p, y_x)) + self.assertTrue(mx.array_equal(s_p, s_x)) + + def test_packed_matches_ops_reference(self): + if mx.default_device() != mx.gpu: + raise unittest.SkipTest("gated delta kernels are GPU only") + + q, k, v, g, beta, state = self._inputs(1, 64, 16, 32, 128, 128, mx.bfloat16) + y_p, s_p = gated_delta_kernel(q, k, v, g, beta, state, None) + y_r, s_r = gated_delta_ops(q, k, v, g, beta, state, None) + mx.eval(y_p, s_p, y_r, s_r) + self.assertLess(_rel_l2(y_p, y_r), 2e-3) + self.assertLess(_rel_l2(s_p, s_r), 2e-3) + + def test_explicit_tree_matches_simd_sum_kernel(self): + # On all current Apple GPUs simd_sum lowers to the same ascending + # butterfly the comparator writes out, so the packed kernel is also + # bit-identical to the pre-existing kernel. If this ever fails on a + # new device or toolchain, the packed default should be revisited + # (its contract vs the comparator still holds). + if mx.default_device() != mx.gpu: + raise unittest.SkipTest("gated delta kernels are GPU only") + args = self._inputs(1, 257, 16, 32, 128, 128, mx.bfloat16) + y_x, s_x = gated_delta_kernel_xtree(*args, None) + y_u, s_u = gated_delta_kernel_unpacked(*args, None) + mx.eval(y_x, s_x, y_u, s_u) + self.assertTrue(mx.array_equal(y_x, y_u)) + self.assertTrue(mx.array_equal(s_x, s_u)) + + def test_masked_generic_matches_ops_reference(self): + if mx.default_device() != mx.gpu: + raise unittest.SkipTest("gated delta kernels are GPU only") + q, k, v, g, beta, state = self._inputs(2, 33, 8, 16, 128, 128, mx.bfloat16) + mask = mx.arange(33)[None] < mx.array([[29], [17]]) + y_k, s_k = gated_delta_kernel(q, k, v, g, beta, state, mask) + y_r, s_r = gated_delta_ops(q, k, v, g, beta, state, mask) + mx.eval(y_k, s_k, y_r, s_r) + # Outputs at padded positions are unspecified (the kernel zeros + # them, the ops reference does not); compare valid positions only. + valid = mask[..., None, None] + y_k = mx.where(valid, y_k, 0) + y_r = mx.where(valid, y_r, 0) + self.assertLess(_rel_l2(y_k, y_r), 2e-3) + self.assertLess(_rel_l2(s_k, s_r), 2e-3) + + def test_vector_gate_generic_matches_ops_reference(self): + if mx.default_device() != mx.gpu: + raise unittest.SkipTest("gated delta kernels are GPU only") + q, k, v, _, beta, state = self._inputs(1, 65, 4, 8, 128, 128, mx.bfloat16) + g = mx.exp(-mx.random.uniform(shape=(1, 65, 8, 128)) * 0.2).astype(mx.float32) + mx.eval(g) + y_k, s_k = gated_delta_kernel(q, k, v, g, beta, state, None) + y_r, s_r = gated_delta_ops(q, k, v, g, beta, state, None) + mx.eval(y_k, s_k, y_r, s_r) + self.assertLess(_rel_l2(y_k, y_r), 2e-3) + self.assertLess(_rel_l2(s_k, s_r), 2e-3) + + def test_small_head_dim_generic_matches_ops_reference(self): + if mx.default_device() != mx.gpu: + raise unittest.SkipTest("gated delta kernels are GPU only") + q, k, v, g, beta, state = self._inputs(1, 65, 4, 8, 64, 64, mx.bfloat16) + y_k, s_k = gated_delta_kernel(q, k, v, g, beta, state, None) + y_r, s_r = gated_delta_ops(q, k, v, g, beta, state, None) + mx.eval(y_k, s_k, y_r, s_r) + self.assertLess(_rel_l2(y_k, y_r), 2e-3) + self.assertLess(_rel_l2(s_k, s_r), 2e-3) + + def test_unsupported_shapes_fall_back(self): + if mx.default_device() != mx.gpu: + raise unittest.SkipTest("gated delta kernels are GPU only") + + # Dk != 128 must take the general kernel and stay exact against it. + q, k, v, g, beta, state = self._inputs(1, 32, 4, 8, 64, 64, mx.bfloat16) + y_p, s_p = gated_delta_kernel(q, k, v, g, beta, state, None) + y_u, s_u = gated_delta_kernel_unpacked(q, k, v, g, beta, state, None) + mx.eval(y_p, s_p, y_u, s_u) + self.assertTrue(mx.array_equal(y_p, y_u)) + self.assertTrue(mx.array_equal(s_p, s_u)) + + # A padding mask also forces the general kernel. + q, k, v, g, beta, state = self._inputs(1, 16, 16, 32, 128, 128, mx.bfloat16) + mask = mx.ones((1, 16), dtype=mx.bool_) + y_p, s_p = gated_delta_kernel(q, k, v, g, beta, state, mask) + y_u, s_u = gated_delta_kernel_unpacked(q, k, v, g, beta, state, mask) + mx.eval(y_p, s_p, y_u, s_u) + self.assertTrue(mx.array_equal(y_p, y_u)) + self.assertTrue(mx.array_equal(s_p, s_u)) + + +if __name__ == "__main__": + unittest.main()