nn/recurrent: run the gated-delta step in one launch

Decode-length scans spend more time in launch gaps than math: the q/k
norms, decay gate, and recurrence each dispatched separately per layer.
Fuse the step into one Metal kernel over the activated conv output,
with per-token boundary states available from the same pass. The graph
implementation remains as the fallback and contract-miss path, and pins
the kernel bit-for-bit in the parity test.

x/mlxrunner/cache/recurrent_test.go

@@ -89,52 +89,28 @@ func TestRecurrentCachePaddedRoundTrip(t *testing.T) {
 	const L = 4
 	const qLen = 2
 
-	// Use distinct values for the real prefix and the padded tail so
-	// we can detect any leak from padded positions into the result.
-	makeQKV := func(seed float32, T int) (q, k, v *mlx.Array) {
-		mkLast := func(off float32, T, n, d int) *mlx.Array {
-			vals := make([]float32, 1*T*n*d)
-			for i := range vals {
-				vals[i] = off + 0.05*float32(i)
-			}
-			return mlx.FromValues(vals, 1, T, n, d)
+	// Distinct values for the real prefix and large junk in the padded
+	// tail so any leak from padded positions is visible.
+	const packedDim = 2*headKDim + numVHeads*headVDim
+	mkPacked := func(seed float32, T int) (packed, ba *mlx.Array) {
+		pv := make([]float32, T*packedDim)
+		bv := make([]float32, T*2*numVHeads)
+		for i := range pv {
+			pv[i] = seed + 0.05*float32(i)
 		}
-		q = mkLast(seed, T, 1, headKDim)
-		k = mkLast(seed+0.1, T, 1, headKDim)
-		v = mkLast(seed+0.2, T, numVHeads, headVDim)
-		return
-	}
-	makeGB := func(seed float32, T int) (g, beta *mlx.Array) {
-		gVals := make([]float32, 1*T*numVHeads)
-		bVals := make([]float32, 1*T*numVHeads)
-		for i := range gVals {
-			gVals[i] = seed + 0.01*float32(i)
-			bVals[i] = seed - 0.02*float32(i)
+		for i := range bv {
+			bv[i] = seed - 0.02*float32(i)
 		}
-		g = mlx.FromValues(gVals, 1, T, numVHeads)
-		beta = mlx.FromValues(bVals, 1, T, numVHeads)
-		return
+		return mlx.FromValues(pv, 1, T, packedDim), mlx.FromValues(bv, 1, T, 2*numVHeads)
 	}
-	makeQKVPadded := func() (q, k, v *mlx.Array) {
-		qReal, kReal, vReal := makeQKV(0.3, qLen)
-		// Distinct, large junk values in the padded tail to surface
-		// any leak (real outputs are O(1)).
-		qPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, 1, headKDim), 99)
-		kPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, 1, headKDim), 99)
-		vPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, numVHeads, headVDim), 99)
-		q = mlx.Concatenate([]*mlx.Array{qReal, qPad}, 1)
-		k = mlx.Concatenate([]*mlx.Array{kReal, kPad}, 1)
-		v = mlx.Concatenate([]*mlx.Array{vReal, vPad}, 1)
-		return
-	}
-	makeGBPadded := func() (g, beta *mlx.Array) {
-		gReal, betaReal := makeGB(0.1, qLen)
-		gPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, numVHeads), 99)
-		betaPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, numVHeads), 99)
-		g = mlx.Concatenate([]*mlx.Array{gReal, gPad}, 1)
-		beta = mlx.Concatenate([]*mlx.Array{betaReal, betaPad}, 1)
-		return
+	mkPackedPadded := func() (packed, ba *mlx.Array) {
+		pReal, baReal := mkPacked(0.3, qLen)
+		pPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, packedDim), 99)
+		baPad := mlx.AddScalar(mlx.Zeros(mlx.DTypeFloat32, 1, L-qLen, 2*numVHeads), 99)
+		return mlx.Concatenate([]*mlx.Array{pReal, pPad}, 1), mlx.Concatenate([]*mlx.Array{baReal, baPad}, 1)
 	}
+	dtBias := mlx.FromValues([]float32{0.3}, numVHeads)
+	aExp := mlx.FromValues([]float32{0.12}, numVHeads)
 
 	// The conv input dimension must match the cache's convDim.
 	mkConvInput := func(seed float32, T int) *mlx.Array {
@@ -170,16 +146,13 @@ func TestRecurrentCachePaddedRoundTrip(t *testing.T) {
 		_, convStates := nn.CausalConv1D(b, convInput, conv, convTail,
 			nn.WithRecurrentHistory(history))
 
-		var q, k, v, g, beta *mlx.Array
+		var packed, ba *mlx.Array
 		if T == L {
-			q, k, v = makeQKVPadded()
-			g, beta = makeGBPadded()
+			packed, ba = mkPackedPadded()
 		} else {
-			q, k, v = makeQKV(0.3, T)
-			g, beta = makeGB(0.1, T)
+			packed, ba = mkPacked(0.3, T)
 		}
-		_, deltaStates := nn.GatedDelta(b, q, k, v, g, beta,
-			nn.WithRecurrentHistory(history))
+		_, deltaStates := nn.GatedDelta(b, packed, ba, dtBias, aExp, nn.WithRecurrentHistory(history))
 
 		c.Put(b, convStates, deltaStates)
 		return convStates[len(convStates)-1], deltaStates[len(deltaStates)-1]

x/mlxrunner/mlx/gated_delta.go

@@ -1,19 +1,21 @@
 package mlx
 
-var gatedDelta = &gpuKernel{
-	name:    "gated_delta_step",
+import "math"
+
+var gatedDeltaRecurrenceKernel = &gpuKernel{
+	name:    "gated_delta_recurrence",
 	inputs:  []string{"q", "k", "v", "g", "beta", "state_in", "T"},
 	outputs: []string{"y", "state_out"},
-	metal:   gpuSource{source: gatedDeltaMetalKernelSource},
-	cuda:    gpuSource{source: gatedDeltaCUDAKernelSource},
+	metal:   gpuSource{source: gatedDeltaRecurrenceMetalSource},
+	cuda:    gpuSource{source: gatedDeltaRecurrenceCUDASource},
 	fallback: func(launch gpuLaunch) []*Array {
 		in := launch.inputs
-		y, state := gatedDeltaFallback(in[0], in[1], in[2], in[3], in[4], in[5])
+		y, state := gatedDeltaRecurrenceGraph(in[0], in[1], in[2], in[3], in[4], in[5])
 		return []*Array{y, state}
 	},
 }
 
-const gatedDeltaMetalKernelSource = `
+const gatedDeltaRecurrenceMetalSource = `
 auto n = thread_position_in_grid.z;
 auto b_idx = n / Hv;
 auto hv_idx = n % Hv;
@@ -81,7 +83,7 @@ for (int i = 0; i < n_per_t; ++i) {
 }
 `
 
-const gatedDeltaCUDAKernelSource = `
+const gatedDeltaRecurrenceCUDASource = `
 auto tid_x = threadIdx.x;
 auto tid_y = threadIdx.y;
 auto grid_y = blockIdx.y * blockDim.y + tid_y;
@@ -161,18 +163,18 @@ for (int i = 0; i < n_per_t; ++i) {
 }
 `
 
-// gatedDeltaDims are the batch and head geometry of one scan, recovered
-// from the input shapes.
-type gatedDeltaDims struct {
+// gatedDeltaRecurrenceDims are the batch and head geometry of one scan,
+// recovered from the input shapes.
+type gatedDeltaRecurrenceDims struct {
 	B, T, Hk, Dk, Hv, Dv int
 }
 
-// resolveGatedDeltaDims validates the inputs against the GPU kernels'
-// contract and recovers the launch geometry. ok=false routes to the graph
-// fallback: shapes that disagree, Dk not a multiple of the 32-lane simd
-// width, or mixed input dtypes.
-func resolveGatedDeltaDims(q, k, v, g, beta, state *Array) (gatedDeltaDims, bool) {
-	var dims gatedDeltaDims
+// resolveGatedDeltaRecurrenceDims validates the inputs against the GPU
+// kernels' contract and recovers the launch geometry. ok=false routes to
+// the graph implementation: shapes that disagree, Dk not a multiple of the
+// 32-lane simd width, or mixed input dtypes.
+func resolveGatedDeltaRecurrenceDims(q, k, v, g, beta, state *Array) (gatedDeltaRecurrenceDims, bool) {
+	var dims gatedDeltaRecurrenceDims
 	if q == nil || k == nil || v == nil || g == nil || beta == nil || state == nil {
 		return dims, false
 	}
@@ -216,7 +218,7 @@ func repeatHeadsForGatedDelta(x *Array, repeatFactor int) *Array {
 	return Reshape(x, int32(shape[0]), int32(shape[1]), int32(shape[2]*repeatFactor), int32(shape[3]))
 }
 
-func gatedDeltaFallback(q, k, v, g, beta, state *Array) (y, nextState *Array) {
+func gatedDeltaRecurrenceGraph(q, k, v, g, beta, state *Array) (y, nextState *Array) {
 	if q == nil || k == nil || v == nil || g == nil || beta == nil || state == nil {
 		return nil, nil
 	}
@@ -290,43 +292,18 @@ func gatedDeltaFallback(q, k, v, g, beta, state *Array) (y, nextState *Array) {
 	return Concatenate(outs, 1), nextState
 }
 
-// FastGatedDelta runs the recurrent update operation.
-//
-// When mask is non-nil, it must be a [B, T] bool tensor identifying real
-// (true) vs. padded (false) positions in q/k/v/g/beta. Padded positions
-// are substituted with neutral values (q=k=v=beta=0, g=1) so each padded
-// kernel iteration is a no-op — state passes through unchanged and the
-// final state equals the state after the last real token of each row.
-//
-// Inputs that fit the GPU kernels' contract run there (CUDA or Metal, with
-// the graph implementation covering boxes where neither can run); anything
-// else runs the graph implementation directly.
-func FastGatedDelta(q, k, v, g, beta, state, mask *Array) (y, nextState *Array) {
-	// TODO: handle this more efficiently with a masked kernel (MLX-LM has one).
-	if mask != nil {
-		B := int32(mask.Dim(0))
-		T := int32(mask.Dim(1))
-		m4 := Reshape(mask, B, T, 1, 1)
-		m3 := Reshape(mask, B, T, 1)
-		zeroQ := FromValue(float32(0)).AsType(q.DType())
-		zeroK := FromValue(float32(0)).AsType(k.DType())
-		zeroV := FromValue(float32(0)).AsType(v.DType())
-		zeroBeta := FromValue(float32(0)).AsType(beta.DType())
-		oneG := FromValue(float32(1)).AsType(g.DType())
-		q = Where(m4, q, zeroQ)
-		k = Where(m4, k, zeroK)
-		v = Where(m4, v, zeroV)
-		beta = Where(m3, beta, zeroBeta)
-		g = Where(m3, g, oneG)
-	}
-
-	if dims, ok := resolveGatedDeltaDims(q, k, v, g, beta, state); ok {
-		outs := gatedDelta.run(gpuLaunch{
+// gatedDeltaRecurrence runs the scan. Inputs that fit the GPU kernels'
+// contract run there (CUDA or Metal, with the graph implementation covering
+// boxes where neither can run); anything else runs the graph implementation
+// directly.
+func gatedDeltaRecurrence(q, k, v, g, beta, state *Array) (y, nextState *Array) {
+	if dims, ok := resolveGatedDeltaRecurrenceDims(q, k, v, g, beta, state); ok {
+		outs := gatedDeltaRecurrenceKernel.run(gpuLaunch{
 			dtypes: []gpuDTypeArg{{"InT", q.DType()}, {"StT", state.DType()}},
 			ints:   []gpuIntArg{{"Dk", dims.Dk}, {"Dv", dims.Dv}, {"Hk", dims.Hk}, {"Hv", dims.Hv}},
 			outputs: []gpuOutputSpec{
-				{"GATED_DELTA_Y", []int32{int32(dims.B), int32(dims.T), int32(dims.Hv), int32(dims.Dv)}, q.DType()},
-				{"GATED_DELTA_STATE", []int32{int32(dims.B), int32(dims.Hv), int32(dims.Dv), int32(dims.Dk)}, state.DType()},
+				{"GATED_DELTA_RECURRENCE_Y", []int32{int32(dims.B), int32(dims.T), int32(dims.Hv), int32(dims.Dv)}, q.DType()},
+				{"GATED_DELTA_RECURRENCE_STATE", []int32{int32(dims.B), int32(dims.Hv), int32(dims.Dv), int32(dims.Dk)}, state.DType()},
 			},
 			grid:        [3]int{32, dims.Dv, dims.B * dims.Hv},
 			threadGroup: [3]int{32, min(dims.Dv, 4), 1},
@@ -335,9 +312,391 @@ func FastGatedDelta(q, k, v, g, beta, state, mask *Array) (y, nextState *Array)
 		return outs[0], outs[1]
 	}
 
-	y, nextState = gatedDeltaFallback(q, k, v, g, beta, state)
+	y, nextState = gatedDeltaRecurrenceGraph(q, k, v, g, beta, state)
 	if y == nil || nextState == nil {
-		panic("mlx.FastGatedDelta: fallback failed (invalid inputs or unsupported shapes)")
+		panic("mlx: gated-delta recurrence: invalid inputs or unsupported shapes")
 	}
 	return y, nextState
 }
+
+// gatedDeltaMaxTokens caps the fused scan length at the current token plus
+// a ten-token draft — several times the depth the EV controller selects in
+// practice. SeqT is a template argument, so every accepted length compiles
+// its own pipeline variant; the cap bounds that set, and longer windows run
+// the same step as graph ops.
+const gatedDeltaMaxTokens = 11
+
+var (
+	gatedDelta = &gpuKernel{
+		name:    "gated_delta",
+		inputs:  []string{"packed", "ba", "dt_bias", "a_exp", "qk_scale", "state_in"},
+		outputs: []string{"y", "state_out"},
+		metal: gpuSource{
+			source: gatedDeltaMetalSource,
+			header: gatedDeltaMetalHeader + "#define GDN_STORE_INTERIOR(index, value)\n",
+		},
+		fallback: func(launch gpuLaunch) []*Array {
+			in := launch.inputs
+			y, end, _ := gatedDeltaGraph(in[0], in[1], in[2], in[3], in[5], false)
+			return []*Array{y, end}
+		},
+	}
+	gatedDeltaStates = &gpuKernel{
+		name:    "gated_delta_states",
+		inputs:  []string{"packed", "ba", "dt_bias", "a_exp", "qk_scale", "state_in"},
+		outputs: []string{"y", "state_out", "state_seq"},
+		metal: gpuSource{
+			source: gatedDeltaMetalSource,
+			header: gatedDeltaMetalHeader + "#define GDN_STORE_INTERIOR(index, value) state_seq[index] = value\n",
+		},
+		fallback: func(launch gpuLaunch) []*Array {
+			in := launch.inputs
+			y, end, interior := gatedDeltaGraph(in[0], in[1], in[2], in[3], in[5], true)
+			for i, s := range interior {
+				interior[i] = ExpandDims(s, 0)
+			}
+			return []*Array{y, end, Concatenate(interior, 0)}
+		},
+	}
+)
+
+// gatedDeltaGraph is the graph implementation of the fused gated-delta step
+// over the kernels' contract domain: q/k RMS norm and scaling, the decay
+// gate, and the (per-token, for captureAll) scan. Geometry is recovered
+// from the packed and state shapes, and the q/k scales are recomputed from
+// Dk with the same bits the kernel input carries.
+func gatedDeltaGraph(packed, ba, dtBias, aExp, state *Array, captureAll bool) (y, nextState *Array, interior []*Array) {
+	B, T := int32(packed.Dim(0)), int32(packed.Dim(1))
+	sd := state.Dims()
+	Hv, Dv, Dk := int32(sd[1]), int32(sd[2]), int32(sd[3])
+	keyDim := (int32(packed.Dim(2)) - Hv*Dv) / 2
+	Hk := keyDim / Dk
+
+	q := SliceStartStop(packed, []int32{0, 0, 0}, []int32{B, T, keyDim})
+	k := SliceStartStop(packed, []int32{0, 0, keyDim}, []int32{B, T, 2 * keyDim})
+	v := SliceStartStop(packed, []int32{0, 0, 2 * keyDim}, []int32{B, T, 2*keyDim + Hv*Dv})
+	q = Reshape(q, B, T, Hk, Dk)
+	k = Reshape(k, B, T, Hk, Dk)
+	v = Reshape(v, B, T, Hv, Dv)
+	invScale := gatedDeltaInvScale(int(Dk))
+	q = MulScalar(RMSNormFn(q, nil, 1e-6), invScale*invScale)
+	k = MulScalar(RMSNormFn(k, nil, 1e-6), invScale)
+
+	beta := SliceStartStop(ba, []int32{0, 0, 0}, []int32{B, T, Hv})
+	alpha := SliceStartStop(ba, []int32{0, 0, Hv}, []int32{B, T, 2 * Hv})
+	decay := Softplus(Add(alpha, dtBias))
+	decay = Mul(decay, aExp)
+	decay = Exp(MulScalar(decay, -1)).AsType(alpha.DType())
+	betaGate := Sigmoid(beta)
+
+	if !captureAll || T == 1 {
+		y, nextState = gatedDeltaRecurrence(q, k, v, decay, betaGate, state)
+		return y, nextState, nil
+	}
+
+	sliceT := func(a *Array, t int32) *Array {
+		dims := a.Dims()
+		start := make([]int32, len(dims))
+		stop := make([]int32, len(dims))
+		for d := range dims {
+			stop[d] = int32(dims[d])
+		}
+		start[1], stop[1] = t, t+1
+		return SliceStartStop(a, start, stop)
+	}
+	outs := make([]*Array, T)
+	for t := range T {
+		outs[t], state = gatedDeltaRecurrence(sliceT(q, t), sliceT(k, t), sliceT(v, t), sliceT(decay, t), sliceT(betaGate, t), state)
+		if t+1 < T {
+			interior = append(interior, state)
+		}
+	}
+	return Concatenate(outs, 1), state, interior
+}
+
+// The fused scan consumes the activated causal-conv output rows [q | k | v]
+// plus the packed [beta | alpha] projection and performs q/k RMS norm and
+// scaling, the decay gate, and the gated-delta recurrence in one launch.
+const gatedDeltaMetalSource = `
+constexpr int SIMDGroups = 4;
+constexpr int ValuesPerSIMD = DvTile / SIMDGroups;
+constexpr int StatePerLane = Dk / 32;
+constexpr int QDim = Hk * Dk;
+
+auto lane = thread_position_in_threadgroup.x;
+auto simd_idx = thread_position_in_threadgroup.y;
+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 value_tile = threadgroup_position_in_grid.y;
+
+threadgroup InT q_cache[Dk];
+threadgroup InT k_cache[Dk];
+threadgroup InT decay_cache[1];
+threadgroup InT beta_cache[1];
+
+float state[ValuesPerSIMD][StatePerLane];
+for (int value_iter = 0; value_iter < ValuesPerSIMD; ++value_iter) {
+  auto dv_idx = value_tile * DvTile + simd_idx + value_iter * SIMDGroups;
+  auto state_offset = (n * Dv + dv_idx) * Dk;
+  for (int i = 0; i < StatePerLane; ++i) {
+    auto d = StatePerLane * lane + i;
+    state[value_iter][i] = state_in[state_offset + d];
+  }
+}
+
+for (int t = 0; t < SeqT; ++t) {
+  if (simd_idx == 0) {
+    auto token = packed + (b_idx * SeqT + t) * PackedDim;
+    float q_values[StatePerLane];
+    float k_values[StatePerLane];
+    float q_squares = 0.0f;
+    float k_squares = 0.0f;
+
+    for (int i = 0; i < StatePerLane; ++i) {
+      auto d = StatePerLane * lane + i;
+      float q_value = static_cast<float>(token[hk_idx * Dk + d]);
+      float k_value = static_cast<float>(token[QDim + hk_idx * Dk + d]);
+      q_values[i] = q_value;
+      k_values[i] = k_value;
+      q_squares += q_value * q_value;
+      k_squares += k_value * k_value;
+    }
+
+    q_squares = simd_sum(q_squares);
+    k_squares = simd_sum(k_squares);
+    float q_inv_rms = metal::precise::rsqrt(q_squares / float(Dk) + 1.0e-6f);
+    float k_inv_rms = metal::precise::rsqrt(k_squares / float(Dk) + 1.0e-6f);
+    InT q_scale = static_cast<InT>(qk_scale[0]);
+    InT k_scale = static_cast<InT>(qk_scale[1]);
+    for (int i = 0; i < StatePerLane; ++i) {
+      auto d = StatePerLane * lane + i;
+      InT q_norm = static_cast<InT>(q_values[i] * q_inv_rms);
+      InT k_norm = static_cast<InT>(k_values[i] * k_inv_rms);
+      q_cache[d] = static_cast<InT>(q_norm * q_scale);
+      k_cache[d] = static_cast<InT>(k_norm * k_scale);
+    }
+
+    if (lane == 0) {
+      auto ba_row = (b_idx * SeqT + t) * 2 * Hv;
+      InT gate_input = static_cast<InT>(ba[ba_row + Hv + hv_idx] + dt_bias[hv_idx]);
+      InT softplus = gdn_logaddexp(gate_input, static_cast<InT>(0));
+      float decay = metal::precise::exp(-static_cast<float>(softplus) * a_exp[hv_idx]);
+      decay_cache[0] = static_cast<InT>(decay);
+      beta_cache[0] = gdn_sigmoid(ba[ba_row + hv_idx]);
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+
+  for (int value_iter = 0; value_iter < ValuesPerSIMD; ++value_iter) {
+    auto dv_idx = value_tile * DvTile + simd_idx + value_iter * SIMDGroups;
+    float projection = 0.0f;
+    for (int i = 0; i < StatePerLane; ++i) {
+      auto d = StatePerLane * lane + i;
+      state[value_iter][i] *= static_cast<float>(decay_cache[0]);
+      projection += state[value_iter][i] * static_cast<float>(k_cache[d]);
+    }
+    projection = simd_sum(projection);
+
+    auto token = packed + (b_idx * SeqT + t) * PackedDim;
+    float v_value = static_cast<float>(token[2 * QDim + hv_idx * Dv + dv_idx]);
+    float delta = (v_value - projection) * static_cast<float>(beta_cache[0]);
+
+    float out = 0.0f;
+    for (int i = 0; i < StatePerLane; ++i) {
+      auto d = StatePerLane * lane + i;
+      state[value_iter][i] += static_cast<float>(k_cache[d]) * delta;
+      out += state[value_iter][i] * static_cast<float>(q_cache[d]);
+    }
+    out = simd_sum(out);
+    if (lane == 0) {
+      y[((b_idx * SeqT + t) * Hv + hv_idx) * Dv + dv_idx] = static_cast<InT>(out);
+    }
+
+    if (t + 1 < SeqT) {
+      auto seq_offset = ((t * threads_per_grid.z + n) * Dv + dv_idx) * Dk;
+      for (int i = 0; i < StatePerLane; ++i) {
+        auto d = StatePerLane * lane + i;
+        GDN_STORE_INTERIOR(seq_offset + d, state[value_iter][i]);
+      }
+    }
+  }
+  threadgroup_barrier(mem_flags::mem_threadgroup);
+}
+
+for (int value_iter = 0; value_iter < ValuesPerSIMD; ++value_iter) {
+  auto dv_idx = value_tile * DvTile + simd_idx + value_iter * SIMDGroups;
+  auto state_offset = (n * Dv + dv_idx) * Dk;
+  for (int i = 0; i < StatePerLane; ++i) {
+    auto d = StatePerLane * lane + i;
+    state_out[state_offset + d] = state[value_iter][i];
+  }
+}
+`
+
+const gatedDeltaMetalHeader = `
+template <typename T>
+T gdn_sigmoid(T x) {
+  auto y = 1 / (1 + metal::exp(metal::abs(x)));
+  return (x < 0) ? y : 1 - y;
+}
+
+template <typename T>
+T gdn_logaddexp(T x, T y) {
+  if (metal::isnan(x) || metal::isnan(y)) {
+    return metal::numeric_limits<T>::quiet_NaN();
+  }
+  constexpr T inf = metal::numeric_limits<T>::infinity();
+  T maxval = metal::max(x, y);
+  T minval = metal::min(x, y);
+  return (minval == -inf || maxval == inf)
+      ? maxval
+      : (maxval + log1p(metal::exp(minval - maxval)));
+}
+`
+
+// gatedDeltaInvScale is the q/k norm scale for key dimension dk; the
+// kernel's qk_scale input and the graph implementation share these bits.
+func gatedDeltaInvScale(dk int) float32 {
+	return float32(1.0 / math.Sqrt(float64(dk)))
+}
+
+// gatedDeltaDims are the batch and head geometry the kernel is instantiated
+// for, recovered from the input shapes.
+type gatedDeltaDims struct {
+	B, Hk, Dk, Hv, Dv, PackedDim, T int
+}
+
+// GatedDelta runs the whole gated-delta step — q/k norms, decay gate, and
+// the scan — in one launch over the activated causal-conv output. packed is
+// [B, T, 2*Hk*Dk + Hv*Dv] with rows packed [q | k | v], ba is [B, T, 2*Hv]
+// packed [beta | alpha] rows, and state is [B, Hv, Dv, Dk]. captureAll
+// additionally emits every interior per-token state. Inputs that fit the
+// one-launch kernels' contract run there; anything else runs the same step
+// as graph ops.
+//
+// When mask is non-nil, it must be a [B, T] bool tensor identifying real
+// (true) vs. padded (false) positions. Padded rows are neutralized before
+// the kernels' own preprocessing: zeroed conv rows make the q/k/v norms
+// emit zeros, and -inf beta/alpha rows yield beta 0 and decay 1, so each
+// padded position is an identity step with zero output, exactly as on the
+// recurrence path.
+func GatedDelta(packed, ba, dtBias, aExp, state, mask *Array, captureAll bool) (y, nextState *Array, interior []*Array) {
+	if mask != nil {
+		mask3 := Reshape(mask, int32(mask.Dim(0)), int32(mask.Dim(1)), 1)
+		zero := FromValue(float32(0)).AsType(packed.DType())
+		packed = Where(mask3, packed, zero)
+		negInf := FromValue(float32(math.Inf(-1))).AsType(ba.DType())
+		ba = Where(mask3, ba, negInf)
+	}
+
+	dims, ok := resolveGatedDeltaDims(packed, ba, dtBias, aExp, state)
+	if !ok {
+		return gatedDeltaGraph(packed, ba, dtBias, aExp, state, captureAll)
+	}
+
+	inv := gatedDeltaInvScale(dims.Dk)
+	qkScale := FromValues([]float32{inv * inv, inv}, 2)
+	dvTile := 32
+	if dims.Dv%dvTile != 0 {
+		// resolveGatedDeltaDims guarantees Dv%16 == 0.
+		dvTile = 16
+	}
+	useAllStates := captureAll && dims.T > 1
+	kernel := gatedDelta
+	outputs := []gpuOutputSpec{
+		{"GATED_DELTA_Y", []int32{int32(dims.B), int32(dims.T), int32(dims.Hv), int32(dims.Dv)}, packed.DType()},
+		{"GATED_DELTA_STATE", []int32{int32(dims.B), int32(dims.Hv), int32(dims.Dv), int32(dims.Dk)}, DTypeFloat32},
+	}
+	if useAllStates {
+		kernel = gatedDeltaStates
+		outputs = append(outputs, gpuOutputSpec{
+			"GATED_DELTA_STATE_SEQ", []int32{int32(dims.T - 1), int32(dims.B), int32(dims.Hv), int32(dims.Dv), int32(dims.Dk)}, DTypeFloat32,
+		})
+	}
+
+	outs := kernel.run(gpuLaunch{
+		dtypes: []gpuDTypeArg{{"InT", packed.DType()}},
+		ints: []gpuIntArg{
+			{"Hk", dims.Hk},
+			{"Dk", dims.Dk},
+			{"Hv", dims.Hv},
+			{"Dv", dims.Dv},
+			{"PackedDim", dims.PackedDim},
+			{"SeqT", dims.T},
+			{"DvTile", dvTile},
+		},
+		outputs:     outputs,
+		grid:        [3]int{32, (dims.Dv / dvTile) * 4, dims.B * dims.Hv},
+		threadGroup: [3]int{32, 4, 1},
+		inputs:      []*Array{packed, ba, dtBias, aExp, qkScale, state},
+	})
+	if useAllStates {
+		interior = sliceGatedDeltaStates(outs[2], dims)
+	}
+	return outs[0], outs[1], interior
+}
+
+func resolveGatedDeltaDims(packed, ba, dtBias, aExp, state *Array) (gatedDeltaDims, bool) {
+	var dims gatedDeltaDims
+	if packed == nil || ba == nil || dtBias == nil || aExp == nil || state == nil {
+		return dims, false
+	}
+	sd := state.Dims()
+	if len(sd) != 4 || sd[0] < 1 {
+		return dims, false
+	}
+	dims.B, dims.Hv, dims.Dv, dims.Dk = sd[0], sd[1], sd[2], sd[3]
+
+	pd := packed.Dims()
+	if len(pd) != 3 || pd[0] != dims.B || pd[1] < 1 || pd[1] > gatedDeltaMaxTokens {
+		return dims, false
+	}
+	dims.T, dims.PackedDim = pd[1], pd[2]
+
+	keyRows := dims.PackedDim - dims.Hv*dims.Dv
+	if keyRows <= 0 || dims.Dk <= 0 || keyRows%(2*dims.Dk) != 0 {
+		return dims, false
+	}
+	dims.Hk = keyRows / (2 * dims.Dk)
+	if dims.Hk <= 0 || dims.Hv%dims.Hk != 0 || dims.Dk%32 != 0 || dims.Dv%16 != 0 {
+		return dims, false
+	}
+
+	if !exactShape(ba, dims.B, dims.T, 2*dims.Hv) ||
+		!exactShape(dtBias, dims.Hv) ||
+		!exactShape(aExp, dims.Hv) {
+		return dims, false
+	}
+	if packed.DType() != DTypeBFloat16 || ba.DType() != DTypeBFloat16 ||
+		dtBias.DType() != DTypeBFloat16 || aExp.DType() != DTypeFloat32 ||
+		state.DType() != DTypeFloat32 {
+		return dims, false
+	}
+	return dims, true
+}
+
+func sliceGatedDeltaStates(stateSeq *Array, dims gatedDeltaDims) []*Array {
+	interior := make([]*Array, dims.T-1)
+	for t := range interior {
+		s := SliceStartStop(stateSeq,
+			[]int32{int32(t), 0, 0, 0, 0},
+			[]int32{int32(t) + 1, int32(dims.B), int32(dims.Hv), int32(dims.Dv), int32(dims.Dk)})
+		interior[t] = Reshape(s, int32(dims.B), int32(dims.Hv), int32(dims.Dv), int32(dims.Dk))
+	}
+	return interior
+}
+
+func exactShape(value *Array, shape ...int) bool {
+	dims := value.Dims()
+	if len(dims) != len(shape) {
+		return false
+	}
+	for i := range shape {
+		if dims[i] != shape[i] {
+			return false
+		}
+	}
+	return true
+}

x/mlxrunner/mlx/gated_delta_test.go

@@ -0,0 +1,242 @@
+package mlx
+
+import (
+	"fmt"
+	"math"
+	"testing"
+)
+
+type gatedDeltaTestGeometry struct {
+	Hk, Dk, Hv, Dv int
+}
+
+func (g gatedDeltaTestGeometry) packedDim() int { return 2*g.Hk*g.Dk + g.Hv*g.Dv }
+
+type gatedDeltaTestInputs struct {
+	packed, ba, dtBias, aExp, state *Array
+}
+
+func gatedDeltaTestInputs36(g gatedDeltaTestGeometry, T int) gatedDeltaTestInputs {
+	return gatedDeltaTestInputs{
+		packed: patternArray(DTypeBFloat16, []int{1, T, g.packedDim()}, 0.01, 0.003, 37, 257),
+		ba:     patternArray(DTypeBFloat16, []int{1, T, 2 * g.Hv}, 0, 0.4, 13, 37),
+		dtBias: patternArray(DTypeBFloat16, []int{g.Hv}, -2.5, 0.08, 5, 29),
+		aExp:   patternArray(DTypeFloat32, []int{g.Hv}, 0.12, 0.002, 3, 19),
+		state:  patternArray(DTypeFloat32, []int{1, g.Hv, g.Dv, g.Dk}, 0, 0.0001, 17, 101),
+	}
+}
+
+// gatedDeltaReference runs the graph implementation the fused kernels fall
+// back to; the test pins the kernels against it bit-for-bit.
+func gatedDeltaReference(in gatedDeltaTestInputs, captureAll bool) (y, nextState *Array, interior []*Array) {
+	return gatedDeltaGraph(in.packed, in.ba, in.dtBias, in.aExp, in.state, captureAll)
+}
+
+// Exactness is per compiler pair: the metallib and the runtime JIT round
+// metal::exp a float ulp apart, which after bf16 rounding leaves rare
+// differing inputs (beta sigmoid at -6.84375); the lattice here avoids them.
+func TestGatedDeltaMatchesGraph(t *testing.T) {
+	skipIfNoMLX(t)
+	withMLXThread(t, func() {
+		testGatedDeltaMatchesGraph(t)
+	})
+}
+
+func testGatedDeltaMatchesGraph(t *testing.T) {
+	geometries := []gatedDeltaTestGeometry{
+		{Hk: 16, Dk: 128, Hv: 32, Dv: 128},
+		{Hk: 4, Dk: 64, Hv: 8, Dv: 32},
+	}
+	for _, g := range geometries {
+		for T := 1; T <= gatedDeltaMaxTokens; T++ {
+			for _, B := range []int{1, 3} {
+				for _, captureAll := range []bool{false, true} {
+					name := fmt.Sprintf("hk%d_dk%d_T%d_B%d_all%v", g.Hk, g.Dk, T, B, captureAll)
+					in := batchGatedDeltaRows(gatedDeltaTestInputs36(g, T), B)
+					refY, refState, refInterior := gatedDeltaReference(in, captureAll)
+					y, state, interior := GatedDelta(in.packed, in.ba, in.dtBias, in.aExp, in.state, nil, captureAll)
+					if err := requireExact("y", y, refY); err != nil {
+						t.Fatalf("%s: %v", name, err)
+					}
+					if err := requireExact("state", state, refState); err != nil {
+						t.Fatalf("%s: %v", name, err)
+					}
+					if captureAll && len(interior) != len(refInterior) {
+						t.Fatalf("%s: interior count = %d, want %d", name, len(interior), len(refInterior))
+					}
+					for i := range interior {
+						if err := requireExact(fmt.Sprintf("interior[%d]", i), interior[i], refInterior[i]); err != nil {
+							t.Fatalf("%s: %v", name, err)
+						}
+					}
+				}
+			}
+		}
+	}
+}
+
+func TestGatedDeltaGraphRouting(t *testing.T) {
+	skipIfNoMLX(t)
+	withMLXThread(t, func() {
+		testGatedDeltaGraphRouting(t)
+	})
+}
+
+// Contract misses — T beyond the kernel cap and a float32 packed input —
+// run the same step as graph ops.
+func testGatedDeltaGraphRouting(t *testing.T) {
+	g := gatedDeltaTestGeometry{Hk: 4, Dk: 64, Hv: 8, Dv: 32}
+
+	check := func(name string, in gatedDeltaTestInputs, captureAll bool) {
+		y, end, interior := GatedDelta(in.packed, in.ba, in.dtBias, in.aExp, in.state, nil, captureAll)
+		refY, refEnd, refInterior := gatedDeltaReference(in, captureAll)
+		if err := requireExact("y", y, refY); err != nil {
+			t.Fatalf("%s: %v", name, err)
+		}
+		if err := requireExact("state", end, refEnd); err != nil {
+			t.Fatalf("%s: %v", name, err)
+		}
+		if len(interior) != len(refInterior) {
+			t.Fatalf("%s: interior count = %d, want %d", name, len(interior), len(refInterior))
+		}
+	}
+
+	check("tooLong", gatedDeltaTestInputs36(g, gatedDeltaMaxTokens+1), true)
+
+	f32 := gatedDeltaTestInputs36(g, 2)
+	f32.packed = f32.packed.AsType(DTypeFloat32)
+	check("float32", f32, false)
+}
+
+// batchGatedDeltaRows widens a single-row input to B rows with distinct
+// per-row contents.
+func batchGatedDeltaRows(row gatedDeltaTestInputs, B int) gatedDeltaTestInputs {
+	if B == 1 {
+		return row
+	}
+	packed := make([]*Array, B)
+	ba := make([]*Array, B)
+	state := make([]*Array, B)
+	for i := range B {
+		r := row
+		if i > 0 {
+			r = scaledGatedDeltaRow(row, 1-0.3*float32(i))
+		}
+		packed[i], ba[i], state[i] = r.packed, r.ba, r.state
+	}
+	return gatedDeltaTestInputs{
+		packed: Concatenate(packed, 0),
+		ba:     Concatenate(ba, 0),
+		dtBias: row.dtBias,
+		aExp:   row.aExp,
+		state:  Concatenate(state, 0),
+	}
+}
+
+func scaledGatedDeltaRow(base gatedDeltaTestInputs, scale float32) gatedDeltaTestInputs {
+	return gatedDeltaTestInputs{
+		packed: MulScalar(base.packed, scale),
+		ba:     MulScalar(base.ba, -scale),
+		dtBias: base.dtBias,
+		aExp:   base.aExp,
+		state:  MulScalar(base.state, scale),
+	}
+}
+
+func TestGatedDeltaBatchedRows(t *testing.T) {
+	skipIfNoMLX(t)
+	withMLXThread(t, func() {
+		testGatedDeltaBatchedRows(t)
+	})
+}
+
+// Each batched row must match its own single-row launch bit-for-bit.
+func testGatedDeltaBatchedRows(t *testing.T) {
+	g := gatedDeltaTestGeometry{Hk: 4, Dk: 64, Hv: 8, Dv: 32}
+	for _, T := range []int{1, 5, 11} {
+		rows := []gatedDeltaTestInputs{gatedDeltaTestInputs36(g, T)}
+		rows = append(rows, scaledGatedDeltaRow(rows[0], 0.7), scaledGatedDeltaRow(rows[0], 0.4))
+
+		packed := Concatenate([]*Array{rows[0].packed, rows[1].packed, rows[2].packed}, 0)
+		ba := Concatenate([]*Array{rows[0].ba, rows[1].ba, rows[2].ba}, 0)
+		state := Concatenate([]*Array{rows[0].state, rows[1].state, rows[2].state}, 0)
+
+		y, end, _ := GatedDelta(packed, ba, rows[0].dtBias, rows[0].aExp, state, nil, false)
+		for i, row := range rows {
+			refY, refEnd, _ := GatedDelta(row.packed, row.ba, row.dtBias, row.aExp, row.state, nil, false)
+			gotY := SliceStartStop(y,
+				[]int32{int32(i), 0, 0, 0},
+				[]int32{int32(i) + 1, int32(T), int32(g.Hv), int32(g.Dv)})
+			gotEnd := SliceStartStop(end,
+				[]int32{int32(i), 0, 0, 0},
+				[]int32{int32(i) + 1, int32(g.Hv), int32(g.Dv), int32(g.Dk)})
+			if err := requireExact("y", gotY, refY); err != nil {
+				t.Fatalf("T=%d row %d: %v", T, i, err)
+			}
+			if err := requireExact("state", gotEnd, refEnd); err != nil {
+				t.Fatalf("T=%d row %d: %v", T, i, err)
+			}
+		}
+	}
+}
+
+func TestGatedDeltaRaggedRows(t *testing.T) {
+	skipIfNoMLX(t)
+	withMLXThread(t, func() {
+		testGatedDeltaRaggedRows(t)
+	})
+}
+
+// A padded row neutralized the way the nn seam does it — zeroed conv rows,
+// ba tail poisoned to -inf — runs identity steps with zero output there: its
+// full output and final state must match a launch of only its real tokens
+// with zeros appended, while the full-length row is unaffected.
+func testGatedDeltaRaggedRows(t *testing.T) {
+	g := gatedDeltaTestGeometry{Hk: 4, Dk: 64, Hv: 8, Dv: 32}
+	const T, realLen = 6, 4
+	row0 := gatedDeltaTestInputs36(g, T)
+	row1 := scaledGatedDeltaRow(row0, 0.7)
+
+	pad := make([]float32, (T-realLen)*2*g.Hv)
+	negInf := float32(math.Inf(-1))
+	for i := range pad {
+		pad[i] = negInf
+	}
+	row1BA := Concatenate([]*Array{
+		SliceStartStop(row1.ba, []int32{0, 0, 0}, []int32{1, realLen, int32(2 * g.Hv)}),
+		FromValues(pad, 1, T-realLen, 2*g.Hv).AsType(DTypeBFloat16),
+	}, 1)
+	row1Packed := Concatenate([]*Array{
+		SliceStartStop(row1.packed, []int32{0, 0, 0}, []int32{1, realLen, int32(g.packedDim())}),
+		Zeros(DTypeBFloat16, 1, T-realLen, g.packedDim()),
+	}, 1)
+
+	packed := Concatenate([]*Array{row0.packed, row1Packed}, 0)
+	ba := Concatenate([]*Array{row0.ba, row1BA}, 0)
+	state := Concatenate([]*Array{row0.state, row1.state}, 0)
+
+	y, end, _ := GatedDelta(packed, ba, row0.dtBias, row0.aExp, state, nil, false)
+
+	ref0Y, ref0End, _ := GatedDelta(row0.packed, row0.ba, row0.dtBias, row0.aExp, row0.state, nil, false)
+	got0Y := SliceStartStop(y, []int32{0, 0, 0, 0}, []int32{1, T, int32(g.Hv), int32(g.Dv)})
+	got0End := SliceStartStop(end, []int32{0, 0, 0, 0}, []int32{1, int32(g.Hv), int32(g.Dv), int32(g.Dk)})
+	if err := requireExact("row0 y", got0Y, ref0Y); err != nil {
+		t.Fatal(err)
+	}
+	if err := requireExact("row0 state", got0End, ref0End); err != nil {
+		t.Fatal(err)
+	}
+
+	prefixPacked := SliceStartStop(row1.packed, []int32{0, 0, 0}, []int32{1, realLen, int32(g.packedDim())})
+	prefixBA := SliceStartStop(row1.ba, []int32{0, 0, 0}, []int32{1, realLen, int32(2 * g.Hv)})
+	ref1Y, ref1End, _ := GatedDelta(prefixPacked, prefixBA, row1.dtBias, row1.aExp, row1.state, nil, false)
+	want1Y := Concatenate([]*Array{ref1Y, Zeros(DTypeBFloat16, 1, T-realLen, g.Hv, g.Dv)}, 1)
+	got1Y := SliceStartStop(y, []int32{1, 0, 0, 0}, []int32{2, T, int32(g.Hv), int32(g.Dv)})
+	got1End := SliceStartStop(end, []int32{1, 0, 0, 0}, []int32{2, int32(g.Hv), int32(g.Dv), int32(g.Dk)})
+	if err := requireExact("row1 y", got1Y, want1Y); err != nil {
+		t.Fatal(err)
+	}
+	if err := requireExact("row1 state", got1End, ref1End); err != nil {
+		t.Fatal(err)
+	}
+}

x/models/nn/recurrent.go

@@ -48,8 +48,7 @@ func WithConvSiLU() RecurrentOption {
 // WithSnapshotSplits requests that the scan run in segments cut at the given
 // offsets within this forward (0 < offset < L), capturing the recurrent state
 // at each boundary. The wrapper returns those per-boundary states to the
-// caller. Offsets must be sorted ascending and strictly interior; out-of-range
-// or duplicate offsets are ignored.
+// caller. Offsets must be sorted ascending, unique, and strictly interior.
 func WithSnapshotSplits(offsets []int) RecurrentOption {
 	return func(c *recurrentConfig) { c.splits = offsets }
 }
@@ -72,7 +71,7 @@ func segmentRanges(splits []int, L int32) []seg {
 
 // sliceSeg slices x to the segment's window [s.start, s.end) along the L axis
 // (axis 1), keeping all other axes whole. Works for any rank — [B, L],
-// [B, L, H], [B, L, H, D] — so the padding mask, gate/beta, and q/k/v all
+// [B, L, H], [B, L, H, D] — so the padding mask and the packed projections
 // slice the same L range and stay aligned. Returns nil when x is nil (the
 // no-padding-mask fast path). Slicing the full-forward mask this way yields
 // the segment's own mask, so masks are built once per forward and reused
@@ -215,57 +214,42 @@ func convStateAt(concat *mlx.Array, queryLens []int32, convTail int, boundary in
 		[]int32{B, boundary + int32(convTail), D})
 }
 
-// GatedDelta wraps mlx.FastGatedDelta with recurrent state management.
-// Reads prior delta state from the supplied option and returns the output
-// and the delta states at each boundary. The boundary states end with the
-// forward-end state. Without WithSnapshotSplits there is one boundary (the
-// end), so states has length 1.
-//
-// Shape conventions:
-//
-//	q:     [B, L, numKeyHeads,   headKDim]
-//	k:     [B, L, numKeyHeads,   headKDim]
-//	v:     [B, L, numValueHeads, headVDim]
-//	state: [B, numValueHeads,    headVDim, headKDim]
-//
-// Prior state comes from exactly one of WithRecurrentHistory (cache
-// path) or WithRecurrentState (no-cache path).
-//
-// When WithSnapshotSplits supplies interior offsets, the scan runs in
-// segments cut at those offsets, threading delta state between them; states
-// holds the delta state at each interior split and at the end. out is
-// identical to the unsegmented scan.
-func GatedDelta(b *batch.Batch, q, k, v, gDecay, beta *mlx.Array, opts ...RecurrentOption) (out *mlx.Array, states []*mlx.Array) {
+// GatedDelta runs the whole gated-delta step over the activated causal-conv
+// output: q/k norms, decay gate, and the scan. convOut rows are packed
+// [q | k | v]; ba is the packed [beta | alpha] projection output. Per-token
+// splits map to the kernels' captureAll shape — one launch emitting every
+// interior state — and any other split pattern composes mlx.GatedDelta per
+// segment, threading the delta state, so each segment runs the fused kernel
+// when it fits. Returns the output and the delta states at each boundary,
+// ending with the forward-end state; without WithSnapshotSplits there is one
+// boundary (the end), so states has length 1.
+func GatedDelta(b *batch.Batch, convOut, ba, dtBias, aExp *mlx.Array, opts ...RecurrentOption) (*mlx.Array, []*mlx.Array) {
 	cfg := resolveRecurrentConfig(opts)
-	var prior *mlx.Array
+	prior := cfg.deltaState
 	if cfg.history != nil {
 		prior = cfg.history.DeltaState()
-	} else {
-		prior = cfg.deltaState
 	}
 
-	L := int32(q.Dim(1))
-	mask := paddingMask(b, L) // built once per forward, sliced per segment
+	L := int32(convOut.Dim(1))
+	mask := paddingMask(b, L)
+
+	// No splits and per-token splits are both a single whole-forward call:
+	// len(splits) == L-1 means the sorted interior offsets are exactly
+	// 1..L-1, the kernels' captureAll shape.
+	if n := len(cfg.splits); n == 0 || n == int(L)-1 {
+		y, end, interior := mlx.GatedDelta(convOut, ba, dtBias, aExp, prior, mask, n > 0)
+		return y, append(interior, end)
+	}
+
 	segs := segmentRanges(cfg.splits, L)
-	if len(segs) <= 1 {
-		out, end := mlx.FastGatedDelta(q, k, v, gDecay, beta, prior, mask)
-		return out, []*mlx.Array{end}
-	}
-
-	// Segmented scan: run each [a,c) piece with the prior segment's state,
-	// recording the delta state at every boundary (interior splits + end).
 	outs := make([]*mlx.Array, 0, len(segs))
-	states = make([]*mlx.Array, 0, len(segs))
+	states := make([]*mlx.Array, 0, len(segs))
 	state := prior
 	for _, seg := range segs {
-		segOut, segState := mlx.FastGatedDelta(
-			sliceSeg(q, seg), sliceSeg(k, seg), sliceSeg(v, seg),
-			sliceSeg(gDecay, seg), sliceSeg(beta, seg),
-			state, sliceSeg(mask, seg),
-		)
-		outs = append(outs, segOut)
-		state = segState
-		states = append(states, segState)
+		var y *mlx.Array
+		y, state, _ = mlx.GatedDelta(sliceSeg(convOut, seg), sliceSeg(ba, seg), dtBias, aExp, state, sliceSeg(mask, seg), false)
+		outs = append(outs, y)
+		states = append(states, state)
 	}
 	return mlx.Concatenate(outs, 1), states
 }

x/models/nn/recurrent_test.go

@@ -8,10 +8,6 @@ import (
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
 
-func ones(dtype mlx.DType, shape ...int) *mlx.Array {
-	return mlx.AddScalar(mlx.Zeros(dtype, shape...), 1)
-}
-
 // lastState returns the forward-end state — the last boundary the recurrent
 // wrappers return.
 func lastState(states []*mlx.Array) *mlx.Array { return states[len(states)-1] }
@@ -142,130 +138,30 @@ func TestCausalConv1DPaddedRowParity(t *testing.T) {
 	}
 }
 
-// TestGatedDeltaDelegatesToKernel checks the wrapper produces the same output
-// and final state as a direct mlx.FastGatedDelta call, for both a zero prior
-// state and a non-zero one (so the wrapper is shown to thread the prior through,
-// not just handle the zero path).
-func TestGatedDeltaDelegatesToKernel(t *testing.T) {
-	skipIfNoMLX(t)
-	B, L, nK, nV, dK, dV := 1, 2, 1, 1, 4, 4
-	q := ones(mlx.DTypeFloat32, B, L, nK, dK)
-	k := ones(mlx.DTypeFloat32, B, L, nK, dK)
-	v := ones(mlx.DTypeFloat32, B, L, nV, dV)
-	gDecay := ones(mlx.DTypeFloat32, B, L, nV)
-	beta := ones(mlx.DTypeFloat32, B, L, nV)
-
-	cases := []struct {
-		name  string
-		prior *mlx.Array
-	}{
-		{"zero", mlx.Zeros(mlx.DTypeFloat32, B, nV, dV, dK)},
-		{"non-zero", mlx.MulScalar(ones(mlx.DTypeFloat32, B, nV, dV, dK), 3)},
-	}
-	for _, tc := range cases {
-		t.Run(tc.name, func(t *testing.T) {
-			outA, statesA := GatedDelta(&batch.Batch{}, q, k, v, gDecay, beta, WithRecurrentState(nil, tc.prior))
-			stateA := lastState(statesA)
-			outB, stateB := mlx.FastGatedDelta(q, k, v, gDecay, beta, tc.prior, nil)
-			mlx.Eval(outA, stateA, outB, stateB)
-
-			gotOut, wantOut := outA.Floats(), outB.Floats()
-			for i := range wantOut {
-				if gotOut[i] != wantOut[i] {
-					t.Fatalf("output[%d]: wrapper=%v direct=%v", i, gotOut[i], wantOut[i])
-				}
-			}
-			gotState, wantState := stateA.Floats(), stateB.Floats()
-			for i := range wantState {
-				if gotState[i] != wantState[i] {
-					t.Fatalf("state[%d]: wrapper=%v direct=%v", i, gotState[i], wantState[i])
-				}
-			}
-		})
-	}
+// gatedDeltaPackedInputs builds deterministic packed conv-output and
+// projection rows plus the per-head parameters for a GatedDelta call.
+func gatedDeltaPackedInputs(B, T, Hk, Dk, Hv, Dv int) (packed, ba, dtBias, aExp *mlx.Array) {
+	packed = fromValues(0.05, B, T, 2*Hk*Dk+Hv*Dv)
+	ba = fromValues(-0.2, B, T, 2*Hv)
+	dtBias = fromValues(0.3, Hv)
+	aExp = fromValues(0.12, Hv)
+	return packed, ba, dtBias, aExp
 }
 
-// TestGatedDeltaPaddedRowParity drives a B=2 batch where row 1 is
-// short (qLen < L). The wrapper must substitute neutral values
-// (q=k=v=beta=0, g=1) at row 1's padded positions so the recurrence
-// is a no-op there — and row 1's final state must equal the state
-// after its last real token. Pinned via parity against a B=1 length-
-// qLen call on the same row.
-func TestGatedDeltaPaddedRowParity(t *testing.T) {
-	skipIfNoMLX(t)
-	L, nK, nV, dK, dV := 4, 1, 1, 4, 4
-	qLenShort := 2
-
-	makeRows := func(seedA, seedB float32, shape ...int) *mlx.Array {
-		// Build a rank-(len(shape)+1) tensor with B=2 rows from two
-		// distinct seeds so the rows are not accidentally identical.
-		n := 1
-		for _, d := range shape {
-			n *= d
-		}
-		vals := make([]float32, 2*n)
-		for i := range n {
-			vals[i] = seedA + 0.1*float32(i)
-		}
-		for i := range n {
-			vals[n+i] = seedB + 0.1*float32(i)
-		}
-		full := append([]int{2}, shape...)
-		return mlx.FromValues(vals, full...)
+// slicePrefix returns rows [lo, hi) of a truncated to the first n positions
+// along axis 1.
+func slicePrefix(a *mlx.Array, lo, hi, n int32) *mlx.Array {
+	dims := a.Dims()
+	start := make([]int32, len(dims))
+	stop := make([]int32, len(dims))
+	start[0], stop[0] = lo, hi
+	for i := 1; i < len(dims); i++ {
+		stop[i] = int32(dims[i])
 	}
-
-	q := makeRows(0.5, -0.5, L, nK, dK)
-	k := makeRows(0.7, -0.7, L, nK, dK)
-	v := makeRows(0.3, -0.3, L, nV, dV)
-	gDecay := makeRows(0.1, -0.1, L, nV)
-	beta := makeRows(0.4, -0.4, L, nV)
-	priorState := makeRows(0.2, -0.2, nV, dV, dK)
-
-	b := &batch.Batch{
-		InputIDs:     mlx.Zeros(mlx.DTypeInt32, 2, L),
-		SeqOffsets:   []int32{0, 0},
-		SeqQueryLens: []int32{int32(L), int32(qLenShort)},
-	}
-	_, states := GatedDelta(b, q, k, v, gDecay, beta, WithRecurrentState(nil, priorState))
-	state := lastState(states)
-	mlx.Eval(state)
-
-	// Reference for row 1: B=1 length-qLenShort call against the
-	// row's real prefix and its prior state slice.
-	row1Slice := func(a *mlx.Array, axisLens ...int32) *mlx.Array {
-		dims := a.Dims()
-		start := make([]int32, len(dims))
-		stop := make([]int32, len(dims))
-		start[0], stop[0] = 1, 2
-		for i := 1; i < len(dims); i++ {
-			stop[i] = int32(dims[i])
-		}
-		// Optionally truncate axis 1 (sequence axis) to qLenShort.
-		if len(axisLens) >= 1 && len(dims) >= 2 {
-			stop[1] = axisLens[0]
-		}
-		return mlx.SliceStartStop(a, start, stop)
-	}
-	q1 := row1Slice(q, int32(qLenShort))
-	k1 := row1Slice(k, int32(qLenShort))
-	v1 := row1Slice(v, int32(qLenShort))
-	gDecay1 := row1Slice(gDecay, int32(qLenShort))
-	beta1 := row1Slice(beta, int32(qLenShort))
-	priorRow1 := row1Slice(priorState)
-
-	_, refState := mlx.FastGatedDelta(q1, k1, v1, gDecay1, beta1, priorRow1, nil)
-	mlx.Eval(refState)
-
-	gotState := state.Floats()
-	wantState := refState.Floats()
-	row1Stride := nV * dV * dK
-	for i := range row1Stride {
-		gotV := gotState[row1Stride+i]
-		wantV := wantState[i]
-		if math.Abs(float64(gotV-wantV)) > 1e-4 {
-			t.Fatalf("row 1 final state[%d]: got %v, want %v", i, gotV, wantV)
-		}
+	if len(dims) >= 2 {
+		stop[1] = n
 	}
+	return mlx.SliceStartStop(a, start, stop)
 }
 
 // floatsClose compares two flat float slices within tolerance.
@@ -281,53 +177,48 @@ func floatsClose(t *testing.T, label string, got, want []float32, tol float64) {
 	}
 }
 
-// TestGatedDeltaSegmentEquivalence checks that running the scan in segments
-// cut at WithSnapshotSplits offsets yields identical output and final state to
-// the single-shot scan, and that each boundary state equals the single-shot
-// state over the corresponding prefix.
+// TestGatedDeltaSegmentEquivalence checks that split forwards match the
+// single-shot call — both the per-token split pattern (the kernels'
+// captureAll shape) and a sparse split (the per-segment composition) — and
+// that each boundary state equals the single-shot state over the
+// corresponding prefix.
 func TestGatedDeltaSegmentEquivalence(t *testing.T) {
 	skipIfNoMLX(t)
-	B, T, Hk, Dk, Hv, Dv := 1, 4, 1, 32, 1, 32
-
-	q := fromValues(0.1, B, T, Hk, Dk)
-	k := fromValues(-0.2, B, T, Hk, Dk)
-	v := fromValues(0.3, B, T, Hv, Dv)
-	gDecay := mlx.MulScalar(ones(mlx.DTypeFloat32, B, T, Hv), 0.9)
-	beta := mlx.MulScalar(ones(mlx.DTypeFloat32, B, T, Hv), 0.5)
+	B, T, Hk, Dk, Hv, Dv := 1, 5, 1, 32, 1, 32
+	packed, ba, dtBias, aExp := gatedDeltaPackedInputs(B, T, Hk, Dk, Hv, Dv)
 	prior := mlx.Zeros(mlx.DTypeFloat32, B, Hv, Dv, Dk)
-
 	full := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{int32(T)}}
 
-	refOut, refStates := GatedDelta(full, q, k, v, gDecay, beta, WithRecurrentState(nil, prior))
+	refOut, refStates := GatedDelta(full, packed, ba, dtBias, aExp, WithRecurrentState(nil, prior))
 	if len(refStates) != 1 {
 		t.Fatalf("unsegmented call returned %d states, want 1", len(refStates))
 	}
 
-	segOut, segStates := GatedDelta(full, q, k, v, gDecay, beta,
-		WithRecurrentState(nil, prior), WithSnapshotSplits([]int{1, 2, 3}))
-	mlx.Eval(refOut, segOut)
-
-	floatsClose(t, "out", segOut.Floats(), refOut.Floats(), 1e-4)
-	// 3 splits + end = 4 boundary states.
-	if len(segStates) != 4 {
-		t.Fatalf("got %d boundary states, want 4", len(segStates))
+	cases := []struct {
+		name   string
+		splits []int
+	}{
+		{"perToken", []int{1, 2, 3, 4}},
+		{"sparse", []int{2}},
 	}
-	mlx.Eval(lastState(segStates), lastState(refStates))
-	floatsClose(t, "final state", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
-
-	// boundary i (offset i+1) must equal a single-shot scan over prefix [0, i+1).
-	for i := range segStates {
-		n := int32(i + 1)
-		pb := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{n}}
-		_, want := GatedDelta(pb,
-			mlx.SliceStartStop(q, []int32{0, 0, 0, 0}, []int32{int32(B), n, int32(Hk), int32(Dk)}),
-			mlx.SliceStartStop(k, []int32{0, 0, 0, 0}, []int32{int32(B), n, int32(Hk), int32(Dk)}),
-			mlx.SliceStartStop(v, []int32{0, 0, 0, 0}, []int32{int32(B), n, int32(Hv), int32(Dv)}),
-			mlx.SliceStartStop(gDecay, []int32{0, 0, 0}, []int32{int32(B), n, int32(Hv)}),
-			mlx.SliceStartStop(beta, []int32{0, 0, 0}, []int32{int32(B), n, int32(Hv)}),
-			WithRecurrentState(nil, prior))
-		mlx.Eval(segStates[i], lastState(want))
-		floatsClose(t, "boundary delta", segStates[i].Floats(), lastState(want).Floats(), 1e-4)
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			segOut, segStates := GatedDelta(full, packed, ba, dtBias, aExp,
+				WithRecurrentState(nil, prior), WithSnapshotSplits(tc.splits))
+			mlx.Eval(refOut, segOut)
+			floatsClose(t, "out", segOut.Floats(), refOut.Floats(), 1e-4)
+			if len(segStates) != len(tc.splits)+1 {
+				t.Fatalf("got %d boundary states, want %d", len(segStates), len(tc.splits)+1)
+			}
+			boundaries := append(append([]int{}, tc.splits...), T)
+			for i, n := range boundaries {
+				_, want, _ := mlx.GatedDelta(
+					slicePrefix(packed, 0, 1, int32(n)), slicePrefix(ba, 0, 1, int32(n)),
+					dtBias, aExp, prior, nil, false)
+				mlx.Eval(segStates[i], want)
+				floatsClose(t, "boundary delta", segStates[i].Floats(), want.Floats(), 1e-4)
+			}
+		})
 	}
 }
 
@@ -372,57 +263,64 @@ func TestCausalConv1DSegmentEquivalence(t *testing.T) {
 	}
 }
 
-// TestGatedDeltaSegmentEquivalenceBatched checks the segmented scan matches the
-// single-shot scan for B>1, including a ragged batch where rows have different
-// real lengths — the per-segment sliced mask must zero each row's padded
-// positions so a short row's boundary state freezes at its real end.
+// TestGatedDeltaSegmentEquivalenceBatched checks split forwards match the
+// single-shot call for B>1 with a ragged batch, for both the per-token and
+// sparse split patterns — the per-segment sliced mask must neutralize each
+// row's padded positions so a short row's boundary state freezes at its
+// real end.
 func TestGatedDeltaSegmentEquivalenceBatched(t *testing.T) {
 	skipIfNoMLX(t)
 	B, T, Hk, Dk, Hv, Dv := 2, 4, 1, 32, 1, 32
-
-	q := fromValues(0.1, B, T, Hk, Dk)
-	k := fromValues(-0.2, B, T, Hk, Dk)
-	v := fromValues(0.3, B, T, Hv, Dv)
-	gDecay := mlx.MulScalar(ones(mlx.DTypeFloat32, B, T, Hv), 0.9)
-	beta := mlx.MulScalar(ones(mlx.DTypeFloat32, B, T, Hv), 0.5)
+	packed, ba, dtBias, aExp := gatedDeltaPackedInputs(B, T, Hk, Dk, Hv, Dv)
 	prior := mlx.Zeros(mlx.DTypeFloat32, B, Hv, Dv, Dk)
 
 	// Row 0 full length T; row 1 ends at 3 (so segment [3,4) is all padding
 	// for row 1).
-	full := &batch.Batch{SeqOffsets: []int32{0, 0}, SeqQueryLens: []int32{int32(T), 3}}
-
-	refOut, refStates := GatedDelta(full, q, k, v, gDecay, beta, WithRecurrentState(nil, prior))
-	segOut, segStates := GatedDelta(full, q, k, v, gDecay, beta,
-		WithRecurrentState(nil, prior), WithSnapshotSplits([]int{1, 2, 3}))
-	mlx.Eval(refOut, segOut, lastState(refStates), lastState(segStates))
-
-	floatsClose(t, "batched out", segOut.Floats(), refOut.Floats(), 1e-4)
-	floatsClose(t, "batched final state", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
-	if len(segStates) != 4 {
-		t.Fatalf("got %d boundary states, want 4", len(segStates))
+	rowReal := []int32{int32(T), 3}
+	full := &batch.Batch{
+		InputIDs:     mlx.Zeros(mlx.DTypeInt32, B, T),
+		SeqOffsets:   []int32{0, 0},
+		SeqQueryLens: rowReal,
 	}
 
-	// Each row's boundary i (offset i+1) must equal a B=1 single-shot scan over
-	// that row's real prefix: row 0 advances the full length, row 1 freezes once
-	// it reaches its real length 3. Per-row B=1 references avoid the ambiguity of
-	// re-declaring a ragged length over a uniform input slice.
-	rowReal := []int32{int32(T), 3}
-	for i := range segStates {
-		for r := range B {
-			n := min(int32(i+1), rowReal[r])
-			lo, hi := int32(r), int32(r)+1
-			rowPrior := mlx.SliceStartStop(prior, []int32{lo, 0, 0, 0}, []int32{hi, int32(Hv), int32(Dv), int32(Dk)})
-			_, want := GatedDelta(&batch.Batch{},
-				mlx.SliceStartStop(q, []int32{lo, 0, 0, 0}, []int32{hi, n, int32(Hk), int32(Dk)}),
-				mlx.SliceStartStop(k, []int32{lo, 0, 0, 0}, []int32{hi, n, int32(Hk), int32(Dk)}),
-				mlx.SliceStartStop(v, []int32{lo, 0, 0, 0}, []int32{hi, n, int32(Hv), int32(Dv)}),
-				mlx.SliceStartStop(gDecay, []int32{lo, 0, 0}, []int32{hi, n, int32(Hv)}),
-				mlx.SliceStartStop(beta, []int32{lo, 0, 0}, []int32{hi, n, int32(Hv)}),
-				WithRecurrentState(nil, rowPrior))
-			gotRow := mlx.SliceStartStop(segStates[i], []int32{lo, 0, 0, 0}, []int32{hi, int32(Hv), int32(Dv), int32(Dk)})
-			mlx.Eval(gotRow, lastState(want))
-			floatsClose(t, "batched boundary delta", gotRow.Floats(), lastState(want).Floats(), 1e-4)
-		}
+	refOut, refStates := GatedDelta(full, packed, ba, dtBias, aExp, WithRecurrentState(nil, prior))
+
+	cases := []struct {
+		name   string
+		splits []int
+	}{
+		{"perToken", []int{1, 2, 3}},
+		{"sparse", []int{2}},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			segOut, segStates := GatedDelta(full, packed, ba, dtBias, aExp,
+				WithRecurrentState(nil, prior), WithSnapshotSplits(tc.splits))
+			mlx.Eval(refOut, segOut, lastState(refStates), lastState(segStates))
+			floatsClose(t, "batched out", segOut.Floats(), refOut.Floats(), 1e-4)
+			floatsClose(t, "batched final state", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
+			if len(segStates) != len(tc.splits)+1 {
+				t.Fatalf("got %d boundary states, want %d", len(segStates), len(tc.splits)+1)
+			}
+
+			// Each row's boundary must equal a B=1 single-shot call over that
+			// row's real prefix: row 0 advances the full length, row 1 freezes
+			// once it reaches its real length.
+			boundaries := append(append([]int{}, tc.splits...), T)
+			for i, bound := range boundaries {
+				for r := range B {
+					n := min(int32(bound), rowReal[r])
+					lo, hi := int32(r), int32(r)+1
+					rowPrior := mlx.SliceStartStop(prior, []int32{lo, 0, 0, 0}, []int32{hi, int32(Hv), int32(Dv), int32(Dk)})
+					_, want, _ := mlx.GatedDelta(
+						slicePrefix(packed, lo, hi, n), slicePrefix(ba, lo, hi, n),
+						dtBias, aExp, rowPrior, nil, false)
+					gotRow := mlx.SliceStartStop(segStates[i], []int32{lo, 0, 0, 0}, []int32{hi, int32(Hv), int32(Dv), int32(Dk)})
+					mlx.Eval(gotRow, want)
+					floatsClose(t, "batched boundary delta", gotRow.Floats(), want.Floats(), 1e-4)
+				}
+			}
+		})
 	}
 }
 

x/models/qwen3_5/qwen3_5.go

@@ -984,10 +984,6 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
 	return nil
 }
 
-func softplus(x *mlx.Array) *mlx.Array {
-	return mlx.Logaddexp(x, mlx.Zeros(x.DType(), x.Dims()...))
-}
-
 func (a *FullAttention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array {
 	qg := a.QProj.Forward(x)
 	qg = mlx.Reshape(qg, B, L, cfg.NumAttentionHeads, cfg.HeadDim*2)
@@ -1035,8 +1031,6 @@ func (g *GatedDeltaNet) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, B,
 	qkv := mlx.SliceStartStop(mixedQKVZ, []int32{0, 0, 0}, []int32{B, L, qkvDim})
 	z := mlx.SliceStartStop(mixedQKVZ, []int32{0, 0, qkvDim}, []int32{B, L, qkvDim + valueDim})
 	z = mlx.Reshape(z, B, L, cfg.LinearNumValueHeads, cfg.LinearValueHeadDim)
-	beta := mlx.SliceStartStop(mixedBA, []int32{0, 0, 0}, []int32{B, L, cfg.LinearNumValueHeads})
-	alpha := mlx.SliceStartStop(mixedBA, []int32{0, 0, cfg.LinearNumValueHeads}, []int32{B, L, 2 * cfg.LinearNumValueHeads})
 	convTail := cfg.LinearConvKernelDim - 1
 	var rc *cache.RecurrentCache
 	opts := make([]nn.RecurrentOption, 0, 3)
@@ -1058,25 +1052,7 @@ func (g *GatedDeltaNet) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, B,
 	}
 
 	convOut, convStates := nn.CausalConv1D(b, qkv, g.Conv1D, int(convTail), opts...)
-
-	q := mlx.SliceStartStop(convOut, []int32{0, 0, 0}, []int32{B, L, keyDim})
-	k := mlx.SliceStartStop(convOut, []int32{0, 0, keyDim}, []int32{B, L, 2 * keyDim})
-	v := mlx.SliceStartStop(convOut, []int32{0, 0, 2 * keyDim}, []int32{B, L, 2*keyDim + valueDim})
-	q = mlx.Reshape(q, B, L, cfg.LinearNumKeyHeads, cfg.LinearKeyHeadDim)
-	k = mlx.Reshape(k, B, L, cfg.LinearNumKeyHeads, cfg.LinearKeyHeadDim)
-	v = mlx.Reshape(v, B, L, cfg.LinearNumValueHeads, cfg.LinearValueHeadDim)
-	invScale := float32(1.0 / math.Sqrt(float64(cfg.LinearKeyHeadDim)))
-	q = mlx.MulScalar(mlx.RMSNormFn(q, nil, 1e-6), invScale*invScale)
-	k = mlx.MulScalar(mlx.RMSNormFn(k, nil, 1e-6), invScale)
-
-	gDecay := softplus(mlx.Add(alpha, g.DtBias))
-	gDecay = mlx.Mul(gDecay, g.AExp)
-	gDecay = mlx.Exp(mlx.MulScalar(gDecay, -1))
-	gDecay = gDecay.AsType(alpha.DType())
-
-	betaGate := mlx.Sigmoid(beta)
-
-	out, deltaStates := nn.GatedDelta(b, q, k, v, gDecay, betaGate, opts...)
+	out, deltaStates := nn.GatedDelta(b, convOut, mixedBA, g.DtBias, g.AExp, opts...)
 	outDType := out.DType()
 	out = mlx.RMSNormFn(out, g.NormWeight, cfg.RMSNormEps)
 	out = mlx.Mul(out.AsType(mlx.DTypeFloat32), mlx.SiLU(z.AsType(mlx.DTypeFloat32))).AsType(outDType)