Skip to content

Add packed gated delta kernel, bitwise-pinned by an explicit-tree comparator #1559

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wyanzhao wants to merge 4 commits into
base: main
Choose a base branch
from
Open

Add packed gated delta kernel, bitwise-pinned by an explicit-tree comparator #1559

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
3afea86
Add packed gated delta kernel, bitwise-pinned by an explicit-tree com…
wyanzhao Jul 10, 2026
7cbe5f5
Format packed gated delta implementation
wyanzhao Jul 13, 2026
3fedd5d
Merge branch 'main' into mlx-lm-pr-packed-gdn
wyanzhao Jul 13, 2026
e1cc9c5
Merge branch 'main' into mlx-lm-pr-packed-gdn
wyanzhao Aug 7, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
323 changes: 320 additions & 3 deletions mlx_lm/models/gated_delta.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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<float>(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<InT>(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<StT>(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<float>(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<float>(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<float>(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<float>(v_[dv_idx]) - kv_mem) *
static_cast<float>(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<float>(k_[e]) * delta;
acc += state[e] * static_cast<float>(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<InT>(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<StT>(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
Expand Down Expand Up @@ -168,31 +400,56 @@ 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,
g: mx.array,
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,
Expand All @@ -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,
Expand Down
Loading