mlxrunner: choose the speculative draft length to maximize throughput

The heuristic schedule grew the draft toward a fixed cap on acceptance alone,
maximizing accepted-tokens-per-step rather than throughput, and on a
steep-forward target it regressed below no speculation. Replace it with an
engine-level controller that drafts the depth maximizing
committed-tokens-per-wallclock from live per-position acceptance and persisted
per-width forward cost, with no draft-length cap; the heuristic schedule and
the OLLAMA_MLX_MTP_* env vars go with it.

x/mlxrunner/model/base/base.go

@@ -45,20 +45,6 @@ type DraftModel interface {
 	LoadWeights(tensors map[string]*mlx.Array) error
 }
 
-// MTPDefaults holds model-provided draft-token defaults for speculative
-// decoding. Environment settings in the runner may override these values.
-type MTPDefaults struct {
-	InitialDraftTokens int
-	MaxDraftTokens     int
-	Enabled            bool
-}
-
-// MTPDefaultsProvider lets a model provide MTP policy defaults from its own
-// config without teaching the runner model-specific shape heuristics.
-type MTPDefaultsProvider interface {
-	MTPDraftDefaults(sample bool) MTPDefaults
-}
-
 // SelfDraft is implemented by models whose draft head ships inline with the
 // target weights; it returns the head, or nil when the checkpoint shipped none.
 type SelfDraft interface {

x/mlxrunner/mtp.go

@@ -5,38 +5,14 @@ import (
 
 	"github.com/ollama/ollama/x/mlxrunner/batch"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
-	"github.com/ollama/ollama/x/mlxrunner/model/base"
 	sampler "github.com/ollama/ollama/x/mlxrunner/sample"
 )
 
-const (
-	mtpDefaultInitialDraftTokens = 4
-	mtpDefaultMaxDraftTokens     = 16
-)
-
 // mtpPendingFlushTokens caps how many committed look-ahead tokens wait in the
 // pending buffer before a batched flush, bounding the pinned hidden states
 // regardless of what else triggers a flush.
 const mtpPendingFlushTokens = 32
 
-func (r *Runner) mtpDefaults(sample bool) base.MTPDefaults {
-	defaults := base.MTPDefaults{
-		InitialDraftTokens: mtpDefaultInitialDraftTokens,
-		MaxDraftTokens:     mtpDefaultMaxDraftTokens,
-		Enabled:            true,
-	}
-	if p, ok := r.Model.(base.MTPDefaultsProvider); ok {
-		defaults = p.MTPDraftDefaults(sample)
-	}
-	if defaults.InitialDraftTokens <= 0 {
-		defaults.InitialDraftTokens = mtpDefaultInitialDraftTokens
-	}
-	if defaults.MaxDraftTokens <= 0 {
-		defaults.MaxDraftTokens = mtpDefaultMaxDraftTokens
-	}
-	return defaults
-}
-
 // mtpDrafter drafts with a model's multi-token-prediction head, fed through
 // the committed-stream reports. The draft KV pairs each slot S with the
 // look-ahead token at S+1 fused with the target hidden at S, so a pair

x/mlxrunner/mtp_test.go

@@ -268,12 +268,14 @@ func TestAcceptMTPDraftsGreedyAcceptAll(t *testing.T) {
 
 	spec := testSpeculationSession(r, caches)
 	current := sampler.Result{Token: mlx.FromValues([]int32{1}, 1)}
-	results, accepted, err := spec.accept(&position, current, candidates)
+	unpin := pinAcceptInputs(current, candidates)
+	defer unpin()
+	results, accepted, observed, err := spec.accept(&position, current, candidates)
 	if err != nil {
 		t.Fatalf("accept: %v", err)
 	}
-	if accepted != 3 {
-		t.Fatalf("accepted = %d, want 3", accepted)
+	if accepted != 3 || observed != 3 {
+		t.Fatalf("accepted = %d, observed = %d, want 3 and 3", accepted, observed)
 	}
 	// The run is the accepted drafts followed by the bonus token (5).
 	if got := resultIDs(results); !slices.Equal(got, []int{2, 3, 4, 5}) {
@@ -302,12 +304,14 @@ func TestAcceptMTPDraftsGreedyMismatch(t *testing.T) {
 
 	spec := testSpeculationSession(r, caches)
 	current := sampler.Result{Token: mlx.FromValues([]int32{1}, 1)}
-	results, accepted, err := spec.accept(&position, current, candidates)
+	unpin := pinAcceptInputs(current, candidates)
+	defer unpin()
+	results, accepted, observed, err := spec.accept(&position, current, candidates)
 	if err != nil {
 		t.Fatalf("accept: %v", err)
 	}
-	if accepted != 1 {
-		t.Fatalf("accepted = %d, want 1", accepted)
+	if accepted != 1 || observed != 2 {
+		t.Fatalf("accepted = %d, observed = %d, want 1 and 2 (a rejection judges the full round)", accepted, observed)
 	}
 	// The run is the one accepted draft (2) followed by the target's own
 	// prediction at the rejection point (3).
@@ -338,12 +342,14 @@ func TestAcceptMTPDraftsGreedyEOS(t *testing.T) {
 
 	spec := testSpeculationSession(r, caches)
 	current := sampler.Result{Token: mlx.FromValues([]int32{1}, 1)}
-	results, accepted, err := spec.accept(&position, current, candidates)
+	unpin := pinAcceptInputs(current, candidates)
+	defer unpin()
+	results, accepted, observed, err := spec.accept(&position, current, candidates)
 	if err != nil {
 		t.Fatalf("accept: %v", err)
 	}
-	if accepted != 2 {
-		t.Fatalf("accepted = %d, want 2 (token + EOS)", accepted)
+	if accepted != 2 || observed != 2 {
+		t.Fatalf("accepted = %d, observed = %d, want 2 and 2 (token + EOS)", accepted, observed)
 	}
 	// The EOS ends generation, so the run is exactly the accepted tokens
 	// with no bonus appended.
@@ -403,20 +409,21 @@ func TestRunMTPDecodeGreedy(t *testing.T) {
 		t.Fatalf("session outputs = %v, want %v", session.outputs, got)
 	}
 
-	// The target writes the caches contiguously: the seed token at offset 1,
-	// then one fused forward validating the current token and the 4 drafts
-	// together at offset 2.
-	wantForwards := []forwardCall{{offset: 1, n: 1}, {offset: 2, n: 5}}
+	// The unprimed drafter parks the first call, so the seed token and the
+	// stretch's in-flight successor decode as pipelined plain forwards; the
+	// resumed round then validates the current token and the 4 drafts in one
+	// fused forward at offset 3.
+	wantForwards := []forwardCall{{offset: 1, n: 1}, {offset: 2, n: 1}, {offset: 3, n: 5}}
 	model := r.Model.(*fakeMTPModel)
 	if !slices.Equal(model.forwards, wantForwards) {
 		t.Fatalf("target forwards = %v, want %v", model.forwards, wantForwards)
 	}
 
 	// Single-position drafting anchors every Draft call in a round at the
-	// last committed position (offset 1 — the current token at offset 2 is
-	// not yet validated), while the extended token advances along the
-	// proposed chain.
-	wantDraft := []draftCall{{1, 2}, {1, 3}, {1, 4}, {1, eos}}
+	// last committed position (offset 2 — the current token at offset 3 is
+	// not yet validated), while the input token advances along the proposed
+	// chain.
+	wantDraft := []draftCall{{2, 3}, {2, 4}, {2, eos}, {2, 0}}
 	if !slices.Equal(draft.calls, wantDraft) {
 		t.Fatalf("draft calls = %v, want %v", draft.calls, wantDraft)
 	}
@@ -445,8 +452,9 @@ func TestRunMTPDecodeSampled(t *testing.T) {
 	}
 	spec := r.spec.open(req, caches)
 	if spec == nil || !spec.enabled {
-		t.Fatalf("newSpeculationSession rejected a sampled request")
+		t.Fatalf("open rejected a sampled request")
 	}
+	pinDraftLimit(spec, 4)
 	d := spec.decoder([]int32{1}, position)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
 		t.Fatalf("decode: %v", err)
@@ -483,10 +491,11 @@ func TestRunMTPDecodeWarmDrafter(t *testing.T) {
 	req := Request{
 		Responses:         ch,
 		Tokens:            []int32{0},
-		CompletionRequest: CompletionRequest{Options: api.Options{Runner: api.Runner{DraftNumPredict: 4}, NumPredict: 20}},
+		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
 	spec := r.spec.open(req, caches)
+	pinDraftLimit(spec, 4)
 	// The prefill chunk's committed report: token 0 at slot 0 with its
 	// hidden row, leaving the drafter ready to propose from slot 1.
 	spec.committed(mlx.FromValues([]int32{0}, 1, 1), oneHotLogits([]int32{1}), 0)
@@ -524,6 +533,65 @@ func TestRunMTPDecodeWarmDrafter(t *testing.T) {
 	}
 }
 
+func TestRunMTPDecodeEOSCutLeavesPositionsUnjudged(t *testing.T) {
+	skipIfNoMLX(t)
+	// An accepted EOS inside the draft ends the round at a terminator, not a
+	// target rejection. The persisted acceptance model records outcomes only
+	// up to the EOS: the positions past it lie beyond where the round stopped,
+	// and folding them in as rejections would bias the depth controller
+	// shallow.
+	const eos int32 = 7
+	predict := map[int32]int32{1: 2, 2: 3, 3: eos, eos: 0}
+	r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
+	r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict})
+
+	caches, _ := newMTPTestCaches(1)
+	session, ch := newMTPTestSession(caches)
+	position := 1
+
+	req := Request{
+		Responses:         ch,
+		Tokens:            []int32{0},
+		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
+		SamplerOpts:       sampler.Options{},
+	}
+	spec := r.spec.open(req, caches)
+	if spec == nil || !spec.enabled {
+		t.Fatalf("want an enabled speculationSession with a depth controller, got %+v", spec)
+	}
+	// Force a four-token first round so the EOS (third draft) is interior;
+	// the controller still records the round's outcomes.
+	spec.limit = 4
+	spec.committed(mlx.FromValues([]int32{0}, 1, 1), oneHotLogits([]int32{1}), 0)
+
+	d := spec.decoder([]int32{1}, position)
+	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
+		t.Fatalf("decode: %v", err)
+	}
+	d.close()
+
+	content, _ := collectResponses(ch)
+	if content != "23" {
+		t.Fatalf("content = %q, want %q", content, "23")
+	}
+	if got := []int32{2, 3, eos}; !slices.Equal(session.outputs, got) {
+		t.Fatalf("session outputs = %v, want %v", session.outputs, got)
+	}
+
+	// The round drafted {2, 3, eos, 0} and the target accepted through the
+	// EOS at position 3; the fourth draft was cut, not judged. Positions
+	// 1..3 record accepts and position 4 is never grown.
+	acc := r.spec.depth.acc
+	if len(acc.seen) != 4 {
+		t.Fatalf("acceptance model grew to %d positions, want 4 (the cut position stays unjudged)", len(acc.seen))
+	}
+	for i := 1; i <= 3; i++ {
+		if acc.seen[i] != 1 || acc.rate[i] != 1 {
+			t.Fatalf("position %d: seen = %d, rate = %v, want 1 and 1", i, acc.seen[i], acc.rate[i])
+		}
+	}
+}
+
 func TestDecodePlain(t *testing.T) {
 	skipIfNoMLX(t)
 	// The same chain with no speculationSession: decode's pipelined loop runs,
@@ -625,24 +693,36 @@ func TestDecodeCancelledMidStream(t *testing.T) {
 	}
 }
 
+// pinDraftLimit fixes an engine's draft length for the whole run: decode
+// tests assert engine mechanics at known widths, so they disable adaptive
+// depth rather than steer what it learns.
+func pinDraftLimit(spec *speculationSession, limit int) {
+	spec.enabled = false
+	spec.limit = limit
+}
+
 // testDecoder builds the decoder TextGenerationPipeline would construct for
-// this request; tests close it explicitly so close-time effects are visible
-// to assertions.
+// this request, with the draft length pinned to a fixed width; tests close it
+// explicitly so close-time effects are visible to assertions.
 func testDecoder(r *Runner, req Request, caches []cache.Cache, seed []int32, position int) decoder {
 	if spec := r.spec.open(req, caches); spec != nil {
+		if spec.enabled {
+			pinDraftLimit(spec, 4)
+		}
 		return spec.decoder(seed, position)
 	}
-	return r.pipelinedDecoder(nil, caches, seed, position)
+	return r.pipelinedDecoder(nil, caches, mlx.FromValues(seed, 1, len(seed)), position)
 }
 
 func TestDecodeKVDraft(t *testing.T) {
 	skipIfNoMLX(t)
 	// A draft with its own KV cache mirroring the target chain
 	// 1->2->3->4->5->6->EOS.
-	// The decode seed catch-up writes the draft pair for the prompt token,
-	// the first proposal comes from the catch-up's held logits (no draft
-	// call), speculative entries are written at advancing slots, and the
-	// post-accept rebuild rewrites the committed range from target hiddens.
+	// The unprimed drafter parks the first call, whose pipelined tokens pair
+	// the prompt token and the first generated token; the first proposal
+	// comes from the catch-up's held logits (no draft call), speculative
+	// entries are written at advancing slots, and the post-accept rebuild
+	// rewrites the committed range from target hiddens.
 	const eos int32 = 7
 	predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: eos, eos: 0}
 	r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
@@ -663,6 +743,7 @@ func TestDecodeKVDraft(t *testing.T) {
 	if spec == nil || !spec.enabled || len(spec.spec.targets) != 1 {
 		t.Fatalf("speculation engine not built around the draft caches")
 	}
+	pinDraftLimit(spec, 4)
 	defer spec.close()
 	d := spec.decoder([]int32{1}, position)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
@@ -693,15 +774,17 @@ func TestDecodeKVDraft(t *testing.T) {
 	// Every committed draft slot S fuses look-ahead token x_{S+1} with the
 	// target hidden at S (whose hot index equals the look-ahead id on this
 	// mirrored chain); speculative steps fuse the head's own projections
-	// (-1), seeded by the catch-up flush's held projection. The first
-	// proposal of the round consumes the catch-up's held logits, so the
-	// four-token draft makes only three head calls before the rebuild.
+	// (-1), seeded by the catch-up flush's held projection. The unprimed
+	// drafter parks the first call, so two pipelined tokens precede the
+	// round and their pairs flush together at the first proposal, which
+	// consumes the catch-up's held logits — the four-token draft makes only
+	// three head calls before the rebuild.
 	wantExtends := []extendCall{
-		{offset: 0, ids: []int32{2}, hiddens: []int32{2}},                             // frontier pair flushed at the first proposal
-		{offset: 1, ids: []int32{3}, hiddens: []int32{-1}},                            // speculative step 2 (held projection)
-		{offset: 2, ids: []int32{4}, hiddens: []int32{-1}},                            // speculative step 3 (projection)
-		{offset: 3, ids: []int32{5}, hiddens: []int32{-1}},                            // speculative step 4 (projection)
-		{offset: 1, ids: []int32{3, 4, 5, 6, eos}, hiddens: []int32{3, 4, 5, 6, eos}}, // committed pairs from the validated run + finish
+		{offset: 0, ids: []int32{2, 3}, hiddens: []int32{2, 3}},                 // parked pairs flushed at the first proposal
+		{offset: 2, ids: []int32{4}, hiddens: []int32{-1}},                      // speculative step 2 (held projection)
+		{offset: 3, ids: []int32{5}, hiddens: []int32{-1}},                      // speculative step 3 (projection)
+		{offset: 4, ids: []int32{6}, hiddens: []int32{-1}},                      // speculative step 4 (projection)
+		{offset: 2, ids: []int32{4, 5, 6, eos}, hiddens: []int32{4, 5, 6, eos}}, // committed pairs from the validated run + finish
 	}
 	if !reflect.DeepEqual(draft.extends, wantExtends) {
 		t.Fatalf("draft extends = %+v, want %+v", draft.extends, wantExtends)
@@ -741,7 +824,7 @@ func TestDecodeKVDraftRejectionRebuildsFromTarget(t *testing.T) {
 		SamplerOpts:       sampler.Options{},
 	}
 	spec := r.spec.open(req, caches)
-	spec.limit = 4
+	pinDraftLimit(spec, 4)
 	defer spec.close()
 	d := spec.decoder([]int32{1}, position)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
@@ -787,10 +870,11 @@ func TestDecodeKVDraftRejectionRebuildsFromTarget(t *testing.T) {
 func TestDecodeMaintainsDraftCacheWithoutDrafting(t *testing.T) {
 	skipIfNoMLX(t)
 	// A request that cannot speculate (logprobs) on a model whose draft has
-	// its own KV cache still maintains it: the pipelined decoder
-	// reports each forwarded token, the pairs wait in the pending list, and
-	// one batched extend at close — its final pair completed by the decoder's
-	// discarded in-flight sample — leaves the draft level with the target.
+	// its own KV cache still maintains it: the speculationSession permanently parks,
+	// so the inner pipelined decoder reports each forwarded token, the pairs
+	// wait in the pending list, and one batched extend at close — its final
+	// pair completed by the decoder's discarded in-flight sample — leaves the
+	// draft level with the target.
 	const eos int32 = 7
 	predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: eos, eos: 0}
 	opts := sampler.Options{Logprobs: true}
@@ -810,7 +894,7 @@ func TestDecodeMaintainsDraftCacheWithoutDrafting(t *testing.T) {
 	}
 	spec := r.spec.open(req, caches)
 	if spec == nil || spec.enabled {
-		t.Fatalf("want a maintain-only speculationSession, got %+v", spec)
+		t.Fatalf("want a permanent-park speculationSession, got %+v", spec)
 	}
 	d := spec.decoder([]int32{1}, position)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
@@ -927,6 +1011,94 @@ func TestCommittedRunBatchesPastFlushCap(t *testing.T) {
 	}
 }
 
+func TestDecodeParkedDraftResume(t *testing.T) {
+	skipIfNoMLX(t)
+	// Leaving a parked stretch: the inner decoder's in-flight sample is
+	// emitted without any new forward, becomes the round's current, and the
+	// next call drafts from it — with the draft pairs contiguous across the
+	// boundary.
+	const eos int32 = 7
+	predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: eos, eos: 0}
+	r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
+	draft := &fakeKVDraft{predict: predict}
+	r.spec = newSpeculation(r, draft)
+
+	caches, _ := newMTPTestCaches(2)
+	req := Request{
+		Tokens:            []int32{1},
+		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
+		SamplerOpts:       sampler.Options{},
+	}
+	spec := r.spec.open(req, caches)
+	if spec == nil || !spec.enabled {
+		t.Fatalf("want a drafting speculationSession, got %+v", spec)
+	}
+	pinDraftLimit(spec, 0)
+	d := spec.decoder([]int32{1}, 0).(*speculativeDecoder)
+
+	// Two parked calls arrive pipelined, one token each.
+	for _, want := range []int{2, 3} {
+		results, err := d.next(20)
+		if err != nil {
+			t.Fatalf("parked next: %v", err)
+		}
+		if got := resultIDs(results); !slices.Equal(got, []int{want}) {
+			t.Fatalf("parked results = %v, want [%d]", got, want)
+		}
+	}
+
+	// The controller picks a positive depth: the next call hands back the
+	// in-flight sample without dispatching another forward.
+	model := r.Model.(*fakeMTPModel)
+	forwards := len(model.forwards)
+	spec.limit = 2
+	results, err := d.next(20)
+	if err != nil {
+		t.Fatalf("resume next: %v", err)
+	}
+	if got := resultIDs(results); !slices.Equal(got, []int{4}) {
+		t.Fatalf("resume results = %v, want [4]", got)
+	}
+	if len(model.forwards) != forwards {
+		t.Fatalf("resume dispatched a forward: %v", model.forwards[forwards:])
+	}
+
+	// The following call runs a fused round drafting from the resumed
+	// current token.
+	results, err = d.next(20)
+	if err != nil {
+		t.Fatalf("round next: %v", err)
+	}
+	if got := resultIDs(results); !slices.Equal(got, []int{5, 6, int(eos)}) {
+		t.Fatalf("round results = %v, want [5 6 %d]", got, eos)
+	}
+	wantForwards := []forwardCall{{0, 1}, {1, 1}, {2, 1}, {3, 3}}
+	if !slices.Equal(model.forwards, wantForwards) {
+		t.Fatalf("target forwards = %v, want %v", model.forwards, wantForwards)
+	}
+
+	// The stretch's reports, the resume pair, and the round's validated run
+	// extend the draft cache contiguously: the catch-up flush at the round's
+	// propose, its speculative entry, and the rebuild from target hiddens at
+	// close.
+	d.close()
+	spec.close()
+	wantExtends := []extendCall{
+		{offset: 0, ids: []int32{2, 3, 4}, hiddens: []int32{2, 3, 4}},
+		{offset: 3, ids: []int32{5}, hiddens: []int32{-1}},
+		{offset: 3, ids: []int32{5, 6, eos}, hiddens: []int32{5, 6, eos}},
+	}
+	if !reflect.DeepEqual(draft.extends, wantExtends) {
+		t.Fatalf("draft extends = %+v, want %+v", draft.extends, wantExtends)
+	}
+	if got, want := caches[0].Offset(), 6; got != want {
+		t.Fatalf("target offset = %d, want %d (EOS never forwarded)", got, want)
+	}
+	if got, want := caches[1].Offset(), 6; got != want {
+		t.Fatalf("draft cache offset = %d, want %d (lockstep with target)", got, want)
+	}
+}
+
 // newMTPTestCaches returns n rewindable fake caches sharing one snapshot
 // tracker, matching the cache.Cache the speculation helpers drive.
 func newMTPTestCaches(n int) ([]cache.Cache, *snapshotTracker) {
@@ -983,3 +1155,12 @@ func scriptedCandidates(r *Runner, tokens []int32) *draftCandidates {
 	d.committed(mlx.FromValues([]int32{0}, 1, 1), mlx.Zeros(mlx.DTypeFloat32, 1, 1, mtpTestVocab), 0)
 	return d.propose(mlx.FromValues([]int32{0}, 1), len(tokens))
 }
+
+// pinAcceptInputs pins the arrays accept's caller must keep alive across
+// accept's internal sweep — current and the candidate tokens — as the decoder
+// and next do in the live engine. It returns the matching unpin.
+func pinAcceptInputs(current sampler.Result, candidates *draftCandidates) func() {
+	arrays := append(current.Arrays(), candidates.tokens)
+	mlx.Pin(arrays...)
+	return func() { mlx.Unpin(arrays...) }
+}

x/mlxrunner/pipeline.go

@@ -92,16 +92,15 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er
 	if spec != nil {
 		d = spec.decoder(seed, position)
 	} else {
-		d = r.pipelinedDecoder(nil, caches, seed, position)
+		d = r.pipelinedDecoder(nil, caches, mlx.FromValues(seed, 1, len(seed)), position)
 	}
 	defer d.close()
 	return r.decode(ctx, request, session, d, promptEval)
 }
 
-// prefill evaluates the prompt's unprocessed tokens in chunks, leaving
-// one token for decode to seed from, and schedules the prompt's periodic
-// snapshots. It returns the seed tokens, the resume position, and the
-// prompt-evaluation duration.
+// prefill evaluates the prompt in chunks, leaving one token for decode to
+// seed from, and schedules the prompt's periodic snapshots. It returns the
+// seed tokens, the resume position, and the prompt-evaluation duration.
 func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *speculationSession) ([]int32, int, time.Duration, error) {
 	start := time.Now()
 	inputs := session.inputs
@@ -280,9 +279,9 @@ type pipelinedDecoder struct {
 	emitted  sampler.Result // last call's result, pinned until the next call
 }
 
-func (r *Runner) pipelinedDecoder(spec *speculationSession, caches []cache.Cache, seed []int32, position int) *pipelinedDecoder {
+func (r *Runner) pipelinedDecoder(spec *speculationSession, caches []cache.Cache, seed *mlx.Array, position int) *pipelinedDecoder {
 	t := &pipelinedDecoder{r: r, spec: spec, caches: caches, position: position}
-	t.sample = t.dispatch(mlx.FromValues(seed, 1, len(seed)))
+	t.sample = t.dispatch(seed)
 	return t
 }
 
@@ -311,18 +310,23 @@ func (t *pipelinedDecoder) next(int) ([]sampler.Result, error) {
 	return []sampler.Result{t.emitted}, nil
 }
 
+// detach ends a parked stretch: it hands the in-flight sample (sampled but
+// never forwarded) and the resume position to the caller, releasing the
+// decoder's emitted pin but not settling the drafter. The caller drafts from
+// the sample next, and that round completes the still-open frontier pair.
+func (t *pipelinedDecoder) detach() (sampler.Result, int) {
+	mlx.Unpin(t.emitted.Arrays()...)
+	return t.sample, t.position
+}
+
 func (t *pipelinedDecoder) close() {
-	// The in-flight sample's forward was never dispatched; its report lets
-	// the drafter settle level with the caches' resting offset.
+	// The in-flight sample's forward was never dispatched; its report settles
+	// the drafter level with the caches' resting offset.
 	t.spec.finish(t.sample.Token)
 	mlx.Unpin(t.emitted.Arrays()...)
 	mlx.Unpin(t.sample.Arrays()...)
 }
 
-func lastLogits(logits *mlx.Array) *mlx.Array {
-	return logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice()).Squeeze(1)
-}
-
 // detokenizer serializes sampled tokens into response chunks, holding bytes
 // whose UTF-8 sequence hasn't completed yet and the logprobs that belong
 // with those bytes so Content and Logprobs stay aligned when a chunk does

x/mlxrunner/speculate.go

@@ -2,11 +2,8 @@ package mlxrunner
 
 import (
 	"fmt"
-	"log/slog"
-	"os"
 	"slices"
-	"strconv"
-	"strings"
+	"time"
 
 	"github.com/ollama/ollama/x/mlxrunner/batch"
 	"github.com/ollama/ollama/x/mlxrunner/cache"
@@ -39,73 +36,13 @@ type drafter interface {
 	close()
 }
 
-type draftSchedule string
-
-const (
-	draftScheduleHeuristic draftSchedule = "heuristic"
-	draftScheduleConstant  draftSchedule = "constant"
-)
-
-type specStats struct {
-	iterations  int
-	drafted     int
-	accepted    int
-	mismatches  int
-	allAccepted int
-	maxDraft    int
-}
-
-type specOptions struct {
-	initialDraftTokens int
-	maxDraftTokens     int
-	draftSchedule      draftSchedule
-}
-
-func (r *Runner) loadSpecOptions(sample bool) specOptions {
-	defaults := r.mtpDefaults(sample)
-
-	opts := specOptions{
-		initialDraftTokens: defaults.InitialDraftTokens,
-		maxDraftTokens:     defaults.MaxDraftTokens,
-		draftSchedule:      draftScheduleConstant,
-	}
-	if v := positiveEnvInt("OLLAMA_MLX_MTP_MAX_DRAFT_TOKENS"); v > 0 {
-		opts.maxDraftTokens = v
-	}
-	if v := positiveEnvInt("OLLAMA_MLX_MTP_INITIAL_DRAFT_TOKENS"); v > 0 {
-		opts.initialDraftTokens = v
-	}
-	if opts.initialDraftTokens > opts.maxDraftTokens {
-		opts.initialDraftTokens = opts.maxDraftTokens
-	}
-	switch schedule := strings.ToLower(strings.TrimSpace(os.Getenv("OLLAMA_MLX_MTP_DRAFT_SCHEDULE"))); schedule {
-	case "", string(draftScheduleConstant):
-		opts.draftSchedule = draftScheduleConstant
-	case string(draftScheduleHeuristic):
-		opts.draftSchedule = draftScheduleHeuristic
-	default:
-		slog.Warn("invalid MTP env setting", "key", "OLLAMA_MLX_MTP_DRAFT_SCHEDULE", "value", schedule)
-	}
-	return opts
-}
-
-func positiveEnvInt(key string) int {
-	raw := os.Getenv(key)
-	if raw == "" {
-		return 0
-	}
-	v, err := strconv.Atoi(raw)
-	if err != nil || v <= 0 {
-		slog.Warn("invalid MTP env setting", "key", key, "value", raw)
-		return 0
-	}
-	return v
-}
-
-// speculation is the persistent speculative-decoding subsystem of a Runner,
-// one per loaded model. It holds the draft model; as the feature grows it
-// also holds the cache partition and the depth state learned across requests.
-// A nil *speculation means the checkpoint ships no draft head.
+// speculation is the persistent speculative-decoding subsystem of a Runner —
+// one per loaded model, holding everything that outlives a request: the draft
+// model, the target-cache partition, and the depth state learned across
+// requests so each request starts at the proven-out depth. Each request opens a
+// short-lived speculationSession cursor over it; nothing here is rebuilt per request. A
+// nil *speculation means the checkpoint ships no draft head, so every request
+// decodes plainly.
 type speculation struct {
 	r     *Runner
 	draft base.DraftModel
@@ -118,6 +55,10 @@ type speculation struct {
 	caches  []cache.Cache
 	draftKV []cache.Cache
 	targets []cache.Cache
+
+	// depth selects each request's draft length and owns the cost/acceptance
+	// models and probe cadence it learns across requests.
+	depth *depthController
 }
 
 // newSpeculation builds the speculative-decoding subsystem for a loaded model,
@@ -126,7 +67,7 @@ func newSpeculation(r *Runner, draft base.DraftModel) *speculation {
 	if draft == nil {
 		return nil
 	}
-	return &speculation{r: r, draft: draft}
+	return &speculation{r: r, draft: draft, depth: newDepthController()}
 }
 
 // bind computes the draft/target cache partition the first time the persistent
@@ -157,17 +98,22 @@ func (s *speculation) bind(caches []cache.Cache) {
 	s.targets = targets
 }
 
-// speculationSession is the speculation cursor for one request over a speculation. A
-// nil speculationSession is a plain decode.
+// speculationSession is the per-request cursor over the persistent speculation:
+// it owns the drafter and runs the validate rounds. A nil session is a plain
+// decode.
 type speculationSession struct {
 	spec    *speculation
 	drafter drafter
-	// enabled selects speculative rounds; a maintain-only engine decodes
-	// plainly while streaming committed runs to keep draft KV prefix-cached.
-	enabled bool
-	opts    specOptions
-	limit   int // current draft length
+	enabled bool // whether this request drafts; false parks (maintain-only)
+	limit   int  // current draft length
 	stats   specStats
+
+	// Cost sampling: each round's wall time (start to next start, spanning the
+	// next emit's sync) is attributed to its draft depth only when the depth
+	// matches the previous round's, since batch-shape transitions inflate it.
+	lastRoundStart time.Time
+	prevDrafts     int
+	roundDrafts    int
 }
 
 // open returns the speculation cursor for this request or nil when the model ships
@@ -178,22 +124,49 @@ func (s *speculation) open(request Request, caches []cache.Cache) *speculationSe
 	}
 	s.bind(caches)
 	d := newMTPDrafter(s)
-	if d == nil {
-		return nil
-	}
 
+	// Logprobs are not yet supported, so a logprobs request keeps a speculationSession
+	// only to maintain a draft cache in lockstep (permanently parked).
 	opts := request.SamplerOpts
-	enabled := s.r.mtpDefaults(opts.Temperature != 0).Enabled &&
-		!opts.Logprobs && opts.TopLogprobs == 0
+	enabled := !opts.Logprobs && opts.TopLogprobs == 0
 
-	specOpts := s.r.loadSpecOptions(opts.Temperature != 0)
-	return &speculationSession{
-		spec:    s,
-		drafter: d,
-		enabled: enabled,
-		opts:    specOpts,
-		limit:   specOpts.initialDraftTokens,
-		stats:   specStats{maxDraft: specOpts.initialDraftTokens},
+	spec := &speculationSession{spec: s, drafter: d, enabled: enabled, prevDrafts: -1, roundDrafts: -1}
+	if enabled {
+		spec.limit = s.depth.scheduled
+		spec.stats.maxDraft = spec.limit
+	}
+	return spec
+}
+
+// beginRound records the previous round's cost sample (its wall time runs to
+// this round's start) and starts timing the new one.
+func (s *speculationSession) beginRound() {
+	now := time.Now()
+	if !s.lastRoundStart.IsZero() && s.roundDrafts >= 0 {
+		s.stats.recordRound(s.roundDrafts)
+		if s.roundDrafts == s.prevDrafts {
+			s.spec.depth.cost.observe(s.roundDrafts, now.Sub(s.lastRoundStart))
+		}
+		s.prevDrafts = s.roundDrafts
+	}
+	s.lastRoundStart = now
+}
+
+// endRound records a completed round's draft depth, proposal outcome, and the
+// controller's next draft length. observed is the leading draft positions the
+// acceptance model learns from: the full round, except an accepted EOS holds out
+// positions past it (a terminator, not a target rejection).
+func (s *speculationSession) endRound(drafted, accepted, observed int) {
+	s.roundDrafts = drafted
+	s.stats.iterations++
+	s.stats.drafted += drafted
+	s.stats.accepted += accepted
+	if s.enabled {
+		if observed > 0 {
+			s.spec.depth.acc.observe(observed, accepted)
+		}
+		s.limit = s.spec.depth.next()
+		s.stats.maxDraft = max(s.stats.maxDraft, s.limit)
 	}
 }
 
@@ -229,131 +202,116 @@ func (s *speculationSession) close() {
 }
 
 // speculativeDecoder decodes one speculative round per call: the engine
-// forwards the current token — emitted by the previous call, so a token
-// that ends generation is never forwarded — fused with the drafter's
-// proposals, and the call returns the round's accepted tokens followed by
-// the engine's next token. The last returned token becomes the next call's
-// current; the seed token primes current and is never returned.
+// forwards the current token (emitted by the previous call, so a token that
+// ends generation is never forwarded) fused with the drafter's proposals,
+// returning the round's accepted tokens followed by the next token. The seed
+// primes current and is never returned. While the engine cannot draft (parked
+// at depth zero, or nothing committed to propose from) calls delegate to an
+// inner pipelined decoder at plain decode speed.
 type speculativeDecoder struct {
 	s        *speculationSession
 	position int
-	current  sampler.Result // emitted (or the seed), not yet forwarded
+	current  sampler.Result    // emitted (or the seed), not yet forwarded
+	inner    *pipelinedDecoder // pipelines plain tokens while parked; nil while drafting
 }
 
-// decoder returns the decoder for this engine's session. A maintain-only
-// engine decodes plainly via the pipelined decoder, which reports every
-// forwarded token to keep draft KV level with the target.
+// decoder returns the decoder for this engine's session. A speculationSession that
+// cannot draft (logprobs) has no depth controller and permanently parks,
+// running the inner pipelined decoder whose reports keep the draft KV level.
 func (s *speculationSession) decoder(seed []int32, position int) decoder {
-	if !s.enabled {
-		return s.spec.r.pipelinedDecoder(s, s.spec.caches, seed, position)
-	}
 	current := sampler.Result{Token: mlx.FromValues(seed, len(seed))}
 	mlx.Pin(current.Arrays()...)
 	return &speculativeDecoder{s: s, position: position, current: current}
 }
 
 func (st *speculativeDecoder) next(remaining int) ([]sampler.Result, error) {
-	results, next, err := st.s.round(&st.position, st.current, remaining)
-	if err != nil {
-		return nil, err
-	}
-	if next.Token != nil {
-		results = append(results, next)
+	// Route: end a parked stretch by emitting the inner sample, draft on a
+	// positive length and a primed drafter, else decode parked.
+	var results []sampler.Result
+	if s := st.s; st.inner != nil && s.limit > 0 {
+		results = st.resume()
+	} else {
+		s.beginRound()
+		var candidates *draftCandidates
+		if s.limit > 0 {
+			// A round emits the accepted drafts plus one more token (the bonus
+			// or residual), so cap the draft one below the remaining budget to
+			// land that extra token within it rather than overshooting. At
+			// remaining 1 the cap is 0 and the last token decodes plainly.
+			candidates = s.drafter.propose(st.current.Token, min(s.limit, remaining-1))
+		}
+		var accepted, observed int
+		var err error
+		if candidates == nil {
+			results, err = st.park(remaining)
+		} else {
+			// candidates stays pinned across accept's internal sweep and the
+			// draft-count read below; accept pins only its own intermediates.
+			mlx.Pin(candidates.tokens)
+			defer mlx.Unpin(candidates.tokens)
+			results, accepted, observed, err = st.s.accept(&st.position, st.current, candidates)
+		}
+		if err != nil {
+			return nil, err
+		}
+		drafted := 0
+		if candidates != nil {
+			drafted = candidates.tokens.Dim(1)
+		}
+		s.endRound(drafted, accepted, observed)
 	}
 
-	last := results[len(results)-1]
-	mlx.Pin(last.Arrays()...)
-	mlx.Unpin(st.current.Arrays()...)
-	st.current = last
-	mlx.AsyncEval(st.current.Arrays()...)
+	st.advance(results[len(results)-1])
 	return results, nil
 }
 
+// advance retires the last returned token as the next call's current, pinned
+// across the sweeps the next call runs before reading it. Nothing is forced here.
+func (st *speculativeDecoder) advance(next sampler.Result) {
+	mlx.Pin(next.Arrays()...)
+	mlx.Unpin(st.current.Arrays()...)
+	st.current = next
+}
+
+// resume ends a parked stretch: the inner decoder's in-flight sample (sampled
+// but never forwarded) is exactly the current token a drafting round expects,
+// so emit it and let the next call draft from it.
+func (st *speculativeDecoder) resume() []sampler.Result {
+	next, position := st.inner.detach()
+	st.position = position
+	st.inner = nil
+	// No round spans this call, so the next beginRound attributes no cost.
+	st.s.roundDrafts = -1
+	// detach handed over the pin; advance re-pins, no sweep in between.
+	mlx.Unpin(next.Arrays()...)
+	return []sampler.Result{next}
+}
+
+// park decodes one pipelined plain token while the engine cannot draft. Each
+// is a depth-0 round in the controller's accounting, and the inner decoder's
+// reports keep the drafter primed and maintained.
+func (st *speculativeDecoder) park(remaining int) ([]sampler.Result, error) {
+	s := st.s
+	if st.inner == nil {
+		st.inner = s.spec.r.pipelinedDecoder(s, s.spec.caches, st.current.Token.ExpandDims(-1), st.position)
+	}
+	return st.inner.next(remaining)
+}
+
 func (st *speculativeDecoder) close() {
-	// Generation always ends with a final token that was emitted but never
-	// forwarded; its report lets the drafter settle level with the caches'
-	// resting offset.
-	st.s.finish(st.current.Token)
+	if st.inner != nil {
+		// Ended while parked: the inner decoder's close settles the drafter
+		// with its in-flight sample.
+		st.inner.close()
+	} else {
+		// The final token was emitted but never forwarded; its report settles
+		// the drafter level with the caches' resting offset.
+		st.s.finish(st.current.Token)
+	}
 	mlx.Unpin(st.current.Arrays()...)
 	st.s.logStats()
 }
 
-// round runs one speculative decode round: fuse the current token with the
-// drafter's proposals in one target forward, validate the proposals, and
-// return the accepted run and the bonus or resampled next token.
-func (s *speculationSession) round(position *int, current sampler.Result, remaining int) (results []sampler.Result, next sampler.Result, err error) {
-	r := s.spec.r
-	s.stats.iterations++
-
-	// A round emits the accepted drafts plus one more token (the bonus or
-	// residual), so cap the draft one below the remaining budget to land that
-	// extra token within it rather than overshooting. At remaining 1 the cap is
-	// 0 and the last token decodes plainly.
-	maxDraft := min(s.limit, remaining-1)
-	candidates := s.drafter.propose(current.Token, maxDraft)
-	if candidates == nil {
-		// Nothing to validate: a plain single-token round.
-		tok := tokenInput(current.Token)
-		hidden := r.Model.Forward(&batch.Batch{
-			InputIDs:     tok,
-			SeqOffsets:   []int32{int32(*position)},
-			SeqQueryLens: []int32{1},
-		}, s.spec.caches)
-		*position++
-		s.drafter.committed(tok, hidden, *position-1)
-		return nil, r.Sampler.Sample([]int{pipelineSlot}, lastLogits(r.Model.Unembed(hidden))), nil
-	}
-
-	draftCount := candidates.tokens.Dim(1)
-	candidateArrays := candidates.Arrays()
-	mlx.Pin(candidateArrays...)
-	mlx.Sweep()
-	defer mlx.Unpin(candidateArrays...)
-	s.stats.drafted += draftCount
-
-	results, accepted, err := s.accept(position, current, candidates)
-	if err != nil {
-		return nil, sampler.Result{}, err
-	}
-	// accept folds the bonus token into its results; surface it back out as
-	// the next step's current.
-	if len(results) > accepted {
-		next = results[len(results)-1]
-		results = results[:accepted]
-	}
-
-	s.stats.accepted += accepted
-	if accepted == draftCount {
-		s.stats.allAccepted++
-	} else {
-		s.stats.mismatches++
-	}
-	if s.opts.draftSchedule == draftScheduleHeuristic {
-		if accepted == draftCount {
-			s.limit = min(s.opts.maxDraftTokens, s.limit+2)
-		} else {
-			s.limit = max(1, s.limit-1)
-		}
-		s.stats.maxDraft = max(s.stats.maxDraft, s.limit)
-	}
-	return results, next, nil
-}
-
-// logStats reports the per-request speculation summary.
-func (s *speculationSession) logStats() {
-	acceptance := 0.0
-	if s.stats.drafted > 0 {
-		acceptance = float64(s.stats.accepted) / float64(s.stats.drafted)
-	}
-	avgDraft := 0.0
-	avgAccepted := 0.0
-	if s.stats.iterations > 0 {
-		avgDraft = float64(s.stats.drafted) / float64(s.stats.iterations)
-		avgAccepted = float64(s.stats.accepted) / float64(s.stats.iterations)
-	}
-	slog.Info("speculative decode stats", "drafted", s.stats.drafted, "accepted", s.stats.accepted, "acceptance", acceptance, "iterations", s.stats.iterations, "avg_draft", avgDraft, "avg_accepted", avgAccepted, "mismatches", s.stats.mismatches, "all_accepted", s.stats.allAccepted, "max_draft", s.stats.maxDraft, "draft_schedule", s.opts.draftSchedule)
-}
-
 // draftCandidates is one round's draft tokens and the proposal distribution
 // each was sampled from, weighed against the target during acceptance.
 type draftCandidates struct {
@@ -417,20 +375,18 @@ func commitSpeculation(caches []cache.Cache, accepted, draftCount, before int) {
 	}
 }
 
-// accept accepts the longest draft prefix that survives rejection sampling
-// against the target model. At temperature 0 the distributions are point
-// masses, so acceptance reduces to argmax-match. It returns the accepted
-// drafts followed by the target's own next token — the residual at the
-// rejection point, or the bonus past a fully accepted run — so a round
-// yields accepted+1 tokens, except when an accepted EOS ends generation,
-// where the run stops at the EOS with no continuation. The accepted run is
-// reported back to the drafter with its hidden states.
+// accept accepts the longest draft prefix that survives rejection sampling,
+// returning the accepted drafts followed by the target's own next token
+// (residual at a rejection, bonus past a full run), except an accepted EOS
+// ends the run with no continuation. observed is the leading positions the
+// acceptance model learns from, capped at the EOS (a terminator, not a target
+// rejection). NumPredict is the decode loop's to enforce, so a token past the
+// budget is left for decode to drop, not cut here.
 //
-// The NumPredict budget is the decode loop's to enforce; the drafter is
-// already capped to the remaining budget, so an accepted run never
-// overshoots, and a legitimate token past the budget is left for decode to
-// drop rather than cut here.
-func (s *speculationSession) accept(position *int, current sampler.Result, candidates *draftCandidates) (results []sampler.Result, accepted int, err error) {
+// The caller keeps current and the candidate tokens pinned across the call,
+// since accept sweeps before its eval and reads both afterward; accept pins
+// only the intermediates it produces.
+func (s *speculationSession) accept(position *int, current sampler.Result, candidates *draftCandidates) (results []sampler.Result, accepted, observed int, err error) {
 	r := s.spec.r
 	before := *position
 	draftCount := candidates.tokens.Dim(1)
@@ -471,6 +427,15 @@ func (s *speculationSession) accept(position *int, current sampler.Result, candi
 	// point is known.
 	residualTokens := r.Sampler.SampleDistribution(pipelineSlot, targetDist.SliceRows(0, draftCount).ResidualAgainst(draftDist))
 	bonusToken := r.sampleTokenAt(targetDist, draftCount)
+
+	// Pin the arrays read after the eval, then sweep so the draft proposal
+	// chain and this validation forward's intermediates are freed as the eval
+	// consumes them, the way the plain decode dispatch sweeps before its eval.
+	// current and the candidate tokens stay pinned by the caller across the call.
+	live := []*mlx.Array{hiddenSeq, acceptedMask, residualTokens, bonusToken}
+	mlx.Pin(live...)
+	defer mlx.Unpin(live...)
+	mlx.Sweep()
 	mlx.Eval(candidates.tokens, acceptedMask, residualTokens, bonusToken)
 
 	draftIDs := candidates.tokens.Ints()
@@ -482,11 +447,13 @@ func (s *speculationSession) accept(position *int, current sampler.Result, candi
 		accepted++
 	}
 	if accepted > draftCount {
-		return nil, 0, fmt.Errorf("speculation validation accepted %d tokens for %d draft tokens", accepted, draftCount)
+		return nil, 0, 0, fmt.Errorf("speculation validation accepted %d tokens for %d draft tokens", accepted, draftCount)
 	}
+	observed = draftCount
 
 	// Find where an accepted EOS ends the run, before committing, so the cut
-	// is known while the per-token rollback snapshots still exist.
+	// is known while the per-token rollback snapshots still exist. The EOS also
+	// caps observed: positions past it are held out, not logged as rejections.
 	commitIDs := make([]int32, 0, accepted+1)
 	keep := accepted
 	done := false
@@ -495,6 +462,7 @@ func (s *speculationSession) accept(position *int, current sampler.Result, candi
 		if r.Tokenizer.IsEOS(int32(id)) {
 			done = true
 			accepted = i + 1
+			observed = accepted
 			// Leave the EOS's own state uncommitted: the next sequence won't
 			// contain this EOS, and recurrent state can't drop a folded-in
 			// token, so committing it would carry the caches past the
@@ -507,11 +475,9 @@ func (s *speculationSession) accept(position *int, current sampler.Result, candi
 	commit(keep)
 	*position = before + 1 + keep
 
-	// Report the validated run — current plus the kept drafts, with the
-	// fused hidden state at each token's own slot — to the drafter before
-	// returning, so even a cancelled emission leaves the drafter's state
-	// describing exactly what the target caches hold. A done generation's
-	// final token is never committed; it reaches the drafter through finish.
+	// Report the validated run (current plus kept drafts) to the drafter before
+	// returning, so a cancelled emission still leaves it matching the caches. A
+	// done generation's final token is uncommitted; it reaches finish instead.
 	runIDs := append([]int32{int32(current.Token.Int())}, commitIDs[:keep]...)
 	s.drafter.committed(
 		mlx.FromValues(runIDs, 1, len(runIDs)),
@@ -521,7 +487,7 @@ func (s *speculationSession) accept(position *int, current sampler.Result, candi
 	results = draftResults(draftIDs[:accepted])
 	if done {
 		r.Sampler.Commit(pipelineSlot, commitIDs)
-		return results, accepted, nil
+		return results, accepted, observed, nil
 	}
 
 	var nextID int32
@@ -534,7 +500,7 @@ func (s *speculationSession) accept(position *int, current sampler.Result, candi
 	r.Sampler.Commit(pipelineSlot, commitIDs)
 
 	results = append(results, sampler.Result{Token: mlx.FromValues([]int32{nextID}, 1)})
-	return results, accepted, nil
+	return results, accepted, observed, nil
 }
 
 func (r *Runner) sampleAcceptedMask(targetDist, draftDist sampler.Distribution, draftTokens *mlx.Array) *mlx.Array {
@@ -557,16 +523,3 @@ func draftResults(ids []int) []sampler.Result {
 	}
 	return results
 }
-
-func tokenInput(token *mlx.Array) *mlx.Array {
-	switch token.NumDims() {
-	case 0:
-		return token.Reshape(1, 1)
-	case 1:
-		return token.ExpandDims(-1)
-	case 2:
-		return token
-	default:
-		panic(fmt.Sprintf("token must be rank 0, 1, or 2, got rank %d", token.NumDims()))
-	}
-}

x/mlxrunner/speculate_depth.go

@@ -0,0 +1,302 @@
+package mlxrunner
+
+import (
+	"fmt"
+	"slices"
+	"strings"
+	"time"
+)
+
+// depthProbeInterval is the base cadence at which the controller drafts one past its
+// selection to refresh the next position up; without it a shallow controller never
+// notices a deeper draft becoming worthwhile.
+const depthProbeInterval = 4
+
+// depthProbeIntervalMax caps the probe backoff. Probes that keep changing nothing
+// double the interval up to this cap, trading slower re-engagement — worst case
+// plain decode speed, never below — for fewer probes.
+const depthProbeIntervalMax = 512
+
+// depthController drafts argmax_N EV(N) where EV(N) = committed(N) / cost(N), over
+// depths from 0 (plain decode) up to one past the frontier. The depth-0 floor lets
+// it stop speculating when no draft pays; the frontier ceiling keeps it from scoring
+// a depth on the optimistic inherited rate, so it climbs outward one position at a
+// time as acceptance is measured rather than leaping deep.
+//
+// It holds the depth-selection state learned across requests — the target forward's
+// per-depth cost curve, the drafts' per-position acceptance rates, the probe cadence,
+// and the depth scheduled for the next round — persisted on the speculation so a fresh
+// request starts at the proven-out depth instead of re-ramping from shallow.
+type depthController struct {
+	cost   *costModel
+	acc    *acceptanceModel
+	probed bool // the previous round drafted a probe
+
+	// scheduled is the draft depth next() chose for the upcoming round, carried
+	// across requests so a new request's first round consumes it instead of
+	// recomputing at the boundary. Its zero value is the depth-0 warmup start.
+	scheduled int
+
+	// Probe cadence, persisted so a backed-off request need not restart at the base.
+	probeInterval int // rounds between probes; backs off while probes change nothing
+	probeSince    int // rounds since the cadence was last calibrated
+	lastSelected  int // the selection the cadence was calibrated against
+}
+
+func newDepthController() *depthController {
+	return &depthController{cost: newCostModel(), acc: newAcceptanceModel(), probeInterval: depthProbeInterval}
+}
+
+func (c *depthController) frontier() int { return c.acc.frontier() }
+
+// next returns the draft depth for the upcoming step: the EV-optimal depth (capped
+// at frontier+1), except periodically it probes one past the selection to refresh
+// the next position up. The probe stays within the frontier window. The cadence
+// doubles toward its cap while probes change nothing and resets on any selection
+// change, giving the new selection a full interval to settle. The chosen depth is
+// recorded in c.scheduled for the next request's open to consume.
+func (c *depthController) next() (depth int) {
+	defer func() { c.scheduled = depth }()
+	sel := c.selected()
+	if sel != c.lastSelected {
+		c.probeInterval = depthProbeInterval
+		c.probeSince = 0
+		c.lastSelected = sel
+	} else if c.probed {
+		c.probeInterval = min(c.probeInterval*2, depthProbeIntervalMax)
+	}
+	c.probed = false
+
+	c.probeSince++
+	if c.probeSince >= c.probeInterval {
+		c.probeSince = 0
+		if probe := min(sel+1, c.frontier()+1); probe != sel {
+			c.probed = true
+			return probe
+		}
+	}
+	// Seed a clean cost sample for every depth in [0, frontier+1] before judging by
+	// EV. Cost is only recorded when a round's draft depth matches the prior round's
+	// (the same-depth gate dropping batch-shape transitions), so a depth never dwelt
+	// at has no clean sample; draft the shallowest such depth to dwell there. Without
+	// this the controller stays at the one depth it can sit at without a transition
+	// (depth 0) and never learns that drafting pays on a deep-optimum model.
+	if seed := c.costSeedDepth(); seed >= 0 {
+		return seed
+	}
+	return sel
+}
+
+// costSeedDepth is the shallowest depth in [0, frontier+1] with no clean
+// cost sample, or -1 if all are sampled; bounding to frontier+1 keeps cost-seeding
+// from outrunning the acceptance frontier.
+func (c *depthController) costSeedDepth() int {
+	limit := c.frontier() + 1
+	for n := 0; n <= limit; n++ {
+		if !c.cost.sampled(n) {
+			return n
+		}
+	}
+	return -1
+}
+
+// selected returns the EV-optimal draft depth without mutating probe state, the
+// argmax over [0, frontier+1]. The frontier bound keeps the inherited optimistic
+// rate from making ever-deeper depths look best; the depth-0 floor lets it stop
+// speculating. Returns 0 until the cost model can compare depths.
+func (c *depthController) selected() int {
+	if !c.cost.ready() {
+		return 0
+	}
+	limit := c.frontier() + 1
+	best, bestEV := 0, c.acc.expectedCommitted(0)/c.cost.cost(0)
+	for n := 1; n <= limit; n++ {
+		if ev := c.acc.expectedCommitted(n) / c.cost.cost(n); ev > bestEV {
+			best, bestEV = n, ev
+		}
+	}
+	return best
+}
+
+// costEWMAAlpha weights the newest cost sample. Fixed-depth cost is low-variance,
+// so a responsive alpha converges in a few visits while smoothing scheduler jitter.
+const costEWMAAlpha = 0.3
+
+// costModel is the target-forward cost for validating N drafts — a fused batch of
+// 1 current token plus N drafts — as an EWMA per visited draft depth read by
+// piecewise-linear interpolation between samples (flat beyond the extremes).
+// Interpolation assumes no curve shape, so a steep compute-bound or flat
+// bandwidth-bound forward is represented as measured. Cost is static within a run,
+// learned from decode steps that already sync the forward, so there is no startup
+// probe.
+type costModel struct {
+	ewma   map[int]float64 // EWMA of measured ms by draft depth
+	depths []int           // sampled draft depths, sorted; updated as new depths arrive
+}
+
+func newCostModel() *costModel {
+	return &costModel{ewma: map[int]float64{}}
+}
+
+// costClampFraction bounds one sample's EWMA innovation. A host stall (cache trim,
+// backpressure) can inflate a sample severalfold; unclamped it can flip the EV
+// comparison against plain decode, and once the controller stops parking it stops
+// resampling depth 0, so the error never heals. A genuine cost change still
+// converges because every sample keeps pushing toward it.
+const costClampFraction = 0.25
+
+// observe folds one forward's wall time into the draft depth's EWMA, clamping the
+// innovation so one stall-inflated sample cannot move it far.
+func (m *costModel) observe(drafts int, dt time.Duration) {
+	if drafts < 0 || dt <= 0 {
+		return
+	}
+	ms := float64(dt) / float64(time.Millisecond)
+	if v, ok := m.ewma[drafts]; ok {
+		limit := costClampFraction * v
+		m.ewma[drafts] = v + costEWMAAlpha*max(-limit, min(limit, ms-v))
+	} else {
+		m.ewma[drafts] = ms
+		i, _ := slices.BinarySearch(m.depths, drafts)
+		m.depths = slices.Insert(m.depths, i, drafts)
+	}
+}
+
+// ready reports whether two distinct depths have been sampled, so a slope exists.
+func (m *costModel) ready() bool { return len(m.ewma) >= 2 }
+
+func (m *costModel) sampled(drafts int) bool {
+	_, ok := m.ewma[drafts]
+	return ok
+}
+
+// cost returns the estimated forward wall time (ms) for validating drafts tokens:
+// a piecewise-linear interpolation of the per-depth EWMAs, clamping to the nearest
+// sample outside the sampled range (the curve beyond is unknown).
+func (m *costModel) cost(drafts int) float64 {
+	ds := m.depths
+	if len(ds) == 0 {
+		return 0
+	}
+	if drafts <= ds[0] {
+		return m.ewma[ds[0]]
+	}
+	if drafts >= ds[len(ds)-1] {
+		return m.ewma[ds[len(ds)-1]]
+	}
+	for i := 1; i < len(ds); i++ {
+		if drafts <= ds[i] {
+			lo, hi := ds[i-1], ds[i]
+			t := float64(drafts-lo) / float64(hi-lo)
+			return m.ewma[lo] + t*(m.ewma[hi]-m.ewma[lo])
+		}
+	}
+	return m.ewma[ds[len(ds)-1]]
+}
+
+// sampleString reports the EWMA ms per visited draft depth for diagnostics.
+func (m *costModel) sampleString() string {
+	ds := m.depths
+	parts := make([]string, 0, len(ds))
+	for _, d := range ds {
+		parts = append(parts, fmt.Sprintf("%d:%.0fms", d, m.ewma[d]))
+	}
+	return strings.Join(parts, " ")
+}
+
+// acceptanceEWMAAlpha weights the newest acceptance outcome. Acceptance drifts with
+// content, so an EWMA follows the drift instead of being anchored by early tokens.
+const acceptanceEWMAAlpha = 0.1
+
+// acceptanceMinSamples is how many reaches a position needs before its rate is
+// trusted; since the search reaches one past the frontier, it also gates how fast
+// the frontier advances. Set near the EWMA's memory (~1/alpha); larger only slows
+// the ramp, since each depth is re-measured as it is drafted.
+const acceptanceMinSamples = 10
+
+// acceptanceModel learns the per-position conditional acceptance rate — the chance
+// position i is accepted given every draft before it already was — as an EWMA, shared
+// across requests so a fresh request keeps the proven-out frontier. Drift is handled
+// by the EWMA forgetting, not by discarding the estimate.
+type acceptanceModel struct {
+	rate []float64 // EWMA acceptance rate given the prefix survived; index i is position i
+	seen []int     // times position i was reached (warmup gate for rate)
+}
+
+func newAcceptanceModel() *acceptanceModel {
+	return &acceptanceModel{rate: []float64{0}, seen: []int{0}}
+}
+
+func (a *acceptanceModel) grow(i int) {
+	for len(a.seen) <= i {
+		a.rate = append(a.rate, 0)
+		a.seen = append(a.seen, 0)
+	}
+}
+
+// observe folds a step's outcome into each reached position's EWMA. A position is
+// reached only when the prefix before it survived (accepted >= i-1), and is accepted
+// iff accepted >= i; updating only the surviving prefix avoids diluting deeper
+// positions the step never reached.
+func (a *acceptanceModel) observe(drafted, accepted int) {
+	for i := 1; i <= drafted; i++ {
+		if accepted < i-1 {
+			break // prefix did not survive to position i; deeper positions unreached
+		}
+		a.grow(i)
+		outcome := 0.0
+		if accepted >= i {
+			outcome = 1.0
+		}
+		if a.seen[i] == 0 {
+			a.rate[i] = outcome
+		} else {
+			a.rate[i] += acceptanceEWMAAlpha * (outcome - a.rate[i])
+		}
+		a.seen[i]++
+	}
+}
+
+// acceptance returns the rate position i is accepted given its prefix survived. An under-sampled
+// position inherits the deepest trusted rate rather than zero, so the controller
+// keeps exploring deeper instead of locking shallow on noise.
+func (a *acceptanceModel) acceptance(i int) float64 {
+	if i >= 1 && i < len(a.seen) && a.seen[i] >= acceptanceMinSamples {
+		return a.rate[i]
+	}
+	// Inherit the deepest trusted rate; fall back to optimistic 1 if none yet.
+	for j := i - 1; j >= 1; j-- {
+		if j < len(a.seen) && a.seen[j] >= acceptanceMinSamples {
+			return a.rate[j]
+		}
+	}
+	return 1
+}
+
+// expectedCommitted returns expected committed tokens at depth N: the current token,
+// which always commits, plus the expected number of accepted drafts — each draft
+// position contributes the probability its whole prefix was accepted, the running
+// product of the per-position acceptance rates summed over positions.
+func (a *acceptanceModel) expectedCommitted(n int) float64 {
+	total, prod := 1.0, 1.0
+	for k := 1; k <= n; k++ {
+		prod *= a.acceptance(k)
+		total += prod
+	}
+	return total
+}
+
+// frontier is the deepest position with a trusted acceptance rate. The controller
+// never selects beyond frontier+1, so the selection grows outward one position at a
+// time instead of jumping deep on the inherited optimistic rate.
+func (a *acceptanceModel) frontier() int {
+	f := 0
+	for i := 1; i < len(a.seen); i++ {
+		if a.seen[i] >= acceptanceMinSamples {
+			f = i
+		} else {
+			break
+		}
+	}
+	return f
+}

x/mlxrunner/speculate_depth_test.go

@@ -0,0 +1,298 @@
+package mlxrunner
+
+import (
+	"testing"
+	"time"
+)
+
+// driveDepthController runs the controller for n steps against a regime mirroring
+// the real loop: next() picks a draft depth D; the regime reports the forward time
+// at fused width D+1, fed to the cost model keyed by draft depth D, and how many of
+// the D drafts are accepted as a prefix (fed to the controller's acceptance
+// estimator). It returns the depth the controller selects after the run. draw(pos,
+// step) is the per-position acceptance outcome; a draft of length D accepts the
+// leading run of positions whose draws succeed.
+func driveDepthController(c *depthController, n int, fwdMS func(width int) float64, draw func(pos, step int) bool) int {
+	for step := range n {
+		d := c.next()
+		c.cost.observe(d, time.Duration(fwdMS(d+1)*float64(time.Millisecond)))
+		accepted := 0
+		for i := 1; i <= d; i++ {
+			if draw(i, step) {
+				accepted++
+			} else {
+				break
+			}
+		}
+		c.acc.observe(d, accepted)
+	}
+	return c.selected()
+}
+
+// deterministicDraw turns a per-position acceptance rate into a reproducible,
+// independent Bernoulli accept/reject per (pos, step). It uses a splitmix64-style
+// avalanche so adjacent positions and steps decorrelate — a plain linear hash leaves
+// the prefix acceptance correlated across positions and biases the recovered rates.
+func deterministicDraw(acceptance func(pos int) float64) func(pos, step int) bool {
+	return func(pos, step int) bool {
+		x := uint64(step)*0x9E3779B97F4A7C15 + uint64(pos)*0xD1B54A32D192ED03
+		x ^= x >> 30
+		x *= 0xBF58476D1CE4E5B9
+		x ^= x >> 27
+		x *= 0x94D049BB133111EB
+		x ^= x >> 31
+		frac := float64(x%1_000_000) / 1_000_000.0
+		return frac < acceptance(pos)
+	}
+}
+
+// Per-position acceptance rates (probability draft position i is accepted given the
+// prefix survived), defined directly as valid probabilities that decay with depth —
+// the model property both quants share. bf16 and nvfp4 use the same acceptance curve;
+// only their cost curves differ, which is what shifts the optimum. The expected
+// accepted prefix at depth N is Σ_{k<=N} ∏_{i<=k} a_i, which decays as the product
+// shrinks, matching the measured diminishing returns.
+func decayAcceptance(pos int) float64 {
+	rates := []float64{0, 0.90, 0.82, 0.72, 0.60, 0.48, 0.40, 0.34, 0.30, 0.27, 0.25, 0.23, 0.22, 0.21, 0.20, 0.20, 0.20}
+	if pos < 1 || pos >= len(rates) {
+		return 0.18
+	}
+	return rates[pos]
+}
+
+func bf16Acceptance(pos int) float64  { return decayAcceptance(pos) }
+func nvfp4Acceptance(pos int) float64 { return decayAcceptance(pos) }
+
+// bandwidthBoundCost models a near-flat forward (deep optimum); computeBoundCost a
+// steep one (shallow optimum). width = 1 + drafts.
+func bandwidthBoundCost(width int) float64 { return 210.0 + 2.6*float64(width) }
+func computeBoundCost(width int) float64   { return 30.0 + 19.0*float64(width) }
+
+func TestDepthControllerSelectsDeepWhenForwardFlat(t *testing.T) {
+	c := newDepthController()
+	selected := driveDepthController(c, 600, bandwidthBoundCost, deterministicDraw(bf16Acceptance))
+	// A near-flat forward makes extra width nearly free, so committed/cost keeps
+	// rising until acceptance saturates: the EV peak sits deep (~6 for this curve).
+	if selected < 5 {
+		t.Errorf("expected a deep selection on a near-flat-cost forward, got %d", selected)
+	}
+}
+
+func TestDepthControllerSelectsShallowWhenForwardSteep(t *testing.T) {
+	c := newDepthController()
+	selected := driveDepthController(c, 600, computeBoundCost, deterministicDraw(nvfp4Acceptance))
+	// A steep forward makes each extra validated token expensive, so committed/cost
+	// peaks shallow.
+	if selected > 4 {
+		t.Errorf("expected a shallow selection on a steep-cost forward, got %d", selected)
+	}
+}
+
+func TestDepthControllerSelectsZeroWhenDraftNotWorthIt(t *testing.T) {
+	c := newDepthController()
+	cost := c.cost
+	// A steep forward (each extra column ~19ms over a 30ms base) paired with poor
+	// acceptance: even the first draft is accepted under half the time. One draft
+	// costs cost(1)/cost(0) = 49/30 = 1.63x the plain step but commits only
+	// 1 + 0.4 = 1.4 tokens, so plain decode (depth 0) wins. The controller must be able
+	// to stop speculating, not be floored at drafting one token.
+	poor := func(int) float64 { return 0.4 }
+	selected := driveDepthController(c, 600, computeBoundCost, deterministicDraw(poor))
+	if selected != 0 {
+		t.Errorf("expected depth 0 (plain decode) when no draft pays for itself, got %d", selected)
+	}
+	// The depth-0 (plain-decode) cost must actually have been measured, or depth 0 was
+	// chosen on a phantom cost rather than a real comparison.
+	if !cost.sampled(0) {
+		t.Error("depth-0 cost was never sampled, so depth 0 was not a real EV comparison")
+	}
+}
+
+func TestDepthControllerExploresBeforeCostReady(t *testing.T) {
+	c := newDepthController()
+	cost := c.cost
+	// Before any cost data, warmup walks the floor upward from depth 0 to seed depths
+	// from the shallowest end — never jumping deep. The first draft is 0 (depth-0,
+	// plain-decode cost); once that depth is recorded the next is 1. Those are the two
+	// depths the controller first compares at.
+	first := c.next()
+	if cost.ready() {
+		t.Fatal("cost model should not be ready before any observation")
+	}
+	if first != 0 {
+		t.Errorf("first warmup draft = %d, want 0 (depth-0 cost first, never deep)", first)
+	}
+	cost.observe(first, 24*time.Millisecond)
+	second := c.next()
+	if second != 1 {
+		t.Errorf("second warmup draft = %d, want 1 (depth-1 cost next)", second)
+	}
+	// After two distinct depths are observed the model becomes usable.
+	cost.observe(second, 31*time.Millisecond)
+	if !cost.ready() {
+		t.Error("cost model should be ready after two distinct depths")
+	}
+}
+
+func TestDepthControllerClimbsOutwardWithinFrontier(t *testing.T) {
+	// The bug this guards: on a near-flat-cost forward the inherited optimistic
+	// acceptance (1 for unmeasured positions) makes deep depths look best, so an
+	// unbounded argmax jumps straight to a deep depth and drafts a dozen tokens that are
+	// almost all rejected. The frontier bound forbids drafting beyond one past the
+	// deepest position with trusted data, so the depth can only climb outward as
+	// fast as acceptance is actually measured. Assert that invariant holds at every
+	// step, even with always-accept (the most deep-favoring case) and flat cost.
+	c := newDepthController()
+	cost := c.cost
+	always := func(int) float64 { return 1.0 }
+	draw := deterministicDraw(always)
+	for step := range 400 {
+		d := c.next()
+		if d > c.frontier()+1 {
+			t.Fatalf("step %d drafted depth %d, want <= frontier+1 = %d", step, d, c.frontier()+1)
+		}
+		cost.observe(d, time.Duration(bandwidthBoundCost(d+1)*float64(time.Millisecond)))
+		accepted := 0
+		for i := 1; i <= d; i++ {
+			if draw(i, step) {
+				accepted++
+			} else {
+				break
+			}
+		}
+		c.acc.observe(d, accepted)
+	}
+}
+
+func TestCostModelBoundsSingleSampleInfluence(t *testing.T) {
+	m := newCostModel()
+	for range 20 {
+		m.observe(0, 16*time.Millisecond)
+	}
+	base := m.cost(0)
+	// One stall-inflated sample may move the estimate by at most alpha times
+	// the clamp fraction.
+	m.observe(0, 49*time.Millisecond)
+	if bound := base * (1 + costEWMAAlpha*costClampFraction) * 1.001; m.cost(0) > bound {
+		t.Errorf("one outlier moved cost(0) from %.2fms to %.2fms, beyond the clamp bound %.2fms", base, m.cost(0), bound)
+	}
+	// A genuine cost change is rate-limited, not rejected: repeated samples at
+	// the new level must converge to it.
+	for range 60 {
+		m.observe(0, 32*time.Millisecond)
+	}
+	if got := m.cost(0); got < 31 || got > 33 {
+		t.Errorf("estimate did not converge to a persistent cost change: got %.2fms, want ~32ms", got)
+	}
+}
+
+func TestDepthControllerStaysParkedThroughHostStall(t *testing.T) {
+	c := newDepthController()
+	cost := c.cost
+	// A regime where plain decode wins (steep width cost, mediocre acceptance):
+	// the controller settles parked at depth 0.
+	steep := func(width int) float64 { return 16.0 + 15.0*float64(width-1) }
+	mediocre := func(int) float64 { return 0.6 }
+	if selected := driveDepthController(c, 600, steep, deterministicDraw(mediocre)); selected != 0 {
+		t.Fatalf("expected the controller to park at depth 0, got %d", selected)
+	}
+	// A host stall between two parked rounds lands in a single depth-0 cost
+	// sample at ~3x the real round time. Its influence must stay bounded so the
+	// parked baseline remains credible and the controller does not start
+	// drafting against a phantom plain-decode cost.
+	cost.observe(0, 49*time.Millisecond)
+	if selected := c.selected(); selected != 0 {
+		t.Errorf("a single stall-inflated depth-0 sample un-parked the controller: selected %d with cost(0)=%.1fms", selected, cost.cost(0))
+	}
+}
+
+// driveCountingProbes is driveDepthController with a count of the rounds that
+// drafted (the parked regime's probes).
+func driveCountingProbes(c *depthController, n int, fwdMS func(width int) float64, draw func(pos, step int) bool) int {
+	probes := 0
+	for step := range n {
+		d := c.next()
+		if d > 0 {
+			probes++
+		}
+		c.cost.observe(d, time.Duration(fwdMS(d+1)*float64(time.Millisecond)))
+		accepted := 0
+		for i := 1; i <= d; i++ {
+			if draw(i, step) {
+				accepted++
+			} else {
+				break
+			}
+		}
+		c.acc.observe(d, accepted)
+	}
+	return probes
+}
+
+func TestDepthControllerProbeBackoff(t *testing.T) {
+	c := newDepthController()
+	// A regime where plain decode wins: the controller parks and its probes
+	// keep changing nothing, so the cadence must back off to the cap and the
+	// probe count over a long stretch must follow the doubling series, not the
+	// base cadence.
+	steep := func(width int) float64 { return 16.0 + 15.0*float64(width-1) }
+	mediocre := deterministicDraw(func(int) float64 { return 0.6 })
+	driveDepthController(c, 600, steep, mediocre)
+	if c.selected() != 0 {
+		t.Fatalf("expected the controller to park at depth 0, got %d", c.selected())
+	}
+	probes := driveCountingProbes(c, 4096, steep, mediocre)
+	if c.probeInterval != depthProbeIntervalMax {
+		t.Errorf("probe interval = %d, want backed off to %d", c.probeInterval, depthProbeIntervalMax)
+	}
+	if probes > 10 {
+		t.Errorf("%d probes in 4096 parked rounds, want the backed-off handful", probes)
+	}
+
+	// The cadence is the controller's own persistent state, so a fresh stretch
+	// starts at the backed-off interval: a stretch shorter than it never probes.
+	c.probeSince = 0
+	if probes := driveCountingProbes(c, 400, steep, mediocre); probes != 0 {
+		t.Errorf("%d probes within the first backed-off interval, want 0", probes)
+	}
+
+	// When drafting turns worthwhile — cheap wide forwards and high acceptance —
+	// the probes eventually pay off, and the selection change must snap the
+	// cadence back to the base interval.
+	flat := func(width int) float64 { return 16.0 + 0.5*float64(width-1) }
+	good := deterministicDraw(func(int) float64 { return 0.95 })
+	if selected := driveDepthController(c, 8192, flat, good); selected == 0 {
+		t.Fatal("controller never re-engaged after the regime turned draft-favorable")
+	}
+	if c.probeInterval != depthProbeInterval {
+		t.Errorf("probe interval = %d after the selection changed, want reset to %d", c.probeInterval, depthProbeInterval)
+	}
+}
+
+func TestDepthControllerTracksAcceptanceDrift(t *testing.T) {
+	c := newDepthController()
+	// Flat cost throughout, so depth is governed by acceptance alone. First phase:
+	// deep acceptance is poor, so the optimum is shallow. Second phase: acceptance
+	// stays high deep, so the optimum moves deeper. The EWMA acceptance rate must follow the
+	// shift — a cumulative all-time mean would stay anchored to the first phase.
+	shallowFavoring := func(pos int) float64 {
+		if pos <= 2 {
+			return 0.9
+		}
+		return 0.2 // deep rarely accepted
+	}
+	deepFavoring := func(int) float64 { return 0.95 } // deep readily accepted
+	driveDepthController(c, 400, bandwidthBoundCost, deterministicDraw(shallowFavoring))
+	shallow := c.selected()
+	if shallow > 4 {
+		t.Fatalf("expected shallow selection while deep acceptance is poor, got %d", shallow)
+	}
+	// The probe cadence backed off while phase-one probes changed nothing, so
+	// recovery runs at backed-off probe spacing until the first selection
+	// change snaps it back — the drive must span several backed-off intervals.
+	driveDepthController(c, 4096, bandwidthBoundCost, deterministicDraw(deepFavoring))
+	deep := c.selected()
+	if deep <= shallow {
+		t.Errorf("expected selection to deepen after acceptance rose (was %d, now %d)", shallow, deep)
+	}
+}

x/mlxrunner/speculate_stats.go

@@ -0,0 +1,84 @@
+package mlxrunner
+
+import (
+	"context"
+	"fmt"
+	"log/slog"
+	"strings"
+)
+
+type specStats struct {
+	iterations int
+	drafted    int
+	accepted   int
+	maxDraft   int
+	// chosen is the draft depth picked each round, in order; split into time
+	// buckets it distinguishes a ramp that holds from one that thrashes shallow.
+	chosen []int
+}
+
+func (s *specStats) recordRound(depth int) {
+	if !slog.Default().Enabled(context.TODO(), slog.LevelDebug) {
+		return
+	}
+	s.chosen = append(s.chosen, depth)
+}
+
+// depthBuckets is how many equal time slices depthOverTime splits a run into.
+const depthBuckets = 8
+
+// depthOverTime reports per-bucket mean/max chosen depth across up to depthBuckets
+// equal time buckets, e.g. "0.3/1 2.1/3 4.8/5 5.0/6".
+func (s *specStats) depthOverTime() string {
+	if len(s.chosen) == 0 {
+		return ""
+	}
+	buckets := min(depthBuckets, len(s.chosen))
+	parts := make([]string, 0, buckets)
+	for b := range buckets {
+		lo := b * len(s.chosen) / buckets
+		hi := (b + 1) * len(s.chosen) / buckets
+		sum, mx := 0, 0
+		for _, d := range s.chosen[lo:hi] {
+			sum += d
+			mx = max(mx, d)
+		}
+		parts = append(parts, fmt.Sprintf("%.1f/%d", float64(sum)/float64(hi-lo), mx))
+	}
+	return strings.Join(parts, " ")
+}
+
+func (s *speculationSession) logStats() {
+	if !s.enabled || !slog.Default().Enabled(context.TODO(), slog.LevelDebug) {
+		return
+	}
+	acceptance := 0.0
+	if s.stats.drafted > 0 {
+		acceptance = float64(s.stats.accepted) / float64(s.stats.drafted)
+	}
+	avgDraft := 0.0
+	avgAccepted := 0.0
+	if s.stats.iterations > 0 {
+		avgDraft = float64(s.stats.drafted) / float64(s.stats.iterations)
+		avgAccepted = float64(s.stats.accepted) / float64(s.stats.iterations)
+	}
+	slog.Debug("speculative decode stats", "iterations", s.stats.iterations, "drafted", s.stats.drafted, "accepted", s.stats.accepted, "acceptance", fmt.Sprintf("%.2f", acceptance), "avg_draft", fmt.Sprintf("%.2f", avgDraft), "max_draft", s.stats.maxDraft, "avg_accepted", fmt.Sprintf("%.2f", avgAccepted), "depth_over_time", s.stats.depthOverTime())
+
+	// Log learned acceptance over the trusted positions [1, frontier] and
+	// expected throughput over the searched window [0, frontier+1]; deeper
+	// depths have no data of their own.
+	d := s.spec.depth
+	frontier := d.frontier()
+	rates := make([]string, 0, frontier)
+	for n := 1; n <= frontier; n++ {
+		rates = append(rates, fmt.Sprintf("%d:%.2f", n, d.acc.acceptance(n)))
+	}
+	limit := frontier + 1
+	tps := make([]string, 0, limit+1)
+	if d.cost.ready() {
+		for n := 0; n <= limit; n++ {
+			tps = append(tps, fmt.Sprintf("%d:%.1f", n, 1000*d.acc.expectedCommitted(n)/d.cost.cost(n)))
+		}
+	}
+	slog.Debug("speculation depth controller", "cost", d.cost.sampleString(), "acceptance", strings.Join(rates, " "), "expected_tps", strings.Join(tps, " "), "probe_interval", d.probeInterval)
+}

x/models/gemma4/gemma4.go

@@ -28,10 +28,7 @@ func init() {
 }
 
 // Compile-time interface checks.
-var (
-	_ base.Model               = (*Model)(nil)
-	_ base.MTPDefaultsProvider = (*Model)(nil)
-)
+var _ base.Model = (*Model)(nil)
 
 // RopeParams holds per-layer-type RoPE settings.
 type RopeParams struct {
@@ -489,24 +486,6 @@ func (m *Model) EnableCompile() bool {
 	return true
 }
 
-func (m *Model) MTPDraftDefaults(_ bool) base.MTPDefaults {
-	defaults := base.MTPDefaults{
-		InitialDraftTokens: 4,
-		MaxDraftTokens:     16,
-		Enabled:            true,
-	}
-	if m == nil || m.TextConfig == nil {
-		return defaults
-	}
-	switch {
-	case !m.EnableMoeBlock && m.HiddenSize == 5376 && m.NumHiddenLayers == 60:
-		defaults.InitialDraftTokens = 14
-	case m.EnableMoeBlock && m.HiddenSize == 2816 && m.NumHiddenLayers == 30:
-		defaults.InitialDraftTokens = 8
-	}
-	return defaults
-}
-
 func resolveWeightPrefix(tensors map[string]*mlx.Array) string {
 	for _, prefix := range []string{"", "language_model.", "model.language_model."} {
 		if tensors[prefix+"embed_tokens.weight"] != nil {

x/models/gemma4/gemma4_test.go

@@ -561,54 +561,6 @@ func TestLayerTypeDetection(t *testing.T) {
 	}
 }
 
-func TestMTPDraftDefaults(t *testing.T) {
-	tests := []struct {
-		name        string
-		cfg         *TextConfig
-		wantInitial int
-		wantMax     int
-	}{
-		{
-			name:        "nil config",
-			wantInitial: 4,
-			wantMax:     16,
-		},
-		{
-			name:        "31b bf16",
-			cfg:         &TextConfig{HiddenSize: 5376, NumHiddenLayers: 60},
-			wantInitial: 14,
-			wantMax:     16,
-		},
-		{
-			name:        "31b quantized",
-			cfg:         &TextConfig{HiddenSize: 5376, NumHiddenLayers: 60, QuantBits: 4},
-			wantInitial: 14,
-			wantMax:     16,
-		},
-		{
-			name:        "26b-a4b moe",
-			cfg:         &TextConfig{HiddenSize: 2816, NumHiddenLayers: 30, EnableMoeBlock: true},
-			wantInitial: 8,
-			wantMax:     16,
-		},
-		{
-			name:        "generic default",
-			cfg:         &TextConfig{HiddenSize: 2560, NumHiddenLayers: 42, HiddenSizePerLayer: 256},
-			wantInitial: 4,
-			wantMax:     16,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			got := (&Model{TextConfig: tt.cfg}).MTPDraftDefaults(false)
-			if got.InitialDraftTokens != tt.wantInitial || got.MaxDraftTokens != tt.wantMax || !got.Enabled {
-				t.Fatalf("MTPDraftDefaults() = %+v, want initial=%d max=%d enabled=true", got, tt.wantInitial, tt.wantMax)
-			}
-		})
-	}
-}
-
 func TestNewCachesOmitsSharedKVLayers(t *testing.T) {
 	m := &Model{
 		Layers: []*DecoderLayer{