mlxrunner: construct per-model state at load

The caches, the speculation binding, and the drafter were each built
lazily inside the first request: begin constructed the caches, and open
bound the cache partition and made a fresh drafter every time. All of
it is a property of the loaded model, so build it once at load.
newPrefixCache replaces the lazy construction in begin, speculation
binds when it is created, and the drafter splits the way speculation
does: a persistent mtpDrafter constructed at load opens each request's
mtpDraftSession, whose constructor syncs the pairing cursor to the
draft caches' restored offset.

x/mlxrunner/mtp.go

@@ -13,12 +13,35 @@ import (
 // regardless of what else triggers a flush.
 const mtpPendingFlushTokens = 32
 
-// 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
-// completes only when the next token arrives.
+// mtpDrafter drafts with a model's multi-token-prediction head. Constructed
+// at load, it opens each request's drafting session.
 type mtpDrafter struct {
 	spec *speculation
+}
+
+func newMTPDrafter(s *speculation) *mtpDrafter {
+	return &mtpDrafter{spec: s}
+}
+
+// open returns the drafting session for one request, its pairing frontier
+// synced to the draft caches' restored offset.
+func (d *mtpDrafter) open() *mtpDraftSession {
+	s := &mtpDraftSession{drafter: d}
+	if kv := d.spec.draftKV; len(kv) > 0 {
+		// A restored prefix arrives with the draft caches already written;
+		// pairing resumes from their absolute offset.
+		s.committedDraftOffset = kv[0].Offset()
+		s.frontier = s.committedDraftOffset
+	}
+	return s
+}
+
+// mtpDraftSession runs one request's drafting, 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
+// completes only when the next token arrives.
+type mtpDraftSession struct {
+	drafter *mtpDrafter
 
 	// frontier is the slot after the last reported token; frontierHidden is
 	// the pinned target hidden at frontier-1, fused into the next pair.
@@ -41,22 +64,9 @@ type mtpDrafter struct {
 	heldProjected *mlx.Array
 }
 
-// newMTPDrafter returns the MTP drafter cursor for this request, syncing its
-// pairing frontier to the draft caches' restored offset.
-func newMTPDrafter(s *speculation) *mtpDrafter {
-	d := &mtpDrafter{spec: s}
-	if len(s.draftKV) > 0 {
-		// A restored prefix arrives with the draft caches already written;
-		// pairing resumes from their absolute offset.
-		d.committedDraftOffset = s.draftKV[0].Offset()
-		d.frontier = d.committedDraftOffset
-	}
-	return d
-}
-
-func (d *mtpDrafter) committed(tokens, hiddens *mlx.Array, position int) {
+func (d *mtpDraftSession) committed(tokens, hiddens *mlx.Array, position int) {
 	n := tokens.Dim(1)
-	if len(d.spec.draftKV) > 0 {
+	if len(d.drafter.spec.draftKV) > 0 {
 		// The pair at slot S fuses token[S+1] with hidden[S], so a run pairs its
 		// tokens with its own hiddens shifted one slot back: the first writable
 		// token takes the carried frontier hidden, each later token the row
@@ -96,22 +106,22 @@ func (d *mtpDrafter) committed(tokens, hiddens *mlx.Array, position int) {
 // same rather than level here. Regenerating hidden[S] to re-pair the boundary
 // on the next request is a separate, re-prefill-bound concern for recurrent
 // targets.
-func (d *mtpDrafter) finish(current *mlx.Array) {
-	if len(d.spec.draftKV) == 0 {
+func (d *mtpDraftSession) finish(current *mlx.Array) {
+	if len(d.drafter.spec.draftKV) == 0 {
 		return
 	}
 	d.settle(current)
 }
 
 // settle completes any open frontier pair with current, then flushes.
-func (d *mtpDrafter) settle(current *mlx.Array) {
+func (d *mtpDraftSession) settle(current *mlx.Array) {
 	if d.frontierHidden != nil && d.frontier-1 == d.committedDraftOffset+d.pendingCount {
 		d.queueCacheWrites(current.ExpandDims(-1), d.frontierHidden)
 	}
 	d.flush()
 }
 
-func (d *mtpDrafter) close() {
+func (d *mtpDraftSession) close() {
 	d.flush()
 	d.setFrontierHidden(nil)
 	d.setHeld(nil, nil)
@@ -121,7 +131,7 @@ func (d *mtpDrafter) close() {
 // fused with their target hiddens — flushing once the buffer reaches the token
 // cap so the pinned hiddens stay bounded. flush coalesces the buffered writes
 // into one head forward, so a contiguous run lands in a single draft-cache extend.
-func (d *mtpDrafter) queueCacheWrites(tokens, hiddens *mlx.Array) {
+func (d *mtpDraftSession) queueCacheWrites(tokens, hiddens *mlx.Array) {
 	mlx.Pin(tokens, hiddens)
 	d.pendingTokens = append(d.pendingTokens, tokens)
 	d.pendingHiddens = append(d.pendingHiddens, hiddens)
@@ -134,11 +144,12 @@ func (d *mtpDrafter) queueCacheWrites(tokens, hiddens *mlx.Array) {
 // flush writes the pending pairs to the draft caches in one head forward,
 // dropping speculative entries past the committed range first and holding
 // the last row's logits and projected hidden for the next proposal chain.
-func (d *mtpDrafter) flush() {
+func (d *mtpDraftSession) flush() {
 	if len(d.pendingTokens) == 0 {
 		return
 	}
-	for _, c := range d.spec.draftKV {
+	spec := d.drafter.spec
+	for _, c := range spec.draftKV {
 		if c.Offset() > d.committedDraftOffset {
 			if !c.Restore(nil, d.committedDraftOffset) {
 				panic(fmt.Sprintf("mtp: draft cache rewind to %d failed", d.committedDraftOffset))
@@ -148,19 +159,19 @@ func (d *mtpDrafter) flush() {
 
 	ids := mlx.Concatenate(d.pendingTokens, 1)
 	hiddens := mlx.Concatenate(d.pendingHiddens, 1)
-	hidden, projected := d.spec.draft.Draft(&batch.Batch{
+	hidden, projected := spec.draft.Draft(&batch.Batch{
 		InputIDs:     ids,
 		SeqOffsets:   []int32{int32(d.committedDraftOffset)},
 		SeqQueryLens: []int32{int32(ids.Dim(1))},
 		Hidden:       hiddens,
-	}, d.spec.caches)
+	}, spec.caches)
 	d.setHeld(lastHiddenRow(hidden), lastHiddenRow(projected))
 	d.committedDraftOffset += ids.Dim(1)
 
 	// Force the draft writes: a session that never drafts would otherwise
 	// leave the flush chain unevaluated, pinning every hidden until close.
-	state := make([]*mlx.Array, 0, 2*len(d.spec.draftKV))
-	for _, c := range d.spec.draftKV {
+	state := make([]*mlx.Array, 0, 2*len(spec.draftKV))
+	for _, c := range spec.draftKV {
 		state = append(state, c.State()...)
 	}
 	mlx.AsyncEval(state...)
@@ -171,14 +182,14 @@ func (d *mtpDrafter) flush() {
 	d.pendingCount = 0
 }
 
-func (d *mtpDrafter) setFrontierHidden(h *mlx.Array) {
+func (d *mtpDraftSession) setFrontierHidden(h *mlx.Array) {
 	mlx.Pin(h)
 	mlx.Unpin(d.frontierHidden)
 	d.frontierHidden = h
 }
 
 // setHeld replaces the held flush outputs, pinned until the next flush or close.
-func (d *mtpDrafter) setHeld(hidden, projected *mlx.Array) {
+func (d *mtpDraftSession) setHeld(hidden, projected *mlx.Array) {
 	mlx.Pin(hidden, projected)
 	mlx.Unpin(d.heldHidden, d.heldProjected)
 	d.heldHidden, d.heldProjected = hidden, projected
@@ -188,13 +199,14 @@ func (d *mtpDrafter) setHeld(hidden, projected *mlx.Array) {
 // A head with draft caches settles the frontier pair first, so its first step
 // reuses the held frontier row with no head call; a cacheless head re-attends
 // the target caches read-only, anchored at the last committed slot.
-func (d *mtpDrafter) propose(current *mlx.Array, maxTokens int) *draftCandidates {
+func (d *mtpDraftSession) propose(current *mlx.Array, maxTokens int) *draftCandidates {
 	if maxTokens <= 0 || d.frontierHidden == nil {
 		return nil
 	}
-	r := d.spec.r
+	spec := d.drafter.spec
+	r := spec.r
 
-	if len(d.spec.draftKV) > 0 {
+	if len(spec.draftKV) > 0 {
 		d.settle(current)
 		if d.heldHidden == nil {
 			return nil
@@ -208,7 +220,7 @@ func (d *mtpDrafter) propose(current *mlx.Array, maxTokens int) *draftCandidates
 
 	for i := range maxTokens {
 		var hidden, projected *mlx.Array
-		if i == 0 && len(d.spec.draftKV) > 0 {
+		if i == 0 && len(spec.draftKV) > 0 {
 			// The settle flush already produced the frontier row; reuse it
 			// instead of re-running the head.
 			hidden, projected = d.heldHidden, d.heldProjected
@@ -219,18 +231,18 @@ func (d *mtpDrafter) propose(current *mlx.Array, maxTokens int) *draftCandidates
 			// head stays at the last committed slot every step, re-attending
 			// the committed prefix read-only ("single-position").
 			pos := d.frontier - 1
-			if len(d.spec.draftKV) > 0 {
+			if len(spec.draftKV) > 0 {
 				pos = d.frontier - 1 + i
 			}
-			hidden, projected = d.spec.draft.Draft(&batch.Batch{
+			hidden, projected = spec.draft.Draft(&batch.Batch{
 				InputIDs:     lastToken,
 				SeqOffsets:   []int32{int32(pos)},
 				SeqQueryLens: []int32{1},
 				Hidden:       lastHidden,
-			}, d.spec.caches)
+			}, spec.caches)
 		}
 		// Unembed only the row being sampled, never the batch.
-		stepLogits := d.spec.draft.Unembed(hidden).Squeeze(1)
+		stepLogits := spec.draft.Unembed(hidden).Squeeze(1)
 		lastHidden = projected
 		// The chain's earlier drafts ride along as the row's history, so
 		// penalties shape proposals the same way they shape validation.

x/mlxrunner/mtp_test.go

@@ -221,6 +221,7 @@ func mtpTestRunner(t *testing.T, predict map[int32]int32, eos []int32, opts samp
 		Model:     &fakeMTPModel{predict: predict, tok: tok},
 		Tokenizer: tok,
 		Sampler:   sampler.New(4096),
+		cache:     &prefixCache{},
 	}
 	r.Sampler.Add(pipelineSlot, opts, nil)
 	t.Cleanup(func() { r.Sampler.Remove(pipelineSlot) })
@@ -377,9 +378,9 @@ func TestRunMTPDecodeGreedy(t *testing.T) {
 	r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
 	// The draft mirrors the target chain, so every drafted token is accepted.
 	draft := &fakeMTPDraft{predict: predict}
-	r.spec = newSpeculation(r, draft)
-
 	caches, _ := newMTPTestCaches(1)
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, draft)
 	session, ch := newMTPTestSession(caches)
 	position := 1 // one prefill token already processed
 
@@ -438,9 +439,9 @@ func TestRunMTPDecodeSampled(t *testing.T) {
 	const eos int32 = 7
 	predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: eos, eos: 0}
 	r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{Temperature: 1, Seed: 42, UseSeed: true})
-	r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict})
-
 	caches, _ := newMTPTestCaches(1)
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict})
 	session, ch := newMTPTestSession(caches)
 	position := 1
 
@@ -482,9 +483,9 @@ func TestRunMTPDecodeWarmDrafter(t *testing.T) {
 	predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: eos, eos: 0}
 	r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
 	draft := &fakeMTPDraft{predict: predict}
-	r.spec = newSpeculation(r, draft)
-
 	caches, _ := newMTPTestCaches(1)
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, draft)
 	session, ch := newMTPTestSession(caches)
 	position := 1 // one prefill token already processed
 
@@ -543,9 +544,9 @@ func TestRunMTPDecodeEOSCutLeavesPositionsUnjudged(t *testing.T) {
 	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)
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict})
 	session, ch := newMTPTestSession(caches)
 	position := 1
 
@@ -648,9 +649,9 @@ func TestDecodeCancelledMidStream(t *testing.T) {
 	const eos int32 = 7
 	predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: eos}
 	r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
-	r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict})
-
 	caches, tr := newMTPTestCaches(1)
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict})
 	session := &cacheSession{caches: caches}
 	ch := make(chan CompletionResponse) // unbuffered: every send must rendezvous
 
@@ -727,9 +728,9 @@ func TestDecodeKVDraft(t *testing.T) {
 	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) // caches[0] target, caches[1] draft KV
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, draft)
 	session, ch := newMTPTestSession(caches)
 	position := 0
 
@@ -811,9 +812,9 @@ func TestDecodeKVDraftRejectionRebuildsFromTarget(t *testing.T) {
 	// the target chain); once the target corrects 3->4 the next proposal
 	// re-aligns on the shared chain.
 	draft := &fakeKVDraft{predict: map[int32]int32{1: 2, 2: 3, 3: 6, 6: 0, 4: 5, 5: eos, eos: 0}}
-	r.spec = newSpeculation(r, draft)
-
 	caches, _ := newMTPTestCaches(2) // caches[0] target, caches[1] draft KV
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, draft)
 	session, ch := newMTPTestSession(caches)
 	position := 0
 
@@ -880,9 +881,9 @@ func TestDecodeMaintainsDraftCacheWithoutDrafting(t *testing.T) {
 	opts := sampler.Options{Logprobs: true}
 	r := mtpTestRunner(t, predict, []int32{eos}, opts)
 	draft := &fakeKVDraft{predict: predict}
-	r.spec = newSpeculation(r, draft)
-
 	caches, _ := newMTPTestCaches(2)
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, draft)
 	session, ch := newMTPTestSession(caches)
 	position := 0
 
@@ -943,9 +944,9 @@ func TestFlushLevelsDraftCacheWithPrefill(t *testing.T) {
 	predict := map[int32]int32{2: 3, 3: 4, 4: 5}
 	r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
 	draft := &fakeKVDraft{predict: predict}
-	r.spec = newSpeculation(r, draft)
-
 	caches, _ := newMTPTestCaches(2)
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, draft)
 	req := Request{
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
@@ -983,9 +984,9 @@ func TestCommittedRunBatchesPastFlushCap(t *testing.T) {
 
 	r := mtpTestRunner(t, predict, []int32{0}, sampler.Options{})
 	draft := &fakeKVDraft{predict: predict}
-	r.spec = newSpeculation(r, draft)
-
 	caches, _ := newMTPTestCaches(2)
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, draft)
 	req := Request{
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
@@ -1021,9 +1022,9 @@ func TestDecodeParkedDraftResume(t *testing.T) {
 	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)
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, draft)
 	req := Request{
 		Tokens:            []int32{1},
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
@@ -1123,7 +1124,7 @@ func newMTPTestSession(caches []cache.Cache) (*cacheSession, chan CompletionResp
 func testSpeculationSession(r *Runner, caches []cache.Cache) *speculationSession {
 	if r.spec != nil {
 		r.spec.bind(caches)
-		return &speculationSession{spec: r.spec, drafter: newMTPDrafter(r.spec)}
+		return &speculationSession{spec: r.spec, drafter: r.spec.drafter.open()}
 	}
 	s := &speculation{r: r, caches: caches, targets: caches}
 	return &speculationSession{spec: s, drafter: nopDrafter{}}
@@ -1151,7 +1152,7 @@ func scriptedCandidates(r *Runner, tokens []int32) *draftCandidates {
 		prev = tok
 	}
 	s := &speculation{r: r, draft: &fakeMTPDraft{predict: chain}}
-	d := &mtpDrafter{spec: s}
+	d := (&mtpDrafter{spec: s}).open()
 	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))
 }

x/mlxrunner/pipeline.go

@@ -71,7 +71,7 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er
 
 	inputs := request.Tokens
 
-	session := r.cache.begin(r.Model, inputs)
+	session := r.cache.begin(inputs)
 	defer session.close()
 	caches := session.caches
 

x/mlxrunner/prefix_cache.go

@@ -62,18 +62,17 @@ type cacheSession struct {
 	pendingSnapshots []pendingSnapshot
 }
 
-func (c *prefixCache) ensureCaches(m base.Model) {
-	if len(c.caches) != 0 {
-		return
-	}
+func newPrefixCache(m base.Model) *prefixCache {
+	c := &prefixCache{}
 	if cacheFactory, ok := m.(interface{ NewCaches() []cache.Cache }); ok {
 		c.caches = cacheFactory.NewCaches()
-		return
+		return c
 	}
 	c.caches = make([]cache.Cache, m.NumLayers())
 	for i := range c.caches {
 		c.caches[i] = cache.NewKVCache()
 	}
+	return c
 }
 
 func (c *prefixCache) ensureRoot() {
@@ -87,8 +86,7 @@ func (c *prefixCache) ensureRoot() {
 
 // begin prepares caches for a new request. It finds the nearest
 // matching cache or creates new caches if none match.
-func (c *prefixCache) begin(m base.Model, inputs []int32) *cacheSession {
-	c.ensureCaches(m)
+func (c *prefixCache) begin(inputs []int32) *cacheSession {
 	c.ensureRoot()
 
 	matchPath, matched := findBestMatch(c.root, inputs)

x/mlxrunner/prefix_cache_test.go

@@ -479,7 +479,7 @@ type requestResult struct {
 func simulateRequest(t *testing.T, pc *prefixCache, inputs, generated []int32, userSnapshotAt ...int) requestResult {
 	t.Helper()
 
-	session := pc.begin(nil, inputs)
+	session := pc.begin(inputs)
 	var snapshotOffsets []int
 	for _, at := range userSnapshotAt {
 		if at > 0 {
@@ -923,7 +923,7 @@ func TestSnapshotBeyondPrefillSkipped(t *testing.T) {
 		pc := env.pc
 		inputs := []int32{1, 2, 3, 4, 5}
 
-		session := pc.begin(nil, inputs)
+		session := pc.begin(inputs)
 		// Request a reachable snapshot at 3 and one at len(inputs), which a
 		// prefill that stops one token short never crosses.
 		session.schedulePrefillSnapshots([]int{3, len(inputs)})
@@ -953,7 +953,7 @@ func TestPrefillSnapshotsDiscardedOnCancel(t *testing.T) {
 		pc := env.pc
 		inputs := []int32{1, 2, 3, 4, 5}
 
-		session := pc.begin(nil, inputs)
+		session := pc.begin(inputs)
 		session.schedulePrefillSnapshots([]int{3})
 		// Cross offset 3 so the caches capture it, then close the session as a
 		// canceled prefill would, before the captures are attached to the trie.

x/mlxrunner/runner.go

@@ -38,7 +38,7 @@ type Runner struct {
 	Tokenizer     *tokenizer.Tokenizer
 	Requests      chan Request
 	Sampler       *sample.Sampler
-	cache         prefixCache
+	cache         *prefixCache
 	contextLength int
 	mlxThread     *mlxthread.Thread
 	// spec is the speculative-decoding subsystem. Nil when the model ships no
@@ -105,6 +105,7 @@ func (r *Runner) Load(modelName string) error {
 	r.Model = m
 	r.Tokenizer = m.Tokenizer()
 	r.contextLength = m.MaxContextLength()
+	r.cache = newPrefixCache(m)
 	r.Sampler = sample.New(r.contextLength)
 	r.spec = newSpeculation(r, draftModel)
 

x/mlxrunner/speculate.go

@@ -56,6 +56,10 @@ type speculation struct {
 	draftKV []cache.Cache
 	targets []cache.Cache
 
+	// drafter is the persistent half of the MTP drafting machinery; each
+	// request's session comes from drafter.open.
+	drafter *mtpDrafter
+
 	// depth selects each request's draft length and owns the cost/acceptance
 	// models and probe cadence it learns across requests.
 	depth *depthController
@@ -67,7 +71,10 @@ func newSpeculation(r *Runner, draft base.DraftModel) *speculation {
 	if draft == nil {
 		return nil
 	}
-	return &speculation{r: r, draft: draft, depth: newDepthController()}
+	s := &speculation{r: r, draft: draft, depth: newDepthController()}
+	s.bind(r.cache.caches)
+	s.drafter = newMTPDrafter(s)
+	return s
 }
 
 // bind computes the draft/target cache partition the first time the persistent
@@ -123,7 +130,7 @@ func (s *speculation) open(request Request, caches []cache.Cache) *speculationSe
 		return nil
 	}
 	s.bind(caches)
-	d := newMTPDrafter(s)
+	d := s.drafter.open()
 
 	// Logprobs are not yet supported, so a logprobs request keeps a speculationSession
 	// only to maintain a draft cache in lockstep (permanently parked).