diff --git a/x/mlxrunner/cache/recurrent_test.go b/x/mlxrunner/cache/recurrent_test.go index 035bde6e0..a0d7fa0bc 100644 --- a/x/mlxrunner/cache/recurrent_test.go +++ b/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 diff --git a/x/models/nn/recurrent.go b/x/models/nn/recurrent.go index 255d49f54..24fb32c76 100644 --- a/x/models/nn/recurrent.go +++ b/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 } diff --git a/x/models/nn/recurrent_test.go b/x/models/nn/recurrent_test.go index 7e8eb6f86..8d19b9e5c 100644 --- a/x/models/nn/recurrent_test.go +++ b/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)}) diff --git a/x/models/qwen3_5/qwen3_5.go b/x/models/qwen3_5/qwen3_5.go index dfca63e24..1dcdfe819 100644 --- a/x/models/qwen3_5/qwen3_5.go +++ b/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