mlxrunner: restore System One scoring

Carry the MLX implementation and its tests from the original System One
commit. The API/compiler and llama-server implementation already exist on
the release branch, so retain those versions. Adapt the optional model test
to the renamed decision package and its explicit rendering/context setup.

(cherry picked from commit 26e546ddb3)

mlxrunner/client_score.go

@@ -0,0 +1,46 @@
+package mlxrunner
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+
+	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/llm"
+)
+
+var _ llm.Scorer = (*Client)(nil)
+
+func (c *Client) Score(ctx context.Context, input llm.ScoreRequest) (llm.ScoreResponse, error) {
+	var result llm.ScoreResponse
+	data, err := json.Marshal(input)
+	if err != nil {
+		return result, err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/v1/score", c.port), bytes.NewReader(data))
+	if err != nil {
+		return result, err
+	}
+	req.Header.Set("Content-Type", "application/json")
+	res, err := c.client.Do(req)
+	if err != nil {
+		if errMsg := c.status.LastError(); errMsg != "" {
+			return result, fmt.Errorf("mlx runner failed: %s", errMsg)
+		}
+		return result, err
+	}
+	defer res.Body.Close()
+	if res.StatusCode != http.StatusOK {
+		body, err := io.ReadAll(res.Body)
+		if err != nil {
+			return result, err
+		}
+		return result, api.StatusError{StatusCode: res.StatusCode, ErrorMessage: strings.TrimSpace(string(body))}
+	}
+	err = json.NewDecoder(res.Body).Decode(&result)
+	return result, err
+}

mlxrunner/client_score_test.go

@@ -0,0 +1,77 @@
+package mlxrunner
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"io"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"reflect"
+	"strings"
+	"testing"
+
+	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/llm"
+)
+
+func TestClientScore(t *testing.T) {
+	input := llm.ScoreRequest{MaxTokens: 2048, Rows: []llm.ScoreRow{{Prompt: "prompt", Candidates: []string{"A", "B"}}}}
+	for _, status := range []int{http.StatusOK, http.StatusBadRequest} {
+		srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if r.URL.Path != "/v1/score" || r.Method != "POST" {
+				t.Errorf("wrong endpoint: %s %s", r.Method, r.URL.Path)
+			}
+			var got llm.ScoreRequest
+			if err := json.NewDecoder(r.Body).Decode(&got); err != nil || !reflect.DeepEqual(got, input) {
+				t.Errorf("request changed: %+v, %v", got, err)
+			}
+			if status != http.StatusOK {
+				http.Error(w, "prompt too long", status)
+				return
+			}
+			_ = json.NewEncoder(w).Encode(llm.ScoreResponse{Logits: [][]float32{{-1, 2}}, InputTokens: 3})
+		}))
+		client := &Client{port: srv.Listener.Addr().(*net.TCPAddr).Port, client: srv.Client()}
+		result, err := client.Score(context.Background(), input)
+		srv.Close()
+		if status == http.StatusOK {
+			if err != nil || result.InputTokens != 3 || !reflect.DeepEqual(result.Logits, [][]float32{{-1, 2}}) {
+				t.Fatalf("unexpected result: %+v, %v", result, err)
+			}
+		} else {
+			var statusErr api.StatusError
+			if !errors.As(err, &statusErr) || statusErr.StatusCode != status || statusErr.ErrorMessage != "prompt too long" {
+				t.Fatalf("status error not preserved: %v", err)
+			}
+		}
+	}
+}
+
+func TestClientScoreTransportError(t *testing.T) {
+	for _, message := range []string{"", "MLX: failed to allocate memory"} {
+		t.Run(message, func(t *testing.T) {
+			srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				conn, _, err := w.(http.Hijacker).Hijack()
+				if err != nil {
+					t.Error(err)
+					return
+				}
+				conn.Close()
+			}))
+			defer srv.Close()
+			status := &llm.StatusWriter{}
+			status.SetLastError(message)
+			client := &Client{port: srv.Listener.Addr().(*net.TCPAddr).Port, client: srv.Client(), status: status}
+			_, err := client.Score(t.Context(), llm.ScoreRequest{})
+			if message == "" {
+				if !errors.Is(err, io.EOF) {
+					t.Fatalf("transport error lost: %v", err)
+				}
+			} else if !llm.IsOutOfMemory(err) || !strings.Contains(err.Error(), message) {
+				t.Fatalf("runner OOM diagnostic lost: %v", err)
+			}
+		})
+	}
+}

mlxrunner/model/qwen3_5/score.go

@@ -0,0 +1,23 @@
+package qwen3_5
+
+import (
+	"github.com/ollama/ollama/mlx"
+	"github.com/ollama/ollama/mlxrunner/nn"
+)
+
+// UnembedCandidates scores one hidden position. For dense heads, gather before
+// casting so FP32 accumulation neither rounds logits to BF16 nor allocates an
+// FP32 copy of the full vocabulary head. This also covers tied dense embeddings.
+func (m *Model) UnembedCandidates(hidden, candidates *mlx.Array) *mlx.Array {
+	if head, ok := m.LMHead.(*nn.Linear); ok {
+		weight := head.Weight.TakeAxis(candidates, 0).AsType(mlx.DTypeFloat32)
+		logits := hidden.AsType(mlx.DTypeFloat32).Matmul(weight.Transpose(1, 0)).Reshape(-1)
+		if head.Bias != nil {
+			logits = logits.Add(head.Bias.TakeAxis(candidates, 0).AsType(mlx.DTypeFloat32))
+		}
+		return logits
+	}
+	// Retain existing quantization, bias, and global-scale handling. Selected-row
+	// dequantization can be added separately when its formats are validated.
+	return m.Unembed(hidden).Reshape(-1).TakeAxis(candidates, 0)
+}

mlxrunner/model/qwen3_5/score_test.go

@@ -0,0 +1,44 @@
+package qwen3_5
+
+import (
+	"math"
+	"testing"
+
+	"github.com/ollama/ollama/mlx"
+	"github.com/ollama/ollama/mlx/mlxtest"
+	"github.com/ollama/ollama/mlxrunner/nn"
+)
+
+func TestUnembedCandidates(t *testing.T) {
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		mlx.Scoped(func() {
+			values := make([]float32, 7*8)
+			for i := range values {
+				values[i] = float32(math.Sin(float64(i)))
+			}
+			weight := mlx.FromValues(values, 7, 8).AsType(mlx.DTypeBFloat16)
+			hidden := mlx.FromValues([]float32{0.7, 1.3, -2.9, 3.1, 0.06, 1.01, -5.03, 2.17}, 1, 1, 8).AsType(mlx.DTypeBFloat16)
+			ids := mlx.FromValues([]int32{6, 2, 0}, 3)
+			bias := mlx.FromValues([]float32{0.1, 1.2, 2.3, 3.4, 4.5, 5.6, 6.7}, 7).AsType(mlx.DTypeBFloat16)
+			for _, b := range []*mlx.Array{nil, bias} {
+				m := &Model{LMHead: nn.NewLinear(weight, b)}
+				got := m.UnembedCandidates(hidden, ids)
+				full := hidden.AsType(mlx.DTypeFloat32).Matmul(weight.AsType(mlx.DTypeFloat32).Transpose(1, 0)).Reshape(-1)
+				if b != nil {
+					full = full.Add(b.AsType(mlx.DTypeFloat32))
+				}
+				want := full.TakeAxis(ids, 0)
+				mlx.Eval(got, want)
+				if got.DType() != mlx.DTypeFloat32 || got.Dim(0) != 3 {
+					t.Fatal("candidate logits must retain FP32 precision")
+				}
+				actual, expected := got.Floats(), want.Floats()
+				for i, v := range actual {
+					if math.Abs(float64(v-expected[i])) > 1e-6 {
+						t.Fatalf("candidate %d: selected=%g full=%g", i, v, expected[i])
+					}
+				}
+			}
+		})
+	})
+}

mlxrunner/score.go

@@ -0,0 +1,214 @@
+package mlxrunner
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"net/http"
+	"slices"
+
+	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/llm"
+	"github.com/ollama/ollama/mlx"
+	"github.com/ollama/ollama/mlx/mlxthread"
+	"github.com/ollama/ollama/mlxrunner/batch"
+	"github.com/ollama/ollama/mlxrunner/cache"
+)
+
+type scoreRow struct {
+	tokens     []int32
+	candidates []int32
+}
+
+// Models may project just the requested vocabulary rows at higher precision.
+type candidateUnembedder interface {
+	UnembedCandidates(hidden, candidates *mlx.Array) *mlx.Array
+}
+
+func (r *Runner) scoreHandler(w http.ResponseWriter, req *http.Request) {
+	var input llm.ScoreRequest
+	if err := json.NewDecoder(req.Body).Decode(&input); err != nil {
+		http.Error(w, err.Error(), http.StatusBadRequest)
+		return
+	}
+	result, err := mlxthread.Call(req.Context(), r.mlxThread, func() (llm.ScoreResponse, error) {
+		return r.score(req.Context(), input)
+	})
+	if err != nil {
+		status := http.StatusInternalServerError
+		var statusErr api.StatusError
+		if errors.As(err, &statusErr) {
+			status = statusErr.StatusCode
+		}
+		http.Error(w, err.Error(), status)
+		return
+	}
+	w.Header().Set("Content-Type", "application/json")
+	_ = json.NewEncoder(w).Encode(result)
+}
+
+func (r *Runner) score(ctx context.Context, input llm.ScoreRequest) (llm.ScoreResponse, error) {
+	var result llm.ScoreResponse
+	badRequest := func(format string, args ...any) (llm.ScoreResponse, error) {
+		return result, api.StatusError{StatusCode: http.StatusBadRequest, ErrorMessage: fmt.Sprintf(format, args...)}
+	}
+	if len(input.Rows) < 1 || len(input.Rows) > 64 {
+		return badRequest("scoring requires 1–64 prompts")
+	}
+	if input.MaxTokens < 1 || input.MaxTokens > r.contextLength {
+		return badRequest("max_tokens must be between 1 and %d", r.contextLength)
+	}
+	rows := make([]scoreRow, len(input.Rows))
+	for i, row := range input.Rows {
+		if err := ctx.Err(); err != nil {
+			return result, err
+		}
+		tokens := r.Tokenizer.Encode(row.Prompt, false)
+		if len(tokens) == 0 || len(tokens) > input.MaxTokens {
+			return badRequest("prompt %d has %d tokens; expected 1–%d (input is never truncated)", i, len(tokens), input.MaxTokens)
+		}
+		if len(row.Candidates) < 1 || len(row.Candidates) > 26 {
+			return badRequest("prompt %d requires 1–26 candidates", i)
+		}
+		rows[i].tokens = tokens
+		result.InputTokens += len(tokens)
+		for _, candidate := range row.Candidates {
+			joined := r.Tokenizer.Encode(row.Prompt+candidate, false)
+			_, special := r.Tokenizer.GetSpecialToken(candidate)
+			if special || len(joined) != len(tokens)+1 || !slices.Equal(joined[:len(tokens)], tokens) {
+				return badRequest("candidate %q must append exactly one ordinary token to prompt %d", candidate, i)
+			}
+			id := joined[len(tokens)]
+			if id < 0 || int(id) >= r.Tokenizer.VocabSize() || slices.Contains(rows[i].candidates, id) {
+				return badRequest("prompt %d has invalid or duplicate candidate tokens", i)
+			}
+			rows[i].candidates = append(rows[i].candidates, id)
+		}
+	}
+	logits, err := r.scoreRows(ctx, rows, scorePrefixLength(rows))
+	result.Logits = logits
+	return result, err
+}
+
+// Leave at least one token for each branch, including identical prompts. Its
+// final hidden state predicts the answer; generation's one-token-short prefill
+// cannot supply that state.
+func scorePrefixLength(rows []scoreRow) int {
+	if len(rows) < 2 {
+		return 0
+	}
+	n := len(rows[0].tokens) - 1
+	for _, row := range rows[1:] {
+		n = min(n, len(row.tokens)-1)
+		for i := 0; i < n; i++ {
+			if row.tokens[i] != rows[0].tokens[i] {
+				n = i
+				break
+			}
+		}
+	}
+	return max(0, n)
+}
+
+// scoreRows runs only on the MLX worker. Caches belong to this request and never
+// enter the generation prefix trie. Restore rewinds KV storage and reinstates
+// both convolution and recurrent state before each suffix.
+func (r *Runner) scoreRows(ctx context.Context, rows []scoreRow, prefix int) ([][]float32, error) {
+	caches := r.Model.NewCaches()
+	snapshots := make([]cache.Snapshot, len(caches))
+	defer func() {
+		// Close lazy snapshots before freeing their backing caches.
+		for _, snapshot := range snapshots {
+			if snapshot != nil {
+				snapshot.Close()
+			}
+		}
+		for _, c := range caches {
+			if c != nil {
+				c.Free()
+			}
+		}
+	}()
+	if prefix > 0 {
+		if _, err := r.scoreForward(ctx, rows[0].tokens[:prefix], 0, caches, nil); err != nil {
+			return nil, err
+		}
+		for i, c := range caches {
+			if c != nil {
+				snapshots[i] = c.Snapshot(0)
+				if snapshots[i] == nil {
+					return nil, fmt.Errorf("cache %d cannot snapshot the scoring prefix", i)
+				}
+			}
+		}
+	}
+	result := make([][]float32, len(rows))
+	for i, row := range rows {
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
+		if i > 0 {
+			for j, c := range caches {
+				if c == nil {
+					continue
+				}
+				if prefix == 0 {
+					c.Free()
+				} else if !c.Restore(snapshots[j], prefix) {
+					return nil, fmt.Errorf("cache %d cannot restore the scoring prefix", j)
+				}
+			}
+		}
+		logits, err := r.scoreForward(ctx, row.tokens[prefix:], prefix, caches, row.candidates)
+		if err != nil {
+			return nil, err
+		}
+		result[i] = logits
+	}
+	return result, ctx.Err()
+}
+
+func (r *Runner) scoreForward(ctx context.Context, tokens []int32, offset int, caches []cache.Cache, candidates []int32) ([]float32, error) {
+	var logits []float32
+	for len(tokens) > 0 {
+		if err := ctx.Err(); err != nil {
+			return nil, err
+		}
+		n := min(prefillChunkSize(), len(tokens))
+		mlx.Scoped(func() {
+			output := mlx.ScopedArrays(func() []*mlx.Array {
+				hidden, _ := r.Model.Forward(&batch.Batch{
+					InputIDs:     mlx.FromValues(tokens[:n], 1, n),
+					SeqOffsets:   []int32{int32(offset)},
+					SeqQueryLens: []int32{int32(n)},
+				}, caches)
+				if n != len(tokens) || len(candidates) == 0 {
+					return nil
+				}
+				last := hidden.Slice(mlx.Slice(), mlx.Slice(n-1, n), mlx.Slice())
+				ids := mlx.FromValues(candidates, len(candidates))
+				var selected *mlx.Array
+				if m, ok := r.Model.(candidateUnembedder); ok {
+					selected = m.UnembedCandidates(last, ids)
+				} else {
+					selected = r.Model.Unembed(last).Reshape(-1).TakeAxis(ids, 0)
+				}
+				return []*mlx.Array{selected.AsType(mlx.DTypeFloat32)}
+			})
+			state := slices.Clone(output)
+			for _, c := range caches {
+				if c != nil {
+					state = append(state, c.State()...)
+				}
+			}
+			mlx.Eval(state...)
+			if len(output) > 0 {
+				logits = output[0].Floats()
+			}
+		})
+		tokens = tokens[n:]
+		offset += n
+	}
+	return logits, ctx.Err()
+}

mlxrunner/score_test.go

@@ -0,0 +1,258 @@
+package mlxrunner
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"math"
+	"os"
+	"slices"
+	"testing"
+
+	"github.com/ollama/ollama/api"
+	"github.com/ollama/ollama/decision"
+	"github.com/ollama/ollama/llm"
+	"github.com/ollama/ollama/mlx"
+	"github.com/ollama/ollama/mlx/mlxtest"
+	"github.com/ollama/ollama/mlxrunner/batch"
+	"github.com/ollama/ollama/mlxrunner/cache"
+	"github.com/ollama/ollama/model/renderers"
+)
+
+// This small stateful model uses real KV and recurrent caches. Its outputs
+// depend on both recurrent tensors, all KV values, and absolute positions, so a
+// missing restore or a branch that overwrites the prefix changes the answer.
+type scoringTestModel struct {
+	textOnlyModel
+	caches      []cache.Cache
+	positions   []int
+	lengths     []int
+	projections int
+	cancel      context.CancelFunc
+	cancelAfter int
+}
+
+func (m *scoringTestModel) NewCaches() []cache.Cache {
+	m.caches = []cache.Cache{cache.NewKVCache(), cache.NewRecurrentCache(1, 1, 1, 1, 1), nil}
+	return m.caches
+}
+
+func (m *scoringTestModel) Forward(b *batch.Batch, caches []cache.Cache) (*mlx.Array, *mlx.Array) {
+	n, offset := b.InputIDs.Dim(1), int(b.SeqOffsets[0])
+	m.positions = append(m.positions, offset)
+	m.lengths = append(m.lengths, n)
+	x := b.InputIDs.AsType(mlx.DTypeFloat32)
+	kv := caches[0].(*cache.KVCache).Update(b, x.Reshape(1, 1, n, 1), x.Multiply(x).Reshape(1, 1, n, 1))
+	rc := caches[1].(*cache.RecurrentCache)
+	history := rc.Get(b, mlx.DTypeFloat32)
+	positions := make([]float32, n)
+	for i := range positions {
+		positions[i] = float32(offset + i + 1)
+	}
+	conv := history.ConvState().Add(x.SumAxis(1, true).Reshape(1, 1, 1))
+	delta := history.DeltaState().Add(x.Multiply(mlx.FromValues(positions, 1, n)).SumAxis(1, true).Reshape(1, 1, 1, 1))
+	rc.Put(b, []*mlx.Array{conv}, []*mlx.Array{delta})
+	values := mlx.Concatenate([]*mlx.Array{
+		kv.V().SumAxis(2, true).Reshape(1), conv.Reshape(1), delta.Reshape(1),
+		x.Slice(mlx.Slice(), mlx.Slice(n-1, n)).Reshape(1),
+		mlx.FromValues([]float32{float32(offset + n), 1, 2, 3}, 4),
+	}, 0).Reshape(1, 1, 8)
+	hidden := mlx.Zeros(mlx.DTypeFloat32, 1, n, 8).Add(values)
+	if m.cancel != nil && len(m.positions) == m.cancelAfter {
+		m.cancel()
+	}
+	return hidden, hidden
+}
+
+func (m *scoringTestModel) Unembed(hidden *mlx.Array) *mlx.Array {
+	if hidden.Dim(1) != 1 {
+		panic("scoring projected more than the final position")
+	}
+	m.projections++
+	return hidden
+}
+
+func TestScoreSharedPrefix(t *testing.T) {
+	for _, tt := range []struct {
+		name   string
+		tokens [][]int32
+	}{
+		{"branch lengths and repeated restore", [][]int32{{1, 2, 3, 4}, {1, 2, 5}, {1, 2, 3, 6, 7}, {1, 2, 3, 4}}},
+		{"identical prompts", [][]int32{{1, 2, 3}, {1, 2, 3}}},
+		{"empty prefix", [][]int32{{1, 2}, {2, 1, 3}}},
+		{"one prompt", [][]int32{{1, 2, 3, 4}}},
+		{"one token", [][]int32{{1}, {1}, {2}}},
+		{"chunked prefix", [][]int32{append(slices.Repeat([]int32{1}, prefillChunkSize()+3), 2), append(slices.Repeat([]int32{1}, prefillChunkSize()+3), 3)}},
+	} {
+		mlxtest.RunSubtest(t, tt.name, func(t *mlxtest.T) {
+			m := &scoringTestModel{}
+			r := &Runner{Model: m}
+			rows := make([]scoreRow, len(tt.tokens))
+			for i, tokens := range tt.tokens {
+				rows[i] = scoreRow{tokens: tokens, candidates: []int32{2, 0, 1, 3, 4}}
+			}
+			independent, err := r.scoreRows(context.Background(), rows, 0)
+			if err != nil {
+				t.Fatal(err)
+			}
+			prefix := scorePrefixLength(rows)
+			// Repeated requests must release cache state and isolate branches.
+			for range 3 {
+				m.positions, m.lengths, m.projections = nil, nil, 0
+				got, err := r.scoreRows(context.Background(), rows, prefix)
+				if err != nil {
+					t.Fatal(err)
+				}
+				for i := range got {
+					if !slices.Equal(got[i], independent[i]) {
+						t.Fatalf("row %d shared %v != independent %v", i, got[i], independent[i])
+					}
+				}
+				for _, c := range m.caches {
+					if c != nil && c.Offset() != 0 {
+						t.Fatal("request left its caches alive")
+					}
+					if c != nil {
+						for _, state := range c.State() {
+							if state != nil {
+								t.Fatal("request retained cache arrays")
+							}
+						}
+					}
+				}
+				if m.projections != len(rows) {
+					t.Fatalf("got %d projections for %d rows", m.projections, len(rows))
+				}
+				total := 0
+				for _, n := range m.lengths {
+					total += n
+				}
+				want := prefix
+				for _, row := range rows {
+					want += len(row.tokens) - prefix
+				}
+				if total != want {
+					t.Fatalf("forwarded %d tokens, want %d", total, want)
+				}
+			}
+		})
+	}
+}
+
+func TestScoreCancellation(t *testing.T) {
+	for _, cancelAfter := range []int{1, 2} {
+		mlxtest.Run(t, func(t *mlxtest.T) {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+			m := &scoringTestModel{cancel: cancel, cancelAfter: cancelAfter}
+			r := &Runner{Model: m}
+			rows := []scoreRow{{tokens: []int32{1, 2, 3}, candidates: []int32{0}}, {tokens: []int32{1, 2, 4}, candidates: []int32{0}}}
+			_, err := r.scoreRows(ctx, rows, 2)
+			if !errors.Is(err, context.Canceled) {
+				t.Fatalf("got %v, want context cancellation", err)
+			}
+			if len(m.positions) != cancelAfter {
+				t.Fatal("cancellation did extra forwards")
+			}
+			for _, c := range m.caches {
+				if c != nil && c.Offset() != 0 {
+					t.Fatal("cancelled request retained cache state")
+				}
+			}
+		})
+	}
+}
+
+func TestScoreValidation(t *testing.T) {
+	tok := newTestTokenizer(t, []int32{7})
+	r := &Runner{Tokenizer: tok, contextLength: 10}
+	for _, input := range []llm.ScoreRequest{
+		{}, {MaxTokens: 10}, {MaxTokens: 11, Rows: []llm.ScoreRow{{Prompt: "1", Candidates: []string{"2"}}}},
+		{MaxTokens: 1, Rows: []llm.ScoreRow{{Prompt: "12", Candidates: []string{"3"}}}},
+		{MaxTokens: 10, Rows: []llm.ScoreRow{{Prompt: "", Candidates: []string{"1"}}}},
+		{MaxTokens: 10, Rows: []llm.ScoreRow{{Prompt: "1"}}},
+		{MaxTokens: 10, Rows: []llm.ScoreRow{{Prompt: "1", Candidates: []string{"23"}}}},
+		{MaxTokens: 10, Rows: []llm.ScoreRow{{Prompt: "1", Candidates: []string{"2", "2"}}}},
+	} {
+		_, err := r.score(context.Background(), input)
+		var status api.StatusError
+		if !errors.As(err, &status) || status.StatusCode != 400 {
+			t.Fatalf("expected validation error before model access, got %v", err)
+		}
+	}
+}
+
+// Optional full-model check against independently prefilling each question.
+func TestNimbleModel(t *testing.T) {
+	name := os.Getenv("OLLAMA_NIMBLE_MODEL")
+	if name == "" {
+		t.Skip("set OLLAMA_NIMBLE_MODEL to an imported Qwen3.5/Nimble model")
+	}
+	var req decision.Request
+	if err := json.Unmarshal([]byte(`{
+		"model":"nimble", "state":"I was charged twice. Please refund the duplicate.",
+		"questions":{
+			"department":{"type":"choice","instructions":"Which department?","criteria":{"billing":"Payments","technical":"Software"}},
+			"refund":{"type":"noul","instructions":"Refund requested?"},
+			"urgency":{"type":"score","instructions":"How urgent?","criteria":["Routine","Urgent","Emergency"]}
+		}
+	}`), &req); err != nil {
+		t.Fatal(err)
+	}
+	c, err := decision.Compile(req)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := c.Render(func(messages []api.Message) (string, error) {
+		return renderers.RenderWithRenderer("qwen3.5", messages, nil, &api.ThinkValue{Value: false})
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	mlxtest.Run(t, func(t *mlxtest.T) {
+		r := &Runner{}
+		if err := r.Load(name); err != nil {
+			t.Fatal(err)
+		}
+		defer r.Close()
+		c.Request.MaxTokens = r.contextLength
+		shared, err := r.score(context.Background(), c.Request)
+		if err != nil {
+			t.Fatal(err)
+		}
+		probabilities := func(logits []float32) []float64 {
+			p := make([]float64, len(logits))
+			peak := slices.Max(logits)
+			var sum float64
+			for i, x := range logits {
+				p[i] = math.Exp(float64(x - peak))
+				sum += p[i]
+			}
+			for i := range p {
+				p[i] /= sum
+			}
+			return p
+		}
+		for i, row := range c.Request.Rows {
+			single, err := r.score(context.Background(), llm.ScoreRequest{Rows: []llm.ScoreRow{row}, MaxTokens: c.Request.MaxTokens})
+			if err != nil {
+				t.Fatal(err)
+			}
+			// Different prefill shapes can choose different kernels, and this
+			// backbone retains its native precision, as do quantized heads.
+			// Check the resulting decisions/probabilities;
+			// log exact logits to make numerical differences visible.
+			p, q := probabilities(shared.Logits[i]), probabilities(single.Logits[0])
+			for j := range p {
+				if diff := math.Abs(p[j] - q[j]); diff > 0.01 {
+					t.Errorf("row %d candidate %d probability difference=%g", i, j, diff)
+				}
+			}
+			if slices.Index(p, slices.Max(p)) != slices.Index(q, slices.Max(q)) {
+				t.Errorf("row %d changed its winning candidate", i)
+			}
+			t.Logf("row %d: shared=%v independent=%v", i, shared.Logits[i], single.Logits[0])
+		}
+	})
+}

mlxrunner/server.go

@@ -87,6 +87,7 @@ func Execute(args []string) error {
 	)
 
 	mux := http.NewServeMux()
+	mux.HandleFunc("POST /v1/score", runner.scoreHandler)
 	mux.HandleFunc("GET /v1/status", func(w http.ResponseWriter, r *http.Request) {
 		if err := json.NewEncoder(w).Encode(statusResponse{
 			Status:        0,