mlxrunner: feed media features to the model during prefill

Each media item's features are encoded lazily when a prefill chunk
first overlaps its expansion and stay pinned until the expansion is
fully evaluated. A chunk never ends strictly inside an atomic
expansion: a bidirectional run's early rows attend its later keys, so
its first evaluation must cover the whole run in one forward. Items
marked Causal are exempt and split at any boundary.

Draft models need the same request state — reference MTP drafters
embed prompt tokens with the image features merged in, and an M-RoPE
drafter cannot compute positions without the request's layout — so the
layout is stamped on every forward, target and draft alike, and the
MTP session holds feature rows across its deferred flush. The dflash
drafter ignores media: its context rows are target hiddens.

x/mlxrunner/batch/batch.go

@@ -21,11 +21,39 @@ type Batch struct {
 	// It is nil for ordinary forward passes.
 	Hidden *mlx.Array
 
+	// Media lists a row's media items, on prefill forwards and draft
+	// forwards that embed prompt tokens; items outside the query range
+	// ride featureless. Nil at decode and for text-only requests.
+	Media []MediaItem
+
+	// Layout carries each row's opaque layout state from PrepareMedia,
+	// identical on every forward of the request; the runner never reads
+	// it. Nil entries derive nothing from layout.
+	Layout []any
+
 	// Memo is per-forward memoization used to cache results, such as masks,
 	// which are often the same across layers.
 	Memo Memo
 }
 
+// MediaItem is one media occurrence in a row's sequence. The runner
+// knows only where the expansion was spliced; which positions bear
+// features is the model's, derived from Pos and Opaque.
+type MediaItem struct {
+	// Seq is the batch row the item belongs to.
+	Seq int
+
+	// Pos is the absolute sequence position of the expansion's first token.
+	Pos int
+
+	// Features is the item's whole feature-row array, attached only while
+	// the item's token range overlaps this forward's query range.
+	Features *mlx.Array
+
+	// Opaque is the item's PreparedMedia.Opaque, round-tripped untouched.
+	Opaque any
+}
+
 type Memo struct {
 	entries map[any]any
 }

x/mlxrunner/dflash.go

@@ -28,8 +28,8 @@ func newDFlashDrafter(s *speculation, draft base.BlockDraft) *dflashDrafter {
 func (d *dflashDrafter) draftLimit() int { return d.blockSize - 1 }
 
 // open returns a session synced to the draft caches' restored offset.
-func (d *dflashDrafter) open() draftSession {
-	s := &dflashDraftSession{drafter: d}
+func (d *dflashDrafter) open(layout []any) draftSession {
+	s := &dflashDraftSession{drafter: d, layout: layout}
 	if kv := d.spec.draftKV; len(kv) > 0 {
 		s.ctxOffset = kv[0].Offset()
 	}
@@ -41,6 +41,7 @@ func (d *dflashDrafter) open() draftSession {
 // after the last reported token.
 type dflashDraftSession struct {
 	drafter *dflashDrafter
+	layout  []any
 
 	// ctxOffset is the slot after the last feature row written; pendingCount
 	// rows are buffered past it.
@@ -53,7 +54,9 @@ type dflashDraftSession struct {
 	blockOutstanding bool
 }
 
-func (d *dflashDraftSession) committed(tokens, features *mlx.Array, position int) {
+// committed ignores the media manifest: a context row derives from the
+// target hidden at its slot, which already carries any image content.
+func (d *dflashDraftSession) committed(tokens, features *mlx.Array, position int, _ []batch.MediaItem) {
 	n := tokens.Dim(1)
 	// Skip leading rows the session already has (a restored prefix). A run
 	// that starts past the frontier would leave a gap, which is a bug.
@@ -121,6 +124,7 @@ func (d *dflashDraftSession) flush() {
 	spec.draft.Forward(&batch.Batch{
 		SeqOffsets: []int32{int32(offset)},
 		Hidden:     features,
+		Layout:     d.layout,
 	}, spec.targets, spec.draftKV)
 
 	// Force the cache writes: a session that never drafts would otherwise
@@ -165,6 +169,7 @@ func (d *dflashDraftSession) propose(current *mlx.Array, maxTokens int) *draftCa
 		SeqOffsets:   []int32{int32(offset)},
 		SeqQueryLens: []int32{int32(n + 1)},
 		Hidden:       features,
+		Layout:       d.layout,
 	}, spec.targets, spec.draftKV)
 
 	// Row i predicts the token at its own position, so the anchor row is

x/mlxrunner/dflash_test.go

@@ -103,7 +103,7 @@ func newBlockTestSession(t *testing.T, predict map[int32]int32, blockSize int) (
 	draft := &fakeBlockDraft{predict: predict, blockSize: blockSize, maskToken: 6, draftCaches: caches[1:]}
 	r.cache.caches = caches
 	r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
-	return r, draft, r.spec.drafter.open().(*dflashDraftSession), caches
+	return r, draft, r.spec.drafter.open(nil).(*dflashDraftSession), caches
 }
 
 // draftTokensOf reads the draft cache's fed token stream.
@@ -122,7 +122,7 @@ func TestDFlashCommittedBuffersPastFlushCap(t *testing.T) {
 	for i := range ids {
 		ids[i] = int32(i % mtpTestVocab)
 	}
-	session.committed(mlx.FromValues(ids, 1, n), oneHotLogits(ids), 0)
+	session.committed(mlx.FromValues(ids, 1, n), oneHotLogits(ids), 0, nil)
 	if got := len(draft.calls); got != 1 {
 		t.Fatalf("draft calls after cap-sized run = %d, want 1", got)
 	}
@@ -133,7 +133,7 @@ func TestDFlashCommittedBuffersPastFlushCap(t *testing.T) {
 	// A run below the cap only buffers; settle writes it through, skipping
 	// the leading rows the flush already covered.
 	tail := []int32{1, 2, 3}
-	session.committed(mlx.FromValues(tail, 1, 3), oneHotLogits(tail), n-1)
+	session.committed(mlx.FromValues(tail, 1, 3), oneHotLogits(tail), n-1, nil)
 	if got := len(draft.calls); got != 1 {
 		t.Fatalf("draft calls after buffered run = %d, want 1 (buffered)", got)
 	}
@@ -151,14 +151,14 @@ func TestDFlashCommittedGapPanics(t *testing.T) {
 	skipIfNoMLX(t)
 	_, _, session, _ := newBlockTestSession(t, nil, 4)
 
-	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0)
+	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
 	defer func() {
 		if recover() == nil {
 			t.Fatalf("committed run past the frontier did not panic")
 		}
 	}()
 	// The frontier is at slot 1; a run starting at 3 leaves slot 1..2 unfed.
-	session.committed(mlx.FromValues([]int32{4}, 1, 1), oneHotLogits([]int32{4}), 3)
+	session.committed(mlx.FromValues([]int32{4}, 1, 1), oneHotLogits([]int32{4}), 3, nil)
 }
 
 func TestDFlashRestoredPrefixResumes(t *testing.T) {
@@ -172,7 +172,7 @@ func TestDFlashRestoredPrefixResumes(t *testing.T) {
 	// A restored prefix arrives with the draft caches already written.
 	restored := []int32{1, 2, 3, 4, 5}
 	caches[1].(*fakeRewindableCache).feed(restored)
-	session := r.spec.drafter.open().(*dflashDraftSession)
+	session := r.spec.drafter.open(nil).(*dflashDraftSession)
 	if session.ctxOffset != len(restored) {
 		t.Fatalf("ctxOffset = %d, want %d (synced to restored offset)", session.ctxOffset, len(restored))
 	}
@@ -180,7 +180,7 @@ func TestDFlashRestoredPrefixResumes(t *testing.T) {
 	// The resumed prefill's run overlaps the restore point; only the rows
 	// past the frontier are buffered and written.
 	run := []int32{2, 3, 0, 1}
-	session.committed(mlx.FromValues(run, 1, 4), oneHotLogits(run), 3)
+	session.committed(mlx.FromValues(run, 1, 4), oneHotLogits(run), 3, nil)
 	session.settle(nil)
 	want := blockCall{offset: 5, ctx: []int32{0, 1}}
 	if got := draft.calls[0]; got.offset != want.offset || !slices.Equal(got.ctx, want.ctx) || got.block != nil {
@@ -201,7 +201,7 @@ func TestDFlashProposeBounds(t *testing.T) {
 	if session.propose(current, 4) != nil {
 		t.Fatalf("propose with no context did not decline")
 	}
-	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0)
+	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
 	if session.propose(current, 0) != nil {
 		t.Fatalf("propose with no budget did not decline")
 	}
@@ -225,7 +225,7 @@ func TestDFlashBlockRewoundBeforeContextWrites(t *testing.T) {
 	predict := map[int32]int32{2: 3, 3: 4, 4: 5}
 	_, draft, session, caches := newBlockTestSession(t, predict, 4)
 
-	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0)
+	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
 	if session.propose(mlx.FromValues([]int32{2}, 1), 3) == nil {
 		t.Fatalf("propose declined")
 	}
@@ -237,7 +237,7 @@ func TestDFlashBlockRewoundBeforeContextWrites(t *testing.T) {
 	// The next round's report rewinds the block before appending context, so
 	// the accepted tokens' rows land at their true slots.
 	run := []int32{2, 3, 4}
-	session.committed(mlx.FromValues(run, 1, 3), oneHotLogits(run), 1)
+	session.committed(mlx.FromValues(run, 1, 3), oneHotLogits(run), 1, nil)
 	session.settle(nil)
 	if got, want := draftTokensOf(caches), []int32{1, 2, 3, 4}; !slices.Equal(got, want) {
 		t.Fatalf("draft cache after settle = %v, want %v (block rewound)", got, want)
@@ -256,7 +256,7 @@ func TestDFlashCloseDrainsOutstandingBlock(t *testing.T) {
 	predict := map[int32]int32{2: 3, 3: 4, 4: 5}
 	_, _, session, caches := newBlockTestSession(t, predict, 4)
 
-	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0)
+	session.committed(mlx.FromValues([]int32{1}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
 	if session.propose(mlx.FromValues([]int32{2}, 1), 3) == nil {
 		t.Fatalf("propose declined")
 	}
@@ -287,7 +287,7 @@ func TestDecodeBlockDraft(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	spec := r.spec.open(req)
+	spec := r.spec.open(req, nil)
 	if spec == nil || !spec.enabled {
 		t.Fatalf("open rejected a block-draft request")
 	}

x/mlxrunner/media.go

@@ -10,12 +10,16 @@ import (
 	"strconv"
 
 	"github.com/ollama/ollama/llm"
+	"github.com/ollama/ollama/x/mlxrunner/batch"
+	"github.com/ollama/ollama/x/mlxrunner/mlx"
 	"github.com/ollama/ollama/x/mlxrunner/model/base"
 )
 
 var imgTagPattern = regexp.MustCompile(`\[img-(\d+)\]`)
 
-// mediaItem is one media occurrence in a request's token stream.
+// mediaItem is one media occurrence in a request's token stream: the
+// absolute position and length of its placeholder expansion, its trie-key
+// fold value, and the prepared item.
 type mediaItem struct {
 	pos    int
 	length int
@@ -23,9 +27,9 @@ type mediaItem struct {
 	item   *base.PreparedItem
 }
 
-// foldValue derives an item's trie-key substitute: a hash of the raw
-// bytes and preprocessing dims, with bit 31 forced so it can never equal
-// a token ID.
+// foldValue derives the trie-key substitute for a media item: a hash of the
+// raw bytes and the preprocessing dims (which pin the feature geometry for
+// given bytes), with bit 31 forced so it can never equal a token ID.
 func foldValue(data []byte, dims []int) uint32 {
 	h := fnv.New64a()
 	h.Write(data)
@@ -38,6 +42,133 @@ func foldValue(data []byte, dims []int) uint32 {
 	return (uint32(sum>>32) ^ uint32(sum)) | 1<<31
 }
 
+// requestMedia manages one request's media features: encoded on first
+// use, released when the expansion is fully evaluated. A nil
+// *requestMedia is a text-only request; every method is nil-safe.
+type requestMedia struct {
+	model    base.MediaModel
+	items    []mediaItem
+	inputLen int
+
+	// manifest is the request-scoped batch view of items; Features is
+	// toggled in place so every batch shares the same slice.
+	manifest []batch.MediaItem
+	features []*mlx.Array // parallel to items; nil until encoded
+
+	// layout is the request's one-row Batch.Layout, shared by every batch
+	// like the manifest; nil when the model returned no layout.
+	layout []any
+}
+
+func (r *Runner) openMedia(request Request) *requestMedia {
+	if len(request.MediaItems) == 0 {
+		return nil
+	}
+	m := &requestMedia{
+		model:    r.Model.(base.MediaModel),
+		items:    request.MediaItems,
+		inputLen: len(request.Tokens),
+		manifest: make([]batch.MediaItem, len(request.MediaItems)),
+		features: make([]*mlx.Array, len(request.MediaItems)),
+	}
+	if request.Layout != nil {
+		m.layout = []any{request.Layout}
+	}
+	for i, item := range m.items {
+		m.manifest[i] = batch.MediaItem{Pos: item.pos, Opaque: item.item.Opaque}
+	}
+	return m
+}
+
+// rowLayout returns the request's per-row Batch.Layout value.
+func (m *requestMedia) rowLayout() []any {
+	if m == nil {
+		return nil
+	}
+	return m.layout
+}
+
+func (item *mediaItem) atomic() bool { return !item.item.Causal }
+
+// extendChunk keeps a chunk from ending strictly inside an atomic
+// expansion: cut before one starting inside the chunk, else grow to its
+// end, clipped one short of the prompt to preserve the decode seed.
+func (m *requestMedia) extendChunk(pos, n int) int {
+	if m == nil {
+		return n
+	}
+	end := pos + n
+	for i := range m.items {
+		item := &m.items[i]
+		if !item.atomic() {
+			continue
+		}
+		if item.pos < end && end < item.pos+item.length {
+			if item.pos > pos {
+				return item.pos - pos
+			}
+			return min(item.pos+item.length, m.inputLen-1) - pos
+		}
+	}
+	return n
+}
+
+// batchMedia returns the manifest for chunk [pos, pos+n), encoding and
+// pinning each item's features on first overlap; nothing evaluates here —
+// the consuming forward pulls the encoder.
+func (m *requestMedia) batchMedia(pos, n int) []batch.MediaItem {
+	if m == nil {
+		return nil
+	}
+	for i, item := range m.items {
+		if item.pos >= pos+n || item.pos+item.length <= pos {
+			continue
+		}
+		if m.features[i] == nil {
+			data := mlx.FromValues(item.item.MediaData, item.item.Dims...)
+			m.features[i] = m.model.EncodeMedia(item.item, data)
+			mlx.Pin(m.features[i])
+			// The upload copied the pixels; free them here — release never
+			// passes the end of an expansion reaching the prompt's last token.
+			item.item.MediaData = nil
+		}
+		m.manifest[i].Features = m.features[i]
+	}
+	return m.manifest
+}
+
+// release frees what items fully evaluated or restored at position pos no
+// longer need: the pinned features and the preprocessed pixel buffer.
+func (m *requestMedia) release(pos int) {
+	if m == nil {
+		return
+	}
+	for i, item := range m.items {
+		if item.pos+item.length <= pos {
+			item.item.MediaData = nil
+			if m.features[i] != nil {
+				mlx.Unpin(m.features[i])
+				m.features[i] = nil
+				m.manifest[i].Features = nil
+			}
+		}
+	}
+}
+
+// close unpins whatever remains when the pipeline exits.
+func (m *requestMedia) close() {
+	if m == nil {
+		return
+	}
+	for i, f := range m.features {
+		if f != nil {
+			mlx.Unpin(f)
+			m.features[i] = nil
+			m.manifest[i].Features = nil
+		}
+	}
+}
+
 // expandMedia tokenizes the [img-N]-tagged prompt into segments, expands
 // them in a single PrepareMedia call, and validates the authored items
 // before keying cache identity on them.

x/mlxrunner/media_test.go

@@ -3,6 +3,9 @@ package mlxrunner
 import (
 	"slices"
 	"testing"
+
+	"github.com/ollama/ollama/x/mlxrunner/mlx"
+	"github.com/ollama/ollama/x/mlxrunner/model/base"
 )
 
 func TestEffectiveKeyTokens(t *testing.T) {
@@ -23,6 +26,100 @@ func TestEffectiveKeyTokens(t *testing.T) {
 	}
 }
 
+func TestExtendChunk(t *testing.T) {
+	m := &requestMedia{
+		items: []mediaItem{
+			{pos: 10, length: 4, item: &base.PreparedItem{}},
+			{pos: 40, length: 8, item: &base.PreparedItem{Causal: true}},
+			{pos: 96, length: 4, item: &base.PreparedItem{}},
+		},
+		inputLen: 100,
+	}
+
+	cases := []struct{ pos, n, want int }{
+		{0, 10, 10}, // ends at the expansion start: not inside
+		{0, 12, 10}, // expansion starts inside: cut so it begins the next chunk
+		{0, 14, 14}, // ends at the expansion end: not inside
+		{10, 2, 4},  // chunk starts at the expansion: extend to its end
+		{12, 1, 2},  // resume mid-expansion: extend to its end
+		{38, 6, 6},  // causal expansion: ending inside is legal
+		{42, 4, 4},  // causal expansion at chunk start: no extension
+		{90, 7, 6},  // trailing expansion starts inside: cut before it
+		{96, 2, 3},  // trailing expansion at chunk start: clip one short of the prompt
+	}
+	for _, c := range cases {
+		if got := m.extendChunk(c.pos, c.n); got != c.want {
+			t.Errorf("extendChunk(%d, %d) = %d, want %d", c.pos, c.n, got, c.want)
+		}
+	}
+
+	var nilMedia *requestMedia
+	if got := nilMedia.extendChunk(0, 12); got != 12 {
+		t.Errorf("nil extendChunk = %d, want 12", got)
+	}
+}
+
+// encodeCountingModel counts EncodeMedia calls and returns a real array so
+// the pin/release lifecycle runs against live handles.
+type encodeCountingModel struct {
+	stubMediaModel
+	calls *int
+}
+
+func (m encodeCountingModel) EncodeMedia(item *base.PreparedItem, data *mlx.Array) *mlx.Array {
+	*m.calls++
+	return mlx.Zeros(mlx.DTypeFloat32, item.Range[1]-item.Range[0], 4)
+}
+
+func TestBatchMediaLifecycle(t *testing.T) {
+	skipIfNoMLX(t)
+
+	calls := 0
+	prepared := &base.PreparedItem{
+		Range:     [2]int{2, 6},
+		MediaData: []float32{1, 2},
+		Dims:      []int{2},
+		Opaque:    7,
+	}
+	r := &Runner{Model: encodeCountingModel{calls: &calls}}
+	request := Request{
+		Tokens:     make([]int32, 8),
+		MediaItems: []mediaItem{{pos: 2, length: 4, item: prepared}},
+	}
+
+	m := r.openMedia(request)
+	if m == nil {
+		t.Fatal("openMedia returned nil for a media request")
+	}
+	if m.manifest[0].Pos != 2 || m.manifest[0].Opaque != 7 {
+		t.Fatalf("manifest = %+v", m.manifest[0])
+	}
+
+	if items := m.batchMedia(0, 2); items[0].Features != nil || calls != 0 {
+		t.Fatal("non-overlapping chunk encoded features")
+	}
+	if items := m.batchMedia(0, 4); items[0].Features == nil || calls != 1 {
+		t.Fatalf("overlap did not encode once (calls=%d)", calls)
+	}
+	if items := m.batchMedia(4, 2); items[0].Features == nil || calls != 1 {
+		t.Fatalf("second overlap re-encoded (calls=%d)", calls)
+	}
+
+	m.release(4)
+	if m.manifest[0].Features == nil {
+		t.Fatal("release dropped features before the expansion was evaluated")
+	}
+	m.release(6)
+	if m.manifest[0].Features != nil {
+		t.Fatal("release kept features past the expansion end")
+	}
+	m.close()
+
+	if r.openMedia(Request{Tokens: make([]int32, 8)}) != nil {
+		t.Fatal("openMedia returned non-nil for a text-only request")
+	}
+}
+
 // Two prompts that differ only in their image diverge at the expansion's
 // first key — one position earlier under bigram packing — and prompts with
 // the same image share keys through the whole expansion.

x/mlxrunner/mtp.go

@@ -2,6 +2,7 @@ package mlxrunner
 
 import (
 	"fmt"
+	"slices"
 
 	"github.com/ollama/ollama/x/mlxrunner/batch"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
@@ -35,8 +36,8 @@ func (d *mtpDrafter) draftLimit() int { return 0 }
 
 // open returns the drafting session for one request, its pairing frontier
 // synced to the draft caches' restored offset.
-func (d *mtpDrafter) open() draftSession {
-	s := &mtpDraftSession{drafter: d}
+func (d *mtpDrafter) open(layout []any) draftSession {
+	s := &mtpDraftSession{drafter: d, layout: layout}
 	if kv := d.spec.draftKV; len(kv) > 0 {
 		// A restored prefix arrives with the draft caches already written;
 		// pairing resumes from their absolute offset.
@@ -52,6 +53,7 @@ func (d *mtpDrafter) open() draftSession {
 // completes only when the next token arrives.
 type mtpDraftSession struct {
 	drafter *mtpDrafter
+	layout  []any
 
 	// frontier is the slot after the last reported token; frontierHidden is
 	// the pinned target hidden at frontier-1, fused into the next pair.
@@ -72,10 +74,17 @@ type mtpDraftSession struct {
 	// reuses them without a head call.
 	heldHidden    *mlx.Array
 	heldAuxHidden *mlx.Array
+
+	// pendingMedia holds manifest rows the deferred flush may still embed,
+	// pinned since prefill releases them after the target's chunk;
+	// lastDelivered marks each row's newest delivered end.
+	pendingMedia  map[int]batch.MediaItem
+	lastDelivered map[int]int
 }
 
-func (d *mtpDraftSession) committed(tokens, hiddens *mlx.Array, position int) {
+func (d *mtpDraftSession) committed(tokens, hiddens *mlx.Array, position int, media []batch.MediaItem) {
 	n := tokens.Dim(1)
+	d.captureMedia(media, position+n)
 	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
@@ -104,6 +113,59 @@ func (d *mtpDraftSession) committed(tokens, hiddens *mlx.Array, position int) {
 	d.setFrontierHidden(lastHiddenRow(hiddens))
 }
 
+// captureMedia pins the run's feature-bearing rows for the deferred
+// flush, which embeds them after prefill has released the features. A row
+// spanning chunks arrives once per chunk.
+func (d *mtpDraftSession) captureMedia(media []batch.MediaItem, end int) {
+	if len(d.drafter.spec.draftKV) == 0 {
+		return
+	}
+	for _, item := range media {
+		if item.Features == nil {
+			continue
+		}
+		if _, ok := d.pendingMedia[item.Pos]; !ok {
+			if d.pendingMedia == nil {
+				d.pendingMedia = make(map[int]batch.MediaItem)
+				d.lastDelivered = make(map[int]int)
+			}
+			mlx.Pin(item.Features)
+			d.pendingMedia[item.Pos] = item
+		}
+		d.lastDelivered[item.Pos] = end
+	}
+}
+
+// flushMedia returns the held rows for a flush batch and drops rows the
+// flush finishes: fully delivered below embedEnd means never embedded
+// again.
+func (d *mtpDraftSession) flushMedia(embedEnd int) []batch.MediaItem {
+	if len(d.pendingMedia) == 0 {
+		return nil
+	}
+	manifest := make([]batch.MediaItem, 0, len(d.pendingMedia))
+	for _, item := range d.pendingMedia {
+		manifest = append(manifest, item)
+	}
+	slices.SortFunc(manifest, func(a, b batch.MediaItem) int { return a.Pos - b.Pos })
+	for pos, last := range d.lastDelivered {
+		if last <= embedEnd {
+			mlx.Unpin(d.pendingMedia[pos].Features)
+			delete(d.pendingMedia, pos)
+			delete(d.lastDelivered, pos)
+		}
+	}
+	return manifest
+}
+
+func (d *mtpDraftSession) closeMedia() {
+	for pos, item := range d.pendingMedia {
+		mlx.Unpin(item.Features)
+		delete(d.pendingMedia, pos)
+		delete(d.lastDelivered, pos)
+	}
+}
+
 // settle completes any open frontier pair with next — the token after the
 // last committed slot — and flushes, leveling the draft caches with the
 // target.
@@ -119,6 +181,7 @@ func (d *mtpDraftSession) settle(next *mlx.Array) {
 
 func (d *mtpDraftSession) close() {
 	d.flush()
+	d.closeMedia()
 	d.setFrontierHidden(nil)
 	d.setHeld(nil, nil)
 }
@@ -155,11 +218,15 @@ func (d *mtpDraftSession) flush() {
 
 	ids := mlx.Concatenate(d.pendingTokens, 1)
 	hiddens := mlx.Concatenate(d.pendingHiddens, 1)
+	// The pair at slot S embeds the look-ahead token S+1, so this flush
+	// embeds prompt tokens up to committedDraftOffset+len+1.
 	hidden, auxHidden := spec.draft.Forward(&batch.Batch{
 		InputIDs:     ids,
 		SeqOffsets:   []int32{int32(d.committedDraftOffset)},
 		SeqQueryLens: []int32{int32(ids.Dim(1))},
 		Hidden:       hiddens,
+		Media:        d.flushMedia(d.committedDraftOffset + ids.Dim(1) + 1),
+		Layout:       d.layout,
 	}, spec.targets, spec.draftKV)
 	d.setHeld(lastHiddenRow(hidden), lastHiddenRow(auxHidden))
 	d.committedDraftOffset += ids.Dim(1)
@@ -235,6 +302,7 @@ func (d *mtpDraftSession) propose(current *mlx.Array, maxTokens int) *draftCandi
 				SeqOffsets:   []int32{int32(pos)},
 				SeqQueryLens: []int32{1},
 				Hidden:       lastHidden,
+				Layout:       d.layout,
 			}, spec.targets, spec.draftKV)
 		}
 		// Unembed only the row being sampled, never the batch.

x/mlxrunner/mtp_test.go

@@ -50,6 +50,9 @@ type fakeMTPModel struct {
 	tok     *tokenizer.Tokenizer
 	// forwards records each Forward call so tests can assert contiguous writes.
 	forwards []forwardCall
+	// layouts records each Forward's Batch.Layout so tests can assert the
+	// request's layout state reaches every target forward.
+	layouts [][]any
 }
 
 type forwardCall struct {
@@ -61,6 +64,7 @@ func (m *fakeMTPModel) Forward(b *batch.Batch, caches []cache.Cache) (hidden, au
 	mlx.Eval(b.InputIDs)
 	ids := b.InputIDs.Ints()
 	m.forwards = append(m.forwards, forwardCall{offset: b.SeqOffsets[0], n: int32(len(ids))})
+	m.layouts = append(m.layouts, b.Layout)
 	for i, c := range caches {
 		if i >= m.NumLayers() {
 			break
@@ -96,6 +100,9 @@ var _ base.Model = (*fakeMTPModel)(nil)
 // fakeMTPDraft is a cacheless draft that extends b.InputIDs through predict;
 // a map (not a step counter) keeps drafting consistent regardless of batching.
 type fakeMTPDraft struct {
+	// layouts records each Forward's Batch.Layout, like fakeMTPModel's.
+	layouts [][]any
+
 	predict map[int32]int32
 	// calls records each Draft call so tests can assert the position convention.
 	calls []draftCall
@@ -111,6 +118,7 @@ func (d *fakeMTPDraft) LoadWeights(map[string]*mlx.Array) error { return nil }
 func (d *fakeMTPDraft) NewCaches() []cache.Cache { return nil }
 
 func (d *fakeMTPDraft) Forward(b *batch.Batch, _, _ []cache.Cache) (hidden, auxHidden *mlx.Array) {
+	d.layouts = append(d.layouts, b.Layout)
 	mlx.Eval(b.InputIDs)
 	prev := int32(b.InputIDs.Ints()[0])
 	d.calls = append(d.calls, draftCall{position: b.SeqOffsets[0], from: prev})
@@ -135,12 +143,14 @@ type fakeKVDraft struct {
 }
 
 // extendCall is one recorded Forward call: the absolute slot of the first
-// entry written, the look-ahead token ids, and the hot index of each fused
-// hidden row (-1 for the head's own aux hidden).
+// entry written, the look-ahead token ids, the hot index of each fused
+// hidden row (-1 for the head's own aux hidden), and the Pos of each media
+// row the batch carried.
 type extendCall struct {
 	offset  int32
 	ids     []int32
 	hiddens []int32
+	media   []int
 }
 
 func (d *fakeKVDraft) LoadWeights(map[string]*mlx.Array) error { return nil }
@@ -166,7 +176,11 @@ func (d *fakeKVDraft) Forward(b *batch.Batch, _, draftCaches []cache.Cache) (hid
 			}
 		}
 	}
-	d.extends = append(d.extends, extendCall{offset: b.SeqOffsets[0], ids: ids, hiddens: hot})
+	var media []int
+	for _, item := range b.Media {
+		media = append(media, item.Pos)
+	}
+	d.extends = append(d.extends, extendCall{offset: b.SeqOffsets[0], ids: ids, hiddens: hot, media: media})
 
 	if rc, ok := draftCaches[0].(*fakeRewindableCache); ok {
 		rc.feed(ids)
@@ -452,7 +466,7 @@ func TestRunMTPDecodeSampled(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{Temperature: 1, Seed: 42, UseSeed: true},
 	}
-	spec := r.spec.open(req)
+	spec := r.spec.open(req, nil)
 	if spec == nil || !spec.enabled {
 		t.Fatalf("open rejected a sampled request")
 	}
@@ -496,11 +510,11 @@ func TestRunMTPDecodeWarmDrafter(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	spec := r.spec.open(req)
+	spec := r.spec.open(req, nil)
 	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)
+	spec.committed(mlx.FromValues([]int32{0}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
 
 	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
@@ -557,14 +571,14 @@ func TestRunMTPDecodeEOSCutLeavesPositionsUnjudged(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	spec := r.spec.open(req)
+	spec := r.spec.open(req, nil)
 	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)
+	spec.committed(mlx.FromValues([]int32{0}, 1, 1), oneHotLogits([]int32{1}), 0, nil)
 
 	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
@@ -695,6 +709,54 @@ func TestDecodeCancelledMidStream(t *testing.T) {
 	}
 }
 
+func TestLayoutRidesEveryForward(t *testing.T) {
+	skipIfNoMLX(t)
+	// The request's opaque layout state must reach every forward: parked
+	// pipelined dispatches, the fused verification forward, and the draft
+	// model's own forwards alike.
+	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{})
+	caches, _ := newMTPTestCaches(1)
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict}, caches, nil)
+	session, ch := newMTPTestSession(caches)
+
+	req := Request{
+		Responses:         ch,
+		Tokens:            []int32{0},
+		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
+		SamplerOpts:       sampler.Options{},
+	}
+	spec := r.spec.open(req, []any{"layout"})
+	pinDraftLimit(spec, 4)
+	d := spec.decoder(mlx.FromValues([]int32{1}, 1), 1)
+	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
+		t.Fatalf("decode: %v", err)
+	}
+	d.close()
+
+	model := r.Model.(*fakeMTPModel)
+	if len(model.layouts) < 3 {
+		t.Fatalf("forwards recorded = %d, want parked dispatches and a fused round", len(model.layouts))
+	}
+	for i, l := range model.layouts {
+		if len(l) != 1 || l[0] != "layout" {
+			t.Fatalf("forward %d layout = %v", i, l)
+		}
+	}
+
+	draft := r.spec.draft.(*fakeMTPDraft)
+	if len(draft.layouts) == 0 {
+		t.Fatal("no draft forwards recorded")
+	}
+	for i, l := range draft.layouts {
+		if len(l) != 1 || l[0] != "layout" {
+			t.Fatalf("draft forward %d layout = %v", i, l)
+		}
+	}
+}
+
 // 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.
@@ -707,13 +769,13 @@ func pinDraftLimit(spec *speculationSession, limit int) {
 // 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); spec != nil {
+	if spec := r.spec.open(req, nil); spec != nil {
 		if spec.enabled {
 			pinDraftLimit(spec, 4)
 		}
 		return spec.decoder(mlx.FromValues(seed, len(seed)), position)
 	}
-	return r.pipelinedDecoder(nil, caches, mlx.FromValues(seed, 1, len(seed)), position)
+	return r.pipelinedDecoder(nil, caches, mlx.FromValues(seed, 1, len(seed)), position, nil)
 }
 
 func TestDecodeKVDraft(t *testing.T) {
@@ -742,7 +804,7 @@ func TestDecodeKVDraft(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	spec := r.spec.open(req)
+	spec := r.spec.open(req, nil)
 	if spec == nil || !spec.enabled || len(spec.spec.targets) != 1 {
 		t.Fatalf("speculation engine not built around the draft caches")
 	}
@@ -827,7 +889,7 @@ func TestDecodeKVDraftRejectionRebuildsFromTarget(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	spec := r.spec.open(req)
+	spec := r.spec.open(req, nil)
 	pinDraftLimit(spec, 4)
 	defer spec.close()
 	d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
@@ -897,7 +959,7 @@ func TestDecodeMaintainsDraftCacheWithoutDrafting(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       opts,
 	}
-	spec := r.spec.open(req)
+	spec := r.spec.open(req, nil)
 	if spec == nil || spec.enabled {
 		t.Fatalf("want a permanent-park speculationSession, got %+v", spec)
 	}
@@ -955,12 +1017,12 @@ func TestSettleLevelsDraftCacheWithPrefill(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	spec := r.spec.open(req)
+	spec := r.spec.open(req, nil)
 	defer spec.close()
 
 	// The prompt's only chunk: tokens 1..4 at slots 0..3 with their hiddens;
 	// token 5 is the seed.
-	spec.committed(mlx.FromValues([]int32{1, 2, 3, 4}, 1, 4), oneHotLogits([]int32{1, 2, 3, 4}), 0)
+	spec.committed(mlx.FromValues([]int32{1, 2, 3, 4}, 1, 4), oneHotLogits([]int32{1, 2, 3, 4}), 0, nil)
 	spec.settle(mlx.FromValues([]int32{5}, 1))
 
 	if got := caches[1].Offset(); got != 4 {
@@ -974,6 +1036,46 @@ func TestSettleLevelsDraftCacheWithPrefill(t *testing.T) {
 	}
 }
 
+func TestFlushMediaHeldUntilEmbedded(t *testing.T) {
+	skipIfNoMLX(t)
+	// The deferred flush embeds prompt tokens after prefill has released the
+	// media features, so the session holds delivered feature rows itself:
+	// each flush carries the held rows, a row is dropped once the flush's
+	// embed frontier passes its delivered end, and a row a later chunk
+	// redelivers is held again.
+	predict := map[int32]int32{}
+	r := mtpTestRunner(t, predict, []int32{7}, sampler.Options{})
+	draft := &fakeKVDraft{predict: predict}
+	caches, _ := newMTPTestCaches(2)
+	draft.draftCaches = caches[1:]
+	r.cache.caches = caches
+	r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
+
+	d := r.spec.drafter.open(nil).(*mtpDraftSession)
+	item := batch.MediaItem{Pos: 1, Features: mlx.Zeros(mlx.DTypeFloat32, 4, 4)}
+
+	// Chunk [0, 4) delivers the item; the flush embeds through the delivered
+	// end, so the row rides that flush and is dropped.
+	d.committed(mlx.FromValues([]int32{1, 2, 3, 4}, 1, 4), oneHotLogits([]int32{1, 2, 3, 4}), 0, []batch.MediaItem{item})
+	d.flush()
+	// A chunk with no delivery flushes without media.
+	d.committed(mlx.FromValues([]int32{5, 6}, 1, 2), oneHotLogits([]int32{5, 6}), 4, nil)
+	d.flush()
+	// A redelivered row (one arrives per chunk it spans) is held and
+	// flushed again.
+	d.committed(mlx.FromValues([]int32{2, 3}, 1, 2), oneHotLogits([]int32{2, 3}), 6, []batch.MediaItem{item})
+	d.flush()
+	d.close()
+
+	got := make([][]int, len(draft.extends))
+	for i, e := range draft.extends {
+		got[i] = e.media
+	}
+	if want := [][]int{{1}, nil, {1}}; !reflect.DeepEqual(got, want) {
+		t.Fatalf("flush media = %v, want %v", got, want)
+	}
+}
+
 func TestCommittedRunBatchesPastFlushCap(t *testing.T) {
 	skipIfNoMLX(t)
 	// A committed run longer than the pending-flush cap still writes the draft
@@ -997,12 +1099,12 @@ func TestCommittedRunBatchesPastFlushCap(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	spec := r.spec.open(req)
+	spec := r.spec.open(req, nil)
 	defer spec.close()
 
 	// One prefill-sized chunk: n tokens at slots 0..n-1 with their hiddens.
 	// The run crosses the flush cap, so the write happens inside committed.
-	spec.committed(mlx.FromValues(tokens, 1, n), oneHotLogits(tokens), 0)
+	spec.committed(mlx.FromValues(tokens, 1, n), oneHotLogits(tokens), 0, nil)
 
 	if got := len(draft.extends); got != 1 {
 		t.Fatalf("draft extends = %d calls, want 1 batched extend", got)
@@ -1039,7 +1141,7 @@ func TestRestoredPrefixRewritesBoundaryPair(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	spec := r.spec.open(req)
+	spec := r.spec.open(req, nil)
 	pinDraftLimit(spec, 4)
 	d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0)
 	if err := r.decode(context.Background(), req, session, d, 0); err != nil {
@@ -1061,8 +1163,8 @@ func TestRestoredPrefixRewritesBoundaryPair(t *testing.T) {
 			t.Fatal("restore to 5 failed")
 		}
 	}
-	spec = r.spec.open(req)
-	spec.committed(mlx.FromValues([]int32{6, 1}, 1, 2), oneHotLogits([]int32{eos, 2}), 5)
+	spec = r.spec.open(req, nil)
+	spec.committed(mlx.FromValues([]int32{6, 1}, 1, 2), oneHotLogits([]int32{eos, 2}), 5, nil)
 	spec.close()
 
 	last := draft.extends[len(draft.extends)-1]
@@ -1093,7 +1195,7 @@ func TestDecodeParkedDraftResume(t *testing.T) {
 		CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
 		SamplerOpts:       sampler.Options{},
 	}
-	spec := r.spec.open(req)
+	spec := r.spec.open(req, nil)
 	if spec == nil || !spec.enabled {
 		t.Fatalf("want a drafting speculationSession, got %+v", spec)
 	}
@@ -1186,7 +1288,7 @@ func newMTPTestSession(caches []cache.Cache) (*cacheSession, chan CompletionResp
 // no-op drafter, since the engine requires one.
 func testSpeculationSession(r *Runner, caches []cache.Cache) *speculationSession {
 	if r.spec != nil {
-		return &speculationSession{spec: r.spec, drafter: r.spec.drafter.open()}
+		return &speculationSession{spec: r.spec, drafter: r.spec.drafter.open(nil)}
 	}
 	s := &speculation{r: r, targets: caches}
 	return &speculationSession{spec: s, drafter: nopDrafter{}}
@@ -1196,10 +1298,10 @@ func testSpeculationSession(r *Runner, caches []cache.Cache) *speculationSession
 // directly and never propose.
 type nopDrafter struct{}
 
-func (nopDrafter) propose(*mlx.Array, int) *draftCandidates { return nil }
-func (nopDrafter) committed(_, _ *mlx.Array, _ int)         {}
-func (nopDrafter) settle(*mlx.Array)                        {}
-func (nopDrafter) close()                                   {}
+func (nopDrafter) propose(*mlx.Array, int) *draftCandidates              { return nil }
+func (nopDrafter) committed(_, _ *mlx.Array, _ int, _ []batch.MediaItem) {}
+func (nopDrafter) settle(*mlx.Array)                                     {}
+func (nopDrafter) close()                                                {}
 
 // scriptedCandidates builds draft candidates by running the real drafter
 // against a fake whose prediction chain, starting from seed token 0, yields
@@ -1213,8 +1315,8 @@ func scriptedCandidates(r *Runner, tokens []int32) *draftCandidates {
 		prev = tok
 	}
 	s := &speculation{r: r, draft: &fakeMTPDraft{predict: chain}}
-	d := (&mtpDrafter{spec: s}).open()
-	d.committed(mlx.FromValues([]int32{0}, 1, 1), mlx.Zeros(mlx.DTypeFloat32, 1, 1, mtpTestVocab), 0)
+	d := (&mtpDrafter{spec: s}).open(nil)
+	d.committed(mlx.FromValues([]int32{0}, 1, 1), mlx.Zeros(mlx.DTypeFloat32, 1, 1, mtpTestVocab), 0, nil)
 	return d.propose(mlx.FromValues([]int32{0}, 1), len(tokens))
 }
 

x/mlxrunner/pipeline.go

@@ -97,12 +97,15 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er
 	defer session.close()
 	caches := session.caches
 
+	media := r.openMedia(request)
+	defer media.close()
+
 	// Built before prefill so a drafter with draft caches follows the prompt
 	// through prefill alongside the target.
-	spec := r.spec.open(request)
+	spec := r.spec.open(request, media.rowLayout())
 	defer spec.close()
 
-	seed, position, promptEval, err := r.prefill(ctx, session, spec)
+	seed, position, promptEval, err := r.prefill(ctx, session, spec, media)
 	if err != nil {
 		return err
 	}
@@ -114,7 +117,7 @@ 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.ExpandDims(-1), position)
+		d = r.pipelinedDecoder(nil, caches, seed.ExpandDims(-1), position, media.rowLayout())
 	}
 	defer d.close()
 	return r.decode(ctx, request, session, d, promptEval)
@@ -123,7 +126,7 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er
 // 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 token, the resume position, and the prompt-evaluation duration.
-func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *speculationSession) (*mlx.Array, int, time.Duration, error) {
+func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *speculationSession, media *requestMedia) (*mlx.Array, int, time.Duration, error) {
 	start := time.Now()
 	inputs := session.inputs
 	tokens := session.remaining
@@ -159,20 +162,30 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *specu
 
 	total, processed := len(tokens), 0
 	position := len(inputs) - len(tokens)
+	// Free restored items' buffers now: on a full cache hit the loop never runs.
+	media.release(position)
 	for total-processed > 1 {
 		if err := ctx.Err(); err != nil {
 			return nil, 0, 0, err
 		}
 
 		n := min(prefillChunk, total-processed-1)
+		n = media.extendChunk(position, n)
 
 		chunkIDs := mlx.FromValues(tokens[processed:processed+n], 1, n)
+		manifest := media.batchMedia(position, n)
 		_, auxHidden := r.Model.Forward(&batch.Batch{
 			InputIDs:     chunkIDs,
 			SeqOffsets:   []int32{int32(position)},
 			SeqQueryLens: []int32{int32(n)},
+			Media:        manifest,
+			Layout:       media.rowLayout(),
 		}, caches)
-		spec.committed(chunkIDs, auxHidden, position)
+		spec.committed(chunkIDs, auxHidden, position, manifest)
+		// Unpin finished items before the sweep: the consuming forward's
+		// graph retains their data until evaluation, so the buffers die with
+		// this chunk's eval instead of surviving into the next chunk.
+		media.release(position + n)
 		mlx.Sweep()
 		materializeCaches()
 		processed += n
@@ -311,12 +324,13 @@ type pipelinedDecoder struct {
 	// drafter at close, keeping a non-drafting session's draft KV level.
 	spec     *speculationSession
 	caches   []cache.Cache
+	layout   []any // the request's per-row layout state, stamped on every forward
 	position int
 	sample   sampler.Result // in flight: sampled, not yet forwarded
 }
 
-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}
+func (r *Runner) pipelinedDecoder(spec *speculationSession, caches []cache.Cache, seed *mlx.Array, position int, layout []any) *pipelinedDecoder {
+	t := &pipelinedDecoder{r: r, spec: spec, caches: caches, layout: layout, position: position}
 	t.sample = t.dispatch(seed)
 	return t
 }
@@ -329,8 +343,9 @@ func (t *pipelinedDecoder) dispatch(token *mlx.Array) sampler.Result {
 		InputIDs:     token,
 		SeqOffsets:   []int32{int32(t.position)},
 		SeqQueryLens: []int32{int32(token.Dim(1))},
+		Layout:       t.layout,
 	}, t.caches)
-	t.spec.committed(token, auxHidden, t.position)
+	t.spec.committed(token, auxHidden, t.position, nil)
 	t.position += token.Dim(1)
 	logits := r.Model.Unembed(hidden)
 	next := r.Sampler.Sample([]int{pipelineSlot}, logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice()).Squeeze(1))

x/mlxrunner/speculate.go

@@ -21,8 +21,11 @@ type draftSession interface {
 	// committed reports a run of tokens committed to the target caches:
 	// tokens[i] sits at slot position+i and hiddens row i is the target
 	// hidden state at that slot. Runs arrive in slot order — prefill
-	// chunks, the decode seed, then each round's validated tokens.
-	committed(tokens, hiddens *mlx.Array, position int)
+	// chunks, the decode seed, then each round's validated tokens. media is
+	// the run's manifest (feature-bearing for items the run overlaps), valid
+	// only for the call; a session that defers its forward pins what it
+	// keeps. Nil outside prefill.
+	committed(tokens, hiddens *mlx.Array, position int, media []batch.MediaItem)
 
 	// settle completes any open frontier pair with next — the token after
 	// the last committed slot — and writes buffered reports through,
@@ -59,9 +62,9 @@ type speculation struct {
 }
 
 // drafter is the per-model half of a drafting implementation, opening each
-// request's drafting session.
+// request's drafting session with the request's per-row layout state.
 type drafter interface {
-	open() draftSession
+	open(layout []any) draftSession
 
 	// draftLimit is the deepest draft this drafter can produce, 0 when nothing
 	// bounds it. A depth past it is never measured, so the depth search must
@@ -91,8 +94,9 @@ func newSpeculation(r *Runner, draft base.DraftModel, targets, draftKV []cache.C
 type speculationSession struct {
 	spec    *speculation
 	drafter draftSession
-	enabled bool // whether this request drafts; false parks (maintain-only)
-	limit   int  // current draft length
+	enabled bool  // whether this request drafts; false parks (maintain-only)
+	limit   int   // current draft length
+	layout  []any // the request's per-row layout state, stamped on every target forward
 	stats   specStats
 
 	// Cost sampling: each round's wall time (start to next start, spanning the
@@ -105,18 +109,18 @@ type speculationSession struct {
 
 // open returns the speculation cursor for this request or nil when the model ships
 // no draft head (a nil receiver), which decodes plainly.
-func (s *speculation) open(request Request) *speculationSession {
+func (s *speculation) open(request Request, layout []any) *speculationSession {
 	if s == nil {
 		return nil
 	}
-	d := s.drafter.open()
+	d := s.drafter.open(layout)
 
 	// 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 := !opts.Logprobs && opts.TopLogprobs == 0
 
-	spec := &speculationSession{spec: s, drafter: d, enabled: enabled, prevDrafts: -1, roundDrafts: -1}
+	spec := &speculationSession{spec: s, drafter: d, layout: layout, enabled: enabled, prevDrafts: -1, roundDrafts: -1}
 	if enabled {
 		spec.limit = s.depth.scheduled
 	}
@@ -155,11 +159,11 @@ func (s *speculationSession) endRound(drafted, accepted, observed int) {
 	}
 }
 
-func (s *speculationSession) committed(tokens, hiddens *mlx.Array, position int) {
+func (s *speculationSession) committed(tokens, hiddens *mlx.Array, position int, media []batch.MediaItem) {
 	if s == nil {
 		return
 	}
-	s.drafter.committed(tokens, hiddens, position)
+	s.drafter.committed(tokens, hiddens, position, media)
 }
 
 // settle completes the drafter's open frontier pair with next and writes
@@ -272,7 +276,7 @@ func (st *speculativeDecoder) resume() []sampler.Result {
 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.targets, st.current.Token.ExpandDims(-1), st.position)
+		st.inner = s.spec.r.pipelinedDecoder(s, s.spec.targets, st.current.Token.ExpandDims(-1), st.position, s.layout)
 	}
 	return st.inner.next(remaining)
 }
@@ -398,6 +402,7 @@ func (s *speculationSession) accept(position *int, current sampler.Result, candi
 		InputIDs:     current.Token.ExpandDims(-1).Concatenate(1, candidates.tokens),
 		SeqOffsets:   []int32{int32(before)},
 		SeqQueryLens: []int32{int32(draftCount + 1)},
+		Layout:       s.layout,
 	}, s.spec.targets)
 
 	// Row i of the fused hidden is the state after the token at before+i, so
@@ -470,7 +475,7 @@ func (s *speculationSession) accept(position *int, current sampler.Result, candi
 	s.drafter.committed(
 		mlx.FromValues(runIDs, 1, len(runIDs)),
 		auxHiddenSeq.Slice(mlx.Slice(), mlx.Slice(0, len(runIDs)), mlx.Slice()),
-		before)
+		before, nil)
 
 	results = draftResults(draftIDs[:accepted])
 	if done {