From 3afea8633e9bffd2ea8b9d63783746c71d680888 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Fri, 10 Jul 2026 01:21:48 -0700 Subject: [PATCH 1/2] Add packed gated delta kernel, bitwise-pinned by an explicit-tree comparator The gated delta prefill kernel is latency-bound: each timestep performs two full 32-lane SIMD reductions per value row, and with Dk=128 each lane holds only four state elements. This packs eight independent value rows into a SIMD-group: four lanes own each row and each lane keeps 32 contiguous state elements in registers. Numerics contract: the packed kernel reproduces, bitwise and by construction on any device, an explicit-tree comparator kernel that writes out the ascending butterfly (shuffle_xor 1,2,4,8,16) instead of calling simd_sum; the tests pin packed == comparator with array_equal across shapes, dtypes and lengths. On all current Apple GPUs simd_sum lowers to this same tree (a canary test asserts comparator == existing kernel and would flag any future device or toolchain where the lowering differs, and metal_kernel compiles with MathMode::Safe so the written order is preserved). On the validated M5 toolchain the packed kernel is therefore also bit-identical to the existing kernel, including full-model logits. Blast radius is limited to the shapes the packed kernel serves (Dk == 128, Dv % 8 == 0, scalar gating, no padding mask; generic in batch, head counts and element type). Masked, vector-gate and other-shape paths keep the pre-existing simd_sum kernels byte-for-byte, and the MLX_GDN_PACKED=0 kill-switch routes back to the pre-existing kernel as well. New ops-reference tests cover the masked, vector-gate and Dk=64 generic paths. Timings on M5 Max, Qwen3.6-35B-A3B shapes (B=1, Hk=16, Hv=32, D=128, T=2048): ~1.8-2.0x per layer across measurement rounds (e.g. 2.54-2.56ms -> 1.34-1.38ms on layers 0/20/38); end-to-end prefill 1.064x (median of six saturated adjacent A/B pairs, 5/6 positive). --- mlx_lm/models/gated_delta.py | 328 ++++++++++++++++++++++++++++++++++- tests/test_gated_delta.py | 176 +++++++++++++++++++ 2 files changed, 501 insertions(+), 3 deletions(-) create mode 100644 tests/test_gated_delta.py diff --git a/mlx_lm/models/gated_delta.py b/mlx_lm/models/gated_delta.py index fa6a2ed3f..d04fc6345 100644 --- a/mlx_lm/models/gated_delta.py +++ b/mlx_lm/models/gated_delta.py @@ -1,3 +1,4 @@ +import os from functools import partial from typing import Optional, Tuple @@ -5,6 +6,14 @@ 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): return mx.exp(-mx.exp(A_log.astype(mx.float32)) * nn.softplus(a + dt_bias)) @@ -115,12 +124,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 +401,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 +409,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,10 +462,74 @@ 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 ) diff --git a/tests/test_gated_delta.py b/tests/test_gated_delta.py new file mode 100644 index 000000000..e920c9b0e --- /dev/null +++ b/tests/test_gated_delta.py @@ -0,0 +1,176 @@ +# 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() From 7cbe5f50c5475768800cb0d804351f595f453637 Mon Sep 17 00:00:00 2001 From: Yanzhao Wang Date: Mon, 13 Jul 2026 13:56:05 -0700 Subject: [PATCH 2/2] Format packed gated delta implementation --- mlx_lm/models/gated_delta.py | 9 ++------- tests/test_gated_delta.py | 4 +--- 2 files changed, 3 insertions(+), 10 deletions(-) diff --git a/mlx_lm/models/gated_delta.py b/mlx_lm/models/gated_delta.py index d04fc6345..b1cd782ea 100644 --- a/mlx_lm/models/gated_delta.py +++ b/mlx_lm/models/gated_delta.py @@ -5,7 +5,6 @@ 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, @@ -514,9 +513,7 @@ def gated_delta_kernel_unpacked( 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 - ) + return _gated_delta_kernel_impl(q, k, v, g, beta, state, mask, allow_packed=False) def gated_delta_kernel( @@ -528,9 +525,7 @@ def gated_delta_kernel( 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 - ) + return _gated_delta_kernel_impl(q, k, v, g, beta, state, mask, allow_packed=True) def gated_delta_ops( diff --git a/tests/test_gated_delta.py b/tests/test_gated_delta.py index e920c9b0e..da7565d0c 100644 --- a/tests/test_gated_delta.py +++ b/tests/test_gated_delta.py @@ -130,9 +130,7 @@ 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 - ) + 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)