nn/recurrent: derive conv boundary states from a single conv pass

The MTP validation forward schedules a snapshot at every drafted token, which
made CausalConv1D re-run the depthwise conv once per segment to recover each
boundary's conv tail. A conv boundary state is just the trailing convTail
input positions, so run the conv once over the whole window and slice each
boundary tail from the shared buffer, removing the per-token conv launches.

x/mlxrunner/cache/recurrent_test.go

@@ -152,6 +152,9 @@ func TestRecurrentCachePaddedRoundTrip(t *testing.T) {
 		return mlx.FromValues(vals, convDim, convTail+1)
 	}
 	weight := mkWeight(0.2)
+	// Build the depthwise causal Conv1d as the model does at load time: the
+	// [C, K] kernel becomes [C, K, 1] and the conv is grouped per channel.
+	conv := nn.NewConv1d(mlx.ExpandDims(weight, 2), nil, 1, 0, 1, convDim)
 
 	runForward := func(c *RecurrentCache, b *batch.Batch, T int) (*mlx.Array, *mlx.Array) {
 		var convInput *mlx.Array
@@ -164,7 +167,7 @@ func TestRecurrentCachePaddedRoundTrip(t *testing.T) {
 		}
 
 		history := c.Get(b, mlx.DTypeFloat32)
-		_, convStates := nn.CausalConv1D(b, convInput, nil, weight, convTail,
+		_, convStates := nn.CausalConv1D(b, convInput, conv, convTail,
 			nn.WithRecurrentHistory(history))
 
 		var q, k, v, g, beta *mlx.Array

x/models/nn/recurrent.go

@@ -1,6 +1,8 @@
 package nn
 
 import (
+	"slices"
+
 	"github.com/ollama/ollama/x/mlxrunner/batch"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
 )
@@ -47,8 +49,6 @@ func WithSnapshotSplits(offsets []int) RecurrentOption {
 // seg is a half-open token range [start, end) within a forward.
 type seg struct{ start, end int32 }
 
-func (s seg) len() int32 { return s.end - s.start }
-
 // segmentRanges expands the interior cut offsets into consecutive [a,c) ranges
 // covering [0, L). Cuts are assumed sorted, deduped, and strictly interior; an
 // empty slice yields a single {0, L} segment.
@@ -62,21 +62,9 @@ func segmentRanges(splits []int, L int32) []seg {
 	return append(segs, seg{start, L})
 }
 
-// segmentLens returns each row's real query length within segment
-// [s.start, s.end): the portion of the row's full real length that falls
-// inside the segment, clamped to the segment width. Rows that already ended
-// before the segment contribute 0; rows extending past it are full.
-func segmentLens(b *batch.Batch, s seg) []int32 {
-	lens := make([]int32, len(b.SeqQueryLens))
-	for i := range lens {
-		lens[i] = min(max(b.SeqQueryLens[i]-s.start, 0), s.len())
-	}
-	return lens
-}
-
 // 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, conv/gate/beta, and q/k/v all
+// [B, L, H], [B, L, H, D] — so the padding mask, gate/beta, and q/k/v all
 // 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
@@ -117,12 +105,7 @@ func resolveRecurrentConfig(opts []RecurrentOption) recurrentConfig {
 
 // CausalConv1D runs a depthwise causal 1D convolution with recurrent
 // state management. Prepends the prior conv state along axis 1 and runs
-// the conv.
-//
-// Conv selection: when conv is non-nil (a full nn.Conv1d layer), it
-// runs through conv.Forward. Otherwise weight is treated as the bare
-// depthwise kernel [C, K] and the fallback manual implementation runs.
-// Exactly one of conv or weight should be non-nil.
+// conv.Forward over the combined window.
 //
 // Shapes: input [B, L, D]; prior state [B, convTail, D]; output
 // [B, L, D] (the causal conv strips the prepended state).
@@ -132,10 +115,12 @@ func resolveRecurrentConfig(opts []RecurrentOption) recurrentConfig {
 //
 // Returns the output and the conv states at each boundary, ending with the
 // forward-end conv tail. Without WithSnapshotSplits there is one boundary (the
-// end), so states has length 1. With splits, the conv runs in segments and
-// states holds the conv tail at each interior split and at the end; out is
-// identical to the unsegmented conv.
-func CausalConv1D(b *batch.Batch, input *mlx.Array, conv *Conv1d, weight *mlx.Array, convTail int, opts ...RecurrentOption) (out *mlx.Array, states []*mlx.Array) {
+// end), so states has length 1. With splits, the conv still runs as a single
+// pass over the whole window and each boundary's conv tail is sliced out of the
+// shared input buffer (a boundary state is purely the trailing convTail input
+// positions, so no extra conv launch is needed); out is identical to the
+// unsegmented conv.
+func CausalConv1D(b *batch.Batch, input *mlx.Array, conv *Conv1d, convTail int, opts ...RecurrentOption) (out *mlx.Array, states []*mlx.Array) {
 	cfg := resolveRecurrentConfig(opts)
 	var prior *mlx.Array
 	if cfg.history != nil {
@@ -144,101 +129,57 @@ func CausalConv1D(b *batch.Batch, input *mlx.Array, conv *Conv1d, weight *mlx.Ar
 		prior = cfg.convState
 	}
 
+	// concat is [prior(convTail); input]; each boundary tail is sliced from it
+	// below.
 	L := int32(input.Dim(1))
-	mask := paddingMask(b, L) // built once per forward, sliced per segment
-	segs := segmentRanges(cfg.splits, L)
-	if len(segs) <= 1 {
-		out, end := causalConv1DSegment(input, prior, mask, b.SeqQueryLens, conv, weight, convTail)
-		return out, []*mlx.Array{end}
-	}
-
-	// Segmented conv: each piece prepends the prior piece's conv tail,
-	// recording the conv tail at every boundary (interior splits + end).
-	outs := make([]*mlx.Array, 0, len(segs))
-	states = make([]*mlx.Array, 0, len(segs))
-	state := prior
-	for _, s := range segs {
-		segOut, segNext := causalConv1DSegment(sliceSeg(input, s), state, sliceSeg(mask, s), segmentLens(b, s), conv, weight, convTail)
-		outs = append(outs, segOut)
-		state = segNext
-		states = append(states, segNext)
-	}
-	return mlx.Concatenate(outs, 1), states
-}
-
-// causalConv1DSegment convolves one segment: prepend prior, convolve, return
-// (output, conv tail). mask is the [B, L] padding mask for input (nil if no
-// row is padded); queryLens is each row's real query length (used to gather
-// the per-row conv tail when masking is in play).
-func causalConv1DSegment(input, prior *mlx.Array, mask *mlx.Array, queryLens []int32, conv *Conv1d, weight *mlx.Array, convTail int) (out, nextConv *mlx.Array) {
-	if mask != nil {
+	if mask := paddingMask(b, L); mask != nil {
 		zero := mlx.FromValue(float32(0)).AsType(input.DType())
 		input = mlx.Where(mlx.ExpandDims(mask, 2), input, zero)
 	}
-
 	concat := mlx.Concatenate([]*mlx.Array{prior, input}, 1)
-	if conv != nil {
-		out = conv.Forward(concat)
-	} else {
-		out = depthwiseCausalConv1d(concat, weight, int32(input.Dim(1)))
-	}
+	out = conv.Forward(concat)
 
+	// Snapshot the conv tail at each segment boundary: interior splits in
+	// ascending order, then the forward end at L. Each boundary at offset O
+	// captures input positions [O-convTail, O) — the trailing convTail rows of
+	// the window through token O.
+	segs := segmentRanges(cfg.splits, L)
+	states = make([]*mlx.Array, 0, len(segs))
+	for _, s := range segs {
+		states = append(states, convStateAt(concat, b.SeqQueryLens, convTail, s.end))
+	}
+	return out, states
+}
+
+// convStateAt returns the conv state to cache at boundary: the trailing convTail
+// input positions ending at boundary, clamped per row to the row's real length so
+// a padded row freezes at its real end rather than capturing padding. The prior
+// prefixed in concat shifts those positions to columns [boundary, boundary+convTail).
+func convStateAt(concat *mlx.Array, queryLens []int32, convTail int, boundary int32) *mlx.Array {
 	B := int32(concat.Dim(0))
-	total := int32(concat.Dim(1))
 	D := int32(concat.Dim(2))
 
-	// Gather the tail from each of the non-padded sequence ends
-	if mask != nil && convTail > 0 {
+	// A row shorter than boundary ends its window at its own real length (inputs
+	// are right-padded), so when any row falls short we gather per row instead of
+	// one shared slice. boundary itself is still batch-wide — per-sequence
+	// boundaries (real batching) are a future change to this gather and the callers.
+	clamped := slices.ContainsFunc(queryLens, func(q int32) bool { return boundary > q })
+
+	if clamped && convTail > 0 {
 		offsets := make([]int32, int(B)*convTail)
-
 		for i := range int(B) {
-			end := queryLens[i]
-
+			end := min(boundary, queryLens[i])
 			for k := range convTail {
 				offsets[i*convTail+k] = end + int32(k)
 			}
 		}
-
 		positions := mlx.NewArrayInt32(offsets, []int32{B, int32(convTail), 1})
-		nextConv = mlx.TakeAlongAxis(concat, positions, 1)
-	} else {
-		nextConv = mlx.SliceStartStop(concat,
-			[]int32{0, total - int32(convTail), 0},
-			[]int32{B, total, D})
+		return mlx.TakeAlongAxis(concat, positions, 1)
 	}
 
-	return out, nextConv
-}
-
-// depthwiseCausalConv1d implements a depthwise 1D causal convolution
-// manually as a sum of kernel-offset multiplies. x has shape
-// [B, inLen, C], weight has shape [C, K]; output has shape [B, outLen, C]
-// where outLen = inLen - K + 1 (the caller passes outLen to avoid the
-// subtraction). Used as the fallback path in CausalConv1D when no
-// full Conv1d layer is configured.
-func depthwiseCausalConv1d(x, w *mlx.Array, outLen int32) *mlx.Array {
-	if x == nil || w == nil {
-		return nil
-	}
-	if w.NumDims() != 2 {
-		return nil
-	}
-	B := int32(x.Dim(0))
-	C := int32(w.Dim(0))
-	K := int32(w.Dim(1))
-	var out *mlx.Array
-	for i := range K {
-		seg := mlx.SliceStartStop(x, []int32{0, i, 0}, []int32{B, i + outLen, C})
-		wi := mlx.SliceStartStop(w, []int32{0, i}, []int32{C, i + 1})
-		wi = mlx.Reshape(wi, 1, 1, C)
-		term := mlx.Mul(seg, wi)
-		if out == nil {
-			out = term
-		} else {
-			out = mlx.Add(out, term)
-		}
-	}
-	return out
+	return mlx.SliceStartStop(concat,
+		[]int32{0, boundary, 0},
+		[]int32{B, boundary + int32(convTail), D})
 }
 
 // GatedDelta wraps mlx.FastGatedDelta with recurrent state management.
@@ -337,14 +278,7 @@ type paddingMaskInputs struct {
 func (in paddingMaskInputs) build() *mlx.Array {
 	B := len(in.batch.SeqQueryLens)
 
-	needed := false
-	for i := range B {
-		if in.batch.SeqQueryLens[i] < in.L {
-			needed = true
-			break
-		}
-	}
-	if !needed {
+	if !slices.ContainsFunc(in.batch.SeqQueryLens, func(q int32) bool { return q < in.L }) {
 		return nil
 	}
 

x/models/nn/recurrent_test.go

@@ -30,84 +30,11 @@ func fromValues(seed float32, shape ...int) *mlx.Array {
 	return mlx.FromValues(vals, shape...)
 }
 
-// depthwiseCausalRef is a Go-side reference for the depthwise causal
-// 1D conv fallback. concat is [B, total, C], weight is [C, K], output
-// is [B, total-K+1, C]. Used to anchor the wrapper's parity tests.
-func depthwiseCausalRef(concat, weight *mlx.Array) []float32 {
-	mlx.Eval(concat, weight)
-	cVals := concat.Floats()
-	wVals := weight.Floats()
-	B := concat.Dim(0)
-	total := concat.Dim(1)
-	C := concat.Dim(2)
-	K := weight.Dim(1)
-	outLen := total - K + 1
-	out := make([]float32, B*outLen*C)
-	for bi := range B {
-		for q := range outLen {
-			for c := range C {
-				var sum float32
-				for k := range K {
-					x := cVals[bi*total*C+(q+k)*C+c]
-					w := wVals[c*K+k]
-					sum += x * w
-				}
-				out[bi*outLen*C+q*C+c] = sum
-			}
-		}
-	}
-	return out
-}
-
-// TestCausalConv1DParity drives the wrapper with non-trivial prior,
-// input, and weight values, then compares against a direct depthwise-
-// causal-conv reference.
-func TestCausalConv1DParity(t *testing.T) {
-	skipIfNoMLX(t)
-	B, L, D, convTail := 1, 4, 3, 2
-	K := convTail + 1
-
-	input := fromValues(0.5, B, L, D)
-	prior := fromValues(-0.3, B, convTail, D)
-	weight := fromValues(0.2, D, K)
-
-	out, convStates := CausalConv1D(&batch.Batch{}, input, nil, weight, convTail, WithRecurrentState(prior, nil))
-	nextConv := lastState(convStates)
-	mlx.Eval(out, nextConv)
-
-	concat := mlx.Concatenate([]*mlx.Array{prior, input}, 1)
-	want := depthwiseCausalRef(concat, weight)
-	got := out.Floats()
-	if len(got) != len(want) {
-		t.Fatalf("out len = %d, want %d", len(got), len(want))
-	}
-	for i := range want {
-		if math.Abs(float64(got[i]-want[i])) > 1e-5 {
-			t.Fatalf("out[%d]: got %v, want %v", i, got[i], want[i])
-		}
-	}
-
-	// nextConv (no padding) is the trailing convTail rows of concat.
-	mlx.Eval(concat)
-	cVals := concat.Floats()
-	total := concat.Dim(1)
-	wantTail := make([]float32, B*convTail*D)
-	for bi := range B {
-		for k := range convTail {
-			for d := range D {
-				wantTail[bi*convTail*D+k*D+d] = cVals[bi*total*D+(total-convTail+k)*D+d]
-			}
-		}
-	}
-	tail := nextConv.Floats()
-	if len(tail) != len(wantTail) {
-		t.Fatalf("nextConv len = %d, want %d", len(tail), len(wantTail))
-	}
-	for i := range wantTail {
-		if tail[i] != wantTail[i] {
-			t.Fatalf("nextConv[%d]: got %v, want %v", i, tail[i], wantTail[i])
-		}
-	}
+// convFromKernel builds the depthwise causal Conv1d the model constructs at
+// load time from a bare [C, K] kernel, so the wrapper tests drive the same
+// mlx.Conv1d path production runs.
+func convFromKernel(w *mlx.Array) *Conv1d {
+	return NewConv1d(mlx.ExpandDims(w, 2), nil, 1, 0, 1, int32(w.Dim(0)))
 }
 
 // TestCausalConv1DPaddedRowParity drives a B=2 batch with one short
@@ -122,6 +49,7 @@ func TestCausalConv1DPaddedRowParity(t *testing.T) {
 	K := convTail + 1
 
 	weight := fromValues(0.2, D, K)
+	conv := convFromKernel(weight)
 	priorFull := fromValues(0.5, 2, convTail, D)
 	priorShort := mlx.SliceStartStop(priorFull,
 		[]int32{1, 0, 0},
@@ -148,20 +76,20 @@ func TestCausalConv1DPaddedRowParity(t *testing.T) {
 		SeqQueryLens: []int32{int32(L), int32(qLenShort)},
 	}
 
-	out, convStates := CausalConv1D(b, input, nil, weight, convTail, WithRecurrentState(priorFull, nil))
+	out, convStates := CausalConv1D(b, input, conv, convTail, WithRecurrentState(priorFull, nil))
 	nextConv := lastState(convStates)
 	mlx.Eval(out, nextConv)
 
 	// Reference for row 0: B=1 unpadded length-L call.
 	refOut0, refConvStates0 := CausalConv1D(&batch.Batch{},
-		inputFull, nil, weight, convTail,
+		inputFull, conv, convTail,
 		WithRecurrentState(mlx.SliceStartStop(priorFull,
 			[]int32{0, 0, 0},
 			[]int32{1, int32(convTail), int32(D)}), nil))
 	refNextConv0 := lastState(refConvStates0)
 	// Reference for row 1: B=1 unpadded length-qLenShort call.
 	refOut1, refConvStates1 := CausalConv1D(&batch.Batch{},
-		inputShortReal, nil, weight, convTail,
+		inputShortReal, conv, convTail,
 		WithRecurrentState(priorShort, nil))
 	refNextConv1 := lastState(refConvStates1)
 	mlx.Eval(refOut0, refNextConv0, refOut1, refNextConv1)
@@ -413,15 +341,16 @@ func TestCausalConv1DSegmentEquivalence(t *testing.T) {
 	input := fromValues(0.5, B, L, D)
 	prior := fromValues(-0.3, B, convTail, D)
 	weight := fromValues(0.2, D, K)
+	conv := convFromKernel(weight)
 
 	full := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{int32(L)}}
 
-	refOut, refStates := CausalConv1D(full, input, nil, weight, convTail, WithRecurrentState(prior, nil))
+	refOut, refStates := CausalConv1D(full, input, conv, convTail, WithRecurrentState(prior, nil))
 	if len(refStates) != 1 {
 		t.Fatalf("unsegmented call returned %d states, want 1", len(refStates))
 	}
 
-	segOut, segStates := CausalConv1D(full, input, nil, weight, convTail,
+	segOut, segStates := CausalConv1D(full, input, conv, convTail,
 		WithRecurrentState(prior, nil), WithSnapshotSplits([]int{1, 2, 3}))
 	mlx.Eval(refOut, segOut)
 
@@ -437,7 +366,7 @@ func TestCausalConv1DSegmentEquivalence(t *testing.T) {
 		pb := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{n}}
 		_, want := CausalConv1D(pb,
 			mlx.SliceStartStop(input, []int32{0, 0, 0}, []int32{int32(B), n, int32(D)}),
-			nil, weight, convTail, WithRecurrentState(prior, nil))
+			conv, convTail, WithRecurrentState(prior, nil))
 		mlx.Eval(segStates[i], lastState(want))
 		floatsClose(t, "boundary conv", segStates[i].Floats(), lastState(want).Floats(), 1e-4)
 	}
@@ -445,7 +374,8 @@ 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 — segmentLens must clamp each row's per-segment query length.
+// 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.
 func TestGatedDeltaSegmentEquivalenceBatched(t *testing.T) {
 	skipIfNoMLX(t)
 	B, T, Hk, Dk, Hv, Dv := 2, 4, 1, 32, 1, 32
@@ -496,8 +426,10 @@ func TestGatedDeltaSegmentEquivalenceBatched(t *testing.T) {
 	}
 }
 
-// TestCausalConv1DSegmentEquivalenceBatched is the conv analog: segmented vs
-// single-shot for a ragged B>1 batch.
+// TestCausalConv1DSegmentEquivalenceBatched is the conv analog of the gated-delta
+// batched test: boundary tails from the single conv pass vs per-row single-shot
+// references for a ragged B>1 batch, where a short row must freeze its tail at
+// its real end rather than reach into padding.
 func TestCausalConv1DSegmentEquivalenceBatched(t *testing.T) {
 	skipIfNoMLX(t)
 	B, L, D, convTail := 2, 4, 3, 2
@@ -506,11 +438,12 @@ func TestCausalConv1DSegmentEquivalenceBatched(t *testing.T) {
 	input := fromValues(0.5, B, L, D)
 	prior := fromValues(-0.3, B, convTail, D)
 	weight := fromValues(0.2, D, K)
+	conv := convFromKernel(weight)
 
 	full := &batch.Batch{SeqOffsets: []int32{0, 0}, SeqQueryLens: []int32{int32(L), 3}}
 
-	refOut, refStates := CausalConv1D(full, input, nil, weight, convTail, WithRecurrentState(prior, nil))
-	segOut, segStates := CausalConv1D(full, input, nil, weight, convTail,
+	refOut, refStates := CausalConv1D(full, input, conv, convTail, WithRecurrentState(prior, nil))
+	segOut, segStates := CausalConv1D(full, input, conv, convTail,
 		WithRecurrentState(prior, nil), WithSnapshotSplits([]int{1, 2, 3}))
 	mlx.Eval(refOut, segOut, lastState(refStates), lastState(segStates))
 
@@ -532,7 +465,7 @@ func TestCausalConv1DSegmentEquivalenceBatched(t *testing.T) {
 				[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, int32(convTail), int32(D)})
 			rowInput := mlx.SliceStartStop(input,
 				[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, n, int32(D)})
-			_, want := CausalConv1D(&batch.Batch{}, rowInput, nil, weight, convTail,
+			_, want := CausalConv1D(&batch.Batch{}, rowInput, conv, convTail,
 				WithRecurrentState(rowPrior, nil))
 			gotRow := mlx.SliceStartStop(segStates[i],
 				[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, int32(convTail), int32(D)})

x/models/qwen3_5/qwen3_5.go

@@ -127,7 +127,6 @@ type GatedDeltaNet struct {
 	OutProj    nn.LinearLayer
 
 	Conv1D     *nn.Conv1d
-	ConvWeight *mlx.Array
 	NormWeight *mlx.Array
 	DtBias     *mlx.Array
 	ALog       *mlx.Array
@@ -762,31 +761,6 @@ func sanitizeConvWeight(w *mlx.Array) *mlx.Array {
 	return w
 }
 
-func depthwiseConv1dKernelWeight(w *mlx.Array) *mlx.Array {
-	if w == nil {
-		return nil
-	}
-	switch w.NumDims() {
-	case 2:
-		// qwen3.5 manual path stores [C, K]; MLX grouped conv expects [Cout, K, Cin/groups].
-		// For depthwise conv (groups=C), that is [C, K, 1].
-		return mlx.ExpandDims(w, 2)
-	case 3:
-		switch {
-		case w.Dim(2) == 1:
-			// [C, K, 1]
-			return w
-		case w.Dim(1) == 1:
-			// [C, 1, K] -> [C, K, 1]
-			return mlx.Transpose(w, 0, 2, 1)
-		case w.Dim(0) == 1:
-			// [1, K, C] -> [C, K, 1]
-			return mlx.Transpose(w, 2, 1, 0)
-		}
-	}
-	return nil
-}
-
 func shouldShiftNormKey(key string) bool {
 	for _, suffix := range []string{
 		".input_layernorm.weight",
@@ -897,9 +871,9 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
 			lin.InProjBA = linears.Make(layerPrefix + ".linear_attn.in_proj_ba")
 			lin.OutProj = linears.Make(layerPrefix + ".linear_attn.out_proj")
 
-			lin.ConvWeight = sanitizeConvWeight(tensors[layerPrefix+".linear_attn.conv1d.weight"])
-			if lin.ConvWeight == nil {
-				lin.ConvWeight = sanitizeConvWeight(tensors[layerPrefix+".linear_attn.conv1d"])
+			convWeight := sanitizeConvWeight(tensors[layerPrefix+".linear_attn.conv1d.weight"])
+			if convWeight == nil {
+				convWeight = sanitizeConvWeight(tensors[layerPrefix+".linear_attn.conv1d"])
 			}
 			lin.NormWeight, _ = tensorAny(tensors,
 				layerPrefix+".linear_attn.norm.weight",
@@ -922,15 +896,15 @@ func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error {
 			if (!hasSplit && !hasCombined) || lin.OutProj == nil {
 				return fmt.Errorf("layer %d: missing linear attention projections", i)
 			}
-			if lin.ConvWeight == nil || lin.NormWeight == nil || lin.DtBias == nil || lin.ALog == nil || lin.AExp == nil {
+			if convWeight == nil || lin.NormWeight == nil || lin.DtBias == nil || lin.ALog == nil || lin.AExp == nil {
 				return fmt.Errorf("layer %d: missing linear attention state tensors", i)
 			}
-			if lin.ConvWeight.NumDims() != 2 {
-				return fmt.Errorf("layer %d: conv1d weight must be 2D after sanitization, got %dD", i, lin.ConvWeight.NumDims())
-			}
-			if convKernel := depthwiseConv1dKernelWeight(lin.ConvWeight); convKernel != nil {
-				lin.Conv1D = nn.NewConv1d(convKernel, nil, 1, 0, 1, int32(lin.ConvWeight.Dim(0)))
+			if convWeight.NumDims() != 2 {
+				return fmt.Errorf("layer %d: conv1d weight must be 2D after sanitization, got %dD", i, convWeight.NumDims())
 			}
+			// MLX grouped conv expects [Cout, K, Cin/groups]; for depthwise
+			// conv (groups=C) the [C, K] kernel becomes [C, K, 1].
+			lin.Conv1D = nn.NewConv1d(mlx.ExpandDims(convWeight, 2), nil, 1, 0, 1, int32(convWeight.Dim(0)))
 
 			layer.Linear = lin
 		} else {
@@ -1182,7 +1156,7 @@ func (g *GatedDeltaNet) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, B,
 		))
 	}
 
-	convOut, convStates := nn.CausalConv1D(b, qkv, g.Conv1D, g.ConvWeight, int(convTail), opts...)
+	convOut, convStates := nn.CausalConv1D(b, qkv, g.Conv1D, int(convTail), opts...)
 	convOut = mlx.SiLU(convOut)
 
 	keyDim := cfg.LinearNumKeyHeads * cfg.LinearKeyHeadDim