From 3ee5dd7f531cb970dc642f676c8db9b42e11da98 Mon Sep 17 00:00:00 2001 From: ParthSareen Date: Wed, 23 Sep 2026 00:49:30 -0700 Subject: [PATCH] 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 26e546ddb3897b4564f03d3bdbdbefce9281edf8) --- mlxrunner/client_score.go | 46 +++++ mlxrunner/client_score_test.go | 77 ++++++++ mlxrunner/model/qwen3_5/score.go | 23 +++ mlxrunner/model/qwen3_5/score_test.go | 44 +++++ mlxrunner/score.go | 214 +++++++++++++++++++++ mlxrunner/score_test.go | 258 ++++++++++++++++++++++++++ mlxrunner/server.go | 1 + 7 files changed, 663 insertions(+) create mode 100644 mlxrunner/client_score.go create mode 100644 mlxrunner/client_score_test.go create mode 100644 mlxrunner/model/qwen3_5/score.go create mode 100644 mlxrunner/model/qwen3_5/score_test.go create mode 100644 mlxrunner/score.go create mode 100644 mlxrunner/score_test.go diff --git a/mlxrunner/client_score.go b/mlxrunner/client_score.go new file mode 100644 index 000000000..4e0a8b4ae --- /dev/null +++ b/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 +} diff --git a/mlxrunner/client_score_test.go b/mlxrunner/client_score_test.go new file mode 100644 index 000000000..594ef0ff6 --- /dev/null +++ b/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) + } + }) + } +} diff --git a/mlxrunner/model/qwen3_5/score.go b/mlxrunner/model/qwen3_5/score.go new file mode 100644 index 000000000..8cee39d74 --- /dev/null +++ b/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) +} diff --git a/mlxrunner/model/qwen3_5/score_test.go b/mlxrunner/model/qwen3_5/score_test.go new file mode 100644 index 000000000..dc592df26 --- /dev/null +++ b/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]) + } + } + } + }) + }) +} diff --git a/mlxrunner/score.go b/mlxrunner/score.go new file mode 100644 index 000000000..d63fe50b4 --- /dev/null +++ b/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() +} diff --git a/mlxrunner/score_test.go b/mlxrunner/score_test.go new file mode 100644 index 000000000..4abd68690 --- /dev/null +++ b/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]) + } + }) +} diff --git a/mlxrunner/server.go b/mlxrunner/server.go index 0470f599b..36385a849 100644 --- a/mlxrunner/server.go +++ b/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,