Revert "mlxrunner: add DFlash speculative decoding (#16134)"

This reverts commit 98e26b8c37.

The DFlash integration is too invasive to keep at this stage: it
threads DFlash-specific logic through the pipeline, base model
interfaces, and the cache layer. The recurrent cache also now
has qwen3.5 model-specific code. Revert it now and reintroduce
the self-contained, generally-useful pieces (YaRN RoPE DRY-out, draft
architecture autodetection, gated-delta fp32 state) as separate
follow-up commits.

x/create/client/create.go

@@ -243,40 +243,6 @@ 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)
@@ -521,11 +487,12 @@ 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 != "" {
-			draft, err := draftMetadata(opts.Modelfile.Draft)
-			if err != nil {
-				return err
+			configData.Draft = &model.Draft{
+				ModelFormat:  "safetensors",
+				Architecture: "Gemma4AssistantForCausalLM",
+				TensorPrefix: "draft.",
+				Config:       "draft/config.json",
 			}
-			configData.Draft = draft
 		}
 		configJSON, err := json.Marshal(configData)
 		if err != nil {

x/create/client/create_test.go

@@ -544,15 +544,10 @@ 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: draftDir},
+		Modelfile: &ModelfileConfig{Draft: "/tmp/assistant"},
 	}
 
 	writer := newManifestWriter(opts, []string{"completion"}, "gemma4", "gemma4")
@@ -586,9 +581,6 @@ 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) {

x/mlxrunner/cache.go

@@ -38,8 +38,6 @@ 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
@@ -63,22 +61,18 @@ type cacheSession struct {
 	pendingSnapshots []pendingSnapshot
 }
 
-func (c *kvCache) ensureCachesWithFactory(newCaches cacheFactoryFunc) {
+func (c *kvCache) ensureCaches(m base.Model) {
 	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 {
-		return cacheFactory.NewCaches()
+		c.caches = cacheFactory.NewCaches()
+		return
 	}
-	caches := make([]cache.Cache, m.NumLayers())
-	for i := range caches {
-		caches[i] = cache.NewKVCache()
+	c.caches = make([]cache.Cache, m.NumLayers())
+	for i := range c.caches {
+		c.caches[i] = cache.NewKVCache()
 	}
-	return caches
 }
 
 func (c *kvCache) ensureRoot() {
@@ -93,29 +87,15 @@ 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 {
-	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.ensureCaches(m)
 	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 keepSeedToken && matched == len(inputs) && matched > 0 {
+	if matched == len(inputs) && matched > 0 {
 		matchPath, matched = findBestMatch(c.root, inputs[:len(inputs)-1])
 	}
 
@@ -143,9 +123,6 @@ func (c *kvCache) beginWithFactoryLimit(inputs []int32, newCaches cacheFactoryFu
 	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

x/mlxrunner/cache/cache.go

@@ -88,10 +88,6 @@ 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
 		}
@@ -112,8 +108,6 @@ 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
 		}

x/mlxrunner/cache/recurrent.go

@@ -6,19 +6,6 @@ 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]
@@ -66,12 +53,9 @@ 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() != deltaDType ||
+	needDelta := c.deltaState == nil || !c.deltaState.Valid() || c.deltaState.DType() != dtype ||
 		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
@@ -81,7 +65,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(deltaDType, batch, c.numVHeads, c.headVDim, c.headKDim), false)
+		c.deltaState = c.setState(c.deltaState, mlx.Zeros(dtype, batch, c.numVHeads, c.headVDim, c.headKDim), false)
 	}
 }
 
@@ -177,109 +161,3 @@ 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")
-	}
-}

x/mlxrunner/cache/recurrent_test.go

@@ -63,55 +63,8 @@ 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.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)
+	if got := h.DeltaState().DType(); got != mlx.DTypeBFloat16 {
+		t.Fatalf("delta state dtype = %v, want %v", got, mlx.DTypeBFloat16)
 	}
 }
 

x/mlxrunner/cache_test.go

@@ -159,50 +159,6 @@ 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.

x/mlxrunner/dflash.go

@@ -1,715 +0,0 @@
-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
-}

x/mlxrunner/imports.go

@@ -1,7 +1,6 @@
 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"

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<StT>(state[i]);
+  o_state[s_idx] = static_cast<InT>(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<StT>(state[i]);
+  o_state[s_idx] = static_cast<InT>(state[i]);
 }
 `
 
@@ -262,9 +262,8 @@ func gatedDeltaKernel(q, k, v, g, beta, state *Array) (y, nextState *Array, ok b
 		return nil, nil, false
 	}
 
-	inputDType := q.DType()
-	stateDType := state.DType()
-	if k.DType() != inputDType || v.DType() != inputDType || g.DType() != inputDType || beta.DType() != inputDType {
+	dtype := q.DType()
+	if k.DType() != dtype || v.DType() != dtype || g.DType() != dtype || beta.DType() != dtype || state.DType() != dtype {
 		return nil, nil, false
 	}
 
@@ -278,13 +277,7 @@ 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(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 {
+	if C.mlx_fast_metal_kernel_config_add_template_arg_dtype(cfg, cInT, C.mlx_dtype(dtype)) != 0 {
 		gatedDeltaMetalDisabled = true
 		return nil, nil, false
 	}
@@ -308,11 +301,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(inputDType)) != 0 {
+	if C.mlx_fast_metal_kernel_config_add_output_arg(cfg, unsafe.SliceData(yShape), C.size_t(len(yShape)), C.mlx_dtype(dtype)) != 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(stateDType)) != 0 {
+	if C.mlx_fast_metal_kernel_config_add_output_arg(cfg, unsafe.SliceData(stateShape), C.size_t(len(stateShape)), C.mlx_dtype(dtype)) != 0 {
 		gatedDeltaMetalDisabled = true
 		return nil, nil, false
 	}
@@ -523,9 +516,8 @@ func gatedDeltaCUDAKernelApply(q, k, v, g, beta, state *Array) (y, nextState *Ar
 		return nil, nil, false
 	}
 
-	inputDType := q.DType()
-	stateDType := state.DType()
-	if k.DType() != inputDType || v.DType() != inputDType || g.DType() != inputDType || beta.DType() != inputDType {
+	dtype := q.DType()
+	if k.DType() != dtype || v.DType() != dtype || g.DType() != dtype || beta.DType() != dtype || state.DType() != dtype {
 		return nil, nil, false
 	}
 
@@ -539,13 +531,7 @@ 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(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 {
+	if C.mlx_fast_cuda_kernel_config_add_template_arg_dtype(cfg, cInT, C.mlx_dtype(dtype)) != 0 {
 		gatedDeltaCUDADisabled = true
 		return nil, nil, false
 	}
@@ -569,11 +555,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(inputDType)) != 0 {
+	if C.mlx_fast_cuda_kernel_config_add_output_arg(cfg, unsafe.SliceData(yShape), C.size_t(len(yShape)), C.mlx_dtype(dtype)) != 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(stateDType)) != 0 {
+	if C.mlx_fast_cuda_kernel_config_add_output_arg(cfg, unsafe.SliceData(stateShape), C.size_t(len(stateShape)), C.mlx_dtype(dtype)) != 0 {
 		gatedDeltaCUDADisabled = true
 		return nil, nil, false
 	}

x/mlxrunner/model/base/base.go

@@ -57,23 +57,6 @@ 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))

x/mlxrunner/pipeline.go

@@ -12,9 +12,7 @@ import (
 	"github.com/ollama/ollama/llm"
 	"github.com/ollama/ollama/logutil"
 	"github.com/ollama/ollama/x/mlxrunner/batch"
-	"github.com/ollama/ollama/x/mlxrunner/cache"
 	"github.com/ollama/ollama/x/mlxrunner/mlx"
-	"github.com/ollama/ollama/x/mlxrunner/model/base"
 	sampler "github.com/ollama/ollama/x/mlxrunner/sample"
 	"github.com/ollama/ollama/x/tokenizer"
 )
@@ -81,90 +79,24 @@ func (r *Runner) TextGenerationPipeline(ctx context.Context, request Request) er
 	caches := session.caches
 	tokens := session.remaining
 	prefillChunk := prefillChunkSize()
-	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,
-		)
+
+	// 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)
 	}
 
-	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
+	const preThinking = 4
+	if end := len(inputs) - preThinking; end > 0 {
+		session.requestSnapshot(end)
 	}
 
-	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}
-		}
+	materializeCaches := func() {
 		state := make([]*mlx.Array, 0, 2*len(caches))
-		for _, set := range cacheSets {
-			for _, c := range set {
-				if c == nil {
-					continue
-				}
-				state = append(state, c.State()...)
-			}
+		for _, c := range caches {
+			state = append(state, c.State()...)
 		}
 		if len(state) == 0 {
 			return
@@ -172,61 +104,6 @@ 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)
@@ -239,49 +116,37 @@ 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 := nextSnapshotOffset(); snapOffset > 0 {
+		if snapOffset := session.nextPendingSnapshot(); snapOffset > 0 {
 			tokensUntilSnapshot := snapOffset - position
 			if tokensUntilSnapshot > 0 && tokensUntilSnapshot < n {
 				n = tokensUntilSnapshot
 			}
 		}
 
-		b := &batch.Batch{
+		r.Model.Forward(&batch.Batch{
 			InputIDs:     mlx.FromValues(tokens[processed:processed+n], 1, n),
 			SeqOffsets:   []int32{int32(position)},
 			SeqQueryLens: []int32{int32(n)},
-		}
-		if dflashEnabled {
-			_, targetHidden := dflashTarget.ForwardDFlash(b, caches, dflashDraft.TargetLayerIDs())
-			dflashDraft.AppendContext(targetHidden, dflashCaches)
-		} else {
-			r.Model.Forward(b, caches)
-		}
+		}, caches)
 		mlx.Sweep()
-		if dflashEnabled {
-			materializeCaches(caches, dflashCaches)
-		} else {
-			materializeCaches()
-		}
+		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.
-		snapshotReadySessions(position)
+		if snapOffset := session.nextPendingSnapshot(); snapOffset > 0 {
+			if position >= snapOffset {
+				session.snapshot()
+			}
+		}
 
 		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)
 	}

x/mlxrunner/runner.go

@@ -40,7 +40,6 @@ type Runner struct {
 	Requests      chan Request
 	Sampler       *sample.Sampler
 	cache         kvCache
-	dflashCache   kvCache
 	contextLength int
 	mlxThread     *mlxthread.Thread
 }

x/models/dflash/dflash.go

@@ -1,439 +0,0 @@
-// 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)))
-}

x/models/dflash/dflash_test.go

@@ -1,52 +0,0 @@
-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)
-	}
-}

x/models/laguna/laguna.go

@@ -22,45 +22,57 @@ 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    *nn.RopeParameters
-	full    *nn.RopeParameters
-	sliding *nn.RopeParameters
+	flat    *RopeParameters
+	full    *RopeParameters
+	sliding *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              *nn.RopeParameters `json:"rope_parameters"`
-	RopeScaling                 *nn.RopeParameters `json:"rope_scaling"`
-	SWARopeParameters           *nn.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              *RopeParameters `json:"rope_parameters"`
+	RopeScaling                 *RopeParameters `json:"rope_scaling"`
+	SWARopeParameters           *RopeParameters `json:"swa_rope_parameters"`
 
 	QuantGroupSize int                               `json:"-"`
 	QuantBits      int                               `json:"-"`
@@ -158,36 +170,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                 *nn.RopeParameters `json:"rope_scaling"`
-		SWARopeParameters           *nn.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                 *RopeParameters `json:"rope_scaling"`
+		SWARopeParameters           *RopeParameters `json:"swa_rope_parameters"`
 	}
 
 	var raw rawConfig
@@ -301,8 +313,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.TypeName(), "yarn") {
-		cfg.FullRopeFreqs, cfg.FullRopeScale = nn.BuildYarnRopeFreqs(cfg.FullRopeDim, cfg.FullRopeBase, ropeParams)
+	if ropeParams != nil && strings.EqualFold(ropeParams.ropeType(), "yarn") {
+		cfg.FullRopeFreqs, cfg.FullRopeScale = buildYarnRopeFreqs(cfg.FullRopeDim, cfg.FullRopeBase, ropeParams)
 	}
 
 	cfg.SlidingRopeBase = cfg.FullRopeBase
@@ -366,12 +378,12 @@ func (r *ropeConfig) UnmarshalJSON(b []byte) error {
 
 	if raw, ok := probe["full_attention"]; ok {
 		r.nested = true
-		r.full = &nn.RopeParameters{}
+		r.full = &RopeParameters{}
 		if err := json.Unmarshal(raw, r.full); err != nil {
 			return err
 		}
 		if raw = probe["sliding_attention"]; raw != nil {
-			r.sliding = &nn.RopeParameters{}
+			r.sliding = &RopeParameters{}
 			if err := json.Unmarshal(raw, r.sliding); err != nil {
 				return err
 			}
@@ -381,12 +393,12 @@ func (r *ropeConfig) UnmarshalJSON(b []byte) error {
 
 	if raw, ok := probe["global_attention"]; ok {
 		r.nested = true
-		r.full = &nn.RopeParameters{}
+		r.full = &RopeParameters{}
 		if err := json.Unmarshal(raw, r.full); err != nil {
 			return err
 		}
 		if raw = probe["sliding_attention"]; raw != nil {
-			r.sliding = &nn.RopeParameters{}
+			r.sliding = &RopeParameters{}
 			if err := json.Unmarshal(raw, r.sliding); err != nil {
 				return err
 			}
@@ -394,18 +406,18 @@ func (r *ropeConfig) UnmarshalJSON(b []byte) error {
 		return nil
 	}
 
-	r.flat = &nn.RopeParameters{}
+	r.flat = &RopeParameters{}
 	return json.Unmarshal(b, r.flat)
 }
 
-func (r ropeConfig) fullParams() *nn.RopeParameters {
+func (r ropeConfig) fullParams() *RopeParameters {
 	if r.nested {
 		return r.full
 	}
 	return r.flat
 }
 
-func (r ropeConfig) slidingParams() *nn.RopeParameters {
+func (r ropeConfig) slidingParams() *RopeParameters {
 	if !r.nested {
 		return nil
 	}
@@ -440,6 +452,88 @@ 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
@@ -922,8 +1016,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 = 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)
+	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)
 
 	var kv nn.SDPAOption
 	if c != nil {
@@ -940,6 +1034,30 @@ 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)))
 }

x/models/nn/rope.go

@@ -1,129 +0,0 @@
-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)
-}

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.Recurrent
+	var rc *cache.RecurrentCache
 	var rec nn.RecurrentOption
-	if typed, ok := c.(cache.Recurrent); ok {
+	if typed, ok := c.(*cache.RecurrentCache); 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(mlx.DTypeFloat32, int(B), int(cfg.LinearNumValueHeads), int(cfg.LinearValueHeadDim), int(cfg.LinearKeyHeadDim)),
+			mlx.Zeros(x.DType(), int(B), int(cfg.LinearNumValueHeads), int(cfg.LinearValueHeadDim), int(cfg.LinearKeyHeadDim)),
 		)
 	}
 
@@ -1193,9 +1193,6 @@ 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()
@@ -1303,48 +1300,26 @@ 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)
-	if len(captured) == 0 {
-		return out, nil
-	}
-	return out, mlx.Concatenate(captured, -1)
+	return out
 }
 
 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)
 }