From 98e26b8c3723e4b7c8ff9b8985b630b42698153d Mon Sep 17 00:00:00 2001 From: Patrick Devine Date: Thu, 14 May 2026 14:02:34 -0700 Subject: [PATCH] mlxrunner: add DFlash speculative decoding (#16134) This change adds dflash block diffusion speculative decoding to the MLX runner. Included in this change: support for qwen3.6 moe/dense speculative decoding draft model recurrent cache playback RoPE/YaRN changes (DRY out the laguna/dflash MoE YaRN implementation) support for greedy sampling / leviathan/chen sampling --- x/create/client/create.go | 43 +- x/create/client/create_test.go | 10 +- x/mlxrunner/cache.go | 39 +- x/mlxrunner/cache/cache.go | 6 + x/mlxrunner/cache/recurrent.go | 126 ++++- x/mlxrunner/cache/recurrent_test.go | 51 +- x/mlxrunner/cache_test.go | 44 ++ x/mlxrunner/dflash.go | 715 ++++++++++++++++++++++++++++ x/mlxrunner/imports.go | 1 + x/mlxrunner/mlx/gated_delta.go | 38 +- x/mlxrunner/model/base/base.go | 17 + x/mlxrunner/pipeline.go | 179 ++++++- x/mlxrunner/runner.go | 1 + x/models/dflash/dflash.go | 439 +++++++++++++++++ x/models/dflash/dflash_test.go | 52 ++ x/models/laguna/laguna.go | 264 +++------- x/models/nn/rope.go | 129 +++++ x/models/qwen3_5/qwen3_5.go | 33 +- 18 files changed, 1940 insertions(+), 247 deletions(-) create mode 100644 x/mlxrunner/dflash.go create mode 100644 x/models/dflash/dflash.go create mode 100644 x/models/dflash/dflash_test.go create mode 100644 x/models/nn/rope.go diff --git a/x/create/client/create.go b/x/create/client/create.go index d8a1a8f8..a6f1cd0f 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -243,6 +243,40 @@ func appendLayersManifestWriter(next create.ManifestWriter, extra []create.Layer } } +func draftMetadata(draftDir string) (*model.Draft, error) { + configPath := filepath.Join(draftDir, "config.json") + data, err := os.ReadFile(configPath) + if err != nil { + return nil, fmt.Errorf("failed to read draft config %s: %w", configPath, err) + } + + var cfg struct { + Architectures []string `json:"architectures"` + ModelType string `json:"model_type"` + } + if err := json.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("failed to parse draft config %s: %w", configPath, err) + } + + arch := "" + if len(cfg.Architectures) > 0 { + arch = cfg.Architectures[0] + } + if arch == "" { + arch = cfg.ModelType + } + if arch == "" { + return nil, fmt.Errorf("draft architecture not found in %s", configPath) + } + + return &model.Draft{ + ModelFormat: "safetensors", + Architecture: arch, + TensorPrefix: "draft.", + Config: "draft/config.json", + }, nil +} + func createModelFromBaseWithDraft(opts CreateOptions, draftLayers []create.LayerInfo, progressFn func(string)) error { progressFn(fmt.Sprintf("loading base model %s", opts.ModelDir)) baseManifest, err := imagemanifest.LoadManifest(opts.ModelDir) @@ -487,12 +521,11 @@ func newManifestWriter(opts CreateOptions, capabilities []string, parserName, re configData.Parser = resolveParserName(opts.Modelfile, parserName) configData.Renderer = resolveRendererName(opts.Modelfile, rendererName) if opts.Modelfile != nil && opts.Modelfile.Draft != "" { - configData.Draft = &model.Draft{ - ModelFormat: "safetensors", - Architecture: "Gemma4AssistantForCausalLM", - TensorPrefix: "draft.", - Config: "draft/config.json", + draft, err := draftMetadata(opts.Modelfile.Draft) + if err != nil { + return err } + configData.Draft = draft } configJSON, err := json.Marshal(configData) if err != nil { diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index bc8f4d6e..47cd9d8b 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -544,10 +544,15 @@ func TestNewManifestWriter_PopulatesFileTypeFromQuantize(t *testing.T) { func TestNewManifestWriter_PopulatesDraftMetadata(t *testing.T) { t.Setenv("OLLAMA_MODELS", t.TempDir()) + draftDir := t.TempDir() + if err := os.WriteFile(filepath.Join(draftDir, "config.json"), []byte(`{"architectures":["DFlashDraftModel"],"model_type":"qwen3"}`), 0o644); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + opts := CreateOptions{ ModelName: "test-draft", ModelDir: t.TempDir(), - Modelfile: &ModelfileConfig{Draft: "/tmp/assistant"}, + Modelfile: &ModelfileConfig{Draft: draftDir}, } writer := newManifestWriter(opts, []string{"completion"}, "gemma4", "gemma4") @@ -581,6 +586,9 @@ func TestNewManifestWriter_PopulatesDraftMetadata(t *testing.T) { if cfg.Draft.TensorPrefix != "draft." || cfg.Draft.Config != "draft/config.json" { t.Fatalf("Draft = %#v, want draft prefix/config", cfg.Draft) } + if cfg.Draft.Architecture != "DFlashDraftModel" { + t.Fatalf("Draft architecture = %q, want DFlashDraftModel", cfg.Draft.Architecture) + } } func TestSupportsThinking(t *testing.T) { diff --git a/x/mlxrunner/cache.go b/x/mlxrunner/cache.go index c1c53b66..f12c7215 100644 --- a/x/mlxrunner/cache.go +++ b/x/mlxrunner/cache.go @@ -38,6 +38,8 @@ type kvCache struct { pagedOutBytes int64 // total bytes in paged-out snapshots across the trie } +type cacheFactoryFunc func() []cache.Cache + // pendingSnapshot is a snapshot scheduled to be taken during prefill. type pendingSnapshot struct { offset int @@ -61,18 +63,22 @@ type cacheSession struct { pendingSnapshots []pendingSnapshot } -func (c *kvCache) ensureCaches(m base.Model) { +func (c *kvCache) ensureCachesWithFactory(newCaches cacheFactoryFunc) { if len(c.caches) != 0 { return } + c.caches = newCaches() +} + +func newModelCaches(m base.Model) []cache.Cache { if cacheFactory, ok := m.(interface{ NewCaches() []cache.Cache }); ok { - c.caches = cacheFactory.NewCaches() - return + return cacheFactory.NewCaches() } - c.caches = make([]cache.Cache, m.NumLayers()) - for i := range c.caches { - c.caches[i] = cache.NewKVCache() + caches := make([]cache.Cache, m.NumLayers()) + for i := range caches { + caches[i] = cache.NewKVCache() } + return caches } func (c *kvCache) ensureRoot() { @@ -87,15 +93,29 @@ func (c *kvCache) ensureRoot() { // begin prepares caches for a new request. It finds the nearest // matching cache or creates new caches if none match. func (c *kvCache) begin(m base.Model, inputs []int32) *cacheSession { - c.ensureCaches(m) + return c.beginWithFactory(inputs, func() []cache.Cache { return newModelCaches(m) }, "") +} + +func (c *kvCache) beginWithFactory(inputs []int32, newCaches cacheFactoryFunc, logPrefix string) *cacheSession { + return c.beginWithFactoryLimit(inputs, newCaches, logPrefix, -1, true) +} + +func (c *kvCache) beginWithFactoryLimit(inputs []int32, newCaches cacheFactoryFunc, logPrefix string, maxCachedPrefix int, keepSeedToken bool) *cacheSession { + c.ensureCachesWithFactory(newCaches) c.ensureRoot() matchPath, matched := findBestMatch(c.root, inputs) originalMatched := matched + if maxCachedPrefix >= 0 { + maxCachedPrefix = min(maxCachedPrefix, len(inputs)) + if matched > maxCachedPrefix { + matchPath, matched = findBestMatch(c.root, inputs[:maxCachedPrefix]) + } + } // Always keep at least one token to re-evaluate so the // pipeline can seed token generation from it. - if matched == len(inputs) && matched > 0 { + if keepSeedToken && matched == len(inputs) && matched > 0 { matchPath, matched = findBestMatch(c.root, inputs[:len(inputs)-1]) } @@ -123,6 +143,9 @@ func (c *kvCache) begin(m base.Model, inputs []int32) *cacheSession { if prefix == 0 { msg = "cache miss" } + if logPrefix != "" { + msg = logPrefix + " " + msg + } slog.Info(msg, "total", len(inputs), "matched", originalMatched, "cached", prefix, "left", len(remaining)) return session diff --git a/x/mlxrunner/cache/cache.go b/x/mlxrunner/cache/cache.go index d4485df5..913a59e7 100644 --- a/x/mlxrunner/cache/cache.go +++ b/x/mlxrunner/cache/cache.go @@ -88,6 +88,10 @@ func BeginSpeculation(caches []Cache) ([]Cache, *Speculation, bool) { sc := newSpeculativeKVCache(c) specCaches[i] = sc layers[i] = sc + case *RecurrentCache: + sc := newSpeculativeRecurrentCache(c) + specCaches[i] = sc + layers[i] = sc default: return nil, nil, false } @@ -108,6 +112,8 @@ func BeginIsolatedSpeculation(caches []Cache) ([]Cache, bool) { specCaches[i] = newSpeculativeRotatingKVCache(c) case *KVCache: specCaches[i] = newIsolatedKVCache(c) + case *RecurrentCache: + specCaches[i] = newSpeculativeRecurrentCache(c) default: return nil, false } diff --git a/x/mlxrunner/cache/recurrent.go b/x/mlxrunner/cache/recurrent.go index cf87953d..2006de1c 100644 --- a/x/mlxrunner/cache/recurrent.go +++ b/x/mlxrunner/cache/recurrent.go @@ -6,6 +6,19 @@ import ( "github.com/ollama/ollama/x/models/nn" ) +// Recurrent is the contract for caches that back recurrent linear-attention layers. +type Recurrent interface { + Cache + Get(b *batch.Batch, dtype mlx.DType) *nn.RecurrentHistory + Put(b *batch.Batch, newConv, newDelta *mlx.Array) +} + +// RecurrentRecorder records the per-token scan inputs needed to commit an +// accepted prefix after a speculative recurrent forward. +type RecurrentRecorder interface { + Record(qkv, q, k, v, gDecay, beta *mlx.Array) +} + // RecurrentCache stores state for linear-recurrent layers. // // Conv state shape: [B, convTail, convDim] @@ -53,9 +66,12 @@ func (c *RecurrentCache) ensure(batch int, dtype mlx.DType) { batch = 1 } + // Keep the gated-delta recurrent state in float32 even when activations are + // bf16/fp16. The convolution tail stays in the activation dtype. + deltaDType := mlx.DTypeFloat32 needConv := c.convState == nil || !c.convState.Valid() || c.convState.DType() != dtype || c.convState.Dim(0) != batch || c.convState.Dim(1) != c.convTail || c.convState.Dim(2) != c.convDim - needDelta := c.deltaState == nil || !c.deltaState.Valid() || c.deltaState.DType() != dtype || + needDelta := c.deltaState == nil || !c.deltaState.Valid() || c.deltaState.DType() != deltaDType || c.deltaState.Dim(0) != batch || c.deltaState.Dim(1) != c.numVHeads || c.deltaState.Dim(2) != c.headVDim || c.deltaState.Dim(3) != c.headKDim if !needConv && !needDelta { return @@ -65,7 +81,7 @@ func (c *RecurrentCache) ensure(batch int, dtype mlx.DType) { c.convState = c.setState(c.convState, mlx.Zeros(dtype, batch, c.convTail, c.convDim), false) } if needDelta { - c.deltaState = c.setState(c.deltaState, mlx.Zeros(dtype, batch, c.numVHeads, c.headVDim, c.headKDim), false) + c.deltaState = c.setState(c.deltaState, mlx.Zeros(deltaDType, batch, c.numVHeads, c.headVDim, c.headKDim), false) } } @@ -161,3 +177,109 @@ func (c *RecurrentCache) Free() { } func (c *RecurrentCache) Offset() int { return c.offset } + +type speculativeRecurrentCache struct { + speculativeBase + target *RecurrentCache + + start int + + initialConv *mlx.Array + initialDelta *mlx.Array + + qkv, q, k, v, gDecay, beta *mlx.Array + fullConv, fullDelta *mlx.Array + length int +} + +func newSpeculativeRecurrentCache(target *RecurrentCache) *speculativeRecurrentCache { + return &speculativeRecurrentCache{ + speculativeBase: speculativeBase{offset: target.Offset()}, + target: target, + start: target.Offset(), + } +} + +func (c *speculativeRecurrentCache) Get(b *batch.Batch, dtype mlx.DType) *nn.RecurrentHistory { + if c.fullConv != nil && c.fullDelta != nil { + return nn.NewRecurrentHistory(c.fullConv, c.fullDelta) + } + + history := c.target.Get(b, dtype) + if c.initialConv == nil { + c.initialConv = history.ConvState() + } + if c.initialDelta == nil { + c.initialDelta = history.DeltaState() + } + return history +} + +func (c *speculativeRecurrentCache) Record(qkv, q, k, v, gDecay, beta *mlx.Array) { + c.qkv, c.q, c.k, c.v, c.gDecay, c.beta = qkv, q, k, v, gDecay, beta + if qkv != nil { + c.length = qkv.Dim(1) + } +} + +func (c *speculativeRecurrentCache) Put(b *batch.Batch, newConv, newDelta *mlx.Array) { + c.fullConv, c.fullDelta = newConv, newDelta + c.offset += int(b.SeqQueryLens[0]) +} + +func (c *speculativeRecurrentCache) State() []*mlx.Array { + if c.fullConv != nil && c.fullDelta != nil { + return []*mlx.Array{c.fullConv, c.fullDelta} + } + return c.target.State() +} + +func (c *speculativeRecurrentCache) commit(n int) { + if n <= 0 { + return + } + if c.length > 0 && n > c.length { + n = c.length + } + + if c.length > 0 && n == c.length && c.fullConv != nil && c.fullDelta != nil { + c.target.convState = c.target.setState(c.target.convState, c.fullConv, true) + c.target.deltaState = c.target.setState(c.target.deltaState, c.fullDelta, false) + c.target.offset = c.start + n + return + } + + if c.initialConv == nil || c.initialDelta == nil || c.qkv == nil || c.q == nil || c.k == nil || c.v == nil || c.gDecay == nil || c.beta == nil { + return + } + + qkv := sliceSeq(c.qkv, n) + convConcat := mlx.Concatenate([]*mlx.Array{c.initialConv, qkv}, 1) + total := convConcat.Dim(1) + nextConv := convConcat.Slice(mlx.Slice(), mlx.Slice(total-c.target.convTail, total), mlx.Slice()) + + _, delta := mlx.FastGatedDelta( + sliceSeq(c.q, n), + sliceSeq(c.k, n), + sliceSeq(c.v, n), + sliceSeq(c.gDecay, n), + sliceSeq(c.beta, n), + c.initialDelta, + nil, + ) + + c.target.convState = c.target.setState(c.target.convState, nextConv, true) + c.target.deltaState = c.target.setState(c.target.deltaState, delta, false) + c.target.offset = c.start + n +} + +func sliceSeq(a *mlx.Array, n int) *mlx.Array { + switch a.NumDims() { + case 3: + return a.Slice(mlx.Slice(), mlx.Slice(0, n), mlx.Slice()) + case 4: + return a.Slice(mlx.Slice(), mlx.Slice(0, n), mlx.Slice(), mlx.Slice()) + default: + panic("recurrent speculative sequence tensor must be rank 3 or 4") + } +} diff --git a/x/mlxrunner/cache/recurrent_test.go b/x/mlxrunner/cache/recurrent_test.go index 31e58dae..9af93c88 100644 --- a/x/mlxrunner/cache/recurrent_test.go +++ b/x/mlxrunner/cache/recurrent_test.go @@ -63,8 +63,55 @@ func TestRecurrentCacheGetLazyInit(t *testing.T) { if got := h.ConvState().DType(); got != mlx.DTypeBFloat16 { t.Fatalf("conv state dtype = %v, want %v", got, mlx.DTypeBFloat16) } - if got := h.DeltaState().DType(); got != mlx.DTypeBFloat16 { - t.Fatalf("delta state dtype = %v, want %v", got, mlx.DTypeBFloat16) + if got := h.DeltaState().DType(); got != mlx.DTypeFloat32 { + t.Fatalf("delta state dtype = %v, want %v", got, mlx.DTypeFloat32) + } +} + +func TestSpeculativeRecurrentCacheUsesStagedState(t *testing.T) { + skipIfNoMLX(t) + target := NewRecurrentCache(2, 3, 1, 2, 3) + caches, ok := BeginIsolatedSpeculation([]Cache{target}) + if !ok { + t.Fatal("BeginIsolatedSpeculation failed") + } + c := caches[0].(*speculativeRecurrentCache) + b := &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, 1, 1), + SeqOffsets: []int32{0}, + SeqQueryLens: []int32{1}, + } + + c.Get(b, mlx.DTypeFloat32) + + convVals := []float32{1, 2, 3, 4, 5, 6} + deltaVals := []float32{7, 8, 9, 10, 11, 12} + nextConv := mlx.FromValues(convVals, 1, 2, 3) + nextDelta := mlx.FromValues(deltaVals, 1, 1, 2, 3) + c.Put(b, nextConv, nextDelta) + + h := c.Get(b, mlx.DTypeFloat32) + state := c.State() + if len(state) != 2 { + t.Fatalf("State() returned %d arrays, want 2", len(state)) + } + + assertArray := func(name string, got, want *mlx.Array) { + t.Helper() + if got != want { + t.Fatalf("%s = %p, want %p", name, got, want) + } + } + assertArray("history conv", h.ConvState(), nextConv) + assertArray("history delta", h.DeltaState(), nextDelta) + assertArray("state conv", state[0], nextConv) + assertArray("state delta", state[1], nextDelta) + + if got := c.Offset(); got != 1 { + t.Fatalf("speculative offset = %d, want 1", got) + } + if got := target.Offset(); got != 0 { + t.Fatalf("target offset = %d, want 0", got) } } diff --git a/x/mlxrunner/cache_test.go b/x/mlxrunner/cache_test.go index fba0d4fb..612d2af1 100644 --- a/x/mlxrunner/cache_test.go +++ b/x/mlxrunner/cache_test.go @@ -159,6 +159,50 @@ func (c *fakeRewindableCache) Split(snapshot cache.Snapshot, at int) (cache.Snap return p, ch } +func TestKVCacheBeginWithFactoryLimitCapsPrefix(t *testing.T) { + inputs := []int32{1, 2, 3, 4, 5} + tracker := &snapshotTracker{} + var kc kvCache + + factory := func() []cache.Cache { + return []cache.Cache{&fakeRewindableCache{tracker: tracker}} + } + + session := kc.beginWithFactoryLimit(inputs, factory, "test", -1, false) + session.caches[0].(*fakeRewindableCache).feed(inputs) + session.close() + + session = kc.beginWithFactoryLimit(inputs, factory, "test", 3, false) + if got, want := session.caches[0].Offset(), 3; got != want { + t.Fatalf("cache offset = %d, want %d", got, want) + } + if got, want := session.remaining, inputs[3:]; !slices.Equal(got, want) { + t.Fatalf("remaining = %v, want %v", got, want) + } +} + +func TestKVCacheBeginWithFactoryLimitDoesNotKeepSeedToken(t *testing.T) { + inputs := []int32{1, 2, 3, 4, 5} + tracker := &snapshotTracker{} + var kc kvCache + + factory := func() []cache.Cache { + return []cache.Cache{&fakeRewindableCache{tracker: tracker}} + } + + session := kc.beginWithFactoryLimit(inputs, factory, "test", -1, false) + session.caches[0].(*fakeRewindableCache).feed(inputs) + session.close() + + session = kc.beginWithFactoryLimit(inputs, factory, "test", len(inputs), false) + if got, want := session.caches[0].Offset(), len(inputs); got != want { + t.Fatalf("cache offset = %d, want %d", got, want) + } + if len(session.remaining) != 0 { + t.Fatalf("remaining = %v, want empty", session.remaining) + } +} + // fakeSlidingWindowCache models RotatingKVCache semantics: stores the full // token sequence but only the trailing maxSize tokens are "live" in the window. // Once the window fills, live rewind is impossible without a snapshot. diff --git a/x/mlxrunner/dflash.go b/x/mlxrunner/dflash.go new file mode 100644 index 00000000..7f6cfd86 --- /dev/null +++ b/x/mlxrunner/dflash.go @@ -0,0 +1,715 @@ +package mlxrunner + +import ( + "context" + "fmt" + "log/slog" + "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" + sampler "github.com/ollama/ollama/x/mlxrunner/sample" +) + +type dflashStats struct { + iterations int + drafted int + accepted int + mismatches int + allAccepted int + batched int + serial int + targetDuration time.Duration + draftDuration time.Duration + validateDuration time.Duration +} + +type dflashDecodeMode string + +const ( + dflashDecodeDisabled dflashDecodeMode = "" + dflashDecodeGreedy dflashDecodeMode = "greedy" + dflashDecodeSample dflashDecodeMode = "sample" +) + +func (m dflashDecodeMode) enabled() bool { + return m != dflashDecodeDisabled +} + +func newDFlashTargetCaches(m base.Model) []cache.Cache { + if cacheFactory, ok := m.(interface{ NewCaches() []cache.Cache }); ok { + return cacheFactory.NewCaches() + } + caches := make([]cache.Cache, m.NumLayers()) + for i := range caches { + caches[i] = cache.NewKVCache() + } + return caches +} + +func freeCacheSet(caches []cache.Cache) { + for _, c := range caches { + if c != nil { + c.Free() + } + } +} + +func (r *Runner) dflashGate(opts sampler.Options) (dflashDecodeMode, string) { + if r.Draft == nil { + return dflashDecodeDisabled, "no_draft" + } + if _, ok := r.Draft.(base.DFlashDraftModel); !ok { + return dflashDecodeDisabled, "draft_not_dflash" + } + if _, ok := r.Model.(base.DFlashTargetModel); !ok { + return dflashDecodeDisabled, "target_not_dflash" + } + if _, ok := r.Model.(base.MTPEmbeddingModel); !ok { + return dflashDecodeDisabled, "target_embeddings_missing" + } + if opts.Logprobs || opts.TopLogprobs > 0 { + return dflashDecodeDisabled, "logprobs_requested" + } + + if opts.Temperature > 0 || dflashUsesSamplerHistory(opts) { + return dflashDecodeSample, "" + } + + return dflashDecodeGreedy, "" +} + +func dflashUsesSamplerHistory(opts sampler.Options) bool { + if opts.RepeatLastN == 0 { + return false + } + + repeatPenalty := opts.RepeatPenalty + if repeatPenalty <= 0 { + repeatPenalty = 1 + } + return repeatPenalty != 1 || opts.PresencePenalty != 0 || opts.FrequencyPenalty != 0 +} + +func (r *Runner) runGreedyDFlashDecode(ctx context.Context, request Request, session *cacheSession, targetCaches []cache.Cache, draftCaches []cache.Cache, seed []int32, position *int, started time.Time) error { + target := r.Model.(base.DFlashTargetModel) + draft := r.Draft.(base.DFlashDraftModel) + stats := dflashStats{} + slog.Info("DFlash greedy decode enabled", "block_size", draft.BlockSize(), "target_layers", draft.TargetLayerIDs()) + + targetForward := func(token *mlx.Array) (*mlx.Array, *mlx.Array) { + hidden, targetHidden := target.ForwardDFlash(&batch.Batch{ + InputIDs: token, + SeqOffsets: []int32{int32(*position)}, + SeqQueryLens: []int32{int32(token.Dim(1))}, + }, targetCaches, draft.TargetLayerIDs()) + *position += token.Dim(1) + return hidden, targetHidden + } + + t0 := time.Now() + hidden, targetHidden := targetForward(mlx.FromValues(seed, 1, len(seed))) + draft.AppendContext(targetHidden, draftCaches) + current := sampler.Result{Token: greedyTokenFromLogits(r.lastLogits(hidden))} + mlx.Pin(current.Arrays()...) + mlx.Sweep() + mlx.AsyncEval(current.Arrays()...) + stats.targetDuration += time.Since(t0) + 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 + } + + if generated == 0 { + mlx.Eval(current.Arrays()...) + final.PromptEvalDuration = time.Since(now) + now = time.Now() + } + + done, err := r.emitMTPToken(ctx, request, session, &dec, current, &final) + if err != nil { + return err + } + if !done { + generated++ + } + if done || generated >= request.Options.NumPredict { + break + } + + draftCount := min(draft.BlockSize()-1, request.Options.NumPredict-generated) + if draftCount <= 0 { + t0 = time.Now() + hidden, targetHidden := targetForward(mtpTokenInput(current.Token)) + draft.AppendContext(targetHidden, draftCaches) + stats.targetDuration += time.Since(t0) + next := sampler.Result{Token: greedyTokenFromLogits(r.lastLogits(hidden))} + mlx.Pin(next.Arrays()...) + old := current + current = next + mlx.Unpin(old.Arrays()...) + mlx.Sweep() + mlx.AsyncEval(current.Arrays()...) + continue + } + + stats.iterations++ + t0 = time.Now() + draftTokens := r.generateDFlashDrafts(draft, current.Token, draftCaches, draftCount) + mlx.Pin(draftTokens) + mlx.Eval(draftTokens) + stats.draftDuration += time.Since(t0) + stats.drafted += draftCount + + t0 = time.Now() + next, accepted, done, err := r.acceptDFlashDrafts(ctx, request, session, &dec, target, draft, targetCaches, draftCaches, position, current, draftTokens, &final, &generated, &stats) + stats.validateDuration += time.Since(t0) + mlx.Unpin(draftTokens) + if err != nil { + return err + } + stats.accepted += accepted + if accepted == draftCount { + stats.allAccepted++ + } else { + stats.mismatches++ + } + 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("DFlash decode stats", "mode", "greedy", "generated", generated, "drafted", stats.drafted, "accepted", stats.accepted, "acceptance", acceptance, "iterations", stats.iterations, "avg_draft", avgDraft, "avg_accepted", avgAccepted, "batched", stats.batched, "serial", stats.serial, "mismatches", stats.mismatches, "all_accepted", stats.allAccepted, "max_draft", draft.BlockSize()-1, "block_size", draft.BlockSize(), "target_layers", draft.TargetLayerIDs(), "target_duration", stats.targetDuration, "draft_duration", stats.draftDuration, "validate_duration", stats.validateDuration) + + select { + case <-ctx.Done(): + return ctx.Err() + case request.Responses <- final: + return nil + } +} + +func (r *Runner) runSampleDFlashDecode(ctx context.Context, request Request, session *cacheSession, targetCaches []cache.Cache, draftCaches []cache.Cache, seed []int32, position *int, started time.Time) error { + target := r.Model.(base.DFlashTargetModel) + draft := r.Draft.(base.DFlashDraftModel) + stats := dflashStats{} + slog.Info("DFlash sample decode enabled", + "block_size", draft.BlockSize(), + "target_layers", draft.TargetLayerIDs(), + "temperature", request.SamplerOpts.Temperature, + "top_p", request.SamplerOpts.TopP, + "top_k", request.SamplerOpts.TopK, + "min_p", request.SamplerOpts.MinP, + "repeat_penalty", request.SamplerOpts.RepeatPenalty, + "presence_penalty", request.SamplerOpts.PresencePenalty, + "frequency_penalty", request.SamplerOpts.FrequencyPenalty, + ) + + targetForward := func(token *mlx.Array) (*mlx.Array, *mlx.Array) { + hidden, targetHidden := target.ForwardDFlash(&batch.Batch{ + InputIDs: token, + SeqOffsets: []int32{int32(*position)}, + SeqQueryLens: []int32{int32(token.Dim(1))}, + }, targetCaches, draft.TargetLayerIDs()) + *position += token.Dim(1) + return hidden, targetHidden + } + + t0 := time.Now() + hidden, targetHidden := targetForward(mlx.FromValues(seed, 1, len(seed))) + draft.AppendContext(targetHidden, draftCaches) + current := r.Sampler.Sample([]int{pipelineSlot}, r.lastLogits(hidden)) + mlx.Pin(current.Arrays()...) + mlx.Sweep() + mlx.AsyncEval(current.Arrays()...) + stats.targetDuration += time.Since(t0) + 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 + } + + if generated == 0 { + mlx.Eval(current.Arrays()...) + final.PromptEvalDuration = time.Since(now) + now = time.Now() + } + + done, err := r.emitMTPToken(ctx, request, session, &dec, current, &final) + if err != nil { + return err + } + if !done { + generated++ + } + if done || generated >= request.Options.NumPredict { + break + } + + draftCount := min(draft.BlockSize()-1, request.Options.NumPredict-generated) + if draftCount <= 0 { + t0 = time.Now() + hidden, targetHidden := targetForward(mtpTokenInput(current.Token)) + draft.AppendContext(targetHidden, draftCaches) + stats.targetDuration += time.Since(t0) + next := r.Sampler.Sample([]int{pipelineSlot}, r.lastLogits(hidden)) + mlx.Pin(next.Arrays()...) + old := current + current = next + mlx.Unpin(old.Arrays()...) + mlx.Sweep() + mlx.AsyncEval(current.Arrays()...) + continue + } + + stats.iterations++ + t0 = time.Now() + candidates := r.generateDFlashDraftCandidates(draft, current.Token, draftCaches, draftCount) + var candidateArrays []*mlx.Array + if candidates != nil { + draftCount = candidates.tokens.Dim(1) + candidateArrays = candidates.Arrays() + mlx.Pin(candidateArrays...) + mlx.Sweep() + } + stats.draftDuration += time.Since(t0) + stats.drafted += draftCount + + var next sampler.Result + if draftCount == 0 { + t0 = time.Now() + hidden, targetHidden := targetForward(mtpTokenInput(current.Token)) + draft.AppendContext(targetHidden, draftCaches) + stats.targetDuration += time.Since(t0) + next = r.Sampler.Sample([]int{pipelineSlot}, r.lastLogits(hidden)) + } else { + var accepted int + t0 = time.Now() + next, accepted, done, err = r.acceptSampleDFlashDrafts(ctx, request, session, &dec, target, draft, targetCaches, draftCaches, position, current, candidates, &final, &generated, &stats) + stats.validateDuration += time.Since(t0) + mlx.Unpin(candidateArrays...) + if err != nil { + return err + } + stats.accepted += accepted + if accepted == draftCount { + stats.allAccepted++ + } else { + stats.mismatches++ + } + 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("DFlash decode stats", "mode", "sample", "generated", generated, "drafted", stats.drafted, "accepted", stats.accepted, "acceptance", acceptance, "iterations", stats.iterations, "avg_draft", avgDraft, "avg_accepted", avgAccepted, "batched", stats.batched, "serial", stats.serial, "mismatches", stats.mismatches, "all_accepted", stats.allAccepted, "max_draft", draft.BlockSize()-1, "block_size", draft.BlockSize(), "target_layers", draft.TargetLayerIDs(), "target_duration", stats.targetDuration, "draft_duration", stats.draftDuration, "validate_duration", stats.validateDuration) + + select { + case <-ctx.Done(): + return ctx.Err() + case request.Responses <- final: + return nil + } +} + +func (r *Runner) dflashDraftLogits(draft base.DFlashDraftModel, current *mlx.Array, caches []cache.Cache, draftCount int) *mlx.Array { + blockLen := draftCount + 1 + values := make([]int32, blockLen) + values[0] = int32(tokenID(current)) + for i := 1; i < blockLen; i++ { + values[i] = draft.MaskTokenID() + } + block := mlx.FromValues(values, 1, blockLen) + logits := draft.Draft(block, caches) + return logits.Slice(mlx.Slice(), mlx.Slice(1, blockLen), mlx.Slice()) +} + +func (r *Runner) generateDFlashDrafts(draft base.DFlashDraftModel, current *mlx.Array, caches []cache.Cache, draftCount int) *mlx.Array { + logits := r.dflashDraftLogits(draft, current, caches, draftCount) + return logits.Argmax(-1, false).AsType(mlx.DTypeInt32) +} + +type dflashDraftCandidates struct { + tokens *mlx.Array + dist sampler.Distribution +} + +func (c *dflashDraftCandidates) Arrays() []*mlx.Array { + if c == nil { + return nil + } + return append([]*mlx.Array{c.tokens}, c.dist.Arrays()...) +} + +func (r *Runner) generateDFlashDraftCandidates(draft base.DFlashDraftModel, current *mlx.Array, caches []cache.Cache, draftCount int) *dflashDraftCandidates { + if draftCount <= 0 { + return nil + } + + logits := r.dflashDraftLogits(draft, current, caches, draftCount) + draftTokens := make([]*mlx.Array, 0, draftCount) + draftDists := make([]sampler.Distribution, 0, draftCount) + var prefix *mlx.Array + + for i := range draftCount { + rows := logits.Slice(mlx.Slice(), mlx.Slice(0, i+1), mlx.Slice()) + dist := r.Sampler.Distribution(pipelineSlot, rows, prefix).SliceRows(i, i+1) + nextToken := mtpTokenVector(r.Sampler.SampleDistribution(pipelineSlot, dist)) + nextInput := mtpTokenInput(nextToken) + + draftTokens = append(draftTokens, nextInput) + draftDists = append(draftDists, dist) + if prefix == nil { + prefix = nextInput + } else { + prefix = prefix.Concatenate(1, nextInput) + } + } + if len(draftTokens) == 0 { + return nil + } + return &dflashDraftCandidates{ + tokens: mlx.Concatenate(draftTokens, 1), + dist: sampler.ConcatenateDistributions(draftDists), + } +} + +func (r *Runner) acceptDFlashDrafts(ctx context.Context, request Request, session *cacheSession, dec *decoder, target base.DFlashTargetModel, draft base.DFlashDraftModel, targetCaches []cache.Cache, draftCaches []cache.Cache, position *int, current sampler.Result, draftTokens *mlx.Array, final *CompletionResponse, generated *int, stats *dflashStats) (sampler.Result, int, bool, error) { + specCaches, spec, ok := cache.BeginSpeculation(targetCaches) + if !ok { + stats.serial++ + return r.acceptDFlashDraftsSerial(ctx, request, session, dec, target, draft, targetCaches, draftCaches, position, current, draftTokens, final, generated) + } + stats.batched++ + return r.acceptDFlashDraftsBatched(ctx, request, session, dec, target, draft, specCaches, spec, draftCaches, position, current, draftTokens, final, generated) +} + +func (r *Runner) acceptDFlashDraftsBatched(ctx context.Context, request Request, session *cacheSession, dec *decoder, target base.DFlashTargetModel, draft base.DFlashDraftModel, specCaches []cache.Cache, spec *cache.Speculation, draftCaches []cache.Cache, position *int, current sampler.Result, draftTokens *mlx.Array, final *CompletionResponse, generated *int) (sampler.Result, int, bool, error) { + before := *position + draftCount := draftTokens.Dim(1) + verifyInput := mtpTokenInput(current.Token).Concatenate(1, draftTokens) + hiddenSeq, targetHiddenSeq := target.ForwardDFlash(&batch.Batch{ + InputIDs: verifyInput, + SeqOffsets: []int32{int32(before)}, + SeqQueryLens: []int32{int32(verifyInput.Dim(1))}, + }, specCaches, draft.TargetLayerIDs()) + + selectedTokens := r.Model.Unembed(hiddenSeq).Argmax(-1, false).AsType(mlx.DTypeInt32) + mlx.Eval(draftTokens, selectedTokens) + + draftIDs := draftTokens.Ints() + selectedIDs := selectedTokens.Ints() + if len(selectedIDs) < draftCount+1 { + spec.Commit(0) + return sampler.Result{}, 0, false, fmt.Errorf("dflash validation produced %d tokens for %d draft tokens", len(selectedIDs), draftCount) + } + + accepted := 0 + for i, id := range draftIDs { + if selectedIDs[i] != id { + break + } + accepted++ + if r.Tokenizer.IsEOS(int32(id)) { + break + } + } + + commitN := accepted + 1 + spec.Commit(0) + + done := false + for _, id := range draftIDs[:accepted] { + if *generated >= request.Options.NumPredict { + done = true + break + } + res := sampler.Result{Token: mlx.FromValues([]int32{int32(id)}, 1)} + var err error + done, err = r.emitMTPToken(ctx, request, session, dec, res, final) + if err != nil { + return sampler.Result{}, accepted, done, err + } + if !done { + (*generated)++ + } + if done { + break + } + } + + spec.Commit(commitN) + *position = before + commitN + draft.AppendContext(targetHiddenSeq.Slice(mlx.Slice(), mlx.Slice(0, commitN), mlx.Slice()), draftCaches) + + if done || *generated >= request.Options.NumPredict { + return sampler.Result{}, accepted, true, nil + } + + nextIndex := accepted + if nextIndex >= len(selectedIDs) { + nextIndex = len(selectedIDs) - 1 + } + return sampler.Result{Token: mlx.FromValues([]int32{int32(selectedIDs[nextIndex])}, 1)}, accepted, false, nil +} + +func (r *Runner) acceptDFlashDraftsSerial(ctx context.Context, request Request, session *cacheSession, dec *decoder, target base.DFlashTargetModel, draft base.DFlashDraftModel, targetCaches []cache.Cache, draftCaches []cache.Cache, position *int, current sampler.Result, draftTokens *mlx.Array, final *CompletionResponse, generated *int) (sampler.Result, int, bool, error) { + targetForward := func(token *mlx.Array) *mlx.Array { + hidden, targetHidden := target.ForwardDFlash(&batch.Batch{ + InputIDs: token, + SeqOffsets: []int32{int32(*position)}, + SeqQueryLens: []int32{int32(token.Dim(1))}, + }, targetCaches, draft.TargetLayerIDs()) + *position += token.Dim(1) + draft.AppendContext(targetHidden, draftCaches) + return r.lastLogits(hidden) + } + + logits := targetForward(mtpTokenInput(current.Token)) + accepted := 0 + for _, id := range draftTokens.Ints() { + selected := greedyTokenFromLogits(logits) + mlx.Eval(selected) + selectedID := tokenID(selected) + if selectedID != id { + return sampler.Result{Token: mlx.FromValues([]int32{int32(selectedID)}, 1)}, accepted, false, nil + } + + res := sampler.Result{Token: mlx.FromValues([]int32{int32(id)}, 1)} + done, err := r.emitMTPToken(ctx, request, session, dec, res, final) + if err != nil { + return sampler.Result{}, accepted, done, err + } + accepted++ + if !done { + (*generated)++ + } + if done || *generated >= request.Options.NumPredict { + return sampler.Result{}, accepted, true, nil + } + + logits = targetForward(mtpTokenInput(res.Token)) + } + + return sampler.Result{Token: greedyTokenFromLogits(logits)}, accepted, false, nil +} + +func (r *Runner) acceptSampleDFlashDrafts(ctx context.Context, request Request, session *cacheSession, dec *decoder, target base.DFlashTargetModel, draft base.DFlashDraftModel, targetCaches []cache.Cache, draftCaches []cache.Cache, position *int, current sampler.Result, candidates *dflashDraftCandidates, final *CompletionResponse, generated *int, stats *dflashStats) (sampler.Result, int, bool, error) { + specCaches, spec, ok := cache.BeginSpeculation(targetCaches) + if !ok { + stats.serial++ + return r.acceptSampleDFlashDraftsSerial(ctx, request, session, dec, target, draft, targetCaches, draftCaches, position, current, candidates, final, generated) + } + stats.batched++ + return r.acceptSampleDFlashDraftsBatched(ctx, request, session, dec, target, draft, specCaches, spec, draftCaches, position, current, candidates, final, generated) +} + +func (r *Runner) acceptSampleDFlashDraftsBatched(ctx context.Context, request Request, session *cacheSession, dec *decoder, target base.DFlashTargetModel, draft base.DFlashDraftModel, specCaches []cache.Cache, spec *cache.Speculation, draftCaches []cache.Cache, position *int, current sampler.Result, candidates *dflashDraftCandidates, final *CompletionResponse, generated *int) (sampler.Result, int, bool, error) { + before := *position + draftCount := candidates.tokens.Dim(1) + verifyInput := mtpTokenInput(current.Token).Concatenate(1, candidates.tokens) + hiddenSeq, targetHiddenSeq := target.ForwardDFlash(&batch.Batch{ + InputIDs: verifyInput, + SeqOffsets: []int32{int32(before)}, + SeqQueryLens: []int32{int32(verifyInput.Dim(1))}, + }, specCaches, draft.TargetLayerIDs()) + + targetDist := r.Sampler.Distribution(pipelineSlot, r.Model.Unembed(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 { + spec.Commit(0) + return sampler.Result{}, 0, false, fmt.Errorf("dflash 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 + } + } + + commitN := accepted + 1 + spec.Commit(0) + + for _, id := range draftIDs[:accepted] { + if *generated >= request.Options.NumPredict { + done = true + break + } + res := sampler.Result{Token: mlx.FromValues([]int32{int32(id)}, 1)} + var err error + done, err = r.emitMTPToken(ctx, request, session, dec, res, final) + if err != nil { + return sampler.Result{}, accepted, done, err + } + if !done { + (*generated)++ + } + if done { + break + } + } + + spec.Commit(commitN) + *position = before + commitN + draft.AppendContext(targetHiddenSeq.Slice(mlx.Slice(), mlx.Slice(0, commitN), mlx.Slice()), draftCaches) + + if done || *generated >= request.Options.NumPredict { + 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) acceptSampleDFlashDraftsSerial(ctx context.Context, request Request, session *cacheSession, dec *decoder, target base.DFlashTargetModel, draft base.DFlashDraftModel, targetCaches []cache.Cache, draftCaches []cache.Cache, position *int, current sampler.Result, candidates *dflashDraftCandidates, final *CompletionResponse, generated *int) (sampler.Result, int, bool, error) { + targetForward := func(token *mlx.Array) *mlx.Array { + hidden, targetHidden := target.ForwardDFlash(&batch.Batch{ + InputIDs: token, + SeqOffsets: []int32{int32(*position)}, + SeqQueryLens: []int32{int32(token.Dim(1))}, + }, targetCaches, draft.TargetLayerIDs()) + *position += token.Dim(1) + draft.AppendContext(targetHidden, draftCaches) + return r.lastLogits(hidden) + } + + mlx.Eval(candidates.tokens) + draftIDs := candidates.tokens.Ints() + logits := targetForward(mtpTokenInput(current.Token)) + accepted := 0 + + for i, id := range draftIDs { + targetDist := r.Sampler.Distribution(pipelineSlot, logits, nil) + draftDist := candidates.dist.SliceRows(i, i+1) + draftToken := mlx.FromValues([]int32{int32(id)}, 1) + acceptedMask := r.mtpSampleAcceptedMask(targetDist, draftDist, draftToken) + mlx.Eval(acceptedMask) + + if acceptedMask.Ints()[0] == 0 { + nextToken := mtpTokenVector(r.Sampler.SampleDistribution(pipelineSlot, targetDist.ResidualAgainst(draftDist))) + mlx.Eval(nextToken) + r.Sampler.Commit(pipelineSlot, []int32{int32(tokenID(nextToken))}) + return sampler.Result{Token: nextToken}, accepted, false, nil + } + + accepted++ + r.Sampler.Commit(pipelineSlot, []int32{int32(id)}) + res := sampler.Result{Token: mlx.FromValues([]int32{int32(id)}, 1)} + done, err := r.emitMTPToken(ctx, request, session, dec, res, final) + if err != nil { + return sampler.Result{}, accepted, done, err + } + if !done { + (*generated)++ + } + if done || *generated >= request.Options.NumPredict { + return sampler.Result{}, accepted, true, nil + } + + logits = targetForward(mtpTokenInput(res.Token)) + } + + targetDist := r.Sampler.Distribution(pipelineSlot, logits, nil) + nextToken := mtpTokenVector(r.Sampler.SampleDistribution(pipelineSlot, targetDist)) + mlx.Eval(nextToken) + r.Sampler.Commit(pipelineSlot, []int32{int32(tokenID(nextToken))}) + return sampler.Result{Token: nextToken}, accepted, false, nil +} diff --git a/x/mlxrunner/imports.go b/x/mlxrunner/imports.go index 316691d9..a0eede39 100644 --- a/x/mlxrunner/imports.go +++ b/x/mlxrunner/imports.go @@ -1,6 +1,7 @@ package mlxrunner import ( + _ "github.com/ollama/ollama/x/models/dflash" _ "github.com/ollama/ollama/x/models/gemma3" _ "github.com/ollama/ollama/x/models/gemma4" _ "github.com/ollama/ollama/x/models/glm4_moe_lite" diff --git a/x/mlxrunner/mlx/gated_delta.go b/x/mlxrunner/mlx/gated_delta.go index 7034c33a..729af01c 100644 --- a/x/mlxrunner/mlx/gated_delta.go +++ b/x/mlxrunner/mlx/gated_delta.go @@ -83,7 +83,7 @@ for (int t = 0; t < T; ++t) { for (int i = 0; i < n_per_t; ++i) { auto s_idx = n_per_t * dk_idx + i; - o_state[s_idx] = static_cast(state[i]); + o_state[s_idx] = static_cast(state[i]); } ` @@ -163,7 +163,7 @@ for (int t = 0; t < T_val; ++t) { for (int i = 0; i < n_per_t; ++i) { auto s_idx = n_per_t * dk_idx + i; - o_state[s_idx] = static_cast(state[i]); + o_state[s_idx] = static_cast(state[i]); } ` @@ -262,8 +262,9 @@ func gatedDeltaKernel(q, k, v, g, beta, state *Array) (y, nextState *Array, ok b return nil, nil, false } - dtype := q.DType() - if k.DType() != dtype || v.DType() != dtype || g.DType() != dtype || beta.DType() != dtype || state.DType() != dtype { + inputDType := q.DType() + stateDType := state.DType() + if k.DType() != inputDType || v.DType() != inputDType || g.DType() != inputDType || beta.DType() != inputDType { return nil, nil, false } @@ -277,7 +278,13 @@ func gatedDeltaKernel(q, k, v, g, beta, state *Array) (y, nextState *Array, ok b cInT := C.CString("InT") defer C.free(unsafe.Pointer(cInT)) - if C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, cInT, C.mlx_dtype(dtype)) != 0 { + if C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, cInT, C.mlx_dtype(inputDType)) != 0 { + gatedDeltaMetalDisabled = true + return nil, nil, false + } + cStT := C.CString("StT") + defer C.free(unsafe.Pointer(cStT)) + if C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, cStT, C.mlx_dtype(stateDType)) != 0 { gatedDeltaMetalDisabled = true return nil, nil, false } @@ -301,11 +308,11 @@ func gatedDeltaKernel(q, k, v, g, beta, state *Array) (y, nextState *Array, ok b yShape := []C.int{C.int(B), C.int(T), C.int(Hv), C.int(Dv)} stateShape := []C.int{C.int(B), C.int(Hv), C.int(Dv), C.int(Dk)} - if C.mlx_fast_metal_kernel_config_add_output_arg(cfg, unsafe.SliceData(yShape), C.size_t(len(yShape)), C.mlx_dtype(dtype)) != 0 { + if C.mlx_fast_metal_kernel_config_add_output_arg(cfg, unsafe.SliceData(yShape), C.size_t(len(yShape)), C.mlx_dtype(inputDType)) != 0 { gatedDeltaMetalDisabled = true return nil, nil, false } - if C.mlx_fast_metal_kernel_config_add_output_arg(cfg, unsafe.SliceData(stateShape), C.size_t(len(stateShape)), C.mlx_dtype(dtype)) != 0 { + if C.mlx_fast_metal_kernel_config_add_output_arg(cfg, unsafe.SliceData(stateShape), C.size_t(len(stateShape)), C.mlx_dtype(stateDType)) != 0 { gatedDeltaMetalDisabled = true return nil, nil, false } @@ -516,8 +523,9 @@ func gatedDeltaCUDAKernelApply(q, k, v, g, beta, state *Array) (y, nextState *Ar return nil, nil, false } - dtype := q.DType() - if k.DType() != dtype || v.DType() != dtype || g.DType() != dtype || beta.DType() != dtype || state.DType() != dtype { + inputDType := q.DType() + stateDType := state.DType() + if k.DType() != inputDType || v.DType() != inputDType || g.DType() != inputDType || beta.DType() != inputDType { return nil, nil, false } @@ -531,7 +539,13 @@ func gatedDeltaCUDAKernelApply(q, k, v, g, beta, state *Array) (y, nextState *Ar cInT := C.CString("InT") defer C.free(unsafe.Pointer(cInT)) - if C.mlx_fast_cuda_kernel_config_add_template_arg_dtype(cfg, cInT, C.mlx_dtype(dtype)) != 0 { + if C.mlx_fast_cuda_kernel_config_add_template_arg_dtype(cfg, cInT, C.mlx_dtype(inputDType)) != 0 { + gatedDeltaCUDADisabled = true + return nil, nil, false + } + cStT := C.CString("StT") + defer C.free(unsafe.Pointer(cStT)) + if C.mlx_fast_cuda_kernel_config_add_template_arg_dtype(cfg, cStT, C.mlx_dtype(stateDType)) != 0 { gatedDeltaCUDADisabled = true return nil, nil, false } @@ -555,11 +569,11 @@ func gatedDeltaCUDAKernelApply(q, k, v, g, beta, state *Array) (y, nextState *Ar yShape := []C.int{C.int(B), C.int(T), C.int(Hv), C.int(Dv)} stateShape := []C.int{C.int(B), C.int(Hv), C.int(Dv), C.int(Dk)} - if C.mlx_fast_cuda_kernel_config_add_output_arg(cfg, unsafe.SliceData(yShape), C.size_t(len(yShape)), C.mlx_dtype(dtype)) != 0 { + if C.mlx_fast_cuda_kernel_config_add_output_arg(cfg, unsafe.SliceData(yShape), C.size_t(len(yShape)), C.mlx_dtype(inputDType)) != 0 { gatedDeltaCUDADisabled = true return nil, nil, false } - if C.mlx_fast_cuda_kernel_config_add_output_arg(cfg, unsafe.SliceData(stateShape), C.size_t(len(stateShape)), C.mlx_dtype(dtype)) != 0 { + if C.mlx_fast_cuda_kernel_config_add_output_arg(cfg, unsafe.SliceData(stateShape), C.size_t(len(stateShape)), C.mlx_dtype(stateDType)) != 0 { gatedDeltaCUDADisabled = true return nil, nil, false } diff --git a/x/mlxrunner/model/base/base.go b/x/mlxrunner/model/base/base.go index 4f43d74b..a94a7c31 100644 --- a/x/mlxrunner/model/base/base.go +++ b/x/mlxrunner/model/base/base.go @@ -57,6 +57,23 @@ type MTPEmbeddingModel interface { TokenEmbeddings(inputIDs *mlx.Array) *mlx.Array } +// DFlashTargetModel exposes target-layer hidden states for DFlash drafts. +type DFlashTargetModel interface { + ForwardDFlash(b *batch.Batch, caches []cache.Cache, layerIDs []int) (hidden, targetHidden *mlx.Array) +} + +// DFlashDraftModel is a block-diffusion speculative draft model. +type DFlashDraftModel interface { + DraftModel + + TargetLayerIDs() []int + BlockSize() int + MaskTokenID() int32 + NewCaches() []cache.Cache + AppendContext(targetHidden *mlx.Array, caches []cache.Cache) + Draft(inputIDs *mlx.Array, caches []cache.Cache) *mlx.Array +} + var ( mu sync.Mutex registry = make(map[string]func(root *model.Root) (Model, error)) diff --git a/x/mlxrunner/pipeline.go b/x/mlxrunner/pipeline.go index 777653c0..7da3384d 100644 --- a/x/mlxrunner/pipeline.go +++ b/x/mlxrunner/pipeline.go @@ -12,7 +12,9 @@ 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" + "github.com/ollama/ollama/x/mlxrunner/model/base" sampler "github.com/ollama/ollama/x/mlxrunner/sample" "github.com/ollama/ollama/x/tokenizer" ) @@ -79,24 +81,90 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er caches := session.caches tokens := session.remaining prefillChunk := prefillChunkSize() - - // Request periodic snapshots during prefill and near the end of the - // prompt so that long prompts can be partially restored and - // thinking/generation can be retried without full reprocessing. - const snapshotInterval = 8192 - for offset := snapshotInterval; offset < len(inputs); offset += snapshotInterval { - session.requestSnapshot(offset) + dflashMode, dflashDisabledReason := r.dflashGate(request.SamplerOpts) + dflashEnabled := dflashMode.enabled() + var dflashDraft base.DFlashDraftModel + var dflashTarget base.DFlashTargetModel + var dflashCaches []cache.Cache + var dflashSession *cacheSession + if dflashEnabled { + dflashDraft = r.Draft.(base.DFlashDraftModel) + dflashTarget = r.Model.(base.DFlashTargetModel) + targetCachedPrefix := len(inputs) - len(tokens) + dflashSession = r.dflashCache.beginWithFactoryLimit(inputs, dflashDraft.NewCaches, "DFlash draft", targetCachedPrefix, false) + dflashCaches = dflashSession.caches + defer func() { + dflashSession.outputs = append([]int32(nil), session.outputs...) + dflashSession.close() + }() + } else if _, ok := r.Draft.(base.DFlashDraftModel); ok { + slog.Info("DFlash decode disabled", + "reason", dflashDisabledReason, + "temperature", request.SamplerOpts.Temperature, + "top_p", request.SamplerOpts.TopP, + "top_k", request.SamplerOpts.TopK, + "min_p", request.SamplerOpts.MinP, + "repeat_penalty", request.SamplerOpts.RepeatPenalty, + "presence_penalty", request.SamplerOpts.PresencePenalty, + "frequency_penalty", request.SamplerOpts.FrequencyPenalty, + "logprobs", request.SamplerOpts.Logprobs, + "top_logprobs", request.SamplerOpts.TopLogprobs, + ) } - const preThinking = 4 - if end := len(inputs) - preThinking; end > 0 { - session.requestSnapshot(end) + requestPipelineSnapshots := func(s *cacheSession) { + if s == nil { + return + } + // Request periodic snapshots during prefill and near the end of the + // prompt so that long prompts can be partially restored and + // thinking/generation can be retried without full reprocessing. + const snapshotInterval = 8192 + for offset := snapshotInterval; offset < len(inputs); offset += snapshotInterval { + s.requestSnapshot(offset) + } + + const preThinking = 4 + if end := len(inputs) - preThinking; end > 0 { + s.requestSnapshot(end) + } + } + requestPipelineSnapshots(session) + requestPipelineSnapshots(dflashSession) + + nextSnapshotOffset := func() int { + next := session.nextPendingSnapshot() + if dflashSession != nil { + if offset := dflashSession.nextPendingSnapshot(); offset > 0 && (next == 0 || offset < next) { + next = offset + } + } + return next } - materializeCaches := func() { + snapshotReadySessions := func(position int) { + if snapOffset := session.nextPendingSnapshot(); snapOffset > 0 && position >= snapOffset { + session.snapshot() + } + if dflashSession != nil { + if snapOffset := dflashSession.nextPendingSnapshot(); snapOffset > 0 && position >= snapOffset { + dflashSession.snapshot() + } + } + } + + materializeCaches := func(cacheSets ...[]cache.Cache) { + if len(cacheSets) == 0 { + cacheSets = [][]cache.Cache{caches} + } state := make([]*mlx.Array, 0, 2*len(caches)) - for _, c := range caches { - state = append(state, c.State()...) + for _, set := range cacheSets { + for _, c := range set { + if c == nil { + continue + } + state = append(state, c.State()...) + } } if len(state) == 0 { return @@ -104,6 +172,61 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er mlx.Eval(state...) } + if dflashEnabled { + targetCachedPrefix := len(inputs) - len(tokens) + dflashCachedPrefix := len(inputs) - len(dflashSession.remaining) + if targetCachedPrefix > dflashCachedPrefix { + t0 := time.Now() + rebuildCaches := newDFlashTargetCaches(r.Model) + rebuildProcessed := 0 + for targetCachedPrefix-rebuildProcessed > 0 { + if err := ctx.Err(); err != nil { + freeCacheSet(rebuildCaches) + return err + } + n := min(prefillChunk, targetCachedPrefix-rebuildProcessed) + if snapOffset := dflashSession.nextPendingSnapshot(); snapOffset > rebuildProcessed && snapOffset < rebuildProcessed+n { + n = snapOffset - rebuildProcessed + } + start, end := rebuildProcessed, rebuildProcessed+n + b := &batch.Batch{ + InputIDs: mlx.FromValues(inputs[start:end], 1, n), + SeqOffsets: []int32{int32(start)}, + SeqQueryLens: []int32{int32(n)}, + } + _, targetHidden := dflashTarget.ForwardDFlash(b, rebuildCaches, dflashDraft.TargetLayerIDs()) + if end > dflashCachedPrefix { + appendHidden := targetHidden + if start < dflashCachedPrefix { + appendHidden = targetHidden.Slice(mlx.Slice(), mlx.Slice(dflashCachedPrefix-start, n), mlx.Slice()) + } + dflashDraft.AppendContext(appendHidden, dflashCaches) + } + mlx.Sweep() + materializeCaches(rebuildCaches, dflashCaches) + rebuildProcessed = end + if snapOffset := dflashSession.nextPendingSnapshot(); snapOffset > 0 && rebuildProcessed >= snapOffset { + dflashSession.snapshot() + } + mlx.ClearCache() + } + freeCacheSet(rebuildCaches) + slog.Info("DFlash draft cache rebuild", + "target_cached", targetCachedPrefix, + "draft_cached", dflashCachedPrefix, + "rebuilt", targetCachedPrefix-dflashCachedPrefix, + "draft_offset", r.dflashCache.minCacheOffset(), + "duration", time.Since(t0), + ) + } else { + slog.Info("DFlash draft cache restored", + "target_cached", targetCachedPrefix, + "draft_cached", dflashCachedPrefix, + "draft_offset", r.dflashCache.minCacheOffset(), + ) + } + } + now := time.Now() total, processed := len(tokens), 0 position := len(inputs) - len(tokens) @@ -116,37 +239,49 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er // If there's a pending snapshot, split the batch so we can // capture it at the exact offset. - if snapOffset := session.nextPendingSnapshot(); snapOffset > 0 { + if snapOffset := nextSnapshotOffset(); snapOffset > 0 { tokensUntilSnapshot := snapOffset - position if tokensUntilSnapshot > 0 && tokensUntilSnapshot < n { n = tokensUntilSnapshot } } - r.Model.Forward(&batch.Batch{ + b := &batch.Batch{ InputIDs: mlx.FromValues(tokens[processed:processed+n], 1, n), SeqOffsets: []int32{int32(position)}, SeqQueryLens: []int32{int32(n)}, - }, caches) + } + if dflashEnabled { + _, targetHidden := dflashTarget.ForwardDFlash(b, caches, dflashDraft.TargetLayerIDs()) + dflashDraft.AppendContext(targetHidden, dflashCaches) + } else { + r.Model.Forward(b, caches) + } mlx.Sweep() - materializeCaches() + if dflashEnabled { + materializeCaches(caches, dflashCaches) + } else { + materializeCaches() + } processed += n position += n slog.Info("Prompt processing progress", "processed", processed, "total", total) logutil.TraceContext(ctx, "mlx prompt forward", "processed", processed, "total", total, "tokens", n, "memory", mlx.Memory{}) // Create snapshot if we've reached a pending offset. - if snapOffset := session.nextPendingSnapshot(); snapOffset > 0 { - if position >= snapOffset { - session.snapshot() - } - } + snapshotReadySessions(position) mlx.ClearCache() } // Register the sampler after prefill completes. r.Sampler.Add(pipelineSlot, request.SamplerOpts, inputs) + if dflashMode == dflashDecodeGreedy { + return r.runGreedyDFlashDecode(ctx, request, session, caches, dflashCaches, tokens[processed:], &position, now) + } + if dflashMode == dflashDecodeSample { + return r.runSampleDFlashDecode(ctx, request, session, caches, dflashCaches, tokens[processed:], &position, now) + } if r.useGreedyMTP(request.SamplerOpts) { return r.runGreedyMTPDecode(ctx, request, session, caches, tokens[processed:], &position, now) } diff --git a/x/mlxrunner/runner.go b/x/mlxrunner/runner.go index ccde6ff0..9fd6421c 100644 --- a/x/mlxrunner/runner.go +++ b/x/mlxrunner/runner.go @@ -40,6 +40,7 @@ type Runner struct { Requests chan Request Sampler *sample.Sampler cache kvCache + dflashCache kvCache contextLength int mlxThread *mlxthread.Thread } diff --git a/x/models/dflash/dflash.go b/x/models/dflash/dflash.go new file mode 100644 index 00000000..35236155 --- /dev/null +++ b/x/models/dflash/dflash.go @@ -0,0 +1,439 @@ +// Package dflash implements DFlash block-diffusion draft models for MLX. +package dflash + +import ( + "encoding/json" + "fmt" + "math" + "sort" + "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" + "github.com/ollama/ollama/x/mlxrunner/model/base" + "github.com/ollama/ollama/x/models/nn" +) + +func init() { + base.RegisterDraft("DFlashDraftModel", newModel) + base.RegisterDraft("dflash", newModel) +} + +var _ base.DFlashDraftModel = (*Model)(nil) + +type Config struct { + HiddenSize int32 `json:"hidden_size"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + NumAttentionHeads int32 `json:"num_attention_heads"` + NumKeyValueHeads int32 `json:"num_key_value_heads"` + HeadDim int32 `json:"head_dim"` + IntermediateSize int32 `json:"intermediate_size"` + VocabSize int32 `json:"vocab_size"` + RMSNormEps float32 `json:"rms_norm_eps"` + RopeTheta float32 `json:"rope_theta"` + RopeScaling *nn.RopeParameters `json:"rope_scaling"` + RopeParameters *nn.RopeParameters `json:"rope_parameters"` + MaxPositionEmbeddings int32 `json:"max_position_embeddings"` + BlockSizeValue int32 `json:"block_size"` + NumTargetLayers int32 `json:"num_target_layers"` + LayerTypes []string `json:"layer_types"` + SlidingWindow int32 `json:"sliding_window"` + FinalLogitSoftcapping *float32 `json:"final_logit_softcapping"` + DFlash struct { + TargetLayerIDs []int `json:"target_layer_ids"` + MaskTokenID int32 `json:"mask_token_id"` + } `json:"dflash_config"` + + QuantGroupSize int `json:"-"` + QuantBits int `json:"-"` + QuantMode string `json:"-"` + TensorQuant map[string]*model.TensorQuantInfo `json:"-"` + Scale float32 `json:"-"` + RopeFreqs *mlx.Array `json:"-"` + RopeScale float32 `json:"-"` +} + +type Model struct { + FC nn.LinearLayer + HiddenNorm *nn.RMSNorm + Layers []*Layer + Norm *nn.RMSNorm + + target base.Model + targetEmbeddings base.MTPEmbeddingModel + tensorPrefix string + + *Config +} + +type Layer struct { + Attention *Attention + MLP *MLP + InputNorm *nn.RMSNorm + PostAttentionNorm *nn.RMSNorm +} + +type Attention struct { + QProj nn.LinearLayer + KProj nn.LinearLayer + VProj nn.LinearLayer + OProj nn.LinearLayer + QNorm *nn.RMSNorm + KNorm *nn.RMSNorm + + Sliding bool +} + +type MLP struct { + GateProj nn.LinearLayer + UpProj nn.LinearLayer + DownProj nn.LinearLayer +} + +func parseConfig(data []byte) (Config, error) { + var cfg Config + if err := json.Unmarshal(data, &cfg); err != nil { + return Config{}, fmt.Errorf("parse dflash config: %w", err) + } + if cfg.HiddenSize <= 0 { + return Config{}, fmt.Errorf("invalid hidden_size: %d", cfg.HiddenSize) + } + if cfg.NumHiddenLayers <= 0 { + return Config{}, fmt.Errorf("invalid num_hidden_layers: %d", cfg.NumHiddenLayers) + } + if cfg.NumAttentionHeads <= 0 { + return Config{}, fmt.Errorf("invalid num_attention_heads: %d", cfg.NumAttentionHeads) + } + if cfg.NumKeyValueHeads <= 0 { + cfg.NumKeyValueHeads = cfg.NumAttentionHeads + } + if cfg.HeadDim <= 0 { + if cfg.HiddenSize%cfg.NumAttentionHeads != 0 { + return Config{}, fmt.Errorf("hidden_size (%d) must be divisible by num_attention_heads (%d)", cfg.HiddenSize, cfg.NumAttentionHeads) + } + cfg.HeadDim = cfg.HiddenSize / cfg.NumAttentionHeads + } + if cfg.RMSNormEps == 0 { + cfg.RMSNormEps = 1e-6 + } + if cfg.RopeTheta == 0 { + ropeParams := cfg.RopeParameters + if ropeParams == nil { + ropeParams = cfg.RopeScaling + } + if ropeParams != nil && ropeParams.RopeTheta > 0 { + cfg.RopeTheta = ropeParams.RopeTheta + } + } + if cfg.RopeTheta == 0 { + cfg.RopeTheta = 1000000 + } + if cfg.BlockSizeValue <= 0 { + return Config{}, fmt.Errorf("invalid block_size: %d", cfg.BlockSizeValue) + } + if len(cfg.DFlash.TargetLayerIDs) == 0 { + return Config{}, fmt.Errorf("dflash_config.target_layer_ids is required") + } + if !sort.IntsAreSorted(cfg.DFlash.TargetLayerIDs) { + return Config{}, fmt.Errorf("dflash_config.target_layer_ids must be sorted") + } + if len(cfg.LayerTypes) == 0 { + cfg.LayerTypes = make([]string, cfg.NumHiddenLayers) + for i := range cfg.LayerTypes { + cfg.LayerTypes[i] = "full_attention" + } + } + if len(cfg.LayerTypes) != int(cfg.NumHiddenLayers) { + return Config{}, fmt.Errorf("layer_types length %d does not match num_hidden_layers %d", len(cfg.LayerTypes), cfg.NumHiddenLayers) + } + for i, typ := range cfg.LayerTypes { + switch strings.ToLower(typ) { + case "full_attention": + case "sliding_attention": + if cfg.SlidingWindow <= 0 { + return Config{}, fmt.Errorf("layer %d uses sliding_attention but sliding_window is not set", i) + } + default: + return Config{}, fmt.Errorf("unsupported layer type %q", typ) + } + } + cfg.Scale = float32(1.0 / math.Sqrt(float64(cfg.HeadDim))) + cfg.RopeScale = 1 + ropeParams := cfg.RopeParameters + if ropeParams == nil { + ropeParams = cfg.RopeScaling + } + if ropeParams != nil && strings.EqualFold(ropeParams.TypeName(), "yarn") { + cfg.RopeFreqs, cfg.RopeScale = nn.BuildYarnRopeFreqs(int(cfg.HeadDim), cfg.RopeTheta, ropeParams) + } + return cfg, nil +} + +func newModel(root *model.Root, target base.Model) (base.DraftModel, error) { + if root == nil || root.Draft == nil { + return nil, fmt.Errorf("draft metadata missing") + } + + configPath := root.Draft.Config + if configPath == "" { + configPath = "draft/config.json" + } + configData, err := root.Manifest.ReadConfig(configPath) + if err != nil { + return nil, fmt.Errorf("load dflash config: %w", err) + } + + cfg, err := parseConfig(configData) + if err != nil { + return nil, err + } + if target.NumLayers() < int(cfg.NumTargetLayers) { + return nil, fmt.Errorf("dflash target expects %d layers, target has %d", cfg.NumTargetLayers, target.NumLayers()) + } + for _, layerID := range cfg.DFlash.TargetLayerIDs { + if layerID < 0 || layerID >= target.NumLayers() { + return nil, fmt.Errorf("dflash target layer id %d out of range for %d-layer target", layerID, target.NumLayers()) + } + } + targetEmbeddings, ok := target.(base.MTPEmbeddingModel) + if !ok { + return nil, fmt.Errorf("dflash draft requires target token embeddings, got %T", target) + } + + if qt := root.QuantType(); qt != "" { + cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams(qt) + if gs := root.GroupSize(); gs > 0 { + cfg.QuantGroupSize = gs + } + } else { + cfg.QuantGroupSize, cfg.QuantBits, cfg.QuantMode = model.QuantizationParams("") + } + cfg.TensorQuant = root.AllTensorQuant() + + prefix := root.Draft.TensorPrefix + if prefix == "" { + prefix = "draft." + } + + m := &Model{ + Config: &cfg, + Layers: make([]*Layer, cfg.NumHiddenLayers), + target: target, + targetEmbeddings: targetEmbeddings, + tensorPrefix: prefix, + } + for i := range m.Layers { + m.Layers[i] = &Layer{Attention: &Attention{}, MLP: &MLP{}} + } + return m, nil +} + +func (m *Model) LoadWeights(tensors map[string]*mlx.Array) error { + prefix := m.tensorPrefix + linears := model.NewLinearFactory(tensors, m.QuantGroupSize, m.QuantBits, m.QuantMode, m.TensorQuant) + + m.FC = linears.Make(prefix + "fc") + if m.FC == nil { + return fmt.Errorf("missing dflash fc weight") + } + if w := tensors[prefix+"hidden_norm.weight"]; w != nil { + m.HiddenNorm = nn.NewRMSNorm(w, m.RMSNormEps) + } + if w := tensors[prefix+"norm.weight"]; w != nil { + m.Norm = nn.NewRMSNorm(w, m.RMSNormEps) + } + if m.HiddenNorm == nil || m.Norm == nil { + return fmt.Errorf("missing dflash norm weights") + } + + for i := range m.NumHiddenLayers { + layerPrefix := fmt.Sprintf("%slayers.%d", prefix, i) + layer := &Layer{ + Attention: &Attention{Sliding: strings.ToLower(m.LayerTypes[i]) == "sliding_attention"}, + MLP: &MLP{ + GateProj: linears.Make(layerPrefix + ".mlp.gate_proj"), + UpProj: linears.Make(layerPrefix + ".mlp.up_proj"), + DownProj: linears.Make(layerPrefix + ".mlp.down_proj"), + }, + } + if w := tensors[layerPrefix+".input_layernorm.weight"]; w != nil { + layer.InputNorm = nn.NewRMSNorm(w, m.RMSNormEps) + } + if w := tensors[layerPrefix+".post_attention_layernorm.weight"]; w != nil { + layer.PostAttentionNorm = nn.NewRMSNorm(w, m.RMSNormEps) + } + layer.Attention.QProj = linears.Make(layerPrefix + ".self_attn.q_proj") + layer.Attention.KProj = linears.Make(layerPrefix + ".self_attn.k_proj") + layer.Attention.VProj = linears.Make(layerPrefix + ".self_attn.v_proj") + layer.Attention.OProj = linears.Make(layerPrefix + ".self_attn.o_proj") + if w := tensors[layerPrefix+".self_attn.q_norm.weight"]; w != nil { + layer.Attention.QNorm = nn.NewRMSNorm(w, m.RMSNormEps) + } + if w := tensors[layerPrefix+".self_attn.k_norm.weight"]; w != nil { + layer.Attention.KNorm = nn.NewRMSNorm(w, m.RMSNormEps) + } + + if layer.InputNorm == nil || layer.PostAttentionNorm == nil { + return fmt.Errorf("dflash layer %d: missing layer norms", i) + } + if layer.Attention.QProj == nil || layer.Attention.KProj == nil || layer.Attention.VProj == nil || layer.Attention.OProj == nil { + return fmt.Errorf("dflash layer %d: missing attention projections", i) + } + if layer.Attention.QNorm == nil || layer.Attention.KNorm == nil { + return fmt.Errorf("dflash layer %d: missing attention q/k norms", i) + } + if layer.MLP.GateProj == nil || layer.MLP.UpProj == nil || layer.MLP.DownProj == nil { + return fmt.Errorf("dflash layer %d: missing mlp projections", i) + } + + m.Layers[i] = layer + } + return nil +} + +func (m *Model) TargetLayerIDs() []int { + return append([]int(nil), m.DFlash.TargetLayerIDs...) +} + +func (m *Model) BlockSize() int { + return int(m.BlockSizeValue) +} + +func (m *Model) MaskTokenID() int32 { + return m.DFlash.MaskTokenID +} + +func (m *Model) NewCaches() []cache.Cache { + caches := make([]cache.Cache, len(m.Layers)) + for i, typ := range m.LayerTypes { + if strings.ToLower(typ) == "sliding_attention" { + // RotatingKVCache.View returns maxSize-1 tokens so assistant + // paths can append the current query. DFlash uses that same + // view for target context, so allocate one extra slot to expose + // the draft model's sliding_window-1 context tokens. + caches[i] = cache.NewRotatingKVCache(int(m.SlidingWindow)) + } else { + caches[i] = cache.NewKVCache() + } + } + return caches +} + +func (m *Model) AppendContext(targetHidden *mlx.Array, caches []cache.Cache) { + if targetHidden == nil || targetHidden.Dim(1) == 0 { + return + } + hCtx := m.HiddenNorm.Forward(m.FC.Forward(targetHidden), m.RMSNormEps) + offset := int32(0) + if len(caches) > 0 && caches[0] != nil { + offset = int32(caches[0].Offset()) + } + b := &batch.Batch{ + InputIDs: mlx.Zeros(mlx.DTypeInt32, targetHidden.Dim(0), targetHidden.Dim(1)), + SeqOffsets: []int32{offset}, + SeqQueryLens: []int32{int32(targetHidden.Dim(1))}, + } + positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets)) + for i, layer := range m.Layers { + if i >= len(caches) || caches[i] == nil { + continue + } + layer.Attention.AppendContext(hCtx, b, positions, caches[i], m.Config) + } +} + +func (m *Model) Draft(inputIDs *mlx.Array, caches []cache.Cache) *mlx.Array { + dims := inputIDs.Dims() + B, L := int32(dims[0]), int32(dims[1]) + offset := int32(0) + if len(caches) > 0 && caches[0] != nil { + offset = int32(caches[0].Offset()) + } + b := &batch.Batch{ + InputIDs: inputIDs, + SeqOffsets: []int32{offset}, + SeqQueryLens: []int32{L}, + } + positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets)) + + h := m.targetEmbeddings.TokenEmbeddings(inputIDs) + for i, layer := range m.Layers { + var c cache.Cache + if i < len(caches) { + c = caches[i] + } + h = layer.Forward(h, b, c, positions, B, L, m.Config) + } + logits := m.target.Unembed(m.Norm.Forward(h, m.RMSNormEps)) + if m.FinalLogitSoftcapping != nil { + cap := mlx.FromValue(*m.FinalLogitSoftcapping).AsType(logits.DType()) + logits = mlx.LogitSoftcap(logits, cap) + } + return logits +} + +func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { + h := mlx.Add(x, l.Attention.Forward(l.InputNorm.Forward(x, cfg.RMSNormEps), b, c, positions, B, L, cfg)) + return mlx.Add(h, l.MLP.Forward(l.PostAttentionNorm.Forward(h, cfg.RMSNormEps))) +} + +func (a *Attention) AppendContext(xCtx *mlx.Array, b *batch.Batch, positions *mlx.Array, c cache.Cache, cfg *Config) { + B, L := int32(xCtx.Dim(0)), int32(xCtx.Dim(1)) + k := a.KProj.Forward(xCtx) + v := a.VProj.Forward(xCtx) + + k = mlx.Reshape(k, B, L, cfg.NumKeyValueHeads, cfg.HeadDim) + v = mlx.Reshape(v, B, L, cfg.NumKeyValueHeads, cfg.HeadDim) + k = a.KNorm.Forward(k, cfg.RMSNormEps) + + k = mlx.Transpose(k, 0, 2, 1, 3) + v = mlx.Transpose(v, 0, 2, 1, 3) + k = nn.ScaleRotaryPart(mlx.RoPEWithFreqs(k, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions, cfg.RopeFreqs), int(cfg.HeadDim), cfg.RopeScale) + + c.(cache.Attention).Update(b, k, v) +} + +func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions *mlx.Array, B, L int32, cfg *Config) *mlx.Array { + q := a.QProj.Forward(x) + propK := a.KProj.Forward(x) + propV := a.VProj.Forward(x) + + q = mlx.Reshape(q, B, L, cfg.NumAttentionHeads, cfg.HeadDim) + propK = mlx.Reshape(propK, B, L, cfg.NumKeyValueHeads, cfg.HeadDim) + propV = mlx.Reshape(propV, B, L, cfg.NumKeyValueHeads, cfg.HeadDim) + + q = a.QNorm.Forward(q, cfg.RMSNormEps) + propK = a.KNorm.Forward(propK, cfg.RMSNormEps) + + q = mlx.Transpose(q, 0, 2, 1, 3) + propK = mlx.Transpose(propK, 0, 2, 1, 3) + propV = mlx.Transpose(propV, 0, 2, 1, 3) + + q = nn.ScaleRotaryPart(mlx.RoPEWithFreqs(q, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions, cfg.RopeFreqs), int(cfg.HeadDim), cfg.RopeScale) + propK = nn.ScaleRotaryPart(mlx.RoPEWithFreqs(propK, int(cfg.HeadDim), false, cfg.RopeTheta, 1.0, positions, cfg.RopeFreqs), int(cfg.HeadDim), cfg.RopeScale) + + k, v := propK, propV + if viewer, ok := c.(cache.Viewer); ok { + if history := viewer.View(b); history != nil { + k = history.K().Concatenate(2, propK) + v = history.V().Concatenate(2, propV) + } + } + + mask := nn.AttentionMask{} + if a.Sliding { + mask = nn.CausalMask() + if int(cfg.SlidingWindow) > 0 && k.Dim(2) > int(cfg.SlidingWindow) { + mask = mask.Intersect(nn.SlidingWindowMask(b, k.Dim(2), int(cfg.SlidingWindow), q.DType())) + } + } + out := nn.ScaledDotProductAttention(b, q, cfg.Scale, nn.WithKV(k, v, []int32{int32(k.Dim(2))}), nn.WithMask(mask)) + out = mlx.Reshape(mlx.Transpose(out, 0, 2, 1, 3), B, L, cfg.NumAttentionHeads*cfg.HeadDim) + return a.OProj.Forward(out) +} + +func (m *MLP) Forward(x *mlx.Array) *mlx.Array { + return m.DownProj.Forward(mlx.SwiGLU(m.GateProj.Forward(x), m.UpProj.Forward(x))) +} diff --git a/x/models/dflash/dflash_test.go b/x/models/dflash/dflash_test.go new file mode 100644 index 00000000..3feff6f6 --- /dev/null +++ b/x/models/dflash/dflash_test.go @@ -0,0 +1,52 @@ +package dflash + +import ( + "math" + "testing" + + "github.com/ollama/ollama/x/mlxrunner/mlx" +) + +func TestParseConfigYarnRopeScaling(t *testing.T) { + if err := mlx.CheckInit(); err != nil { + t.Skipf("MLX not available: %v", err) + } + + data := []byte(`{ + "hidden_size": 2048, + "num_hidden_layers": 8, + "num_attention_heads": 32, + "num_key_value_heads": 4, + "head_dim": 128, + "intermediate_size": 6144, + "vocab_size": 248320, + "rms_norm_eps": 0.000001, + "rope_theta": 10000000, + "rope_scaling": { + "beta_fast": 32.0, + "beta_slow": 1.0, + "factor": 64.0, + "original_max_position_embeddings": 4096, + "rope_type": "yarn" + }, + "block_size": 16, + "num_target_layers": 40, + "layer_types": ["full_attention", "full_attention", "full_attention", "full_attention", "full_attention", "full_attention", "full_attention", "full_attention"], + "dflash_config": { + "mask_token_id": 248070, + "target_layer_ids": [1, 10, 19, 28, 37] + } + }`) + + cfg, err := parseConfig(data) + if err != nil { + t.Fatalf("parseConfig failed: %v", err) + } + if cfg.RopeFreqs == nil { + t.Fatalf("RopeFreqs is nil") + } + wantScale := float32(0.1*math.Log(64.0) + 1.0) + if math.Abs(float64(cfg.RopeScale-wantScale)) > 1e-6 { + t.Fatalf("RopeScale = %v, want %v", cfg.RopeScale, wantScale) + } +} diff --git a/x/models/laguna/laguna.go b/x/models/laguna/laguna.go index fea4e3f2..e78edb71 100644 --- a/x/models/laguna/laguna.go +++ b/x/models/laguna/laguna.go @@ -22,57 +22,45 @@ func init() { var _ base.Model = (*Model)(nil) -type RopeParameters struct { - RopeTheta float32 `json:"rope_theta"` - RopeType string `json:"rope_type"` - Type string `json:"type"` - PartialRotaryFactor float32 `json:"partial_rotary_factor"` - Factor float32 `json:"factor"` - OriginalMaxPositionEmbeddings int32 `json:"original_max_position_embeddings"` - BetaFast float32 `json:"beta_fast"` - BetaSlow float32 `json:"beta_slow"` - AttentionFactor float32 `json:"attention_factor"` -} - type gatingMode string type ropeConfig struct { - flat *RopeParameters - full *RopeParameters - sliding *RopeParameters + flat *nn.RopeParameters + full *nn.RopeParameters + sliding *nn.RopeParameters nested bool } type Config struct { - ModelType string `json:"model_type"` - HiddenSize int32 `json:"hidden_size"` - IntermediateSize int32 `json:"intermediate_size"` - MoeIntermediateSize int32 `json:"moe_intermediate_size"` - SharedExpertIntermediate int32 `json:"shared_expert_intermediate_size"` - NumHiddenLayers int32 `json:"num_hidden_layers"` - NumAttentionHeads int32 `json:"num_attention_heads"` - NumAttentionHeadsPerLayer []int32 `json:"num_attention_heads_per_layer"` - NumKeyValueHeads int32 `json:"num_key_value_heads"` - HeadDim int32 `json:"head_dim"` - RMSNormEps float32 `json:"rms_norm_eps"` - VocabSize int32 `json:"vocab_size"` - MaxPositionEmbeddings int32 `json:"max_position_embeddings"` - LayerTypes []string `json:"layer_types"` - SlidingWindow int32 `json:"sliding_window"` - MLPOnlyLayers []int32 `json:"mlp_only_layers"` - DecoderSparseStep int32 `json:"decoder_sparse_step"` - NumExperts int32 `json:"num_experts"` - NumExpertsPerTok int32 `json:"num_experts_per_tok"` - NormTopKProb bool `json:"norm_topk_prob"` - MoeRoutedScalingFactor float32 `json:"moe_routed_scaling_factor"` - MoeApplyRouterWeightOnInput bool `json:"moe_apply_router_weight_on_input"` - Gating string `json:"gating"` - TieWordEmbeddings bool `json:"tie_word_embeddings"` - RopeTheta float32 `json:"rope_theta"` - PartialRotaryFactor float32 `json:"partial_rotary_factor"` - RopeParameters *RopeParameters `json:"rope_parameters"` - RopeScaling *RopeParameters `json:"rope_scaling"` - SWARopeParameters *RopeParameters `json:"swa_rope_parameters"` + ModelType string `json:"model_type"` + HiddenSize int32 `json:"hidden_size"` + IntermediateSize int32 `json:"intermediate_size"` + MoeIntermediateSize int32 `json:"moe_intermediate_size"` + SharedExpertIntermediate int32 `json:"shared_expert_intermediate_size"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + NumAttentionHeads int32 `json:"num_attention_heads"` + NumAttentionHeadsPerLayer []int32 `json:"num_attention_heads_per_layer"` + NumKeyValueHeads int32 `json:"num_key_value_heads"` + HeadDim int32 `json:"head_dim"` + RMSNormEps float32 `json:"rms_norm_eps"` + VocabSize int32 `json:"vocab_size"` + MaxPositionEmbeddings int32 `json:"max_position_embeddings"` + LayerTypes []string `json:"layer_types"` + SlidingWindow int32 `json:"sliding_window"` + MLPOnlyLayers []int32 `json:"mlp_only_layers"` + DecoderSparseStep int32 `json:"decoder_sparse_step"` + NumExperts int32 `json:"num_experts"` + NumExpertsPerTok int32 `json:"num_experts_per_tok"` + NormTopKProb bool `json:"norm_topk_prob"` + MoeRoutedScalingFactor float32 `json:"moe_routed_scaling_factor"` + MoeApplyRouterWeightOnInput bool `json:"moe_apply_router_weight_on_input"` + Gating string `json:"gating"` + TieWordEmbeddings bool `json:"tie_word_embeddings"` + RopeTheta float32 `json:"rope_theta"` + PartialRotaryFactor float32 `json:"partial_rotary_factor"` + RopeParameters *nn.RopeParameters `json:"rope_parameters"` + RopeScaling *nn.RopeParameters `json:"rope_scaling"` + SWARopeParameters *nn.RopeParameters `json:"swa_rope_parameters"` QuantGroupSize int `json:"-"` QuantBits int `json:"-"` @@ -170,36 +158,36 @@ type stackedExpertWeights struct { func parseConfig(configData []byte) (Config, error) { type rawConfig struct { - ModelType string `json:"model_type"` - HiddenSize int32 `json:"hidden_size"` - IntermediateSize int32 `json:"intermediate_size"` - MoeIntermediateSize int32 `json:"moe_intermediate_size"` - SharedExpertIntermediate int32 `json:"shared_expert_intermediate_size"` - NumHiddenLayers int32 `json:"num_hidden_layers"` - NumAttentionHeads int32 `json:"num_attention_heads"` - NumAttentionHeadsPerLayer []int32 `json:"num_attention_heads_per_layer"` - NumKeyValueHeads int32 `json:"num_key_value_heads"` - HeadDim int32 `json:"head_dim"` - RMSNormEps float32 `json:"rms_norm_eps"` - VocabSize int32 `json:"vocab_size"` - MaxPositionEmbeddings int32 `json:"max_position_embeddings"` - LayerTypes []string `json:"layer_types"` - SlidingWindow int32 `json:"sliding_window"` - MLPOnlyLayers []int32 `json:"mlp_only_layers"` - MLPLayerTypes []string `json:"mlp_layer_types"` - DecoderSparseStep int32 `json:"decoder_sparse_step"` - NumExperts int32 `json:"num_experts"` - NumExpertsPerTok int32 `json:"num_experts_per_tok"` - NormTopKProb *bool `json:"norm_topk_prob"` - MoeRoutedScalingFactor float32 `json:"moe_routed_scaling_factor"` - MoeApplyRouterWeightOnInput bool `json:"moe_apply_router_weight_on_input"` - Gating gatingMode `json:"gating"` - TieWordEmbeddings bool `json:"tie_word_embeddings"` - RopeTheta float32 `json:"rope_theta"` - PartialRotaryFactor float32 `json:"partial_rotary_factor"` - RopeParameters ropeConfig `json:"rope_parameters"` - RopeScaling *RopeParameters `json:"rope_scaling"` - SWARopeParameters *RopeParameters `json:"swa_rope_parameters"` + ModelType string `json:"model_type"` + HiddenSize int32 `json:"hidden_size"` + IntermediateSize int32 `json:"intermediate_size"` + MoeIntermediateSize int32 `json:"moe_intermediate_size"` + SharedExpertIntermediate int32 `json:"shared_expert_intermediate_size"` + NumHiddenLayers int32 `json:"num_hidden_layers"` + NumAttentionHeads int32 `json:"num_attention_heads"` + NumAttentionHeadsPerLayer []int32 `json:"num_attention_heads_per_layer"` + NumKeyValueHeads int32 `json:"num_key_value_heads"` + HeadDim int32 `json:"head_dim"` + RMSNormEps float32 `json:"rms_norm_eps"` + VocabSize int32 `json:"vocab_size"` + MaxPositionEmbeddings int32 `json:"max_position_embeddings"` + LayerTypes []string `json:"layer_types"` + SlidingWindow int32 `json:"sliding_window"` + MLPOnlyLayers []int32 `json:"mlp_only_layers"` + MLPLayerTypes []string `json:"mlp_layer_types"` + DecoderSparseStep int32 `json:"decoder_sparse_step"` + NumExperts int32 `json:"num_experts"` + NumExpertsPerTok int32 `json:"num_experts_per_tok"` + NormTopKProb *bool `json:"norm_topk_prob"` + MoeRoutedScalingFactor float32 `json:"moe_routed_scaling_factor"` + MoeApplyRouterWeightOnInput bool `json:"moe_apply_router_weight_on_input"` + Gating gatingMode `json:"gating"` + TieWordEmbeddings bool `json:"tie_word_embeddings"` + RopeTheta float32 `json:"rope_theta"` + PartialRotaryFactor float32 `json:"partial_rotary_factor"` + RopeParameters ropeConfig `json:"rope_parameters"` + RopeScaling *nn.RopeParameters `json:"rope_scaling"` + SWARopeParameters *nn.RopeParameters `json:"swa_rope_parameters"` } var raw rawConfig @@ -313,8 +301,8 @@ func parseConfig(configData []byte) (Config, error) { } cfg.FullRopeDim = clampRopeDim(int(float32(cfg.HeadDim)*fullPartial), int(cfg.HeadDim)) cfg.FullRopeScale = 1 - if ropeParams != nil && strings.EqualFold(ropeParams.ropeType(), "yarn") { - cfg.FullRopeFreqs, cfg.FullRopeScale = buildYarnRopeFreqs(cfg.FullRopeDim, cfg.FullRopeBase, ropeParams) + if ropeParams != nil && strings.EqualFold(ropeParams.TypeName(), "yarn") { + cfg.FullRopeFreqs, cfg.FullRopeScale = nn.BuildYarnRopeFreqs(cfg.FullRopeDim, cfg.FullRopeBase, ropeParams) } cfg.SlidingRopeBase = cfg.FullRopeBase @@ -378,12 +366,12 @@ func (r *ropeConfig) UnmarshalJSON(b []byte) error { if raw, ok := probe["full_attention"]; ok { r.nested = true - r.full = &RopeParameters{} + r.full = &nn.RopeParameters{} if err := json.Unmarshal(raw, r.full); err != nil { return err } if raw = probe["sliding_attention"]; raw != nil { - r.sliding = &RopeParameters{} + r.sliding = &nn.RopeParameters{} if err := json.Unmarshal(raw, r.sliding); err != nil { return err } @@ -393,12 +381,12 @@ func (r *ropeConfig) UnmarshalJSON(b []byte) error { if raw, ok := probe["global_attention"]; ok { r.nested = true - r.full = &RopeParameters{} + r.full = &nn.RopeParameters{} if err := json.Unmarshal(raw, r.full); err != nil { return err } if raw = probe["sliding_attention"]; raw != nil { - r.sliding = &RopeParameters{} + r.sliding = &nn.RopeParameters{} if err := json.Unmarshal(raw, r.sliding); err != nil { return err } @@ -406,18 +394,18 @@ func (r *ropeConfig) UnmarshalJSON(b []byte) error { return nil } - r.flat = &RopeParameters{} + r.flat = &nn.RopeParameters{} return json.Unmarshal(b, r.flat) } -func (r ropeConfig) fullParams() *RopeParameters { +func (r ropeConfig) fullParams() *nn.RopeParameters { if r.nested { return r.full } return r.flat } -func (r ropeConfig) slidingParams() *RopeParameters { +func (r ropeConfig) slidingParams() *nn.RopeParameters { if !r.nested { return nil } @@ -452,88 +440,6 @@ func denseLayers(mlpOnlyLayers []int32, mlpLayerTypes []string) ([]int32, error) return dense, nil } -func (rp *RopeParameters) ropeType() string { - if rp == nil { - return "" - } - if rp.RopeType != "" { - return rp.RopeType - } - return rp.Type -} - -func buildYarnRopeFreqs(dim int, base float32, rp *RopeParameters) (*mlx.Array, float32) { - if rp == nil || dim <= 0 { - return nil, 1 - } - factor := rp.Factor - if factor <= 0 { - factor = 1 - } - attentionFactor := rp.AttentionFactor - if attentionFactor == 0 && factor > 1 { - attentionFactor = float32(0.1*math.Log(float64(factor)) + 1.0) - } else if attentionFactor == 0 { - attentionFactor = 1 - } - if factor <= 1 { - return nil, attentionFactor - } - - originalMax := rp.OriginalMaxPositionEmbeddings - if originalMax <= 0 { - originalMax = 4096 - } - betaFast := rp.BetaFast - if betaFast == 0 { - betaFast = 32 - } - betaSlow := rp.BetaSlow - if betaSlow == 0 { - betaSlow = 1 - } - half := dim / 2 - low, high := yarnCorrectionRange(betaFast, betaSlow, dim, base, originalMax) - freqs := make([]float32, half) - for i := range half { - posFreq := math.Pow(float64(base), float64(2*i)/float64(dim)) - invExtrapolation := 1.0 / posFreq - invInterpolation := 1.0 / (float64(factor) * posFreq) - ramp := yarnRamp(float64(i), low, high) - mask := 1 - ramp - inv := invInterpolation*(1-mask) + invExtrapolation*mask - freqs[i] = float32(1.0 / inv) - } - arr := mlx.FromValues(freqs, half) - mlx.Eval(arr) - return arr, attentionFactor -} - -func yarnCorrectionRange(betaFast, betaSlow float32, dim int, base float32, maxPosition int32) (float64, float64) { - findDim := func(rot float32) float64 { - return float64(dim) * math.Log(float64(maxPosition)/(float64(rot)*2*math.Pi)) / (2 * math.Log(float64(base))) - } - low := math.Floor(findDim(betaFast)) - high := math.Ceil(findDim(betaSlow)) - low = math.Max(low, 0) - high = math.Min(high, float64(dim-1)) - if low == high { - high += 0.001 - } - return low, high -} - -func yarnRamp(i, low, high float64) float64 { - v := (i - low) / (high - low) - if v < 0 { - return 0 - } - if v > 1 { - return 1 - } - return v -} - func clampRopeDim(v, maxDim int) int { if v <= 0 { return maxDim @@ -1016,8 +922,8 @@ func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positio if layer.IsSliding { ropeDim, ropeBase, ropeMSScale, ropeFreqs = cfg.SlidingRopeDim, cfg.SlidingRopeBase, cfg.SlidingRopeScale, nil } - q = scaleRotaryPart(mlx.RoPEWithFreqs(q, ropeDim, false, ropeBase, 1.0, positions, ropeFreqs), ropeDim, ropeMSScale) - k = scaleRotaryPart(mlx.RoPEWithFreqs(k, ropeDim, false, ropeBase, 1.0, positions, ropeFreqs), ropeDim, ropeMSScale) + q = nn.ScaleRotaryPart(mlx.RoPEWithFreqs(q, ropeDim, false, ropeBase, 1.0, positions, ropeFreqs), ropeDim, ropeMSScale) + k = nn.ScaleRotaryPart(mlx.RoPEWithFreqs(k, ropeDim, false, ropeBase, 1.0, positions, ropeFreqs), ropeDim, ropeMSScale) var kv nn.SDPAOption if c != nil { @@ -1034,30 +940,6 @@ func (a *Attention) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positio return a.OProj.Forward(out) } -func scaleRotaryPart(x *mlx.Array, ropeDim int, scale float32) *mlx.Array { - if scale == 1 { - return x - } - dims := x.Dims() - last := dims[len(dims)-1] - if ropeDim >= last { - return mlx.MulScalar(x, scale) - } - start := make([]int32, len(dims)) - stopRot := make([]int32, len(dims)) - stopPass := make([]int32, len(dims)) - startPass := make([]int32, len(dims)) - for i, dim := range dims { - stopRot[i] = int32(dim) - stopPass[i] = int32(dim) - } - stopRot[len(dims)-1] = int32(ropeDim) - startPass[len(dims)-1] = int32(ropeDim) - rot := mlx.MulScalar(mlx.SliceStartStop(x, start, stopRot), scale) - pass := mlx.SliceStartStop(x, startPass, stopPass) - return mlx.Concatenate([]*mlx.Array{rot, pass}, -1) -} - func (m *DenseMLP) Forward(x *mlx.Array, _ *Config) *mlx.Array { return m.DownProj.Forward(mlx.SwiGLU(m.GateProj.Forward(x), m.UpProj.Forward(x))) } diff --git a/x/models/nn/rope.go b/x/models/nn/rope.go new file mode 100644 index 00000000..39d9dab1 --- /dev/null +++ b/x/models/nn/rope.go @@ -0,0 +1,129 @@ +package nn + +import ( + "math" + + "github.com/ollama/ollama/x/mlxrunner/mlx" +) + +// RopeParameters carries common RoPE metadata embedded in model configs. +type RopeParameters struct { + RopeTheta float32 `json:"rope_theta"` + RopeType string `json:"rope_type"` + Type string `json:"type"` + PartialRotaryFactor float32 `json:"partial_rotary_factor"` + Factor float32 `json:"factor"` + OriginalMaxPositionEmbeddings int32 `json:"original_max_position_embeddings"` + BetaFast float32 `json:"beta_fast"` + BetaSlow float32 `json:"beta_slow"` + AttentionFactor float32 `json:"attention_factor"` +} + +// TypeName returns rope_type when present, falling back to type. +func (rp *RopeParameters) TypeName() string { + if rp == nil { + return "" + } + if rp.RopeType != "" { + return rp.RopeType + } + return rp.Type +} + +// BuildYarnRopeFreqs returns YaRN rotary frequencies and the mscale value. +func BuildYarnRopeFreqs(dim int, base float32, rp *RopeParameters) (*mlx.Array, float32) { + if rp == nil || dim <= 0 { + return nil, 1 + } + factor := rp.Factor + if factor <= 0 { + factor = 1 + } + attentionFactor := rp.AttentionFactor + if attentionFactor == 0 && factor > 1 { + attentionFactor = float32(0.1*math.Log(float64(factor)) + 1.0) + } else if attentionFactor == 0 { + attentionFactor = 1 + } + if factor <= 1 { + return nil, attentionFactor + } + + originalMax := rp.OriginalMaxPositionEmbeddings + if originalMax <= 0 { + originalMax = 4096 + } + betaFast := rp.BetaFast + if betaFast == 0 { + betaFast = 32 + } + betaSlow := rp.BetaSlow + if betaSlow == 0 { + betaSlow = 1 + } + half := dim / 2 + low, high := yarnCorrectionRange(betaFast, betaSlow, dim, base, originalMax) + freqs := make([]float32, half) + for i := range half { + posFreq := math.Pow(float64(base), float64(2*i)/float64(dim)) + invExtrapolation := 1.0 / posFreq + invInterpolation := 1.0 / (float64(factor) * posFreq) + ramp := yarnRamp(float64(i), low, high) + mask := 1 - ramp + inv := invInterpolation*(1-mask) + invExtrapolation*mask + freqs[i] = float32(1.0 / inv) + } + arr := mlx.FromValues(freqs, half) + mlx.Eval(arr) + return arr, attentionFactor +} + +func yarnCorrectionRange(betaFast, betaSlow float32, dim int, base float32, maxPosition int32) (float64, float64) { + findDim := func(rot float32) float64 { + return float64(dim) * math.Log(float64(maxPosition)/(float64(rot)*2*math.Pi)) / (2 * math.Log(float64(base))) + } + low := math.Floor(findDim(betaFast)) + high := math.Ceil(findDim(betaSlow)) + low = math.Max(low, 0) + high = math.Min(high, float64(dim-1)) + if low == high { + high += 0.001 + } + return low, high +} + +func yarnRamp(i, low, high float64) float64 { + v := (i - low) / (high - low) + if v < 0 { + return 0 + } + if v > 1 { + return 1 + } + return v +} + +// ScaleRotaryPart applies YaRN's mscale to only the rotated dimensions. +func ScaleRotaryPart(x *mlx.Array, ropeDim int, scale float32) *mlx.Array { + if scale == 1 { + return x + } + dims := x.Dims() + last := dims[len(dims)-1] + if ropeDim >= last { + return mlx.MulScalar(x, scale) + } + start := make([]int32, len(dims)) + stopRot := make([]int32, len(dims)) + stopPass := make([]int32, len(dims)) + startPass := make([]int32, len(dims)) + for i, dim := range dims { + stopRot[i] = int32(dim) + stopPass[i] = int32(dim) + } + stopRot[len(dims)-1] = int32(ropeDim) + startPass[len(dims)-1] = int32(ropeDim) + rot := mlx.MulScalar(mlx.SliceStartStop(x, start, stopRot), scale) + pass := mlx.SliceStartStop(x, startPass, stopPass) + return mlx.Concatenate([]*mlx.Array{rot, pass}, -1) +} diff --git a/x/models/qwen3_5/qwen3_5.go b/x/models/qwen3_5/qwen3_5.go index 47c02ed4..b8f0b3e8 100644 --- a/x/models/qwen3_5/qwen3_5.go +++ b/x/models/qwen3_5/qwen3_5.go @@ -1160,15 +1160,15 @@ func (g *GatedDeltaNet) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, B, }, -1) } convTail := cfg.LinearConvKernelDim - 1 - var rc *cache.RecurrentCache + var rc cache.Recurrent var rec nn.RecurrentOption - if typed, ok := c.(*cache.RecurrentCache); ok { + if typed, ok := c.(cache.Recurrent); ok { rc = typed rec = nn.WithRecurrentHistory(rc.Get(b, x.DType())) } else { rec = nn.WithRecurrentState( mlx.Zeros(x.DType(), int(B), int(convTail), qkv.Dim(2)), - mlx.Zeros(x.DType(), int(B), int(cfg.LinearNumValueHeads), int(cfg.LinearValueHeadDim), int(cfg.LinearKeyHeadDim)), + mlx.Zeros(mlx.DTypeFloat32, int(B), int(cfg.LinearNumValueHeads), int(cfg.LinearValueHeadDim), int(cfg.LinearKeyHeadDim)), ) } @@ -1193,6 +1193,9 @@ func (g *GatedDeltaNet) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, B, gDecay = gDecay.AsType(alpha.DType()) betaGate := mlx.Sigmoid(beta) + if recorder, ok := rc.(cache.RecurrentRecorder); ok { + recorder.Record(qkv, q, k, v, gDecay, betaGate) + } out, state := nn.GatedDelta(b, q, k, v, gDecay, betaGate, rec) outDType := out.DType() @@ -1300,26 +1303,48 @@ func (l *Layer) Forward(x *mlx.Array, b *batch.Batch, c cache.Cache, positions * } func (m *Model) Forward(b *batch.Batch, caches []cache.Cache) *mlx.Array { + out, _ := m.forward(b, caches, nil) + return out +} + +func (m *Model) ForwardDFlash(b *batch.Batch, caches []cache.Cache, layerIDs []int) (hidden, targetHidden *mlx.Array) { + return m.forward(b, caches, layerIDs) +} + +func (m *Model) forward(b *batch.Batch, caches []cache.Cache, captureLayerIDs []int) (*mlx.Array, *mlx.Array) { dims := b.InputIDs.Dims() B, L := int32(dims[0]), int32(dims[1]) positions := mlx.FromValues(b.SeqOffsets, len(b.SeqOffsets)) h := m.EmbedTokens.Forward(b.InputIDs) + captured := make([]*mlx.Array, 0, len(captureLayerIDs)) + nextCapture := 0 for i, layer := range m.Layers { var c cache.Cache if caches != nil && i < len(caches) { c = caches[i] } h = layer.Forward(h, b, c, positions, B, L, m.Config) + if nextCapture < len(captureLayerIDs) && i == captureLayerIDs[nextCapture] { + captured = append(captured, h) + nextCapture++ + } } out := m.Norm.Forward(h, m.RMSNormEps) - return out + if len(captured) == 0 { + return out, nil + } + return out, mlx.Concatenate(captured, -1) } func (m *Model) Unembed(x *mlx.Array) *mlx.Array { return m.LMHead.Forward(x) } +func (m *Model) TokenEmbeddings(inputIDs *mlx.Array) *mlx.Array { + return m.EmbedTokens.Forward(inputIDs) +} + func (m *Model) NumLayers() int { return len(m.Layers) }