mirror of
https://github.com/ollama/ollama.git
synced 2026-08-04 14:56:15 +00:00
d573a2367b
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.
476 lines
18 KiB
Go
476 lines
18 KiB
Go
package nn
|
|
|
|
import (
|
|
"math"
|
|
"testing"
|
|
|
|
"github.com/ollama/ollama/x/mlxrunner/batch"
|
|
"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] }
|
|
|
|
// fromValues builds a tensor with sequentially-numbered float32
|
|
// values so element-by-element parity actually exercises the kernel.
|
|
func fromValues(seed float32, shape ...int) *mlx.Array {
|
|
n := 1
|
|
for _, d := range shape {
|
|
n *= d
|
|
}
|
|
vals := make([]float32, n)
|
|
for i := range vals {
|
|
vals[i] = seed + 0.1*float32(i)
|
|
}
|
|
return mlx.FromValues(vals, shape...)
|
|
}
|
|
|
|
// 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
|
|
// row (qLen<L). For the short row, (a) `out` positions [0..qLen)
|
|
// must equal a B=1 reference at length qLen, (b) `nextConv` for the
|
|
// short row must be the row's last convTail real positions (not the
|
|
// padded tail), (c) the full row must be unaffected.
|
|
func TestCausalConv1DPaddedRowParity(t *testing.T) {
|
|
skipIfNoMLX(t)
|
|
L, D, convTail := 4, 3, 2
|
|
qLenShort := 2
|
|
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},
|
|
[]int32{2, int32(convTail), int32(D)})
|
|
|
|
// Pad row 1 with arbitrary values past qLenShort — the wrapper
|
|
// must zero them before convolving. Distinct values let us catch
|
|
// any leak.
|
|
inputFull := fromValues(1.0, 1, L, D)
|
|
inputShortReal := mlx.FromValues([]float32{
|
|
2.0, 2.1, 2.2,
|
|
2.3, 2.4, 2.5,
|
|
}, 1, qLenShort, D)
|
|
inputShortPad := mlx.FromValues([]float32{
|
|
99, 99, 99,
|
|
99, 99, 99,
|
|
}, 1, L-qLenShort, D)
|
|
inputShortFull := mlx.Concatenate([]*mlx.Array{inputShortReal, inputShortPad}, 1)
|
|
input := mlx.Concatenate([]*mlx.Array{inputFull, inputShortFull}, 0)
|
|
|
|
b := &batch.Batch{
|
|
InputIDs: mlx.Zeros(mlx.DTypeInt32, 2, L),
|
|
SeqOffsets: []int32{0, 0},
|
|
SeqQueryLens: []int32{int32(L), int32(qLenShort)},
|
|
}
|
|
|
|
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, 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, conv, convTail,
|
|
WithRecurrentState(priorShort, nil))
|
|
refNextConv1 := lastState(refConvStates1)
|
|
mlx.Eval(refOut0, refNextConv0, refOut1, refNextConv1)
|
|
|
|
gotOut := out.Floats()
|
|
wantOut0 := refOut0.Floats()
|
|
wantOut1 := refOut1.Floats()
|
|
|
|
for q := range L {
|
|
for d := range D {
|
|
i := q*D + d
|
|
if gotOut[i] != wantOut0[i] {
|
|
t.Fatalf("row 0 out[q=%d,d=%d]: got %v, want %v", q, d, gotOut[i], wantOut0[i])
|
|
}
|
|
}
|
|
}
|
|
for q := range qLenShort {
|
|
for d := range D {
|
|
gotI := L*D + q*D + d
|
|
refI := q*D + d
|
|
if math.Abs(float64(gotOut[gotI]-wantOut1[refI])) > 1e-5 {
|
|
t.Fatalf("row 1 real out[q=%d,d=%d]: got %v, want %v", q, d, gotOut[gotI], wantOut1[refI])
|
|
}
|
|
}
|
|
}
|
|
|
|
// nextConv: row 0 unaffected, row 1 must be the row's real tail
|
|
// (positions [qLenShort - convTail, qLenShort) of the per-row
|
|
// concat, i.e. the last two real input rows in this setup).
|
|
gotTail := nextConv.Floats()
|
|
wantTail0 := refNextConv0.Floats()
|
|
wantTail1 := refNextConv1.Floats()
|
|
for k := range convTail {
|
|
for d := range D {
|
|
i := k*D + d
|
|
if gotTail[i] != wantTail0[i] {
|
|
t.Fatalf("row 0 nextConv[k=%d,d=%d]: got %v, want %v", k, d, gotTail[i], wantTail0[i])
|
|
}
|
|
}
|
|
}
|
|
for k := range convTail {
|
|
for d := range D {
|
|
gotI := convTail*D + k*D + d
|
|
refI := k*D + d
|
|
if gotTail[gotI] != wantTail1[refI] {
|
|
t.Fatalf("row 1 nextConv[k=%d,d=%d]: got %v, want %v (must come from real positions, not the padded tail)",
|
|
k, d, gotTail[gotI], wantTail1[refI])
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// 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])
|
|
}
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
// 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...)
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|
|
}
|
|
|
|
// floatsClose compares two flat float slices within tolerance.
|
|
func floatsClose(t *testing.T, label string, got, want []float32, tol float64) {
|
|
t.Helper()
|
|
if len(got) != len(want) {
|
|
t.Fatalf("%s: len %d, want %d", label, len(got), len(want))
|
|
}
|
|
for i := range want {
|
|
if math.Abs(float64(got[i]-want[i])) > tol {
|
|
t.Fatalf("%s[%d]: got %v, want %v", label, i, got[i], want[i])
|
|
}
|
|
}
|
|
}
|
|
|
|
// 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.
|
|
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)
|
|
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))
|
|
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))
|
|
}
|
|
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)
|
|
}
|
|
}
|
|
|
|
// TestCausalConv1DSegmentEquivalence checks the conv segmented path matches the
|
|
// single-shot conv for output, final conv tail, and each boundary conv state.
|
|
func TestCausalConv1DSegmentEquivalence(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)
|
|
conv := convFromKernel(weight)
|
|
|
|
full := &batch.Batch{SeqOffsets: []int32{0}, SeqQueryLens: []int32{int32(L)}}
|
|
|
|
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, conv, convTail,
|
|
WithRecurrentState(prior, nil), WithSnapshotSplits([]int{1, 2, 3}))
|
|
mlx.Eval(refOut, segOut)
|
|
|
|
floatsClose(t, "conv out", segOut.Floats(), refOut.Floats(), 1e-4)
|
|
if len(segStates) != 4 {
|
|
t.Fatalf("got %d boundary conv states, want 4", len(segStates))
|
|
}
|
|
mlx.Eval(lastState(segStates), lastState(refStates))
|
|
floatsClose(t, "conv final", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
|
|
|
|
for i := range segStates {
|
|
n := int32(i + 1)
|
|
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)}),
|
|
conv, convTail, WithRecurrentState(prior, nil))
|
|
mlx.Eval(segStates[i], lastState(want))
|
|
floatsClose(t, "boundary conv", segStates[i].Floats(), lastState(want).Floats(), 1e-4)
|
|
}
|
|
}
|
|
|
|
// 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.
|
|
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)
|
|
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))
|
|
}
|
|
|
|
// 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)
|
|
}
|
|
}
|
|
}
|
|
|
|
// 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
|
|
K := convTail + 1
|
|
|
|
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, 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))
|
|
|
|
floatsClose(t, "batched conv out", segOut.Floats(), refOut.Floats(), 1e-4)
|
|
floatsClose(t, "batched conv final", lastState(segStates).Floats(), lastState(refStates).Floats(), 1e-4)
|
|
if len(segStates) != 4 {
|
|
t.Fatalf("got %d boundary conv states, want 4", len(segStates))
|
|
}
|
|
|
|
// Each row's boundary i (offset i+1) must equal a B=1 single-shot conv 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(L), 3}
|
|
for i := range segStates {
|
|
for r := range B {
|
|
n := min(int32(i+1), rowReal[r])
|
|
rowPrior := mlx.SliceStartStop(prior,
|
|
[]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, conv, convTail,
|
|
WithRecurrentState(rowPrior, nil))
|
|
gotRow := mlx.SliceStartStop(segStates[i],
|
|
[]int32{int32(r), 0, 0}, []int32{int32(r) + 1, int32(convTail), int32(D)})
|
|
mlx.Eval(gotRow, lastState(want))
|
|
floatsClose(t, "batched boundary conv", gotRow.Floats(), lastState(want).Floats(), 1e-4)
|
|
}
|
|
}
|
|
}
|