CERBERUS/ollama
1
0
Fork 0
mirror of https://github.com/ollama/ollama.git synced 2026-09-22 06:14:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

mlxrunner: let each model declare the cache slots it needs

Browse source 
The runner used to build caches by probing the model for an optional
NewCaches method, with one KV cache per layer as the fallback. A model
with a draft head appended the draft's cache slots to its own list, and
the speculative engine later recovered the two groups by comparing slot
identities, panicking when the lists didn't line up.

NewCaches is now a required method on both the model and the draft, and
each returns only the slots it writes. The runner concatenates the two
lists for the prefix cache and passes them to the speculative engine
separately, so snapshots and rollback apply to the target's slots and
the draft forward receives both groups as arguments. The identity
comparison, its panics, and the per-request rebinding are gone; the two
groups are fixed at load time.
This commit is contained in:
Jesse Gross 2026-08-01 21:53:18 -07:00
parent 2f84872ce0
commit e7fbd528f7
15 changed files with 142 additions and 175 deletions
Download patch file Download diff file Expand all files Collapse all files

41
x/mlxrunner/model/base/base.go
View file

@ -15,36 +15,41 @@ import (
// Model is the interface that model implementations must satisfy.
type Model interface {
// Forward returns the hidden state to unembed and the state a draft
// model conditions on; plain models return the final hidden for both.
Forward(b *batch.Batch, cache []cache.Cache) (hidden, auxHidden *mlx.Array)
Unembed(x *mlx.Array) *mlx.Array
NumLayers() int
Tokenizer() *tokenizer.Tokenizer
MaxContextLength() int
// LoadWeights receives all tensors loaded from the manifest and assigns
// them to model fields. Model-specific logic (MLA absorption, expert
// stacking, quantized layer creation) happens here.
LoadWeights(tensors map[string]*mlx.Array) error
// NewCaches builds the cache slots this model's layers need.
NewCaches() []cache.Cache
// Forward returns the hidden state to unembed and the state a draft
// model conditions on; plain models return the final hidden for both.
Forward(b *batch.Batch, cache []cache.Cache) (hidden, auxHidden *mlx.Array)
Unembed(x *mlx.Array) *mlx.Array
Tokenizer() *tokenizer.Tokenizer
MaxContextLength() int
}
// DraftModel is an auxiliary model alongside a target that proposes speculative
// tokens.
type DraftModel interface {
// Draft fuses b.Hidden (the target hidden state) into its own forward and
// returns the head's hidden plus the aux hidden seeding the next step.
Draft(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden *mlx.Array)
// LoadWeights assigns manifest tensors to the draft model's fields. An
// inline head has nothing to do here; its weights load with the target's.
LoadWeights(tensors map[string]*mlx.Array) error
// NewCaches builds the cache slots this draft model writes, or nil
// when it keeps no KV.
NewCaches() []cache.Cache
// Forward consumes b.Hidden (the draft-conditioning state) and returns
// its hidden plus the aux hidden that seeds the next step. targetCaches
// is read-only, for drafts that attend over the target's history.
Forward(b *batch.Batch, targetCaches, draftCaches []cache.Cache) (hidden, auxHidden *mlx.Array)
// Unembed projects a hidden state to vocabulary logits.
Unembed(x *mlx.Array) *mlx.Array
// DraftCaches selects the draft model's own KV caches from the full
// per-request slice — any subset, or nil when the draft keeps no KV.
DraftCaches(caches []cache.Cache) []cache.Cache
// LoadWeights assigns manifest tensors to the draft head's fields.
LoadWeights(tensors map[string]*mlx.Array) error
}
// SelfDraft is implemented by models whose draft head ships inline with the

8
x/mlxrunner/mtp.go
View file

@ -151,12 +151,12 @@ func (d *mtpDraftSession) flush() {
ids := mlx.Concatenate(d.pendingTokens, 1)
hiddens := mlx.Concatenate(d.pendingHiddens, 1)
hidden, auxHidden := spec.draft.Draft(&batch.Batch{
hidden, auxHidden := spec.draft.Forward(&batch.Batch{
InputIDs: ids,
SeqOffsets: []int32{int32(d.committedDraftOffset)},
SeqQueryLens: []int32{int32(ids.Dim(1))},
Hidden: hiddens,
}, spec.caches)
}, spec.targets, spec.draftKV)
d.setHeld(lastHiddenRow(hidden), lastHiddenRow(auxHidden))
d.committedDraftOffset += ids.Dim(1)
@ -226,12 +226,12 @@ func (d *mtpDraftSession) propose(current *mlx.Array, maxTokens int) *draftCandi
if len(spec.draftKV) > 0 {
pos = d.frontier - 1 + i
}
hidden, auxHidden = spec.draft.Draft(&batch.Batch{
hidden, auxHidden = spec.draft.Forward(&batch.Batch{
InputIDs: lastToken,
SeqOffsets: []int32{int32(pos)},
SeqQueryLens: []int32{1},
Hidden: lastHidden,
}, spec.caches)
}, spec.targets, spec.draftKV)
}
// Unembed only the row being sampled, never the batch.
stepLogits := spec.draft.Unembed(hidden).Squeeze(1)

98
x/mlxrunner/mtp_test.go
View file

@ -82,6 +82,7 @@ func (m *fakeMTPModel) Forward(b *batch.Batch, caches []cache.Cache) (hidden, au
return out, out
}
func (m *fakeMTPModel) NewCaches() []cache.Cache { return nil }
func (m *fakeMTPModel) Unembed(x *mlx.Array) *mlx.Array { return x }
func (m *fakeMTPModel) NumLayers() int { return 1 }
func (m *fakeMTPModel) Tokenizer() *tokenizer.Tokenizer { return m.tok }
@ -107,9 +108,9 @@ type draftCall struct {
func (d *fakeMTPDraft) LoadWeights(map[string]*mlx.Array) error { return nil }
func (d *fakeMTPDraft) DraftCaches([]cache.Cache) []cache.Cache { return nil }
func (d *fakeMTPDraft) NewCaches() []cache.Cache { return nil }
func (d *fakeMTPDraft) Draft(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden *mlx.Array) {
func (d *fakeMTPDraft) Forward(b *batch.Batch, _, _ []cache.Cache) (hidden, auxHidden *mlx.Array) {
mlx.Eval(b.InputIDs)
prev := int32(b.InputIDs.Ints()[0])
d.calls = append(d.calls, draftCall{position: b.SeqOffsets[0], from: prev})
@ -121,18 +122,19 @@ func (d *fakeMTPDraft) Unembed(x *mlx.Array) *mlx.Array { return x }
var _ base.DraftModel = (*fakeMTPDraft)(nil)
// fakeKVDraft is a draft head with its own KV cache: it claims the trailing
// cache slot, writes its input ids there on every Draft call (advancing the
// fakeKVDraft is a draft head that declares a KV cache: it writes its
// input ids there on every Draft call (advancing the
// offset like a real KV write), and records each call's offset, ids, and
// the identity of every fused hidden row. A target hidden row is one-hot
// (its hot index identifies which position it came from); the head's own
// aux hidden is all-zero and records as -1.
type fakeKVDraft struct {
predict map[int32]int32
extends []extendCall
predict map[int32]int32
draftCaches []cache.Cache
extends []extendCall
}
// extendCall is one recorded Draft call: the absolute slot of the first
// extendCall is one recorded Forward call: the absolute slot of the first
// entry written, the look-ahead token ids, and the hot index of each fused
// hidden row (-1 for the head's own aux hidden).
type extendCall struct {
@ -143,11 +145,9 @@ type extendCall struct {
func (d *fakeKVDraft) LoadWeights(map[string]*mlx.Array) error { return nil }
func (d *fakeKVDraft) DraftCaches(caches []cache.Cache) []cache.Cache {
return caches[len(caches)-1:]
}
func (d *fakeKVDraft) NewCaches() []cache.Cache { return d.draftCaches }
func (d *fakeKVDraft) Draft(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden *mlx.Array) {
func (d *fakeKVDraft) Forward(b *batch.Batch, _, draftCaches []cache.Cache) (hidden, auxHidden *mlx.Array) {
mlx.Eval(b.InputIDs, b.Hidden)
rawIDs := b.InputIDs.Ints()
ids := make([]int32, len(rawIDs))
@ -168,7 +168,7 @@ func (d *fakeKVDraft) Draft(b *batch.Batch, caches []cache.Cache) (hidden, auxHi
}
d.extends = append(d.extends, extendCall{offset: b.SeqOffsets[0], ids: ids, hiddens: hot})
if rc, ok := d.DraftCaches(caches)[0].(*fakeRewindableCache); ok {
if rc, ok := draftCaches[0].(*fakeRewindableCache); ok {
rc.feed(ids)
}
@ -381,7 +381,7 @@ func TestRunMTPDecodeGreedy(t *testing.T) {
draft := &fakeMTPDraft{predict: predict}
caches, _ := newMTPTestCaches(1)
r.cache.caches = caches
r.spec = newSpeculation(r, draft)
r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
session, ch := newMTPTestSession(caches)
position := 1 // one prefill token already processed
@ -442,7 +442,7 @@ func TestRunMTPDecodeSampled(t *testing.T) {
r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{Temperature: 1, Seed: 42, UseSeed: true})
caches, _ := newMTPTestCaches(1)
r.cache.caches = caches
r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict})
r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict}, caches, nil)
session, ch := newMTPTestSession(caches)
position := 1
@ -452,7 +452,7 @@ func TestRunMTPDecodeSampled(t *testing.T) {
CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
SamplerOpts: sampler.Options{Temperature: 1, Seed: 42, UseSeed: true},
}
spec := r.spec.open(req, caches)
spec := r.spec.open(req)
if spec == nil || !spec.enabled {
t.Fatalf("open rejected a sampled request")
}
@ -486,7 +486,7 @@ func TestRunMTPDecodeWarmDrafter(t *testing.T) {
draft := &fakeMTPDraft{predict: predict}
caches, _ := newMTPTestCaches(1)
r.cache.caches = caches
r.spec = newSpeculation(r, draft)
r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
session, ch := newMTPTestSession(caches)
position := 1 // one prefill token already processed
@ -496,7 +496,7 @@ func TestRunMTPDecodeWarmDrafter(t *testing.T) {
CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
SamplerOpts: sampler.Options{},
}
spec := r.spec.open(req, caches)
spec := r.spec.open(req)
pinDraftLimit(spec, 4)
// The prefill chunk's committed report: token 0 at slot 0 with its
// hidden row, leaving the drafter ready to propose from slot 1.
@ -547,7 +547,7 @@ func TestRunMTPDecodeEOSCutLeavesPositionsUnjudged(t *testing.T) {
r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
caches, _ := newMTPTestCaches(1)
r.cache.caches = caches
r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict})
r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict}, caches, nil)
session, ch := newMTPTestSession(caches)
position := 1
@ -557,7 +557,7 @@ func TestRunMTPDecodeEOSCutLeavesPositionsUnjudged(t *testing.T) {
CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
SamplerOpts: sampler.Options{},
}
spec := r.spec.open(req, caches)
spec := r.spec.open(req)
if spec == nil || !spec.enabled {
t.Fatalf("want an enabled speculationSession with a depth controller, got %+v", spec)
}
@ -652,7 +652,7 @@ func TestDecodeCancelledMidStream(t *testing.T) {
r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
caches, tr := newMTPTestCaches(1)
r.cache.caches = caches
r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict})
r.spec = newSpeculation(r, &fakeMTPDraft{predict: predict}, caches, nil)
session := &cacheSession{caches: caches}
ch := make(chan CompletionResponse) // unbuffered: every send must rendezvous
@ -707,7 +707,7 @@ func pinDraftLimit(spec *speculationSession, limit int) {
// this request, with the draft length pinned to a fixed width; tests close it
// explicitly so close-time effects are visible to assertions.
func testDecoder(r *Runner, req Request, caches []cache.Cache, seed []int32, position int) decoder {
if spec := r.spec.open(req, caches); spec != nil {
if spec := r.spec.open(req); spec != nil {
if spec.enabled {
pinDraftLimit(spec, 4)
}
@ -728,10 +728,11 @@ func TestDecodeKVDraft(t *testing.T) {
const eos int32 = 7
predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: eos, eos: 0}
r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
draft := &fakeKVDraft{predict: predict}
draft := &fakeKVDraft{predict: predict, draftCaches: nil}
caches, _ := newMTPTestCaches(2) // caches[0] target, caches[1] draft KV
draft.draftCaches = caches[1:]
r.cache.caches = caches
r.spec = newSpeculation(r, draft)
r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
session, ch := newMTPTestSession(caches)
position := 0
@ -741,7 +742,7 @@ func TestDecodeKVDraft(t *testing.T) {
CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
SamplerOpts: sampler.Options{},
}
spec := r.spec.open(req, caches)
spec := r.spec.open(req)
if spec == nil || !spec.enabled || len(spec.spec.targets) != 1 {
t.Fatalf("speculation engine not built around the draft caches")
}
@ -812,10 +813,11 @@ func TestDecodeKVDraftRejectionRebuildsFromTarget(t *testing.T) {
// The draft mirrors the target except at 3, where it proposes 6 (absent from
// the target chain); once the target corrects 3->4 the next proposal
// re-aligns on the shared chain.
draft := &fakeKVDraft{predict: map[int32]int32{1: 2, 2: 3, 3: 6, 6: 0, 4: 5, 5: eos, eos: 0}}
draft := &fakeKVDraft{predict: map[int32]int32{1: 2, 2: 3, 3: 6, 6: 0, 4: 5, 5: eos, eos: 0}, draftCaches: nil}
caches, _ := newMTPTestCaches(2) // caches[0] target, caches[1] draft KV
draft.draftCaches = caches[1:]
r.cache.caches = caches
r.spec = newSpeculation(r, draft)
r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
session, ch := newMTPTestSession(caches)
position := 0
@ -825,7 +827,7 @@ func TestDecodeKVDraftRejectionRebuildsFromTarget(t *testing.T) {
CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
SamplerOpts: sampler.Options{},
}
spec := r.spec.open(req, caches)
spec := r.spec.open(req)
pinDraftLimit(spec, 4)
defer spec.close()
d := spec.decoder(mlx.FromValues([]int32{1}, 1), position)
@ -881,10 +883,11 @@ func TestDecodeMaintainsDraftCacheWithoutDrafting(t *testing.T) {
predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: eos, eos: 0}
opts := sampler.Options{Logprobs: true}
r := mtpTestRunner(t, predict, []int32{eos}, opts)
draft := &fakeKVDraft{predict: predict}
draft := &fakeKVDraft{predict: predict, draftCaches: nil}
caches, _ := newMTPTestCaches(2)
draft.draftCaches = caches[1:]
r.cache.caches = caches
r.spec = newSpeculation(r, draft)
r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
session, ch := newMTPTestSession(caches)
position := 0
@ -894,7 +897,7 @@ func TestDecodeMaintainsDraftCacheWithoutDrafting(t *testing.T) {
CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
SamplerOpts: opts,
}
spec := r.spec.open(req, caches)
spec := r.spec.open(req)
if spec == nil || spec.enabled {
t.Fatalf("want a permanent-park speculationSession, got %+v", spec)
}
@ -943,15 +946,16 @@ func TestSettleLevelsDraftCacheWithPrefill(t *testing.T) {
const eos int32 = 7
predict := map[int32]int32{2: 3, 3: 4, 4: 5}
r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
draft := &fakeKVDraft{predict: predict}
draft := &fakeKVDraft{predict: predict, draftCaches: nil}
caches, _ := newMTPTestCaches(2)
draft.draftCaches = caches[1:]
r.cache.caches = caches
r.spec = newSpeculation(r, draft)
r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
req := Request{
CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
SamplerOpts: sampler.Options{},
}
spec := r.spec.open(req, caches)
spec := r.spec.open(req)
defer spec.close()
// The prompt's only chunk: tokens 1..4 at slots 0..3 with their hiddens;
@ -984,15 +988,16 @@ func TestCommittedRunBatchesPastFlushCap(t *testing.T) {
}
r := mtpTestRunner(t, predict, []int32{0}, sampler.Options{})
draft := &fakeKVDraft{predict: predict}
draft := &fakeKVDraft{predict: predict, draftCaches: nil}
caches, _ := newMTPTestCaches(2)
draft.draftCaches = caches[1:]
r.cache.caches = caches
r.spec = newSpeculation(r, draft)
r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
req := Request{
CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
SamplerOpts: sampler.Options{},
}
spec := r.spec.open(req, caches)
spec := r.spec.open(req)
defer spec.close()
// One prefill-sized chunk: n tokens at slots 0..n-1 with their hiddens.
@ -1022,10 +1027,11 @@ func TestRestoredPrefixRewritesBoundaryPair(t *testing.T) {
const eos int32 = 7
predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: eos, eos: 0}
r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
draft := &fakeKVDraft{predict: predict}
draft := &fakeKVDraft{predict: predict, draftCaches: nil}
caches, _ := newMTPTestCaches(2)
draft.draftCaches = caches[1:]
r.cache.caches = caches
r.spec = newSpeculation(r, draft)
r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
session, ch := newMTPTestSession(caches)
req := Request{
Responses: ch,
@ -1033,7 +1039,7 @@ func TestRestoredPrefixRewritesBoundaryPair(t *testing.T) {
CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
SamplerOpts: sampler.Options{},
}
spec := r.spec.open(req, caches)
spec := r.spec.open(req)
pinDraftLimit(spec, 4)
d := spec.decoder(mlx.FromValues([]int32{1}, 1), 0)
if err := r.decode(context.Background(), req, session, d, 0); err != nil {
@ -1055,7 +1061,7 @@ func TestRestoredPrefixRewritesBoundaryPair(t *testing.T) {
t.Fatal("restore to 5 failed")
}
}
spec = r.spec.open(req, caches)
spec = r.spec.open(req)
spec.committed(mlx.FromValues([]int32{6, 1}, 1, 2), oneHotLogits([]int32{eos, 2}), 5)
spec.close()
@ -1077,16 +1083,17 @@ func TestDecodeParkedDraftResume(t *testing.T) {
const eos int32 = 7
predict := map[int32]int32{1: 2, 2: 3, 3: 4, 4: 5, 5: 6, 6: eos, eos: 0}
r := mtpTestRunner(t, predict, []int32{eos}, sampler.Options{})
draft := &fakeKVDraft{predict: predict}
draft := &fakeKVDraft{predict: predict, draftCaches: nil}
caches, _ := newMTPTestCaches(2)
draft.draftCaches = caches[1:]
r.cache.caches = caches
r.spec = newSpeculation(r, draft)
r.spec = newSpeculation(r, draft, caches[:1], caches[1:])
req := Request{
Tokens: []int32{1},
CompletionRequest: CompletionRequest{Options: api.Options{NumPredict: 20}},
SamplerOpts: sampler.Options{},
}
spec := r.spec.open(req, caches)
spec := r.spec.open(req)
if spec == nil || !spec.enabled {
t.Fatalf("want a drafting speculationSession, got %+v", spec)
}
@ -1179,14 +1186,13 @@ func newMTPTestSession(caches []cache.Cache) (*cacheSession, chan CompletionResp
// no-op drafter, since the engine requires one.
func testSpeculationSession(r *Runner, caches []cache.Cache) *speculationSession {
if r.spec != nil {
r.spec.bind(caches)
return &speculationSession{spec: r.spec, drafter: r.spec.drafter.open()}
}
s := &speculation{r: r, caches: caches, targets: caches}
s := &speculation{r: r, targets: caches}
return &speculationSession{spec: s, drafter: nopDrafter{}}
}
// nopDrafter satisfies drafter for engine tests that supply candidates
// nopDrafter satisfies draftSession for engine tests that supply candidates
// directly and never propose.
type nopDrafter struct{}

2
x/mlxrunner/pipeline.go
View file

@ -77,7 +77,7 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er
// Built before prefill so a drafter with draft caches follows the prompt
// through prefill alongside the target.
spec := r.spec.open(request, caches)
spec := r.spec.open(request)
defer spec.close()
seed, position, promptEval, err := r.prefill(ctx, session, spec)

15
x/mlxrunner/prefix_cache.go
View file

@ -27,7 +27,6 @@ import (
"github.com/ollama/ollama/logutil"
"github.com/ollama/ollama/x/mlxrunner/cache"
"github.com/ollama/ollama/x/mlxrunner/mlx"
"github.com/ollama/ollama/x/mlxrunner/model/base"
)
const maxPagedOutBytes int64 = 8 << 30 // 8 GiB eviction threshold for paged-out snapshot memory
@ -66,17 +65,9 @@ type cacheSession struct {
pendingSnapshots []pendingSnapshot
}
func newPrefixCache(m base.Model) *prefixCache {
c := &prefixCache{}
if cacheFactory, ok := m.(interface{ NewCaches() []cache.Cache }); ok {
c.caches = cacheFactory.NewCaches()
return c
}
c.caches = make([]cache.Cache, m.NumLayers())
for i := range c.caches {
c.caches[i] = cache.NewKVCache()
}
return c
// newPrefixCache manages the given cache slots for the model's life.
func newPrefixCache(caches []cache.Cache) *prefixCache {
return &prefixCache{caches: caches}
}
func (c *prefixCache) ensureRoot() {

16
x/mlxrunner/runner.go
View file

@ -6,12 +6,14 @@ import (
"log/slog"
"net"
"net/http"
"slices"
"strings"
"golang.org/x/sync/errgroup"
"github.com/ollama/ollama/api"
"github.com/ollama/ollama/x/internal/mlxthread"
"github.com/ollama/ollama/x/mlxrunner/cache"
"github.com/ollama/ollama/x/mlxrunner/mlx"
"github.com/ollama/ollama/x/mlxrunner/model"
"github.com/ollama/ollama/x/mlxrunner/model/base"
@ -106,15 +108,25 @@ func (r *Runner) Load(modelName string) error {
r.Model = m
r.Tokenizer = m.Tokenizer()
r.contextLength = m.MaxContextLength()
r.cache = newPrefixCache(m)
caches := m.NewCaches()
draftCaches := newDraftCaches(draftModel)
r.cache = newPrefixCache(slices.Concat(caches, draftCaches))
r.Sampler = sample.New(r.contextLength)
r.spec = newSpeculation(r, draftModel)
r.spec = newSpeculation(r, draftModel, caches, draftCaches)
mlx.EnableCompile()
return nil
}
// newDraftCaches returns nil when the model ships no draft.
func newDraftCaches(draft base.DraftModel) []cache.Cache {
if draft == nil {
return nil
}
return draft.NewCaches()
}
func configureWiredMemory() {
if !mlx.GPUIsAvailable() {
return

51
x/mlxrunner/speculate.go
View file

@ -2,7 +2,6 @@ package mlxrunner
import (
"fmt"
"slices"
"time"
"github.com/ollama/ollama/x/mlxrunner/batch"
@ -44,14 +43,10 @@ type speculation struct {
r *Runner
draft base.DraftModel
// caches is the whole persistent slice, passed to every forward; draftKV
// are the draft head's own caches and targets are the rest — the caches the
// target forward writes, which speculation snapshots and rollback cover.
// Bound the first time the caches exist (the Runner reuses one cache slice
// for its life) and stable thereafter.
caches []cache.Cache
draftKV []cache.Cache
// targets are the model's slots — what speculation snapshots and rolls
// back; draftKV the drafter's. Built at load, stable for the model's life.
targets []cache.Cache
draftKV []cache.Cache
// drafter is the persistent half of the MTP drafting machinery; each
// request's session comes from drafter.open.
@ -64,44 +59,15 @@ type speculation struct {
// 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 {
func newSpeculation(r *Runner, draft base.DraftModel, targets, draftKV []cache.Cache) *speculation {
if draft == nil {
return nil
}
s := &speculation{r: r, draft: draft, depth: newDepthController()}
s.bind(r.cache.caches)
s := &speculation{r: r, draft: draft, targets: targets, draftKV: draftKV, depth: newDepthController()}
s.drafter = newMTPDrafter(s)
return s
}
// bind computes the draft/target cache partition the first time the persistent
// caches exist; later requests reuse the same slice, so it runs once.
func (s *speculation) bind(caches []cache.Cache) {
if s.caches != nil {
if !slices.Equal(s.caches, caches) {
panic("speculation: cache slice changed between requests")
}
return
}
draftKV := s.draft.DraftCaches(caches)
// Partition caches into target slots (everything not in draftKV) in one
// pass. The count check rejects a draft slot that isn't a member of caches.
targets := make([]cache.Cache, 0, len(caches))
for _, c := range caches {
if !slices.Contains(draftKV, c) {
targets = append(targets, c)
}
}
if len(caches)-len(targets) != len(draftKV) {
panic("speculation: DraftCaches must select slots of the cache slice")
}
s.caches = caches
s.draftKV = draftKV
s.targets = targets
}
// speculationSession is the per-request cursor over the persistent speculation:
// it owns the drafter and runs the validate rounds. A nil session is a plain
// decode.
@ -122,11 +88,10 @@ type speculationSession struct {
// open returns the speculation cursor for this request or nil when the model ships
// no draft head (a nil receiver), which decodes plainly.
func (s *speculation) open(request Request, caches []cache.Cache) *speculationSession {
func (s *speculation) open(request Request) *speculationSession {
if s == nil {
return nil
}
s.bind(caches)
d := s.drafter.open()
// Logprobs are not yet supported, so a logprobs request keeps a speculationSession
@ -291,7 +256,7 @@ func (st *speculativeDecoder) resume() []sampler.Result {
func (st *speculativeDecoder) park(remaining int) ([]sampler.Result, error) {
s := st.s
if st.inner == nil {
st.inner = s.spec.r.pipelinedDecoder(s, s.spec.caches, st.current.Token.ExpandDims(-1), st.position)
st.inner = s.spec.r.pipelinedDecoder(s, s.spec.targets, st.current.Token.ExpandDims(-1), st.position)
}
return st.inner.next(remaining)
}
@ -417,7 +382,7 @@ func (s *speculationSession) accept(position *int, current sampler.Result, candi
InputIDs: current.Token.ExpandDims(-1).Concatenate(1, candidates.tokens),
SeqOffsets: []int32{int32(before)},
SeqQueryLens: []int32{int32(draftCount + 1)},
}, s.spec.caches)
}, s.spec.targets)
// Row i of the fused hidden is the state after the token at before+i, so
// the rows already line up with the drafts: row 0 (current's state)

4
x/models/cohere2_moe/cohere2_moe.go
View file

@ -754,10 +754,6 @@ func (m *Model) Unembed(x *mlx.Array) *mlx.Array {
return logits
}
func (m *Model) NumLayers() int {
return len(m.Layers)
}
func (m *Model) MaxContextLength() int {
return int(m.MaxPositionEmbeddings)
}

4
x/models/gemma3/gemma3.go
View file

@ -427,10 +427,6 @@ func (m *Model) Unembed(x *mlx.Array) *mlx.Array {
return m.LMHead.Forward(x)
}
func (m *Model) NumLayers() int {
return len(m.Layers)
}
func (m *Model) MaxContextLength() int {
return int(m.MaxPositionEmbeddings)
}

14
x/models/gemma4/assistant.go
View file

@ -269,22 +269,22 @@ func (m *AssistantModel) precomputeScaledWeights() {
}
}
// DraftCaches returns nil: the assistant keeps no KV and re-attends the
// NewCaches returns nil: the assistant keeps no KV and re-attends the
// target's caches read-only each step.
func (m *AssistantModel) DraftCaches([]cache.Cache) []cache.Cache { return nil }
func (m *AssistantModel) NewCaches() []cache.Cache { return nil }
// Draft is single-position: the head drafts every token as if it sat at the
// Forward is single-position: the head drafts every token as if it sat at the
// last target-seen position, so it anchors RoPE and the mask there — from the
// non-moving target full-attention cache — regardless of the advancing offset
// in b. The input fuses the target's scaled token embedding with b.Hidden.
func (m *AssistantModel) Draft(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden *mlx.Array) {
func (m *AssistantModel) Forward(b *batch.Batch, targetCaches, _ []cache.Cache) (hidden, auxHidden *mlx.Array) {
inputsEmbeds := m.target.TokenEmbeddings(b.InputIDs).Concatenate(-1, b.Hidden)
dims := inputsEmbeds.Dims()
B, L := int32(dims[0]), int32(dims[1])
anchor := int32(0)
if len(caches) > 0 {
if full := caches[len(caches)-1]; full != nil {
if len(targetCaches) > 0 {
if full := targetCaches[len(targetCaches)-1]; full != nil {
anchor = int32(full.Offset() - 1)
}
}
@ -294,7 +294,7 @@ func (m *AssistantModel) Draft(b *batch.Batch, caches []cache.Cache) (hidden, au
SeqQueryLens: []int32{L},
}
sliding, full := m.sharedHistories(ab, caches)
sliding, full := m.sharedHistories(ab, targetCaches)
h := m.PreProjection.Forward(inputsEmbeds)
positions := mlx.FromValues([]int32{anchor}, 1)

4
x/models/gemma4/gemma4.go
View file

@ -1081,10 +1081,6 @@ func suppressTokenLogits(logits, bias *mlx.Array) *mlx.Array {
return logits.Add(bias.AsType(logits.DType()))
}
func (m *Model) NumLayers() int {
return len(m.Layers)
}
func (m *Model) MaxContextLength() int {
return int(m.MaxPositionEmbeddings)
}

10
x/models/glm4_moe_lite/glm4_moe_lite.go
View file

@ -767,8 +767,14 @@ func (m *Model) Unembed(x *mlx.Array) *mlx.Array {
return m.LMHead.Forward(x)
}
// NumLayers returns the number of transformer layers
func (m *Model) NumLayers() int { return len(m.Layers) }
// NewCaches builds a KV cache per layer.
func (m *Model) NewCaches() []cache.Cache {
caches := make([]cache.Cache, len(m.Layers))
for i := range caches {
caches[i] = cache.NewKVCache()
}
return caches
}
// MaxContextLength returns the maximum context length
func (m *Model) MaxContextLength() int { return int(m.MaxPositionEmbeddings) }

4
x/models/llama/llama.go
View file

@ -259,10 +259,6 @@ func (m *Model) Unembed(x *mlx.Array) *mlx.Array {
return m.LMHead.Forward(x)
}
func (m *Model) NumLayers() int {
return len(m.Layers)
}
func (m *Model) MaxContextLength() int {
return int(m.MaxPositionEmbeddings)
}

4
x/models/qwen3/qwen3.go
View file

@ -276,10 +276,6 @@ func (m *Model) Unembed(x *mlx.Array) *mlx.Array {
return m.LMHead.Forward(x)
}
func (m *Model) NumLayers() int {
return len(m.Layers)
}
func (m *Model) MaxContextLength() int {
return int(m.MaxPositionEmbeddings)
}

42
x/models/qwen3_5/qwen3_5.go
View file

@ -27,7 +27,7 @@ func init() {
var (
_ base.Model = (*Model)(nil)
_ base.SelfDraft = (*Model)(nil)
_ base.DraftModel = (*Model)(nil)
_ base.DraftModel = (*mtpDraft)(nil)
)
// RopeParameters carries optional rope metadata embedded under rope_parameters.
@ -102,8 +102,8 @@ type Model struct {
weightPrefix string
}
// MTPHead is the multi-token-prediction draft head; it owns a KV cache appended
// after the per-layer caches and reuses the model's lm_head.
// MTPHead is the multi-token-prediction draft head; it writes one KV cache
// and reuses the model's lm_head.
type MTPHead struct {
Enorm *nn.RMSNorm
Hnorm *nn.RMSNorm
@ -1215,27 +1215,35 @@ func (m *Model) Unembed(x *mlx.Array) *mlx.Array {
return m.LMHead.Forward(x)
}
// DraftCaches returns the MTP head's KV caches: the trailing slots NewCaches
// appended after the per-layer caches.
func (m *Model) DraftCaches(caches []cache.Cache) []cache.Cache {
// mtpDraft is the model viewed as its own draft: the same struct, carrying
// the DraftModel method set for the inline MTP head.
type mtpDraft Model
// NewCaches builds the MTP head's KV cache.
func (m *mtpDraft) NewCaches() []cache.Cache {
if m.MTP == nil {
return nil
}
return caches[len(m.Layers):]
return []cache.Cache{cache.NewKVCache()}
}
// SelfDraft returns the model as its own draft when an MTP head loaded, else
// nil; Draft exists either way, so availability is reported here.
// LoadWeights does nothing: an inline head's weights load with the target's.
func (m *mtpDraft) LoadWeights(map[string]*mlx.Array) error { return nil }
func (m *mtpDraft) Unembed(x *mlx.Array) *mlx.Array { return (*Model)(m).Unembed(x) }
// SelfDraft returns the model's drafting view, or nil when no MTP head was
// loaded. The methods exist either way, so this is what decides availability.
func (m *Model) SelfDraft() base.DraftModel {
if m.MTP == nil {
return nil
}
return m
return (*mtpDraft)(m)
}
// Draft runs one MTP step; the head's hidden also serves as the aux
// Forward runs one MTP step; the head's hidden also serves as the aux
// hidden seeding the next step.
func (m *Model) Draft(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden *mlx.Array) {
func (m *mtpDraft) Forward(b *batch.Batch, _, draftCaches []cache.Cache) (hidden, auxHidden *mlx.Array) {
dims := b.InputIDs.Dims()
B, L := int32(dims[0]), int32(dims[1])
positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets))
@ -1244,8 +1252,7 @@ func (m *Model) Draft(b *batch.Batch, caches []cache.Cache) (hidden, auxHidden *
h := m.MTP.Hnorm.Forward(b.Hidden, m.RMSNormEps)
fused := m.MTP.FC.Forward(emb.Concatenate(-1, h))
c := m.DraftCaches(caches)[0]
out := m.MTP.Layer.Forward(fused, b, c, positions, B, L, m.Config)
out := m.MTP.Layer.Forward(fused, b, draftCaches[0], positions, B, L, m.Config)
hidden = m.MTP.Norm.Forward(out, m.RMSNormEps)
return hidden, hidden
}
@ -1263,8 +1270,7 @@ func (m *Model) Tokenizer() *tokenizer.Tokenizer {
}
func (m *Model) NewCaches() []cache.Cache {
// Reserve a trailing slot for the MTP head's KV cache when present.
caches := make([]cache.Cache, len(m.Layers), len(m.Layers)+1)
caches := make([]cache.Cache, len(m.Layers))
convTail := m.LinearConvKernelDim - 1
convDim := 2*m.LinearNumKeyHeads*m.LinearKeyHeadDim + m.LinearNumValueHeads*m.LinearValueHeadDim
for i, layer := range m.Layers {
@ -1274,9 +1280,5 @@ func (m *Model) NewCaches() []cache.Cache {
caches[i] = cache.NewKVCache()
}
}
// Trailing slot is the MTP head's (DraftCaches); Model.Forward never touches it.
if m.MTP != nil {
caches = append(caches, cache.NewKVCache())
}
return caches
}