mlxrunner: host speculative decoding in the text generation pipeline
The pipeline and the MTP decoder each owned a decode loop with duplicated prefill, budget, and emission handling. Split the pipeline into prefill and decode phases behind a decoder interface, with the decode loop the sole emitter enforcing the NumPredict budget, and split speculation into a generic engine that returns the accepted run and a drafter interface that owns only how proposals are made.
This commit is contained in:
Jesse Gross
@@ -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())
|
||||
}
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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)...)
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
|
||||
513
x/mlxrunner/speculate.go
Normal file
513
x/mlxrunner/speculate.go
Normal file
@@ -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()))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user