diff --git a/x/mlxrunner/mtp.go b/x/mlxrunner/mtp.go index e56d014e..52381f92 100644 --- a/x/mlxrunner/mtp.go +++ b/x/mlxrunner/mtp.go @@ -1,15 +1,6 @@ package mlxrunner import ( - "context" - "fmt" - "log/slog" - "os" - "strconv" - "strings" - "time" - - "github.com/ollama/ollama/x/mlxrunner/batch" "github.com/ollama/ollama/x/mlxrunner/cache" "github.com/ollama/ollama/x/mlxrunner/mlx" "github.com/ollama/ollama/x/mlxrunner/model/base" @@ -21,31 +12,6 @@ const ( mtpDefaultMaxDraftTokens = 16 ) -type mtpDraftSchedule string - -const ( - mtpDraftScheduleHeuristic mtpDraftSchedule = "heuristic" - mtpDraftScheduleConstant mtpDraftSchedule = "constant" -) - -type mtpStats struct { - iterations int - drafted int - accepted int - mismatches int - allAccepted int - maxDraft int - targetDuration time.Duration - draftDuration time.Duration - validateDuration time.Duration -} - -type mtpOptions struct { - initialDraftTokens int - maxDraftTokens int - draftSchedule mtpDraftSchedule -} - func (r *Runner) mtpDefaults(sample bool) base.MTPDefaults { defaults := base.MTPDefaults{ InitialDraftTokens: mtpDefaultInitialDraftTokens, @@ -64,245 +30,83 @@ func (r *Runner) mtpDefaults(sample bool) base.MTPDefaults { return defaults } -func (r *Runner) loadMTPOptions(sample bool) mtpOptions { - defaults := r.mtpDefaults(sample) +// 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. +type mtpDrafter struct { + spec *speculation + draft base.MTPDraftModel + target base.MTPEmbeddingModel + caches []cache.Cache - opts := mtpOptions{ - initialDraftTokens: defaults.InitialDraftTokens, - maxDraftTokens: defaults.MaxDraftTokens, - draftSchedule: mtpDraftScheduleConstant, - } - if v := positiveEnvInt("OLLAMA_MLX_MTP_MAX_DRAFT_TOKENS"); v > 0 { - opts.maxDraftTokens = v - } - if v := positiveEnvInt("OLLAMA_MLX_MTP_INITIAL_DRAFT_TOKENS"); v > 0 { - opts.initialDraftTokens = v - } - if opts.initialDraftTokens > opts.maxDraftTokens { - opts.initialDraftTokens = opts.maxDraftTokens - } - switch schedule := strings.ToLower(strings.TrimSpace(os.Getenv("OLLAMA_MLX_MTP_DRAFT_SCHEDULE"))); schedule { - case "", string(mtpDraftScheduleConstant): - opts.draftSchedule = mtpDraftScheduleConstant - case string(mtpDraftScheduleHeuristic): - opts.draftSchedule = mtpDraftScheduleHeuristic - default: - slog.Warn("invalid MTP env setting", "key", "OLLAMA_MLX_MTP_DRAFT_SCHEDULE", "value", schedule) - } - return opts + // 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 int + frontierHidden *mlx.Array } -func positiveEnvInt(key string) int { - raw := os.Getenv(key) - if raw == "" { - return 0 - } - v, err := strconv.Atoi(raw) - if err != nil || v <= 0 { - slog.Warn("invalid MTP env setting", "key", key, "value", raw) - return 0 - } - return v -} - -// useMTP reports whether the request can run through the MTP speculative -// decode path. Logprobs are not yet supported. -func (r *Runner) useMTP(opts sampler.Options) bool { - if r.Draft == nil { - return false - } - if _, ok := r.Draft.(base.MTPDraftModel); !ok { - return false - } - if _, ok := r.Model.(base.MTPEmbeddingModel); !ok { - return false - } - if !r.mtpDefaults(opts.Temperature != 0).Enabled { - return false - } - if opts.Logprobs || opts.TopLogprobs > 0 { - return false - } - return true -} - -func (r *Runner) runMTPDecode(ctx context.Context, request Request, session *cacheSession, caches []cache.Cache, seed []int32, position *int, started time.Time) error { - targetEmbeddings := r.Model.(base.MTPEmbeddingModel) - draft := r.Draft.(base.MTPDraftModel) - mtpOpts := r.loadMTPOptions(request.SamplerOpts.Temperature != 0) - stats := mtpStats{maxDraft: mtpOpts.initialDraftTokens} - draftLimit := mtpOpts.initialDraftTokens - slog.Info("MTP decode enabled", "initial_draft_tokens", mtpOpts.initialDraftTokens, "max_draft_tokens", mtpOpts.maxDraftTokens, "draft_schedule", mtpOpts.draftSchedule) - - targetForward := func(token *mlx.Array) *mlx.Array { - fwd := r.Model.Forward(&batch.Batch{ - InputIDs: token, - SeqOffsets: []int32{int32(*position)}, - SeqQueryLens: []int32{int32(token.Dim(1))}, - }, caches) - *position += token.Dim(1) - return fwd - } - - hidden := targetForward(mlx.FromValues(seed, 1, len(seed))) - current := r.Sampler.Sample([]int{pipelineSlot}, r.lastLogits(hidden)) - mlx.Pin(current.Arrays()...) - mlx.Sweep() - mlx.AsyncEval(current.Arrays()...) - defer func() { - mlx.Unpin(current.Arrays()...) - }() - - dec := decoder{tokenizer: r.Tokenizer} - final := CompletionResponse{Done: true, PromptEvalCount: len(request.Tokens), DoneReason: 1} - now := started - - generated := 0 - for generated < request.Options.NumPredict { - if err := ctx.Err(); err != nil { - return err - } - - t0 := time.Now() - hidden = targetForward(mtpTokenInput(current.Token)) - baseLogits := r.lastLogits(hidden) - stats.targetDuration += time.Since(t0) - - if generated == 0 { - mlx.Eval(current.Arrays()...) - final.PromptEvalDuration = time.Since(now) - now = time.Now() - } - - done, err := r.emitTokens(ctx, request, session, &dec, []sampler.Result{current}, &final, &generated) - if err != nil { - return err - } - if done { - break - } - - stats.iterations++ - maxDraft := min(draftLimit, request.Options.NumPredict-generated) - t0 = time.Now() - candidates := r.generateMTPDraftCandidates(draft, targetEmbeddings, current.Token, hidden, caches, int32(*position-1), maxDraft) - draftCount := 0 - var candidateArrays []*mlx.Array - if candidates != nil { - draftCount = candidates.tokens.Dim(1) - candidateArrays = append([]*mlx.Array{baseLogits}, candidates.Arrays()...) - mlx.Pin(candidateArrays...) - mlx.Sweep() - } - stats.draftDuration += time.Since(t0) - stats.drafted += draftCount - var next sampler.Result - if draftCount == 0 { - next = r.Sampler.Sample([]int{pipelineSlot}, baseLogits) - } else { - var accepted int - t0 = time.Now() - next, accepted, done, err = r.acceptMTPDrafts(ctx, request, session, &dec, caches, position, baseLogits, candidates, &final, &generated) - stats.validateDuration += time.Since(t0) - mlx.Unpin(candidateArrays...) - if err != nil { - return err - } - stats.accepted += accepted - switch { - case mtpOpts.draftSchedule == mtpDraftScheduleConstant: - case accepted == draftCount: - stats.allAccepted++ - draftLimit = min(mtpOpts.maxDraftTokens, draftLimit+2) - default: - stats.mismatches++ - draftLimit = max(1, draftLimit-1) - } - if mtpOpts.draftSchedule == mtpDraftScheduleConstant { - if accepted == draftCount { - stats.allAccepted++ - } else { - stats.mismatches++ - } - } - stats.maxDraft = max(stats.maxDraft, draftLimit) - if next.Token == nil { - mlx.Sweep() - } - if done || generated >= request.Options.NumPredict { - break - } - } - - mlx.Pin(next.Arrays()...) - old := current - current = next - mlx.Unpin(old.Arrays()...) - mlx.Sweep() - mlx.AsyncEval(current.Arrays()...) - - if generated%256 == 0 { - mlx.ClearCache() - } - } - - final.EvalCount = generated - final.EvalDuration = time.Since(now) - acceptance := 0.0 - if stats.drafted > 0 { - acceptance = float64(stats.accepted) / float64(stats.drafted) - } - avgDraft := 0.0 - avgAccepted := 0.0 - if stats.iterations > 0 { - avgDraft = float64(stats.drafted) / float64(stats.iterations) - avgAccepted = float64(stats.accepted) / float64(stats.iterations) - } - slog.Info("MTP decode stats", "generated", generated, "drafted", stats.drafted, "accepted", stats.accepted, "acceptance", acceptance, "iterations", stats.iterations, "avg_draft", avgDraft, "avg_accepted", avgAccepted, "mismatches", stats.mismatches, "all_accepted", stats.allAccepted, "max_draft", stats.maxDraft, "draft_schedule", mtpOpts.draftSchedule, "target_duration", stats.targetDuration, "draft_duration", stats.draftDuration, "validate_duration", stats.validateDuration) - select { - case <-ctx.Done(): - return ctx.Err() - case request.Responses <- final: +// 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 } -} - -type mtpDraftCandidates struct { - tokens *mlx.Array - // dist is the proposal distribution used to sample each drafted token. - dist sampler.Distribution -} - -func (c *mtpDraftCandidates) Arrays() []*mlx.Array { - if c == nil { + target, ok := s.r.Model.(base.MTPEmbeddingModel) + if !ok { return nil } - return append([]*mlx.Array{c.tokens}, c.dist.Arrays()...) + return &mtpDrafter{spec: s, draft: draft, target: target, caches: caches} } -func (r *Runner) generateMTPDraftCandidates(draft base.MTPDraftModel, target base.MTPEmbeddingModel, token *mlx.Array, hidden *mlx.Array, caches []cache.Cache, position int32, maxDraft int) *mtpDraftCandidates { - if maxDraft <= 0 { +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) + } + d.frontierHidden = h +} + +func (d *mtpDrafter) close() { + if d.frontierHidden != nil { + mlx.Unpin(d.frontierHidden) + d.frontierHidden = nil + } +} + +// 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. +func (d *mtpDrafter) propose(current *mlx.Array, maxTokens int) *draftCandidates { + if maxTokens <= 0 || d.frontierHidden == nil { return nil } + r := d.spec.r - lastToken := mtpTokenInput(token) - lastHidden := hidden - draftTokens := make([]*mlx.Array, 0, maxDraft) - draftDists := make([]sampler.Distribution, 0, maxDraft) + anchor := int32(d.frontier - 1) + lastToken := current.ExpandDims(-1) + lastHidden := d.frontierHidden + draftTokens := make([]*mlx.Array, 0, maxTokens) + draftDists := make([]sampler.Distribution, 0, maxTokens) var prefix *mlx.Array - // Gemma4 assistant MTP is trained as "single-position" drafting: - // keep the RoPE/cache position anchored at the last target-seen token - // while the proposed token and projected hidden state advance. - for range maxDraft { - tokenEmbedding := target.TokenEmbeddings(lastToken) + for range maxTokens { + tokenEmbedding := d.target.TokenEmbeddings(lastToken) inputs := tokenEmbedding.Concatenate(-1, lastHidden) - logits, projected := draft.Draft(inputs, position, caches) - stepLogits := r.lastLogitsFromLogits(logits) + logits, projected := d.draft.Draft(inputs, anchor, d.caches) + stepLogits := lastLogits(logits) dist := r.Sampler.Distribution(pipelineSlot, stepLogits, prefix) nextToken := r.Sampler.SampleDistribution(pipelineSlot, dist) - lastToken = mtpTokenInput(nextToken) + lastToken = nextToken.ExpandDims(-1) lastHidden = projected draftTokens = append(draftTokens, lastToken) draftDists = append(draftDists, dist) @@ -312,248 +116,12 @@ func (r *Runner) generateMTPDraftCandidates(draft base.MTPDraftModel, target bas prefix = prefix.Concatenate(1, lastToken) } } - if len(draftTokens) == 0 { - return nil - } - return &mtpDraftCandidates{ + return &draftCandidates{ tokens: mlx.Concatenate(draftTokens, 1), dist: sampler.ConcatenateDistributions(draftDists), } } -// scheduleSpeculation schedules per-token snapshots at offsets -// [before, before+draftCount) on every cache, so the speculative forward -// captures a rollback point before each draft token's write. -func scheduleSpeculation(caches []cache.Cache, before, draftCount int) { - offsets := make([]int, draftCount) - for i := range offsets { - offsets[i] = before + i - } - for _, c := range caches { - if c != nil { - c.PrepareSnapshots(offsets) - } - } -} - -// commitSpeculation collects the captured snapshots and rolls every cache -// back to before+accepted, keeping only the accepted prefix of the -// speculative forward. The bonus token (at before+draftCount) is never -// committed. On full acceptance no restore is needed. -// -// Rollback tries a live rewind first (Restore(nil)) and falls back to the -// captured snapshot when the live state can't be rewound in place. -func commitSpeculation(caches []cache.Cache, accepted, draftCount, before int) { - target := before + accepted - for _, c := range caches { - if c == nil { - continue - } - snaps := c.TakeSnapshots() - if accepted < draftCount { - // Close the snapshots we won't restore from before restoring: a - // snapshot restore on a wrapped RotatingKVCache copies out every - // outstanding lazy snapshot before it rebuilds the buffer, so - // dropping the unused ones first stops that copy-out from - // materializing snapshots we are about to discard anyway. - for i, s := range snaps { - if s != nil && i != accepted { - s.Close() - snaps[i] = nil - } - } - if !c.Restore(nil, target) && !c.Restore(snaps[accepted], target) { - panic(fmt.Sprintf("mtp: cache restore to %d failed", target)) - } - } - for _, s := range snaps { - if s != nil { - s.Close() - } - } - } -} - -// acceptMTPDrafts accepts the longest draft prefix that survives rejection -// sampling against the target model. At temperature 0 the distributions are -// point masses, so acceptance reduces to argmax-match. -func (r *Runner) acceptMTPDrafts(ctx context.Context, request Request, session *cacheSession, dec *decoder, caches []cache.Cache, position *int, baseLogits *mlx.Array, candidates *mtpDraftCandidates, final *CompletionResponse, generated *int) (sampler.Result, int, bool, error) { - before := *position - draftCount := candidates.tokens.Dim(1) - scheduleSpeculation(caches, before, draftCount) - hiddenSeq := r.Model.Forward(&batch.Batch{ - InputIDs: candidates.tokens, - SeqOffsets: []int32{int32(before)}, - SeqQueryLens: []int32{int32(draftCount)}, - }, caches) - - targetDist := r.Sampler.Distribution(pipelineSlot, r.mtpValidationLogits(baseLogits, hiddenSeq), candidates.tokens) - draftDist := candidates.dist - acceptedMask := r.mtpSampleAcceptedMask(targetDist.SliceRows(0, draftCount), draftDist, candidates.tokens) - mlx.Eval(candidates.tokens, acceptedMask) - - draftIDs := candidates.tokens.Ints() - acceptedFlags := acceptedMask.Ints() - accepted := 0 - for _, ok := range acceptedFlags { - if ok == 0 { - break - } - accepted++ - } - if accepted > draftCount { - // Drain the scheduled snapshots and roll the speculative forward back - // out of the live caches before bailing, so the abandoned drafts don't - // reach the trie via session.close(). - commitSpeculation(caches, 0, draftCount, before) - return sampler.Result{}, 0, false, fmt.Errorf("mtp sample validation accepted %d tokens for %d draft tokens", accepted, draftCount) - } - - commitIDs := make([]int32, 0, accepted+1) - done := false - for i, id := range draftIDs[:accepted] { - commitIDs = append(commitIDs, int32(id)) - if r.Tokenizer.IsEOS(int32(id)) { - done = true - accepted = i + 1 - commitIDs = commitIDs[:accepted] - break - } - } - - commitSpeculation(caches, accepted, draftCount, before) - *position = before + accepted - - emitted, err := r.emitTokens(ctx, request, session, dec, draftResults(draftIDs[:accepted]), final, generated) - if err != nil { - return sampler.Result{}, accepted, emitted || done, err - } - if emitted || done { - r.Sampler.Commit(pipelineSlot, commitIDs) - return sampler.Result{}, accepted, true, nil - } - - var nextToken *mlx.Array - if accepted == draftCount { - nextToken = r.mtpSampleTokenAt(targetDist, draftCount) - } else { - nextToken = r.mtpSampleResidualToken(targetDist, draftDist, accepted) - } - mlx.Eval(nextToken) - nextID := int32(tokenID(nextToken)) - commitIDs = append(commitIDs, nextID) - r.Sampler.Commit(pipelineSlot, commitIDs) - - return sampler.Result{Token: nextToken}, accepted, false, nil -} - -func (r *Runner) mtpSampleAcceptedMask(targetDist, draftDist sampler.Distribution, draftTokens *mlx.Array) *mlx.Array { - p := targetDist.Prob(draftTokens) - q := draftDist.Prob(draftTokens) - acceptP := mlx.Minimum(p.Divide(q), mlx.FromValue(float32(1))) - return r.Sampler.Bernoulli(pipelineSlot, acceptP).AsType(mlx.DTypeInt32) -} - -func (r *Runner) mtpSampleTokenAt(dist sampler.Distribution, index int) *mlx.Array { - return mtpTokenVector(r.Sampler.SampleDistribution(pipelineSlot, dist.SliceRows(index, index+1))) -} - -func (r *Runner) mtpSampleResidualToken(targetDist, draftDist sampler.Distribution, index int) *mlx.Array { - residual := targetDist.SliceRows(index, index+1).ResidualAgainst(draftDist.SliceRows(index, index+1)) - return mtpTokenVector(r.Sampler.SampleDistribution(pipelineSlot, residual)) -} - -func mtpTokenInput(token *mlx.Array) *mlx.Array { - switch token.NumDims() { - case 0: - return token.Reshape(1, 1) - case 1: - return token.ExpandDims(-1) - case 2: - return token - default: - panic(fmt.Sprintf("mtp token must be rank 0, 1, or 2, got rank %d", token.NumDims())) - } -} - -func mtpTokenVector(token *mlx.Array) *mlx.Array { - switch token.NumDims() { - case 0: - return token.Reshape(1) - case 1: - return token - default: - panic(fmt.Sprintf("mtp sampled token must be rank 0 or 1, got rank %d", token.NumDims())) - } -} - -// emitTokens records a run of generated tokens to session.outputs, then streams -// them. A trailing EOS stops generation and is recorded but not streamed. -// Returns whether to stop and any cancellation error. -func (r *Runner) emitTokens(ctx context.Context, request Request, session *cacheSession, dec *decoder, results []sampler.Result, final *CompletionResponse, generated *int) (done bool, err error) { - stream := len(results) - for i, res := range results { - id := int32(tokenID(res.Token)) - session.outputs = append(session.outputs, id) - if r.Tokenizer.IsEOS(id) { - final.DoneReason = 0 - done = true - stream = i - break - } - (*generated)++ - } - if *generated >= request.Options.NumPredict { - done = true - } - - // Record the whole run before streaming any of it: streaming returns early on - // a cancelled context, and a partial stream must not leave the cache ahead of - // session.outputs. - for _, res := range results[:stream] { - resp, ok := dec.decode(res) - if !ok { - continue - } - select { - case <-ctx.Done(): - return done, ctx.Err() - case request.Responses <- resp: - } - } - return done, nil -} - -// draftResults wraps accepted draft token ids as sampler results for emitTokens. -// Accepted drafts carry no logprobs, so only the token id is set. -func draftResults(ids []int) []sampler.Result { - results := make([]sampler.Result, len(ids)) - for i, id := range ids { - results[i] = sampler.Result{Token: mlx.FromValues([]int32{int32(id)}, 1)} - } - return results -} - -func (r *Runner) lastLogits(hidden *mlx.Array) *mlx.Array { - logits := r.Model.Unembed(hidden) - return r.lastLogitsFromLogits(logits) -} - -func (r *Runner) mtpValidationLogits(baseLogits, hiddenSeq *mlx.Array) *mlx.Array { - seqLogits := r.Model.Unembed(hiddenSeq) - return baseLogits.ExpandDims(1).Concatenate(1, seqLogits) -} - -func (r *Runner) lastLogitsFromLogits(logits *mlx.Array) *mlx.Array { - return logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice()).Squeeze(1) -} - -// tokenID reads a single-token array as its host id. It goes through the -// item accessor, which evaluates the array first: raw data reads on a lazy -// array race its evaluation and return garbage. -func tokenID(token *mlx.Array) int { - if token == nil { - return -1 - } - return token.Int() +func lastHiddenRow(hidden *mlx.Array) *mlx.Array { + return hidden.Slice(mlx.Slice(), mlx.Slice(hidden.Dim(1)-1), mlx.Slice()) } diff --git a/x/mlxrunner/mtp_test.go b/x/mlxrunner/mtp_test.go index 59cbf235..c449b3e7 100644 --- a/x/mlxrunner/mtp_test.go +++ b/x/mlxrunner/mtp_test.go @@ -3,11 +3,11 @@ package mlxrunner import ( "context" "encoding/json" + "errors" "fmt" "slices" "strings" "testing" - "time" "github.com/ollama/ollama/api" "github.com/ollama/ollama/x/mlxrunner/batch" @@ -190,6 +190,15 @@ func collectResponses(ch chan CompletionResponse) (content string, final Complet } } +// resultIDs reads the token id of each result. +func resultIDs(results []sampler.Result) []int { + ids := make([]int, 0, len(results)) + for _, res := range results { + ids = append(ids, res.Token.Int()) + } + return ids +} + func TestAcceptMTPDraftsGreedyAcceptAll(t *testing.T) { skipIfNoMLX(t) // Target predicts 1->2->3->4 along the accepted chain; the draft proposed @@ -202,24 +211,20 @@ func TestAcceptMTPDraftsGreedyAcceptAll(t *testing.T) { candidates := scriptedCandidates(r, []int32{2, 3, 4}) baseLogits := oneHotLogits([]int32{2}).Squeeze(1) // target prediction after the seed token 1 - session, ch := newMTPTestSession(caches) position := caches[0].Offset() - final := CompletionResponse{Done: true} - generated := 0 - req := Request{Responses: ch, CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 100}}} - next, accepted, done, err := r.acceptMTPDrafts(context.Background(), req, session, &decoder{tokenizer: r.Tokenizer}, caches, &position, baseLogits, candidates, &final, &generated) + 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) if err != nil { - t.Fatalf("acceptMTPDrafts: %v", err) + t.Fatalf("accept: %v", err) } if accepted != 3 { t.Fatalf("accepted = %d, want 3", accepted) } - if done { - t.Fatalf("done = true, want false") - } - if got := tokenID(next.Token); got != 5 { - t.Fatalf("bonus token = %d, want 5", got) + // The run is the accepted drafts followed by the bonus token (5). + if got := resultIDs(results); !slices.Equal(got, []int{2, 3, 4, 5}) { + t.Fatalf("results = %v, want [2 3 4 5]", got) } if position != 3 { t.Fatalf("position = %d, want 3", position) @@ -227,9 +232,6 @@ func TestAcceptMTPDraftsGreedyAcceptAll(t *testing.T) { if got := caches[0].Offset(); got != 3 { t.Fatalf("cache offset = %d, want 3 (all drafts kept)", got) } - if generated != 3 { - t.Fatalf("generated = %d, want 3", generated) - } } func TestAcceptMTPDraftsGreedyMismatch(t *testing.T) { @@ -244,24 +246,21 @@ func TestAcceptMTPDraftsGreedyMismatch(t *testing.T) { candidates := scriptedCandidates(r, []int32{2, 7}) baseLogits := oneHotLogits([]int32{2}).Squeeze(1) - session, ch := newMTPTestSession(caches) position := caches[0].Offset() - final := CompletionResponse{Done: true} - generated := 0 - req := Request{Responses: ch, CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 100}}} - next, accepted, done, err := r.acceptMTPDrafts(context.Background(), req, session, &decoder{tokenizer: r.Tokenizer}, caches, &position, baseLogits, candidates, &final, &generated) + 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) if err != nil { - t.Fatalf("acceptMTPDrafts: %v", err) + t.Fatalf("accept: %v", err) } if accepted != 1 { t.Fatalf("accepted = %d, want 1", accepted) } - if done { - t.Fatalf("done = true, want false") - } - if got := tokenID(next.Token); got != 3 { - t.Fatalf("bonus token = %d, want 3 (target prediction at rejection)", got) + // The run is the one accepted draft (2) followed by the target's own + // prediction at the rejection point (3). + if got := resultIDs(results); !slices.Equal(got, []int{2, 3}) { + t.Fatalf("results = %v, want [2 3]", got) } if position != 1 { t.Fatalf("position = %d, want 1", position) @@ -274,7 +273,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, no bonus token is produced, and the cache keeps both tokens. + // generation and no bonus token is produced. Both accepted tokens are + // committed, so the caches hold the EOS's KV. const eos int32 = 6 predict := map[int32]int32{1: 2, 2: eos, eos: 0} r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{}) @@ -283,30 +283,30 @@ func TestAcceptMTPDraftsGreedyEOS(t *testing.T) { candidates := scriptedCandidates(r, []int32{2, eos}) baseLogits := oneHotLogits([]int32{2}).Squeeze(1) - session, ch := newMTPTestSession(caches) position := caches[0].Offset() - final := CompletionResponse{Done: true, DoneReason: 1} - generated := 0 - req := Request{Responses: ch, CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 100}}} - next, accepted, done, err := r.acceptMTPDrafts(context.Background(), req, session, &decoder{tokenizer: r.Tokenizer}, caches, &position, baseLogits, candidates, &final, &generated) + 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) if err != nil { - t.Fatalf("acceptMTPDrafts: %v", err) + t.Fatalf("accept: %v", err) } if accepted != 2 { t.Fatalf("accepted = %d, want 2 (token + EOS)", accepted) } - if !done { - t.Fatalf("done = false, want true") + // The EOS ends generation, so the run is exactly the accepted tokens + // with no bonus appended. + if got := resultIDs(results); !slices.Equal(got, []int{2, int(eos)}) { + t.Fatalf("results = %v, want [2 %d]", got, eos) } - if next.Token != nil { - t.Fatalf("bonus token = %d, want none after EOS", tokenID(next.Token)) - } - if final.DoneReason != 0 { - t.Fatalf("DoneReason = %d, want 0 (EOS)", final.DoneReason) + 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", position) + t.Fatalf("position = %d, want 2 (both accepted tokens committed)", position) + } + if got := caches[0].Offset(); got != 2 { + t.Fatalf("cache offset = %d, want 2 (both accepted tokens committed)", got) } } @@ -320,7 +320,7 @@ func TestRunMTPDecodeGreedy(t *testing.T) { r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{}) // The draft mirrors the target chain, so every drafted token is accepted. draft := &fakeMTPDraft{predict: predict} - r.Draft = draft + r.spec = newSpeculation(r, draft) caches, _ := newMTPTestCaches(1) session, ch := newMTPTestSession(caches) @@ -332,9 +332,11 @@ func TestRunMTPDecodeGreedy(t *testing.T) { CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, SamplerOpts: sampler.Options{}, } - if err := r.runMTPDecode(context.Background(), req, session, caches, []int32{1}, &position, time.Now()); err != nil { - t.Fatalf("runMTPDecode: %v", err) + d := testDecoder(r, req, caches, []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 != "234" { @@ -360,9 +362,10 @@ func TestRunMTPDecodeGreedy(t *testing.T) { } // Single-position drafting anchors every Draft call in a round at the - // last target-seen position (the current token, offset 2), while the - // extended token advances along the proposed chain. - wantDraft := []draftCall{{2, 2}, {2, 3}, {2, 4}, {2, eos}} + // last committed position (offset 1 — the current token at offset 2 is + // not yet validated), while the extended token advances along the + // proposed chain. + wantDraft := []draftCall{{1, 2}, {1, 3}, {1, 4}, {1, eos}} if !slices.Equal(draft.calls, wantDraft) { t.Fatalf("draft calls = %v, want %v", draft.calls, wantDraft) } @@ -377,11 +380,7 @@ func TestRunMTPDecodeSampled(t *testing.T) { const eos int32 = 7 predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: eos, eos: 0} r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{Temperature: 1, Seed: 42, UseSeed: true}) - r.Draft = &fakeMTPDraft{predict: predict} - - if !r.useMTP(sampler.Options{Temperature: 1}) { - t.Fatalf("useMTP rejected a sampled request") - } + r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict}) caches, _ := newMTPTestCaches(1) session, ch := newMTPTestSession(caches) @@ -393,9 +392,11 @@ func TestRunMTPDecodeSampled(t *testing.T) { CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, SamplerOpts: sampler.Options{Temperature: 1, Seed: 42, UseSeed: true}, } - if err := r.runMTPDecode(context.Background(), req, session, caches, []int32{1}, &position, time.Now()); err != nil { - t.Fatalf("runMTPDecode: %v", err) + d := testDecoder(r, req, caches, []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 != "234" { @@ -409,6 +410,117 @@ func TestRunMTPDecodeSampled(t *testing.T) { } } +func TestDecodePlain(t *testing.T) { + skipIfNoMLX(t) + // The same chain with no speculationSession: decode's pipelined loop runs, + // dispatching the forward that produces the next token before the + // current one is emitted. + 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) + session, ch := newMTPTestSession(caches) + position := 1 + + req := Request{ + Responses: ch, + Tokens: []int32{0}, + CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, + SamplerOpts: sampler.Options{}, + } + d := testDecoder(r, req, caches, []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 != "234" { + t.Fatalf("content = %q, want %q", content, "234") + } + if final.DoneReason != 0 || final.EvalCount != 3 { + t.Fatalf("final DoneReason = %d, EvalCount = %d, want 0 (EOS) and 3", final.DoneReason, final.EvalCount) + } + if got := []int32{2, 3, 4, eos}; !slices.Equal(session.outputs, got) { + t.Fatalf("session outputs = %v, want %v", session.outputs, got) + } + + // One forward per token: the seed at offset 1, then each sampled token + // in turn — including the ending EOS, whose forward is already in + // flight when the token is checked. Every forwarded token is recorded, + // so the caches rest exactly at the recorded path. + wantForwards := []forwardCall{{1, 1}, {2, 1}, {3, 1}, {4, 1}, {5, 1}} + model := r.Model.(*fakeMTPModel) + if !slices.Equal(model.forwards, wantForwards) { + t.Fatalf("target forwards = %v, want %v", model.forwards, wantForwards) + } +} + +func TestDecodeCancelledMidStream(t *testing.T) { + skipIfNoMLX(t) + // Cancelling while accepted drafts stream must leave the session + // consistent: every token committed to the caches is recorded in + // session.outputs, no speculation snapshot schedule is left pending on + // the shared caches, and every captured snapshot is closed. + const eos int32 = 7 + predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: eos} + r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{}) + r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict}) + + caches, tr := newMTPTestCaches(1) + session := &cacheSession{caches: caches} + ch := make(chan CompletionResponse) // unbuffered: every send must rendezvous + + ctx, cancel := context.WithCancel(context.Background()) + go func() { + <-ch // the first streamed token + cancel() + // Stop reading: the next send blocks until the emit select sees + // the cancelled context. + }() + + position := 0 + req := Request{ + Responses: ch, + Tokens: []int32{1}, + CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}}, + SamplerOpts: sampler.Options{}, + } + d := testDecoder(r, req, caches, []int32{1}, position) + err := r.decode(ctx, req, session, d, 0) + d.close() + if !errors.Is(err, context.Canceled) { + t.Fatalf("decode error = %v, want context.Canceled", err) + } + + // Every committed token is recorded, and the caches never run ahead of + // the record: the final recorded output is the round's next token, + // emitted before it is ever forwarded, so the caches rest one short. + rc := caches[0].(*fakeRewindableCache) + if want := 1 + len(session.outputs) - 1; rc.Offset() != want { + t.Fatalf("cache offset = %d, want %d (seed + recorded outputs %v minus the unforwarded final)", rc.Offset(), want, session.outputs) + } + if rc.pending.offsets != nil { + t.Fatalf("speculation snapshot schedule left pending: %v", rc.pending.offsets) + } + for i, s := range tr.all { + if s.closeCount == 0 { + t.Fatalf("snapshot #%d [%d,%d) leaked: never closed", i, s.from, s.to) + } + } +} + +// testDecoder builds the decoder TextGenerationPipeline would construct for +// this request; tests close it explicitly so close-time effects are visible +// to assertions. +func testDecoder(r *Runner, req Request, caches []cache.Cache, seed []int32, position int) decoder { + if spec := r.spec.open(req, caches); spec != nil { + return spec.decoder(seed, position) + } + return r.pipelinedDecoder(caches, seed, position) +} + // newMTPTestCaches returns n rewindable fake caches sharing one snapshot // tracker, matching the cache.Cache the speculation helpers drive. func newMTPTestCaches(n int) ([]cache.Cache, *snapshotTracker) { @@ -427,20 +539,41 @@ func newMTPTestSession(caches []cache.Cache) (*cacheSession, chan CompletionResp return &cacheSession{caches: caches}, ch } -// scriptedCandidates builds draft candidates by running the real generator -// against a draft whose prediction chain, starting from seed token 0, yields -// exactly the requested tokens. Using the real generator means the proposal -// distributions match what acceptMTPDrafts expects. -func scriptedCandidates(r *Runner, tokens []int32) *mtpDraftCandidates { +// 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. +func testSpeculationSession(r *Runner, caches []cache.Cache) *speculationSession { + s := r.spec + if s == nil { + s = &speculation{r: r} + } + var d drafter = nopDrafter{} + if md := newMTPDrafter(s, caches); md != nil { + d = md + } + return &speculationSession{spec: s, drafter: d, caches: caches} +} + +// nopDrafter satisfies drafter for engine tests that supply candidates +// directly and never propose. +type nopDrafter struct{} + +func (nopDrafter) propose(*mlx.Array, int) *draftCandidates { return nil } +func (nopDrafter) committed(_, _ *mlx.Array, _ int) {} +func (nopDrafter) close() {} + +// scriptedCandidates builds draft candidates by running the real drafter +// against a fake whose prediction chain, starting from seed token 0, yields +// exactly the requested tokens. Using the real drafter means the proposal +// distributions match what the engine's acceptance expects. +func scriptedCandidates(r *Runner, tokens []int32) *draftCandidates { chain := map[int32]int32{} prev := int32(0) for _, tok := range tokens { chain[prev] = tok prev = tok } - draft := &fakeMTPDraft{predict: chain} - target := r.Model.(base.MTPEmbeddingModel) - seed := mlx.FromValues([]int32{0}, 1, 1) - hidden := mlx.Zeros(mlx.DTypeFloat32, 1, 1, mtpTestVocab) - return r.generateMTPDraftCandidates(draft, target, seed, hidden, nil, 0, len(tokens)) + d := &mtpDrafter{spec: &speculation{r: r}, draft: &fakeMTPDraft{predict: chain}, target: r.Model.(base.MTPEmbeddingModel)} + 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 e5720135..3f5f31e2 100644 --- a/x/mlxrunner/pipeline.go +++ b/x/mlxrunner/pipeline.go @@ -12,6 +12,7 @@ import ( "github.com/ollama/ollama/llm" "github.com/ollama/ollama/logutil" "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/cache" "github.com/ollama/ollama/x/mlxrunner/mlx" sampler "github.com/ollama/ollama/x/mlxrunner/sample" "github.com/ollama/ollama/x/tokenizer" @@ -55,12 +56,9 @@ const pipelineSlot = 0 func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) error { mlx.ResetPeakMemory() - var sample, nextSample sampler.Result defer func() { r.Sampler.Remove(pipelineSlot) - mlx.Unpin(sample.Arrays()...) - mlx.Unpin(nextSample.Arrays()...) mlx.Sweep() mlx.ClearCache() @@ -75,9 +73,38 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er session := r.cache.begin(r.Model, inputs) defer session.close() - caches := session.caches + + seed, position, promptEval, err := r.prefill(ctx, session) + if err != nil { + return err + } + + // 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) + } + defer d.close() + return r.decode(ctx, request, session, d, promptEval) +} + +// prefill evaluates the prompt's unprocessed tokens in chunks, leaving +// one token for decode to seed from, and schedules the prompt's periodic +// snapshots. It returns the seed tokens, the resume position, and the +// prompt-evaluation duration. +func (r *Runner) prefill(ctx context.Context, session *cacheSession) ([]int32, int, time.Duration, error) { + start := time.Now() + inputs := session.inputs tokens := session.remaining + caches := session.caches prefillChunk := prefillChunkSize() // Request periodic snapshots during prefill and near the end of the @@ -107,12 +134,11 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er session.schedulePrefillSnapshots(snapshotOffsets) - now := time.Now() total, processed := len(tokens), 0 position := len(inputs) - len(tokens) for total-processed > 1 { if err := ctx.Err(); err != nil { - return err + return nil, 0, 0, err } n := min(prefillChunk, total-processed-1) @@ -135,65 +161,78 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er // Attach the snapshots captured during prefill to the trie. session.attachPrefillSnapshots() - // Register the sampler after prefill completes. - r.Sampler.Add(pipelineSlot, request.SamplerOpts, inputs) - if r.useMTP(request.SamplerOpts) { - return r.runMTPDecode(ctx, request, session, caches, tokens[processed:], &position, now) - } + return tokens[processed:], position, time.Since(start), nil +} - step := func(token *mlx.Array) sampler.Result { - fwd := r.Model.Forward(&batch.Batch{ - InputIDs: token, - SeqOffsets: []int32{int32(position)}, - SeqQueryLens: []int32{int32(token.Dim(1))}, - }, caches) - position += token.Dim(1) - logits := r.Model.Unembed(fwd) - logits = logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice()).Squeeze(1) +// A decoder produces each run of tokens to emit, owning its own dispatch and +// synchronization; the decode loop owns the budget, emission, and +// cancellation. next may return none while its first tokens are in flight. +type decoder interface { + next(remaining int) ([]sampler.Result, error) + close() +} - sample := r.Sampler.Sample([]int{pipelineSlot}, logits) - mlx.Pin(sample.Arrays()...) - mlx.Sweep() - mlx.AsyncEval(sample.Arrays()...) - return sample - } - - sample = step(mlx.FromValues(tokens[processed:], 1, total-processed)) - logutil.TraceContext(ctx, "mlx decode seed", "tokens", total-processed, "memory", mlx.Memory{}) - - dec := decoder{ +// decode drives either decoder and owns where generation stops — at an EOS +// or the NumPredict budget. Every produced token is recorded so the caches +// never rest ahead of session.outputs; tokens past the stop are recorded but +// not streamed or counted. +func (r *Runner) decode(ctx context.Context, request Request, session *cacheSession, d decoder, promptEval time.Duration) error { + detok := detokenizer{ tokenizer: r.Tokenizer, wantLogprobs: request.SamplerOpts.Logprobs, wantTopLogprobs: request.SamplerOpts.TopLogprobs, } - final := CompletionResponse{Done: true, PromptEvalCount: len(inputs), EvalCount: request.Options.NumPredict, DoneReason: 1} - for i := range request.Options.NumPredict { + final := CompletionResponse{Done: true, PromptEvalCount: len(request.Tokens), DoneReason: 1} + final.PromptEvalDuration = promptEval + now := time.Now() + + // Release MLX's cached free buffers every clearCacheInterval tokens so the + // allocator's pool does not grow unbounded over a long generation. + const clearCacheInterval = 256 + + generated := 0 + for generated < request.Options.NumPredict { if err := ctx.Err(); err != nil { return err } - nextSample = step(sample.Token.ExpandDims(-1)) - - if i == 0 { - mlx.Eval(sample.Arrays()...) - final.PromptEvalDuration = time.Since(now) - now = time.Now() + results, err := d.next(request.Options.NumPredict - generated) + if err != nil { + return err } - output := int32(sample.Token.Int()) - session.outputs = append(session.outputs, output) - if i == 0 { - logutil.TraceContext(ctx, "mlx decode first token", "memory", mlx.Memory{}) + // Record the whole run before streaming any of it: a cancelled + // stream returns early and must not leave the caches ahead of + // session.outputs. + done := false + stream := len(results) + for i, res := range results { + // Int evaluates the array before reading it; a raw data read + // on a lazy array races its evaluation and returns garbage. + id := int32(res.Token.Int()) + session.outputs = append(session.outputs, id) + if done { + continue + } + if r.Tokenizer.IsEOS(id) { + final.DoneReason = 0 + done = true + stream = i + continue + } + generated++ + if generated >= request.Options.NumPredict { + done = true + stream = i + 1 + } } - if r.Tokenizer.IsEOS(output) { - final.DoneReason = 0 - final.EvalCount = i - break - } - - if resp, ok := dec.decode(sample); ok { + for _, res := range results[:stream] { + resp, ok := detok.detokenize(res) + if !ok { + continue + } select { case <-ctx.Done(): return ctx.Err() @@ -201,14 +240,16 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er } } - mlx.Unpin(sample.Arrays()...) - sample, nextSample = nextSample, sampler.Result{} + if done { + break + } - if i%256 == 0 { + if generated%clearCacheInterval == 0 { mlx.ClearCache() } } + final.EvalCount = generated final.EvalDuration = time.Since(now) select { case <-ctx.Done(): @@ -218,11 +259,60 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er } } -// decoder serializes sampled tokens into response chunks, holding bytes +// pipelinedDecoder decodes one token per call, one call ahead of emission: +// 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 + 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} + t.sample = t.dispatch(mlx.FromValues(seed, 1, len(seed))) + return t +} + +// dispatch builds one forward-and-sample chain without reading the token's +// value, so it is in flight before the previous token is synchronized. +func (t *pipelinedDecoder) dispatch(token *mlx.Array) sampler.Result { + r := t.r + hidden := r.Model.Forward(&batch.Batch{ + InputIDs: token, + SeqOffsets: []int32{int32(t.position)}, + SeqQueryLens: []int32{int32(token.Dim(1))}, + }, t.caches) + t.position += token.Dim(1) + next := r.Sampler.Sample([]int{pipelineSlot}, lastLogits(r.Model.Unembed(hidden))) + mlx.Pin(next.Arrays()...) + mlx.Sweep() + mlx.AsyncEval(next.Arrays()...) + return next +} + +func (t *pipelinedDecoder) next(int) ([]sampler.Result, error) { + mlx.Unpin(t.emitted.Arrays()...) + t.emitted, t.sample = t.sample, t.dispatch(t.sample.Token.ExpandDims(-1)) + return []sampler.Result{t.emitted}, nil +} + +func (t *pipelinedDecoder) close() { + mlx.Unpin(t.emitted.Arrays()...) + mlx.Unpin(t.sample.Arrays()...) +} + +func lastLogits(logits *mlx.Array) *mlx.Array { + return logits.Slice(mlx.Slice(), mlx.Slice(logits.Dim(1)-1), mlx.Slice()).Squeeze(1) +} + +// detokenizer serializes sampled tokens into response chunks, holding bytes // whose UTF-8 sequence hasn't completed yet and the logprobs that belong // with those bytes so Content and Logprobs stay aligned when a chunk does // flush. -type decoder struct { +type detokenizer struct { tokenizer *tokenizer.Tokenizer buf bytes.Buffer logprobs []llm.Logprob @@ -230,7 +320,7 @@ type decoder struct { wantTopLogprobs int } -func (d *decoder) decode(res sampler.Result) (CompletionResponse, bool) { +func (d *detokenizer) detokenize(res sampler.Result) (CompletionResponse, bool) { output := int32(res.Token.Int()) d.buf.WriteString(d.tokenizer.Decode([]int32{output})) d.logprobs = append(d.logprobs, buildLogprob(res, d.wantLogprobs, d.wantTopLogprobs, d.tokenizer.Decode)...) diff --git a/x/mlxrunner/runner.go b/x/mlxrunner/runner.go index 37866bb2..189bfda9 100644 --- a/x/mlxrunner/runner.go +++ b/x/mlxrunner/runner.go @@ -35,13 +35,15 @@ type Request struct { type Runner struct { Model base.Model - Draft base.DraftModel Tokenizer *tokenizer.Tokenizer Requests chan Request Sampler *sample.Sampler cache kvCache contextLength int mlxThread *mlxthread.Thread + // spec is the speculative-decoding subsystem. Nil when the model ships no + // draft head. + spec *speculation } func (r *Runner) Load(modelName string) error { @@ -69,7 +71,7 @@ func (r *Runner) Load(modelName string) error { return err } - r.Draft = nil + var draftModel base.DraftModel draft, err := base.NewDraft(root, m) if err != nil { return err @@ -78,7 +80,7 @@ func (r *Runner) Load(modelName string) error { if err := draft.LoadWeights(tensors); err != nil { return err } - r.Draft = draft + draftModel = draft } collected := mlx.Collect(m) @@ -101,8 +103,10 @@ func (r *Runner) Load(modelName string) error { r.Tokenizer = m.Tokenizer() r.contextLength = m.MaxContextLength() r.Sampler = sample.New(r.contextLength) + r.spec = newSpeculation(r, draftModel) mlx.EnableCompile() + return nil } diff --git a/x/mlxrunner/speculate.go b/x/mlxrunner/speculate.go new file mode 100644 index 00000000..2b9a46a2 --- /dev/null +++ b/x/mlxrunner/speculate.go @@ -0,0 +1,513 @@ +package mlxrunner + +import ( + "fmt" + "log/slog" + "os" + "strconv" + "strings" + + "github.com/ollama/ollama/x/mlxrunner/batch" + "github.com/ollama/ollama/x/mlxrunner/cache" + "github.com/ollama/ollama/x/mlxrunner/mlx" + "github.com/ollama/ollama/x/mlxrunner/model/base" + sampler "github.com/ollama/ollama/x/mlxrunner/sample" +) + +// drafter proposes speculative tokens for the engine to validate, learning +// the conversation through the committed-stream reports. +type drafter interface { + // propose returns up to maxTokens draft tokens with their proposal + // distributions, or nil to decode this round plainly. + propose(current *mlx.Array, maxTokens int) *draftCandidates + + // 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. + committed(tokens, hiddens *mlx.Array, position int) + + close() +} + +type draftSchedule string + +const ( + draftScheduleHeuristic draftSchedule = "heuristic" + draftScheduleConstant draftSchedule = "constant" +) + +type specStats struct { + iterations int + drafted int + accepted int + mismatches int + allAccepted int + maxDraft int +} + +type specOptions struct { + initialDraftTokens int + maxDraftTokens int + draftSchedule draftSchedule +} + +func (r *Runner) loadSpecOptions(sample bool) specOptions { + defaults := r.mtpDefaults(sample) + + opts := specOptions{ + initialDraftTokens: defaults.InitialDraftTokens, + maxDraftTokens: defaults.MaxDraftTokens, + draftSchedule: draftScheduleConstant, + } + if v := positiveEnvInt("OLLAMA_MLX_MTP_MAX_DRAFT_TOKENS"); v > 0 { + opts.maxDraftTokens = v + } + if v := positiveEnvInt("OLLAMA_MLX_MTP_INITIAL_DRAFT_TOKENS"); v > 0 { + opts.initialDraftTokens = v + } + if opts.initialDraftTokens > opts.maxDraftTokens { + opts.initialDraftTokens = opts.maxDraftTokens + } + switch schedule := strings.ToLower(strings.TrimSpace(os.Getenv("OLLAMA_MLX_MTP_DRAFT_SCHEDULE"))); schedule { + case "", string(draftScheduleConstant): + opts.draftSchedule = draftScheduleConstant + case string(draftScheduleHeuristic): + opts.draftSchedule = draftScheduleHeuristic + default: + slog.Warn("invalid MTP env setting", "key", "OLLAMA_MLX_MTP_DRAFT_SCHEDULE", "value", schedule) + } + return opts +} + +func positiveEnvInt(key string) int { + raw := os.Getenv(key) + if raw == "" { + return 0 + } + v, err := strconv.Atoi(raw) + if err != nil || v <= 0 { + slog.Warn("invalid MTP env setting", "key", key, "value", raw) + return 0 + } + return v +} + +// speculation is the persistent speculative-decoding subsystem of a Runner, +// one per loaded model. It holds the draft model; as the feature grows it +// also holds the cache partition and the depth state learned across requests. +// A nil *speculation means the checkpoint ships no draft head. +type speculation struct { + r *Runner + draft base.DraftModel +} + +// newSpeculation builds the speculative-decoding subsystem for a loaded model, +// or nil when the checkpoint ships no draft head. +func newSpeculation(r *Runner, draft base.DraftModel) *speculation { + if draft == nil { + return nil + } + return &speculation{r: r, draft: draft} +} + +// 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 + 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. +func (s *speculation) open(request Request, caches []cache.Cache) *speculationSession { + if s == nil { + return nil + } + d := newMTPDrafter(s, caches) + 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 + } + + specOpts := s.r.loadSpecOptions(opts.Temperature != 0) + return &speculationSession{ + spec: s, + drafter: d, + caches: caches, + opts: specOpts, + limit: specOpts.initialDraftTokens, + stats: specStats{maxDraft: specOpts.initialDraftTokens}, + } +} + +func (s *speculationSession) committed(tokens, hiddens *mlx.Array, position int) { + if s == nil { + return + } + s.drafter.committed(tokens, hiddens, position) +} + +func (s *speculationSession) close() { + if s == nil { + return + } + s.drafter.close() +} + +// speculativeDecoder decodes one speculative round per call: it forwards +// the current token — emitted by the previous call, so a token that ends +// generation is never forwarded — has the engine draft and validate ahead +// of it, and returns the round's accepted tokens followed by the engine's +// next token. The last returned token becomes the next call's current; the +// seed token primes current and is never returned. +type speculativeDecoder struct { + s *speculationSession + position int + current sampler.Result // emitted (or the seed), not yet forwarded +} + +func (s *speculationSession) decoder(seed []int32, position int) *speculativeDecoder { + current := sampler.Result{Token: mlx.FromValues(seed, len(seed))} + mlx.Pin(current.Arrays()...) + return &speculativeDecoder{s: s, position: position, current: current} +} + +func (st *speculativeDecoder) next(remaining int) ([]sampler.Result, error) { + s := st.s + r := s.spec.r + + hidden := r.Model.Forward(&batch.Batch{ + InputIDs: tokenInput(st.current.Token), + SeqOffsets: []int32{int32(st.position)}, + SeqQueryLens: []int32{1}, + }, s.caches) + st.position++ + + results, next, err := s.round(&st.position, st.current, hidden, remaining) + if err != nil { + return nil, err + } + if next.Token != nil { + results = append(results, next) + } + + last := results[len(results)-1] + mlx.Pin(last.Arrays()...) + mlx.Unpin(st.current.Arrays()...) + st.current = last + mlx.AsyncEval(st.current.Arrays()...) + return results, nil +} + +func (st *speculativeDecoder) close() { + mlx.Unpin(st.current.Arrays()...) + st.s.logStats() +} + +// round runs one speculative decode round for the just-forwarded current +// token and its hidden state: draft candidates after it, validate them +// against the target, and return the accepted run and the bonus or +// resampled next token. +func (s *speculationSession) round(position *int, current sampler.Result, hidden *mlx.Array, remaining int) (results []sampler.Result, next sampler.Result, err error) { + r := s.spec.r + s.stats.iterations++ + + // A round emits the accepted drafts plus one more token (the bonus or + // residual), so cap the draft one below the remaining budget to land that + // extra token within it rather than overshooting. At remaining 1 the cap is + // 0 and the last token decodes plainly. + maxDraft := min(s.limit, remaining-1) + candidates := s.drafter.propose(current.Token, maxDraft) + baseLogits := lastLogits(r.Model.Unembed(hidden)) + if candidates == nil { + s.drafter.committed(tokenInput(current.Token), lastHiddenRow(hidden), *position-1) + return nil, r.Sampler.Sample([]int{pipelineSlot}, baseLogits), nil + } + + draftCount := candidates.tokens.Dim(1) + // hidden survives the sweep alongside the candidates: the post-accept + // report fuses it into the committed stream. + candidateArrays := append([]*mlx.Array{baseLogits, hidden}, candidates.Arrays()...) + mlx.Pin(candidateArrays...) + mlx.Sweep() + defer mlx.Unpin(candidateArrays...) + s.stats.drafted += draftCount + + results, accepted, err := s.accept(position, current, hidden, baseLogits, candidates) + if err != nil { + return nil, sampler.Result{}, err + } + // accept folds the bonus token into its results; surface it back out as + // the next step's current. + if len(results) > accepted { + next = results[len(results)-1] + results = results[:accepted] + } + + s.stats.accepted += accepted + if accepted == draftCount { + s.stats.allAccepted++ + } else { + s.stats.mismatches++ + } + if s.opts.draftSchedule == draftScheduleHeuristic { + if accepted == draftCount { + s.limit = min(s.opts.maxDraftTokens, s.limit+2) + } else { + s.limit = max(1, s.limit-1) + } + s.stats.maxDraft = max(s.stats.maxDraft, s.limit) + } + return results, next, nil +} + +// logStats reports the per-request speculation summary. +func (s *speculationSession) logStats() { + acceptance := 0.0 + if s.stats.drafted > 0 { + acceptance = float64(s.stats.accepted) / float64(s.stats.drafted) + } + avgDraft := 0.0 + avgAccepted := 0.0 + if s.stats.iterations > 0 { + avgDraft = float64(s.stats.drafted) / float64(s.stats.iterations) + avgAccepted = float64(s.stats.accepted) / float64(s.stats.iterations) + } + slog.Info("speculative decode stats", "drafted", s.stats.drafted, "accepted", s.stats.accepted, "acceptance", acceptance, "iterations", s.stats.iterations, "avg_draft", avgDraft, "avg_accepted", avgAccepted, "mismatches", s.stats.mismatches, "all_accepted", s.stats.allAccepted, "max_draft", s.stats.maxDraft, "draft_schedule", s.opts.draftSchedule) +} + +// draftCandidates is one round's draft tokens and the proposal distribution +// each was sampled from, weighed against the target during acceptance. +type draftCandidates struct { + tokens *mlx.Array + dist sampler.Distribution +} + +func (c *draftCandidates) Arrays() []*mlx.Array { + if c == nil { + return nil + } + return append([]*mlx.Array{c.tokens}, c.dist.Arrays()...) +} + +// scheduleSpeculation schedules per-token snapshots at offsets +// [before, before+draftCount) on every cache, so the speculative forward +// captures a rollback point before each draft token's write. +func scheduleSpeculation(caches []cache.Cache, before, draftCount int) { + offsets := make([]int, draftCount) + for i := range offsets { + offsets[i] = before + i + } + for _, c := range caches { + if c != nil { + c.PrepareSnapshots(offsets) + } + } +} + +// commitSpeculation rolls every cache back to before+accepted, keeping only +// the accepted prefix; full acceptance needs no restore. Rollback tries a +// live rewind first (Restore(nil)) and falls back to the captured snapshot. +func commitSpeculation(caches []cache.Cache, accepted, draftCount, before int) { + target := before + accepted + for _, c := range caches { + if c == nil { + continue + } + snaps := c.TakeSnapshots() + if accepted < draftCount { + // Close the snapshots we won't restore from before restoring: a + // snapshot restore on a wrapped RotatingKVCache copies out every + // outstanding lazy snapshot before it rebuilds the buffer, so + // dropping the unused ones first stops that copy-out from + // materializing snapshots we are about to discard anyway. + for i, s := range snaps { + if s != nil && i != accepted { + s.Close() + snaps[i] = nil + } + } + if !c.Restore(nil, target) && !c.Restore(snaps[accepted], target) { + panic(fmt.Sprintf("speculation: cache restore to %d failed", target)) + } + } + for _, s := range snaps { + if s != nil { + s.Close() + } + } + } +} + +// accept accepts the longest draft prefix that survives rejection sampling +// against the target model. At temperature 0 the distributions are point +// masses, so acceptance reduces to argmax-match. It returns the accepted +// drafts followed by the target's own next token — the residual at the +// rejection point, or the bonus past a fully accepted run — so a round +// yields accepted+1 tokens, except when an accepted EOS ends generation, +// where the run stops at the EOS with no continuation. The accepted run is +// reported back to the drafter with its hidden states. +// +// The NumPredict budget is the decode loop's to enforce; the drafter is +// already capped to the remaining budget, so an accepted run never +// overshoots, and a legitimate token past the budget is left for decode to +// drop rather than cut here. +func (s *speculationSession) accept(position *int, current sampler.Result, hidden, baseLogits *mlx.Array, candidates *draftCandidates) (results []sampler.Result, accepted int, err error) { + r := s.spec.r + before := *position + draftCount := candidates.tokens.Dim(1) + scheduleSpeculation(s.caches, before, draftCount) + + // Every exit between schedule and commit must drain the snapshot + // schedule and roll the speculative writes back out of the live caches: + // an undrained schedule panics the next PrepareSnapshots, and + // uncommitted speculative tokens reach the trie through session.close. + committed := false + commit := func(keep int) { + if committed { + return + } + committed = true + commitSpeculation(s.caches, keep, draftCount, before) + } + defer commit(0) + + hiddenSeq := r.Model.Forward(&batch.Batch{ + InputIDs: candidates.tokens, + SeqOffsets: []int32{int32(before)}, + SeqQueryLens: []int32{int32(draftCount)}, + }, s.caches) + + targetDist := r.Sampler.Distribution(pipelineSlot, validationLogits(r, baseLogits, hiddenSeq), candidates.tokens) + draftDist := candidates.dist + acceptedMask := r.sampleAcceptedMask(targetDist.SliceRows(0, draftCount), draftDist, candidates.tokens) + mlx.Eval(candidates.tokens, acceptedMask) + + draftIDs := candidates.tokens.Ints() + acceptedFlags := acceptedMask.Ints() + for _, ok := range acceptedFlags { + if ok == 0 { + break + } + accepted++ + } + if accepted > draftCount { + return nil, 0, fmt.Errorf("speculation validation accepted %d tokens for %d draft tokens", accepted, draftCount) + } + + // Find where an accepted EOS ends the run, before committing, so the cut + // is known while the per-token rollback snapshots still exist. + commitIDs := make([]int32, 0, accepted+1) + done := false + for i, id := range draftIDs[:accepted] { + commitIDs = append(commitIDs, int32(id)) + if r.Tokenizer.IsEOS(int32(id)) { + done = true + accepted = i + 1 + break + } + } + + commit(accepted) + *position = before + accepted + + // Report the validated run — current plus the accepted 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...) + runHiddens := lastHiddenRow(hidden) + if accepted > 0 { + runHiddens = runHiddens.Concatenate(1, hiddenSeq.Slice(mlx.Slice(), mlx.Slice(0, accepted), mlx.Slice())) + } + s.drafter.committed(mlx.FromValues(runIDs, 1, len(runIDs)), runHiddens, before-1) + + results = draftResults(draftIDs[:accepted]) + if done { + r.Sampler.Commit(pipelineSlot, commitIDs) + return results, accepted, nil + } + + var nextToken *mlx.Array + if accepted == draftCount { + nextToken = r.sampleTokenAt(targetDist, draftCount) + } else { + nextToken = r.sampleResidualToken(targetDist, draftDist, accepted) + } + mlx.Eval(nextToken) + nextID := int32(nextToken.Int()) + commitIDs = append(commitIDs, nextID) + r.Sampler.Commit(pipelineSlot, commitIDs) + + results = append(results, sampler.Result{Token: nextToken}) + return results, accepted, nil +} + +// validationLogits stacks the current token's logits ahead of the draft +// positions' logits so row i scores draft i. +func validationLogits(r *Runner, baseLogits, hiddenSeq *mlx.Array) *mlx.Array { + seqLogits := r.Model.Unembed(hiddenSeq) + return baseLogits.ExpandDims(1).Concatenate(1, seqLogits) +} + +func (r *Runner) sampleAcceptedMask(targetDist, draftDist sampler.Distribution, draftTokens *mlx.Array) *mlx.Array { + p := targetDist.Prob(draftTokens) + q := draftDist.Prob(draftTokens) + acceptP := mlx.Minimum(p.Divide(q), mlx.FromValue(float32(1))) + return r.Sampler.Bernoulli(pipelineSlot, acceptP).AsType(mlx.DTypeInt32) +} + +func (r *Runner) sampleTokenAt(dist sampler.Distribution, index int) *mlx.Array { + return r.Sampler.SampleDistribution(pipelineSlot, dist.SliceRows(index, index+1)) +} + +func (r *Runner) sampleResidualToken(targetDist, draftDist sampler.Distribution, index int) *mlx.Array { + residual := targetDist.SliceRows(index, index+1).ResidualAgainst(draftDist.SliceRows(index, index+1)) + return tokenVector(r.Sampler.SampleDistribution(pipelineSlot, residual)) +} + +// draftResults wraps accepted draft ids as sampler results; drafts carry no +// logprobs, so only the token id is set. +func draftResults(ids []int) []sampler.Result { + results := make([]sampler.Result, len(ids)) + for i, id := range ids { + results[i] = sampler.Result{Token: mlx.FromValues([]int32{int32(id)}, 1)} + } + return results +} + +func tokenInput(token *mlx.Array) *mlx.Array { + switch token.NumDims() { + case 0: + return token.Reshape(1, 1) + case 1: + return token.ExpandDims(-1) + case 2: + return token + default: + panic(fmt.Sprintf("token must be rank 0, 1, or 2, got rank %d", token.NumDims())) + } +} + +func tokenVector(token *mlx.Array) *mlx.Array { + switch token.NumDims() { + case 0: + return token.Reshape(1) + case 1: + return token + default: + panic(fmt.Sprintf("sampled token must be rank 0 or 1, got rank %d", token.NumDims())) + } +}