feat(mlx): draft_num_predict per request on MLX: 0 turns drafting off, N caps the depth, no reload

docs/maxusai/vision-suite/README.md

@@ -77,12 +77,17 @@ first is image accounting and the second is generation length.
   `IMAGE_MIN_TOKENS` / `IMAGE_MAX_TOKENS` (fork-only per-request vision budget,
   arch-gated to gemma4 and nemotron_h_omni; unset = build default. Recorded in the
   scores as `req_image_*_tokens` so a control run is identifiable after the fact),
-  `DRAFT_NUM_PREDICT` (MTP / speculative draft depth, `--spec-draft-n-max` on the
-  llama-server side; recorded as `req_draft_num_predict`. **Set it to `0` for an
-  off-arm rather than leaving it unset** — `server/routes.go` zeroes the option
-  unless it was explicitly set, so "unset" and "disabled" are the same request but
-  only one of them is self-describing in the scores. Inert on the MLX runner,
-  which never reads `draft_num_predict` and picks depth adaptively instead),
+  `DRAFT_NUM_PREDICT` (MTP / speculative draft depth; recorded as
+  `req_draft_num_predict`). On llama-server it is `--spec-draft-n-max`, a load-time flag:
+  a change relaunches the runner, `0` turns drafting off, and `server/routes.go` zeroes an
+  unset value when there is no separate draft file. On the MLX runner, since #465, it is
+  read per request with no reload: `0` turns drafting off for that request, `N` caps the
+  adaptive depth, and **unset stays adaptive** (the drafter ships inside MLX tags, so unset
+  maps to −1, not 0). **On MLX, unset and `0` are different arms.** `0` is an off-arm on
+  both runners, but exporting `DRAFT_NUM_PREDICT=0` for a mixed GGUF+MLX run now undrafts
+  the MLX cells too. Builds before #465 ignore it on MLX, so compare
+  `req_draft_num_predict` only within one build. A depth-0 request does not advance the
+  runner's per-load warm-up: in an A/B, send unset requests first),
   `PROMPT_CACHE_RAM` (llama-server's host-RAM prompt cache, ADR 0047. **Default `0`, the
   cache off,** sent on every request since 2026-10-03, because a benchmark's image requests
   pay for the cache's update and get nothing back. A size in MiB or `-1` is sent as given,

docs/maxusai/vision-suite/client.py

@@ -229,8 +229,10 @@ def generate(host, model, prompt, images, num_predict=None, num_ctx=None,
     # routes.go zeroes the option unless it was explicitly set, so "0" is the
     # only way to state "MTP off" rather than "unspecified", and an arm that
     # omitted it could not be distinguished from one that disabled it. Also a
-    # Runner option -- changing it reloads the model. Inert on the MLX runner,
-    # which never reads draft_num_predict and picks depth adaptively instead.
+    # Runner option on llama-server -- changing it reloads the model. On the MLX
+    # runner (since #465) it is read per request, no reload: 0 = no drafting for
+    # this request, N caps the adaptive depth, unset stays adaptive. So on MLX
+    # "unset" and "0" are different arms, and before #465 MLX ignored it.
     if use_env_opts and os.environ.get("DRAFT_NUM_PREDICT") not in (None, ""):
         opts["draft_num_predict"] = int(os.environ["DRAFT_NUM_PREDICT"])
     # llama-server's prompt cache, OFF unless PROMPT_CACHE_RAM says otherwise

mlxrunner/client.go

@@ -163,6 +163,21 @@ type CompletionRequest struct {
 	Logprobs                   bool
 	TopLogprobs                int
 	IncludeIntermediateMetrics bool
+	// DraftLimit is the request's draft_num_predict (requestDraftLimit): nil
+	// keeps the adaptive draft depth, 0 parks the drafter for this request,
+	// N caps the depth at N. A zero-valued request is nil, so drafts as always.
+	DraftLimit *int
+}
+
+// requestDraftLimit maps draft_num_predict onto DraftLimit. The server sends
+// -1 when neither the request nor the Modelfile set it (an MLX model's
+// drafter is part of the model, so unset is not "off").
+func requestDraftLimit(opts api.Options) *int {
+	if opts.DraftNumPredict < 0 {
+		return nil
+	}
+	n := opts.DraftNumPredict
+	return &n
 }
 
 type CompletionResponse struct {
@@ -269,6 +284,7 @@ func (c *Client) Completion(ctx context.Context, req llm.CompletionRequest, fn f
 	}
 	if req.Options != nil {
 		creq.Options = *req.Options
+		creq.DraftLimit = requestDraftLimit(creq.Options)
 	}
 
 	body, err := json.Marshal(creq)

mlxrunner/speculate.go

@@ -98,10 +98,14 @@ func newSpeculation(r *Runner, draft model.DraftModel, targets, draftKV []cache.
 type speculationSession struct {
 	spec    *speculation
 	drafter draftSession
-	enabled bool  // whether this request drafts; false parks (maintain-only)
-	limit   int   // current draft length
-	layout  []any // the request's per-row layout state, stamped on every target forward
-	stats   specStats
+	enabled bool // whether this request drafts; false parks (maintain-only)
+	limit   int  // current draft length
+	// maxLimit is the request's draft_num_predict, the deepest it drafts; -1
+	// leaves the depth to the controller. The controller keeps learning from
+	// a capped request's rounds and schedules for the next request uncapped.
+	maxLimit int
+	layout   []any // the request's per-row layout state, stamped on every target forward
+	stats    specStats
 
 	// Cost sampling: each round's wall time (start to next start, spanning the
 	// next emit's sync) is attributed to its draft depth only when the depth
@@ -150,6 +154,10 @@ func draftingEnabled(request Request) bool {
 		// Logprobs are not yet supported on the speculative path.
 		return false
 	}
+	if request.DraftLimit != nil && *request.DraftLimit == 0 {
+		// draft_num_predict 0: this request does not draft.
+		return false
+	}
 	return request.Grammar == nil || draftUnderGrammar
 }
 
@@ -166,9 +174,13 @@ func (s *speculation) open(request Request, layout []any) *speculationSession {
 	// turned OLLAMA_MLX_DRAFT_UNDER_GRAMMAR off.
 	enabled := draftingEnabled(request)
 
-	spec := &speculationSession{spec: s, drafter: d, layout: layout, enabled: enabled, prevDrafts: -1, roundDrafts: -1}
+	spec := &speculationSession{spec: s, drafter: d, layout: layout, enabled: enabled, prevDrafts: -1, roundDrafts: -1,
+		maxLimit: -1}
+	if request.DraftLimit != nil {
+		spec.maxLimit = *request.DraftLimit
+	}
 	if enabled {
-		spec.limit = s.depth.scheduled
+		spec.limit = spec.capped(s.depth.scheduled)
 	}
 	return spec
 }
@@ -206,8 +218,16 @@ func (s *speculationSession) endRound(drafted, accepted, observed int) {
 		if observed > 0 {
 			s.spec.depth.acc.observe(observed, accepted)
 		}
-		s.limit = s.spec.depth.next()
+		s.limit = s.capped(s.spec.depth.next())
+	}
+}
+
+// capped holds depth to the request's draft_num_predict, if it set one.
+func (s *speculationSession) capped(depth int) int {
+	if s.maxLimit >= 0 {
+		return min(depth, s.maxLimit)
 	}
+	return depth
 }
 
 func (s *speculationSession) committed(tokens, hiddens *mlx.Array, position int, media []batch.MediaItem) {

mlxrunner/speculate_draftlimit_test.go

@@ -0,0 +1,102 @@
+package mlxrunner
+
+import (
+	"testing"
+	"time"
+
+	"github.com/ollama/ollama/api"
+)
+
+// draft_num_predict on MLX (the fork): the request's value reaches the runner as
+// CompletionRequest.DraftLimit. nil keeps the adaptive depth the runner always
+// chose; 0 parks the drafter for that request; N caps the depth at N. None of
+// these reload the model -- the drafter stays loaded either way -- so one loaded
+// drafting tag measures drafted and undrafted generation back to back.
+
+func TestRequestDraftLimit(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		opt  int
+		want *int
+	}{
+		// The server sends -1 for an MLX model whose request and Modelfile
+		// leave draft_num_predict unset (modelOptionsWithEmbeddingBatchDefault).
+		{"unset keeps the adaptive depth", -1, nil},
+		{"0 turns drafting off", 0, new(0)},
+		{"N caps the depth", 6, new(6)},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			got := requestDraftLimit(api.Options{Runner: api.Runner{DraftNumPredict: tc.opt}})
+			if (got == nil) != (tc.want == nil) || (got != nil && *got != *tc.want) {
+				t.Fatalf("requestDraftLimit(%d) = %v, want %v", tc.opt, deref(got), deref(tc.want))
+			}
+		})
+	}
+}
+
+func deref(p *int) any {
+	if p == nil {
+		return nil
+	}
+	return *p
+}
+
+func TestDraftLimitGatesTheRequest(t *testing.T) {
+	for _, tc := range []struct {
+		name  string
+		limit *int
+		want  bool
+	}{
+		{"no limit drafts", nil, true},
+		{"limit 0 parks", new(0), false},
+		{"a positive limit drafts", new(3), true},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			r := Request{}
+			r.DraftLimit = tc.limit
+			if got := draftingEnabled(r); got != tc.want {
+				t.Fatalf("draftingEnabled() = %v, want %v", got, tc.want)
+			}
+		})
+	}
+}
+
+func TestDraftLimitCapsTheDepth(t *testing.T) {
+	s := &speculation{drafter: grammarTestDrafter{}, depth: newDepthController()}
+	s.depth.scheduled = 9
+	// Every depth costs the same, so with full acceptance deeper always pays
+	// and the controller climbs as fast as its acceptance frontier allows.
+	for n := 0; n <= 16; n++ {
+		s.depth.cost.observe(n, time.Millisecond)
+	}
+
+	adaptive := s.open(Request{}, nil)
+	defer adaptive.close()
+	if adaptive.limit != 9 {
+		t.Fatalf("unlimited request opened at depth %d, want the scheduled 9", adaptive.limit)
+	}
+
+	r := Request{}
+	r.DraftLimit = new(2)
+	capped := s.open(r, nil)
+	defer capped.close()
+	if capped.limit != 2 {
+		t.Fatalf("limited request opened at depth %d, want its limit 2", capped.limit)
+	}
+
+	// Every round accepts in full, so the controller keeps reaching deeper; the
+	// capped session must never follow it past its limit, and the uncapped one
+	// must (or this test could not fail).
+	deeper := false
+	for range 200 {
+		capped.endRound(capped.limit, capped.limit, capped.limit)
+		if capped.limit > 2 {
+			t.Fatalf("limited request's depth rose to %d, past its limit 2", capped.limit)
+		}
+		adaptive.endRound(adaptive.limit, adaptive.limit, adaptive.limit)
+		deeper = deeper || adaptive.limit > 2
+	}
+	if !deeper {
+		t.Fatal("the unlimited session never went past depth 2; the cap is untested")
+	}
+}

server/routes.go

@@ -158,6 +158,14 @@ func (s *Server) modelOptionsWithEmbeddingBatchDefault(model *Model, requestOpts
 
 	if model != nil && model.DraftPath == "" && !draftNumPredictSet {
 		opts.DraftNumPredict = 0
+		// An MLX drafter ships inside the model and the runner picks its depth
+		// adaptively, so unset must stay distinguishable from a request's 0,
+		// which turns drafting off: -1 reaches the runner as "no limit"
+		// (mlxrunner.requestDraftLimit). The MLX scheduler reloads only on
+		// num_ctx, so no value here reloads the model.
+		if model.IsMLX() {
+			opts.DraftNumPredict = -1
+		}
 	}
 
 	return opts, nil

server/routes_options_test.go

@@ -272,6 +272,30 @@ func TestModelOptionsDraftNumPredictDefault(t *testing.T) {
 			requestOpts: map[string]any{"draft_num_predict": float64(0)},
 			want:        0,
 		},
+		// MLX drafters ship inside the model and pick their depth adaptively, so
+		// an unset value must stay distinguishable from a request's 0 (off).
+		{
+			name:  "MLX: unset leaves the runner's adaptive depth",
+			model: &Model{Config: model.ConfigV2{ModelFormat: "safetensors"}},
+			want:  -1,
+		},
+		{
+			name:        "MLX: request 0 turns drafting off",
+			model:       &Model{Config: model.ConfigV2{ModelFormat: "safetensors"}},
+			requestOpts: map[string]any{"draft_num_predict": float64(0)},
+			want:        0,
+		},
+		{
+			name:        "MLX: request caps the depth",
+			model:       &Model{Config: model.ConfigV2{ModelFormat: "safetensors"}},
+			requestOpts: map[string]any{"draft_num_predict": float64(6)},
+			want:        6,
+		},
+		{
+			name:  "MLX: a Modelfile value is a cap too",
+			model: &Model{Config: model.ConfigV2{ModelFormat: "safetensors"}, Options: map[string]any{"draft_num_predict": float64(15)}},
+			want:  15,
+		},
 	}
 
 	for _, tt := range tests {