diff --git a/x/mlxrunner/batch/batch.go b/x/mlxrunner/batch/batch.go index 409f5307..011b3110 100644 --- a/x/mlxrunner/batch/batch.go +++ b/x/mlxrunner/batch/batch.go @@ -17,6 +17,10 @@ type Batch struct { // Length equals the batch dimension of InputIDs. SeqQueryLens []int32 + // Hidden is the target hidden state a draft model fuses with its input + // embedding for this step. It is nil for ordinary forward passes. + Hidden *mlx.Array + // Memo is per-forward memoization used to cache results, such as masks, // which are often the same across layers. Memo Memo diff --git a/x/mlxrunner/model/base/base.go b/x/mlxrunner/model/base/base.go index 4f43d74b..9a056956 100644 --- a/x/mlxrunner/model/base/base.go +++ b/x/mlxrunner/model/base/base.go @@ -27,8 +27,21 @@ type Model interface { LoadWeights(tensors map[string]*mlx.Array) error } -// DraftModel is an auxiliary model stored alongside a target model. +// DraftModel is an auxiliary model alongside a target that proposes speculative +// tokens. type DraftModel interface { + // Draft fuses b.Hidden (the target hidden state) into its own forward and + // returns the head's hidden plus the projected hidden seeding the next step. + Draft(b *batch.Batch, caches []cache.Cache) (hidden, projected *mlx.Array) + + // Unembed projects a hidden state to vocabulary logits. + Unembed(x *mlx.Array) *mlx.Array + + // DraftCaches selects the draft model's own KV caches from the full + // per-request slice — any subset, or nil when the draft keeps no KV. + DraftCaches(caches []cache.Cache) []cache.Cache + + // LoadWeights assigns manifest tensors to the draft head's fields. LoadWeights(tensors map[string]*mlx.Array) error } @@ -46,15 +59,10 @@ type MTPDefaultsProvider interface { MTPDraftDefaults(sample bool) MTPDefaults } -// MTPDraftModel is a draft model capable of Gemma-style multi-token -// prediction from target token embeddings, target hidden states, and target KV. -type MTPDraftModel interface { - Draft(inputEmbeds *mlx.Array, position int32, caches []cache.Cache) (logits, hidden *mlx.Array) -} - -// MTPEmbeddingModel exposes the target token embedding path used by MTP drafts. -type MTPEmbeddingModel interface { - TokenEmbeddings(inputIDs *mlx.Array) *mlx.Array +// 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 { + SelfDraft() DraftModel } var ( diff --git a/x/mlxrunner/mtp.go b/x/mlxrunner/mtp.go index 52381f92..ba788d82 100644 --- a/x/mlxrunner/mtp.go +++ b/x/mlxrunner/mtp.go @@ -1,7 +1,9 @@ package mlxrunner import ( - "github.com/ollama/ollama/x/mlxrunner/cache" + "fmt" + + "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" @@ -12,6 +14,11 @@ const ( 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, @@ -30,85 +37,231 @@ func (r *Runner) mtpDefaults(sample bool) base.MTPDefaults { return defaults } -// mtpDrafter drafts with a model's multi-token-prediction head: a small -// head trained to continue the target's hidden states, fed through the -// committed-stream reports. It retains the hidden at the last committed -// slot and tracks the committed frontier, which anchors its drafting. +// 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. type mtpDrafter struct { - spec *speculation - draft base.MTPDraftModel - target base.MTPEmbeddingModel - caches []cache.Cache + spec *speculation - // frontier is the slot after the last committed token; the hidden - // reported for slot frontier-1 is retained (pinned) as the fusion - // input for the next proposal chain. + // frontier is the slot after the last reported token; frontierHidden is + // the pinned target hidden at frontier-1, fused into the next pair. frontier int frontierHidden *mlx.Array + + // committedDraftOffset is the slot after the last pair written to the + // draft caches; later pairs wait pinned in the pending lists until + // flushed. pendingCount is the look-ahead tokens those lists hold, summed + // across the buffered runs. + committedDraftOffset int + pendingTokens []*mlx.Array + pendingHiddens []*mlx.Array + pendingCount int + + // heldHidden is the frontier row's pre-unembed hidden and heldProjected + // its fusion hidden, carried from the last flush so the first proposal + // reuses them without a head call. + heldHidden *mlx.Array + heldProjected *mlx.Array } -// newMTPDrafter returns the MTP drafter for this request's caches, or nil -// when the model carries no MTP head. -func newMTPDrafter(s *speculation, caches []cache.Cache) *mtpDrafter { - draft, ok := s.draft.(base.MTPDraftModel) - if !ok { - return nil +// 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 } - target, ok := s.r.Model.(base.MTPEmbeddingModel) - if !ok { - return nil - } - return &mtpDrafter{spec: s, draft: draft, target: target, caches: caches} + return d } func (d *mtpDrafter) committed(tokens, hiddens *mlx.Array, position int) { - d.frontier = position + tokens.Dim(1) - h := lastHiddenRow(hiddens) - mlx.Pin(h) - if d.frontierHidden != nil { - mlx.Unpin(d.frontierHidden) + n := tokens.Dim(1) + if len(d.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 + // before it. Leading tokens whose slot is already buffered or written + // through (a proposal consumed the run's first token, or a restored + // prefix sits at the run start) are skipped; a slot below the frontier + // is a gap bug. + start := d.committedDraftOffset + d.pendingCount - position + 1 + if start < 0 { + panic(fmt.Sprintf("mtp: committed run at %d leaves a pair gap at %d", position, d.committedDraftOffset+d.pendingCount)) + } + if start < n { + ids := tokens.Slice(mlx.Slice(), mlx.Slice(start, n)) + var h *mlx.Array + if start == 0 { + h = mlx.Concatenate([]*mlx.Array{d.frontierHidden, hiddens.Slice(mlx.Slice(), mlx.Slice(0, n-1), mlx.Slice())}, 1) + } else { + h = hiddens.Slice(mlx.Slice(), mlx.Slice(start-1, n-1), mlx.Slice()) + } + d.queueCacheWrites(ids, h) + } } - d.frontierHidden = h + + d.frontier = position + n + d.setFrontierHidden(lastHiddenRow(hiddens)) +} + +// finish settles the drafter when generation ends: current completes the +// frontier pair, leveling the draft caches with the target's resting offset. +// +// TODO: leveling the draft to the target writes a boundary entry whose +// look-ahead token is outside the stored prefix (here, the never-committed +// current). When a later request restores this prefix and diverges at the +// boundary, that entry is stale and lowers draft acceptance. EAGLE keeps the +// draft one slot behind the target so the unconfirmed boundary entry is never +// written (its bigram partner token[S+1] does not exist yet); we should do the +// 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 { + return + } + d.settle(current) +} + +// settle completes any open frontier pair with current, then flushes. +func (d *mtpDrafter) 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() { - if d.frontierHidden != nil { - mlx.Unpin(d.frontierHidden) - d.frontierHidden = nil + d.flush() + d.setFrontierHidden(nil) + d.setHeld(nil, nil) +} + +// queueCacheWrites buffers completed draft-cache writes — look-ahead tokens +// 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) { + mlx.Pin(tokens, hiddens) + d.pendingTokens = append(d.pendingTokens, tokens) + d.pendingHiddens = append(d.pendingHiddens, hiddens) + d.pendingCount += tokens.Dim(1) + if d.pendingCount >= mtpPendingFlushTokens { + d.flush() } } -// propose drafts a token chain Gemma-style ("single-position"): the head is -// trained to draft every speculative token as if it sat at the last -// committed slot, re-attending the target caches read-only. Each step fuses -// the previous token's target embedding with the hidden chain — the -// reported hidden first, then the head's own projections — and the -// RoPE/cache anchor stays at the last committed slot while the proposed -// tokens advance. +// 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() { + if len(d.pendingTokens) == 0 { + return + } + for _, c := range d.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)) + } + } + } + + ids := mlx.Concatenate(d.pendingTokens, 1) + hiddens := mlx.Concatenate(d.pendingHiddens, 1) + hidden, projected := d.spec.draft.Draft(&batch.Batch{ + InputIDs: ids, + SeqOffsets: []int32{int32(d.committedDraftOffset)}, + SeqQueryLens: []int32{int32(ids.Dim(1))}, + Hidden: hiddens, + }, d.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 = append(state, c.State()...) + } + mlx.AsyncEval(state...) + + mlx.Unpin(d.pendingTokens...) + mlx.Unpin(d.pendingHiddens...) + d.pendingTokens, d.pendingHiddens = nil, nil + d.pendingCount = 0 +} + +func (d *mtpDrafter) 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) { + mlx.Pin(hidden, projected) + mlx.Unpin(d.heldHidden, d.heldProjected) + d.heldHidden, d.heldProjected = hidden, projected +} + +// propose drafts a token chain after the not-yet-validated current token. +// 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 { if maxTokens <= 0 || d.frontierHidden == nil { return nil } r := d.spec.r - anchor := int32(d.frontier - 1) + if len(d.spec.draftKV) > 0 { + d.settle(current) + if d.heldHidden == nil { + return nil + } + } + lastToken := current.ExpandDims(-1) lastHidden := d.frontierHidden - draftTokens := make([]*mlx.Array, 0, maxTokens) draftDists := make([]sampler.Distribution, 0, maxTokens) var prefix *mlx.Array - for range maxTokens { - tokenEmbedding := d.target.TokenEmbeddings(lastToken) - inputs := tokenEmbedding.Concatenate(-1, lastHidden) - logits, projected := d.draft.Draft(inputs, anchor, d.caches) - stepLogits := lastLogits(logits) + for i := range maxTokens { + var hidden, projected *mlx.Array + if i == 0 && len(d.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 + } else { + // A head with draft caches writes each draft token to the next + // draft-cache slot, advancing one per step from the last committed + // slot (the held i==0 step stands in for that slot). A cacheless + // 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 { + pos = d.frontier - 1 + i + } + hidden, projected = d.spec.draft.Draft(&batch.Batch{ + InputIDs: lastToken, + SeqOffsets: []int32{int32(pos)}, + SeqQueryLens: []int32{1}, + Hidden: lastHidden, + }, d.spec.caches) + } + // Unembed only the row being sampled, never the batch. + stepLogits := d.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. dist := r.Sampler.Distribution(pipelineSlot, stepLogits, prefix) nextToken := r.Sampler.SampleDistribution(pipelineSlot, dist) lastToken = nextToken.ExpandDims(-1) - lastHidden = projected - draftTokens = append(draftTokens, lastToken) draftDists = append(draftDists, dist) if prefix == nil { prefix = lastToken @@ -117,7 +270,7 @@ func (d *mtpDrafter) propose(current *mlx.Array, maxTokens int) *draftCandidates } } return &draftCandidates{ - tokens: mlx.Concatenate(draftTokens, 1), + tokens: prefix, dist: sampler.ConcatenateDistributions(draftDists), } } diff --git a/x/mlxrunner/mtp_test.go b/x/mlxrunner/mtp_test.go index c449b3e7..bde9e607 100644 --- a/x/mlxrunner/mtp_test.go +++ b/x/mlxrunner/mtp_test.go @@ -5,6 +5,7 @@ import ( "encoding/json" "errors" "fmt" + "reflect" "slices" "strings" "testing" @@ -60,7 +61,10 @@ func (m *fakeMTPModel) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array mlx.Eval(b.InputIDs) ids := b.InputIDs.Ints() m.forwards = append(m.forwards, forwardCall{offset: b.SeqOffsets[0], n: int32(len(ids))}) - for _, c := range caches { + for i, c := range caches { + if i >= m.NumLayers() { + break + } if rc, ok := c.(*fakeRewindableCache); ok { seg := make([]int32, len(ids)) for i, id := range ids { @@ -85,25 +89,10 @@ func (m *fakeMTPModel) LoadWeights(map[string]*mlx.Array) error { return nil } -// TokenEmbeddings returns a width-1 embedding holding the token id as a float, -// so a draft can recover which token it is extending from inputs[...,0]. -func (m *fakeMTPModel) TokenEmbeddings(inputIDs *mlx.Array) *mlx.Array { - mlx.Eval(inputIDs) - ids := inputIDs.Ints() - data := make([]float32, len(ids)) - for i, id := range ids { - data[i] = float32(id) - } - return mlx.FromValues(data, inputIDs.Dim(0), inputIDs.Dim(1), 1) -} +var _ base.Model = (*fakeMTPModel)(nil) -var ( - _ base.Model = (*fakeMTPModel)(nil) - _ base.MTPEmbeddingModel = (*fakeMTPModel)(nil) -) - -// fakeMTPDraft extends the token in inputEmbeds through predict; a map (not a -// step counter) keeps drafting consistent regardless of batching. +// 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 { predict map[int32]int32 // calls records each Draft call so tests can assert the position convention. @@ -117,17 +106,82 @@ type draftCall struct { func (d *fakeMTPDraft) LoadWeights(map[string]*mlx.Array) error { return nil } -func (d *fakeMTPDraft) Draft(inputEmbeds *mlx.Array, position int32, caches []cache.Cache) (logits, hidden *mlx.Array) { - mlx.Eval(inputEmbeds) - prev := int32(inputEmbeds.Floats()[0]) - d.calls = append(d.calls, draftCall{position: position, from: prev}) +func (d *fakeMTPDraft) DraftCaches([]cache.Cache) []cache.Cache { return nil } + +func (d *fakeMTPDraft) Draft(b *batch.Batch, caches []cache.Cache) (hidden, projected *mlx.Array) { + mlx.Eval(b.InputIDs) + prev := int32(b.InputIDs.Ints()[0]) + d.calls = append(d.calls, draftCall{position: b.SeqOffsets[0], from: prev}) return oneHotLogits([]int32{d.predict[prev]}), mlx.Zeros(mlx.DTypeFloat32, 1, 1, mtpTestVocab) } -var ( - _ base.DraftModel = (*fakeMTPDraft)(nil) - _ base.MTPDraftModel = (*fakeMTPDraft)(nil) -) +// Unembed is the identity: the fake's hidden already is its one-hot logits. +func (d *fakeMTPDraft) Unembed(x *mlx.Array) *mlx.Array { return x } + +var _ base.DraftModel = (*fakeMTPDraft)(nil) + +// fakeKVDraft is a draft head with its own KV cache: it claims the trailing +// cache slot, writes its input ids there on every Draft call (advancing the +// offset like a real KV write), and records each call's offset, ids, and +// the identity of every fused hidden row. A target hidden row is one-hot +// (its hot index identifies which position it came from); the head's own +// projected hidden is all-zero and records as -1. +type fakeKVDraft struct { + predict map[int32]int32 + extends []extendCall +} + +// extendCall is one recorded Draft 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 projected hidden). +type extendCall struct { + offset int32 + ids []int32 + hiddens []int32 +} + +func (d *fakeKVDraft) LoadWeights(map[string]*mlx.Array) error { return nil } + +func (d *fakeKVDraft) DraftCaches(caches []cache.Cache) []cache.Cache { + return caches[len(caches)-1:] +} + +func (d *fakeKVDraft) Draft(b *batch.Batch, caches []cache.Cache) (hidden, projected *mlx.Array) { + mlx.Eval(b.InputIDs, b.Hidden) + rawIDs := b.InputIDs.Ints() + ids := make([]int32, len(rawIDs)) + for i, id := range rawIDs { + ids[i] = int32(id) + } + + hot := make([]int32, b.Hidden.Dim(1)) + flat := b.Hidden.Floats() + for r := range hot { + hot[r] = -1 + for v := range mtpTestVocab { + if flat[r*mtpTestVocab+v] != 0 { + hot[r] = int32(v) + break + } + } + } + d.extends = append(d.extends, extendCall{offset: b.SeqOffsets[0], ids: ids, hiddens: hot}) + + if rc, ok := d.DraftCaches(caches)[0].(*fakeRewindableCache); ok { + rc.feed(ids) + } + + preds := make([]int32, len(ids)) + for i, id := range ids { + preds[i] = d.predict[id] + } + return oneHotLogits(preds), mlx.Zeros(mlx.DTypeFloat32, 1, 1, mtpTestVocab) +} + +// Unembed is the identity: the fake's hidden already is its one-hot logits. +func (d *fakeKVDraft) Unembed(x *mlx.Array) *mlx.Array { return x } + +var _ base.DraftModel = (*fakeKVDraft)(nil) // newTestTokenizer builds a byte-level BPE tokenizer over single-character // tokens "0".."7" with the given EOS ids, so Decode(id) yields that digit and @@ -215,7 +269,7 @@ 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, oneHotLogits([]int32{2}), baseLogits, candidates) + results, accepted, err := spec.accept(&position, current, oneHotLogits([]int32{1}), baseLogits, candidates) if err != nil { t.Fatalf("accept: %v", err) } @@ -250,7 +304,7 @@ 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, oneHotLogits([]int32{2}), baseLogits, candidates) + results, accepted, err := spec.accept(&position, current, oneHotLogits([]int32{1}), baseLogits, candidates) if err != nil { t.Fatalf("accept: %v", err) } @@ -273,8 +327,8 @@ func TestAcceptMTPDraftsGreedyMismatch(t *testing.T) { func TestAcceptMTPDraftsGreedyEOS(t *testing.T) { skipIfNoMLX(t) // The second accepted draft token is EOS: it is recorded but stops - // generation and no bonus token is produced. Both accepted tokens are - // committed, so the caches hold the EOS's KV. + // generation and no bonus token is produced. The EOS's own KV is rolled + // back so the caches rest one token behind the recorded outputs. const eos int32 = 6 predict := map[int32]int32{1: 2, 2: eos, eos: 0} r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{}) @@ -287,7 +341,7 @@ 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, oneHotLogits([]int32{2}), baseLogits, candidates) + results, accepted, err := spec.accept(&position, current, oneHotLogits([]int32{1}), baseLogits, candidates) if err != nil { t.Fatalf("accept: %v", err) } @@ -302,11 +356,11 @@ func TestAcceptMTPDraftsGreedyEOS(t *testing.T) { if len(results) != accepted { t.Fatalf("results has %d tokens, want exactly the %d accepted with no bonus after EOS", len(results), accepted) } - if position != 2 { - t.Fatalf("position = %d, want 2 (both accepted tokens committed)", position) + if position != 1 { + t.Fatalf("position = %d, want 1 (kept draft committed, EOS dropped)", position) } - if got := caches[0].Offset(); got != 2 { - t.Fatalf("cache offset = %d, want 2 (both accepted tokens committed)", got) + if got := caches[0].Offset(); got != 1 { + t.Fatalf("cache offset = %d, want 1 (one behind the recorded outputs, EOS dropped)", got) } } @@ -392,7 +446,11 @@ func TestRunMTPDecodeSampled(t *testing.T) { CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, SamplerOpts: sampler.Options{Temperature: 1, Seed: 42, UseSeed: true}, } - d := testDecoder(r, req, caches, []int32{1}, position) + spec := r.spec.open(req, caches) + if spec == nil || !spec.enabled { + t.Fatalf("newSpeculationSession rejected a sampled request") + } + d := spec.decoder([]int32{1}, position) if err := r.decode(context.Background(), req, session, d, 0); err != nil { t.Fatalf("decode: %v", err) } @@ -518,7 +576,299 @@ func testDecoder(r *Runner, req Request, caches []cache.Cache, seed []int32, pos if spec := r.spec.open(req, caches); spec != nil { return spec.decoder(seed, position) } - return r.pipelinedDecoder(caches, seed, position) + return r.pipelinedDecoder(nil, caches, 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. + 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) // caches[0] target, caches[1] draft KV + session, ch := newMTPTestSession(caches) + position := 0 + + req := Request{ + Responses: ch, + Tokens: []int32{1}, + CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, + SamplerOpts: sampler.Options{}, + } + spec := r.spec.open(req, caches) + if spec == nil || !spec.enabled || len(spec.spec.targets) != 1 { + t.Fatalf("speculation engine not built around the draft caches") + } + defer spec.close() + 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, final := collectResponses(ch) + if content != "23456" { + t.Fatalf("content = %q, want %q", content, "23456") + } + if final.EvalCount != 5 || final.DoneReason != 0 { + t.Fatalf("EvalCount = %d, DoneReason = %d, want 5 and 0", final.EvalCount, final.DoneReason) + } + if got := []int32{2, 3, 4, 5, 6, eos}; !slices.Equal(session.outputs, got) { + t.Fatalf("session outputs = %v, want %v", session.outputs, got) + } + + // All caches rest together at the trie frontier: the prompt token plus + // five generated tokens; the EOS's KV is never committed. + if got := caches[0].Offset(); got != 6 { + t.Fatalf("target offset = %d, want 6", got) + } + if got := caches[1].Offset(); got != 6 { + t.Fatalf("draft cache offset = %d, want 6 (lockstep with target)", got) + } + + // 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. + 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 + } + if !reflect.DeepEqual(draft.extends, wantExtends) { + t.Fatalf("draft extends = %+v, want %+v", draft.extends, wantExtends) + } + + // The committed draft KV holds the look-ahead tokens, one per slot. + if got, want := caches[1].(*fakeRewindableCache).tokens, []int32{2, 3, 4, 5, 6, eos}; !slices.Equal(got, want) { + t.Fatalf("draft cache = %v, want %v", got, want) + } +} + +func TestDecodeKVDraftRejectionRebuildsFromTarget(t *testing.T) { + skipIfNoMLX(t) + // The draft mispredicts mid-chain: it proposes 6 where the target's own + // next token is 4, so the round is accepted only up to the rejection and the + // loop re-proposes from the target's correction. The speculative draft KV + // written for the rejected proposals is rolled back and rewritten from the + // validated run's target hiddens, so the committed draft cache holds the + // target chain with no trace of the rejected tokens. + const eos int32 = 7 + target := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: eos, eos: 0} + r := mtpTestRunner(t, target, []int32{eos}, sampler.Options{}) + // The draft mirrors the target except at 3, where it proposes 6 (absent from + // 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 + session, ch := newMTPTestSession(caches) + position := 0 + + req := Request{ + Responses: ch, + Tokens: []int32{1}, + CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, + SamplerOpts: sampler.Options{}, + } + spec := r.spec.open(req, caches) + spec.limit = 4 + defer spec.close() + 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, final := collectResponses(ch) + if content != "2345" { + t.Fatalf("content = %q, want %q", content, "2345") + } + if final.DoneReason != 0 { + t.Fatalf("DoneReason = %d, want 0 (EOS)", final.DoneReason) + } + if got := []int32{2, 3, 4, 5, eos}; !slices.Equal(session.outputs, got) { + t.Fatalf("session outputs = %v, want %v", session.outputs, got) + } + + // The divergent token was drafted into the speculative KV... + drafted := false + for _, e := range draft.extends { + if slices.Contains(e.ids, 6) { + drafted = true + } + } + if !drafted { + t.Fatalf("draft never proposed the divergent token 6; extends = %+v", draft.extends) + } + // ...but the rejection rolled it back: both caches rest in lockstep at the + // target chain, and the committed draft KV holds only the validated + // look-ahead tokens. + if got := caches[0].Offset(); got != 5 { + t.Fatalf("target offset = %d, want 5", got) + } + if got := caches[1].Offset(); got != 5 { + t.Fatalf("draft cache offset = %d, want 5 (lockstep with target)", got) + } + if got, want := caches[1].(*fakeRewindableCache).tokens, []int32{2, 3, 4, 5, eos}; !slices.Equal(got, want) { + t.Fatalf("draft cache = %v, want %v (rejected proposals rolled back)", got, want) + } +} + +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. + 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} + r := mtpTestRunner(t, predict, []int32{eos}, opts) + draft := &fakeKVDraft{predict: predict} + r.spec = newSpeculation(r, draft) + + caches, _ := newMTPTestCaches(2) + session, ch := newMTPTestSession(caches) + position := 0 + + req := Request{ + Responses: ch, + Tokens: []int32{1}, + CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, + SamplerOpts: opts, + } + spec := r.spec.open(req, caches) + if spec == nil || spec.enabled { + t.Fatalf("want a maintain-only speculationSession, got %+v", spec) + } + 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() + spec.close() // the pipeline defers this before session close + + content, _ := collectResponses(ch) + if content != "23456" { + t.Fatalf("content = %q, want %q", content, "23456") + } + + // The pipelined decoder forwards every emitted token, including the + // ending EOS, so both caches rest at the full recorded path. + if got := caches[0].Offset(); got != 7 { + t.Fatalf("target offset = %d, want 7", got) + } + if got := caches[1].Offset(); got != 7 { + t.Fatalf("draft cache offset = %d, want 7 (lockstep with target)", got) + } + + // A session that never drafts defers every committed pair: the whole + // generation arrives in one batched flush at finish, with the EOS slot's + // pair completed by the in-flight sample (predict[eos] = 0) that decoding + // discarded. + wantExtends := []extendCall{ + {offset: 0, ids: []int32{2, 3, 4, 5, 6, eos, 0}, hiddens: []int32{2, 3, 4, 5, 6, eos, 0}}, + } + if !reflect.DeepEqual(draft.extends, wantExtends) { + t.Fatalf("draft extends = %+v, want %+v", draft.extends, wantExtends) + } + if got, want := caches[1].(*fakeRewindableCache).tokens, []int32{2, 3, 4, 5, 6, eos, 0}; !slices.Equal(got, want) { + t.Fatalf("draft cache = %v, want %v", got, want) + } +} + +func TestFlushLevelsDraftCacheWithPrefill(t *testing.T) { + skipIfNoMLX(t) + // Prefill attaches its scheduled snapshots only at offsets every cache + // has crossed, so the pipeline flushes the drafter's buffered pairs + // first: after a prefill-sized committed report and a flush, the + // draft cache covers every completed pair, one slot behind the target + // (the frontier pair still awaits its look-ahead token). + const eos int32 = 7 + 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) + req := Request{ + CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, + SamplerOpts: sampler.Options{}, + } + spec := r.spec.open(req, caches) + defer spec.close() + + // The prompt's only chunk: tokens 1..4 at slots 0..3 with their hiddens. + spec.committed(mlx.FromValues([]int32{1, 2, 3, 4}, 1, 4), oneHotLogits([]int32{1, 2, 3, 4}), 0) + spec.flush() + + if got := caches[1].Offset(); got != 3 { + t.Fatalf("draft cache offset after flush = %d, want 3 (all completed pairs)", got) + } + wantExtends := []extendCall{ + {offset: 0, ids: []int32{2, 3, 4}, hiddens: []int32{1, 2, 3}}, + } + if !reflect.DeepEqual(draft.extends, wantExtends) { + t.Fatalf("draft extends = %+v, want %+v", draft.extends, wantExtends) + } +} + +func TestCommittedRunBatchesPastFlushCap(t *testing.T) { + skipIfNoMLX(t) + // A committed run longer than the pending-flush cap still writes the draft + // caches in a single head forward: the run's completed pairs coalesce into + // one batched extend at the run's start, rather than splitting at the cap. + const n = mtpPendingFlushTokens + 8 + predict := map[int32]int32{} + tokens := make([]int32, n) + for i := range tokens { + tokens[i] = int32(i%7 + 1) + predict[tokens[i]] = int32((i+1)%7 + 1) + } + + r := mtpTestRunner(t, predict, []int32{0}, sampler.Options{}) + draft := &fakeKVDraft{predict: predict} + r.spec = newSpeculation(r, draft) + + caches, _ := newMTPTestCaches(2) + req := Request{ + CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, + SamplerOpts: sampler.Options{}, + } + spec := r.spec.open(req, caches) + defer spec.close() + + // One prefill-sized chunk: n tokens at slots 0..n-1 with their hiddens. + spec.committed(mlx.FromValues(tokens, 1, n), oneHotLogits(tokens), 0) + spec.flush() + + if got := len(draft.extends); got != 1 { + t.Fatalf("draft extends = %d calls, want 1 batched extend", got) + } + if got, want := len(draft.extends[0].ids), n-1; got != want { + t.Fatalf("batched extend ids = %d, want %d (every completed pair)", got, want) + } + if got := draft.extends[0].offset; got != 0 { + t.Fatalf("batched extend offset = %d, want 0", got) + } + if got := caches[1].Offset(); got != n-1 { + t.Fatalf("draft cache offset = %d, want %d (all completed pairs)", got, n-1) + } } // newMTPTestCaches returns n rewindable fake caches sharing one snapshot @@ -539,19 +889,16 @@ func newMTPTestSession(caches []cache.Cache) (*cacheSession, chan CompletionResp return &cacheSession{caches: caches}, ch } -// testSpeculationSession wires a speculation engine around the runner's drafter and -// caches for tests that drive accept directly. A runner without a draft -// model gets a no-op drafter, since the engine requires one. +// testSpeculationSession wires a speculation engine around a drafter and caches for +// tests that drive accept directly. A runner without a draft model gets a +// no-op drafter, since the engine requires one. func testSpeculationSession(r *Runner, caches []cache.Cache) *speculationSession { - s := r.spec - if s == nil { - s = &speculation{r: r} + if r.spec != nil { + r.spec.bind(caches) + return &speculationSession{spec: r.spec, drafter: newMTPDrafter(r.spec)} } - var d drafter = nopDrafter{} - if md := newMTPDrafter(s, caches); md != nil { - d = md - } - return &speculationSession{spec: s, drafter: d, caches: caches} + s := &speculation{r: r, caches: caches, targets: caches} + return &speculationSession{spec: s, drafter: nopDrafter{}} } // nopDrafter satisfies drafter for engine tests that supply candidates @@ -560,6 +907,8 @@ type nopDrafter struct{} func (nopDrafter) propose(*mlx.Array, int) *draftCandidates { return nil } func (nopDrafter) committed(_, _ *mlx.Array, _ int) {} +func (nopDrafter) finish(*mlx.Array) {} +func (nopDrafter) flush() {} func (nopDrafter) close() {} // scriptedCandidates builds draft candidates by running the real drafter @@ -573,7 +922,8 @@ func scriptedCandidates(r *Runner, tokens []int32) *draftCandidates { chain[prev] = tok prev = tok } - d := &mtpDrafter{spec: &speculation{r: r}, draft: &fakeMTPDraft{predict: chain}, target: r.Model.(base.MTPEmbeddingModel)} + s := &speculation{r: r, draft: &fakeMTPDraft{predict: chain}} + d := &mtpDrafter{spec: s} 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)) } diff --git a/x/mlxrunner/pipeline.go b/x/mlxrunner/pipeline.go index 3f5f31e2..677fc311 100644 --- a/x/mlxrunner/pipeline.go +++ b/x/mlxrunner/pipeline.go @@ -75,7 +75,12 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er defer session.close() caches := session.caches - seed, position, promptEval, err := r.prefill(ctx, session) + // Built before prefill so a drafter with draft caches follows the prompt + // through prefill alongside the target. + spec := r.spec.open(request, caches) + defer spec.close() + + seed, position, promptEval, err := r.prefill(ctx, session, spec) if err != nil { return err } @@ -83,14 +88,11 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er // Register the sampler after prefill completes. r.Sampler.Add(pipelineSlot, request.SamplerOpts, inputs) - spec := r.spec.open(request, caches) - defer spec.close() - var d decoder if spec != nil { d = spec.decoder(seed, position) } else { - d = r.pipelinedDecoder(caches, seed, position) + d = r.pipelinedDecoder(nil, caches, seed, position) } defer d.close() return r.decode(ctx, request, session, d, promptEval) @@ -100,7 +102,7 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er // 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) ([]int32, int, time.Duration, error) { +func (r *Runner) prefill(ctx context.Context, session *cacheSession, spec *speculationSession) ([]int32, int, time.Duration, error) { start := time.Now() inputs := session.inputs tokens := session.remaining @@ -143,11 +145,13 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession) ([]int32, i n := min(prefillChunk, total-processed-1) - r.Model.Forward(&batch.Batch{ - InputIDs: mlx.FromValues(tokens[processed:processed+n], 1, n), + chunkIDs := mlx.FromValues(tokens[processed:processed+n], 1, n) + hidden := r.Model.Forward(&batch.Batch{ + InputIDs: chunkIDs, SeqOffsets: []int32{int32(position)}, SeqQueryLens: []int32{int32(n)}, }, caches) + spec.committed(chunkIDs, hidden, position) mlx.Sweep() materializeCaches() processed += n @@ -158,7 +162,10 @@ func (r *Runner) prefill(ctx context.Context, session *cacheSession) ([]int32, i mlx.ClearCache() } - // Attach the snapshots captured during prefill to the trie. + // Flush before attaching: snapshots attach only at offsets every cache + // has crossed, and a drafter with draft caches keeps buffered pairs that + // would otherwise hold those caches short of the scheduled offsets. + spec.flush() session.attachPrefillSnapshots() return tokens[processed:], position, time.Since(start), nil @@ -263,15 +270,18 @@ func (r *Runner) decode(ctx context.Context, request Request, session *cacheSess // the next token's chain is dispatched before the returned one is // synchronized, so the device runs ahead of host emission. type pipelinedDecoder struct { - r *Runner + r *Runner + // spec, when non-nil, receives every forwarded token and settles its + // drafter at close, keeping a non-drafting session's draft KV level. + spec *speculationSession caches []cache.Cache position int sample sampler.Result // in flight: sampled, not yet forwarded emitted sampler.Result // last call's result, pinned until the next call } -func (r *Runner) pipelinedDecoder(caches []cache.Cache, seed []int32, position int) *pipelinedDecoder { - t := &pipelinedDecoder{r: r, caches: caches, position: position} +func (r *Runner) pipelinedDecoder(spec *speculationSession, caches []cache.Cache, seed []int32, position int) *pipelinedDecoder { + t := &pipelinedDecoder{r: r, spec: spec, caches: caches, position: position} t.sample = t.dispatch(mlx.FromValues(seed, 1, len(seed))) return t } @@ -285,8 +295,10 @@ func (t *pipelinedDecoder) dispatch(token *mlx.Array) sampler.Result { SeqOffsets: []int32{int32(t.position)}, SeqQueryLens: []int32{int32(token.Dim(1))}, }, t.caches) + t.spec.committed(token, hidden, t.position) t.position += token.Dim(1) - next := r.Sampler.Sample([]int{pipelineSlot}, lastLogits(r.Model.Unembed(hidden))) + 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)) mlx.Pin(next.Arrays()...) mlx.Sweep() mlx.AsyncEval(next.Arrays()...) @@ -300,6 +312,9 @@ func (t *pipelinedDecoder) next(int) ([]sampler.Result, error) { } 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. + t.spec.finish(t.sample.Token) mlx.Unpin(t.emitted.Arrays()...) mlx.Unpin(t.sample.Arrays()...) } diff --git a/x/mlxrunner/runner.go b/x/mlxrunner/runner.go index 189bfda9..c76e3e4d 100644 --- a/x/mlxrunner/runner.go +++ b/x/mlxrunner/runner.go @@ -81,6 +81,9 @@ func (r *Runner) Load(modelName string) error { return err } draftModel = draft + } else if sd, ok := m.(base.SelfDraft); ok { + // Inline draft head: already loaded with the target; nil if none shipped. + draftModel = sd.SelfDraft() } collected := mlx.Collect(m) diff --git a/x/mlxrunner/speculate.go b/x/mlxrunner/speculate.go index 2b9a46a2..baf9811e 100644 --- a/x/mlxrunner/speculate.go +++ b/x/mlxrunner/speculate.go @@ -4,6 +4,7 @@ import ( "fmt" "log/slog" "os" + "slices" "strconv" "strings" @@ -23,10 +24,18 @@ type drafter 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 — the decode - // seed, then each round's validated tokens. + // 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) + // finish reports generation ended with current sampled but never + // committed, so the drafter can settle state tracking the target caches. + finish(current *mlx.Array) + + // flush writes any buffered committed reports through to the draft + // caches. A drafter without draft caches has nothing to write. + flush() + close() } @@ -100,6 +109,15 @@ func positiveEnvInt(key string) int { type speculation struct { r *Runner draft base.DraftModel + + // caches is the whole persistent slice, passed to every forward; draftKV + // are the draft head's own caches and targets are the rest — the caches the + // target forward writes, which speculation snapshots and rollback cover. + // Bound the first time the caches exist (the Runner reuses one cache slice + // for its life) and stable thereafter. + caches []cache.Cache + draftKV []cache.Cache + targets []cache.Cache } // newSpeculation builds the speculative-decoding subsystem for a loaded model, @@ -111,41 +129,68 @@ func newSpeculation(r *Runner, draft base.DraftModel) *speculation { return &speculation{r: r, draft: draft} } +// bind computes the draft/target cache partition the first time the persistent +// caches exist; later requests reuse the same slice, so it runs once. +func (s *speculation) bind(caches []cache.Cache) { + if s.caches != nil { + if !slices.Equal(s.caches, caches) { + panic("speculation: cache slice changed between requests") + } + return + } + draftKV := s.draft.DraftCaches(caches) + + // Partition caches into target slots (everything not in draftKV) in one + // pass. The count check rejects a draft slot that isn't a member of caches. + targets := make([]cache.Cache, 0, len(caches)) + for _, c := range caches { + if !slices.Contains(draftKV, c) { + targets = append(targets, c) + } + } + if len(caches)-len(targets) != len(draftKV) { + panic("speculation: DraftCaches must select slots of the cache slice") + } + + s.caches = caches + s.draftKV = draftKV + s.targets = targets +} + // speculationSession is the speculation cursor for one request over a speculation. A // nil speculationSession is a plain decode. type speculationSession struct { spec *speculation drafter drafter - caches []cache.Cache + // 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 stats specStats } -// open returns the speculation cursor for this request, or nil when the model -// carries no draft head or the request cannot speculate. Logprobs are not yet -// supported. A nil receiver (no draft head) decodes plainly. +// 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, caches []cache.Cache) *speculationSession { if s == nil { return nil } - d := newMTPDrafter(s, caches) + s.bind(caches) + d := newMTPDrafter(s) if d == nil { return nil } + opts := request.SamplerOpts - if !s.r.mtpDefaults(opts.Temperature != 0).Enabled { - return nil - } - if opts.Logprobs || opts.TopLogprobs > 0 { - return nil - } + enabled := s.r.mtpDefaults(opts.Temperature != 0).Enabled && + !opts.Logprobs && opts.TopLogprobs == 0 specOpts := s.r.loadSpecOptions(opts.Temperature != 0) return &speculationSession{ spec: s, drafter: d, - caches: caches, + enabled: enabled, opts: specOpts, limit: specOpts.initialDraftTokens, stats: specStats{maxDraft: specOpts.initialDraftTokens}, @@ -159,6 +204,23 @@ func (s *speculationSession) committed(tokens, hiddens *mlx.Array, position int) s.drafter.committed(tokens, hiddens, position) } +// finish reports the end of generation to the drafter: current was sampled +// after the last committed slot and will never be committed. +func (s *speculationSession) finish(current *mlx.Array) { + if s == nil { + return + } + s.drafter.finish(current) +} + +// flush writes the drafter's buffered committed reports to the draft caches. +func (s *speculationSession) flush() { + if s == nil { + return + } + s.drafter.flush() +} + func (s *speculationSession) close() { if s == nil { return @@ -178,7 +240,13 @@ type speculativeDecoder struct { current sampler.Result // emitted (or the seed), not yet forwarded } -func (s *speculationSession) decoder(seed []int32, position int) *speculativeDecoder { +// 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. +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} @@ -192,7 +260,7 @@ func (st *speculativeDecoder) next(remaining int) ([]sampler.Result, error) { InputIDs: tokenInput(st.current.Token), SeqOffsets: []int32{int32(st.position)}, SeqQueryLens: []int32{1}, - }, s.caches) + }, s.spec.caches) st.position++ results, next, err := s.round(&st.position, st.current, hidden, remaining) @@ -212,6 +280,10 @@ func (st *speculativeDecoder) next(remaining int) ([]sampler.Result, error) { } 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) mlx.Unpin(st.current.Arrays()...) st.s.logStats() } @@ -368,7 +440,7 @@ func (s *speculationSession) accept(position *int, current sampler.Result, hidde r := s.spec.r before := *position draftCount := candidates.tokens.Dim(1) - scheduleSpeculation(s.caches, before, draftCount) + scheduleSpeculation(s.spec.targets, before, draftCount) // Every exit between schedule and commit must drain the snapshot // schedule and roll the speculative writes back out of the live caches: @@ -380,7 +452,7 @@ func (s *speculationSession) accept(position *int, current sampler.Result, hidde return } committed = true - commitSpeculation(s.caches, keep, draftCount, before) + commitSpeculation(s.spec.targets, keep, draftCount, before) } defer commit(0) @@ -388,7 +460,7 @@ func (s *speculationSession) accept(position *int, current sampler.Result, hidde InputIDs: candidates.tokens, SeqOffsets: []int32{int32(before)}, SeqQueryLens: []int32{int32(draftCount)}, - }, s.caches) + }, s.spec.caches) targetDist := r.Sampler.Distribution(pipelineSlot, validationLogits(r, baseLogits, hiddenSeq), candidates.tokens) draftDist := candidates.dist @@ -419,18 +491,24 @@ func (s *speculationSession) accept(position *int, current sampler.Result, hidde break } } + // The final token of a generation is recorded and streamed but its KV + // is never committed: the caches rest at the trie frontier. + keep := accepted + if done { + keep-- + } + commit(keep) + *position = before + keep - commit(accepted) - *position = before + accepted - - // Report the validated run — current plus the accepted drafts, with the + // Report the validated run — current plus the kept drafts, with the // 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. - runIDs := append([]int32{int32(current.Token.Int())}, commitIDs...) + // describing exactly what the target caches hold. A done generation's + // final token is never committed; it reaches the drafter through finish. + runIDs := append([]int32{int32(current.Token.Int())}, commitIDs[:keep]...) runHiddens := lastHiddenRow(hidden) - if accepted > 0 { - runHiddens = runHiddens.Concatenate(1, hiddenSeq.Slice(mlx.Slice(), mlx.Slice(0, accepted), mlx.Slice())) + if keep > 0 { + runHiddens = runHiddens.Concatenate(1, hiddenSeq.Slice(mlx.Slice(), mlx.Slice(0, keep), mlx.Slice())) } s.drafter.committed(mlx.FromValues(runIDs, 1, len(runIDs)), runHiddens, before-1) diff --git a/x/models/gemma4/assistant.go b/x/models/gemma4/assistant.go index 20a67560..36b67504 100644 --- a/x/models/gemma4/assistant.go +++ b/x/models/gemma4/assistant.go @@ -12,10 +12,7 @@ import ( "github.com/ollama/ollama/x/models/nn" ) -var ( - _ base.DraftModel = (*AssistantModel)(nil) - _ base.MTPDraftModel = (*AssistantModel)(nil) -) +var _ base.DraftModel = (*AssistantModel)(nil) type AssistantConfig struct { TextConfig TextConfig `json:"text_config"` @@ -37,6 +34,10 @@ type AssistantModel struct { NormScaled *mlx.Array + // target supplies the scaled token embeddings the MTP head fuses with + // the target hidden state. + target *Model + *AssistantConfig tensorPrefix string @@ -145,6 +146,7 @@ func newAssistantModel(root *model.Root, target base.Model) (base.DraftModel, er m := &AssistantModel{ AssistantConfig: &cfg, tensorPrefix: tensorPrefix, + target: targetGemma, Layers: make([]*AssistantLayer, cfg.TextConfig.NumHiddenLayers), TensorQuant: root.AllTensorQuant(), } @@ -267,26 +269,42 @@ func (m *AssistantModel) precomputeScaledWeights() { } } -func (m *AssistantModel) Draft(inputsEmbeds *mlx.Array, position int32, caches []cache.Cache) (logits, hidden *mlx.Array) { +// DraftCaches returns nil: the assistant keeps no KV and re-attends the +// target's caches read-only each step. +func (m *AssistantModel) DraftCaches([]cache.Cache) []cache.Cache { return nil } + +// Draft is single-position: the head drafts every token as if it sat at the +// last target-seen position, so it anchors RoPE and the mask there — from the +// non-moving target full-attention cache — regardless of the advancing offset +// in b. The input fuses the target's scaled token embedding with b.Hidden. +func (m *AssistantModel) Draft(b *batch.Batch, caches []cache.Cache) (hidden, projected *mlx.Array) { + inputsEmbeds := m.target.TokenEmbeddings(b.InputIDs).Concatenate(-1, b.Hidden) dims := inputsEmbeds.Dims() B, L := int32(dims[0]), int32(dims[1]) - b := &batch.Batch{ + + anchor := int32(0) + if len(caches) > 0 { + if full := caches[len(caches)-1]; full != nil { + anchor = int32(full.Offset() - 1) + } + } + ab := &batch.Batch{ InputIDs: mlx.Zeros(mlx.DTypeInt32, int(B), int(L)), - SeqOffsets: []int32{position}, + SeqOffsets: []int32{anchor}, SeqQueryLens: []int32{L}, } - sliding, full := m.sharedHistories(b, caches) + sliding, full := m.sharedHistories(ab, caches) h := m.PreProjection.Forward(inputsEmbeds) - positions := mlx.FromValues([]int32{position}, 1) + positions := mlx.FromValues([]int32{anchor}, 1) for _, layer := range m.Layers { - h = layer.Forward(h, b, positions, B, L, &m.TextConfig, sliding, full) + h = layer.Forward(h, ab, positions, B, L, &m.TextConfig, sliding, full) } hidden = mlx.RMSNormFn(h, m.NormScaled, m.TextConfig.RMSNormEps) - projected := m.PostProjection.Forward(hidden) - return m.unembed(hidden), projected + projected = m.PostProjection.Forward(hidden) + return hidden, projected } func (m *AssistantModel) sharedHistories(b *batch.Batch, caches []cache.Cache) (sliding, full *nn.KVHistory) { @@ -302,7 +320,7 @@ func (m *AssistantModel) sharedHistories(b *batch.Batch, caches []cache.Cache) ( return sliding, full } -func (m *AssistantModel) unembed(hidden *mlx.Array) *mlx.Array { +func (m *AssistantModel) Unembed(hidden *mlx.Array) *mlx.Array { if m.UseOrderedEmbeddings { return m.applyCentroidMasking(hidden) }