diff --git a/README.md b/README.md index e1e727b3a0a..ff0626b8836 100644 --- a/README.md +++ b/README.md @@ -110,7 +110,7 @@ decision and its measurements live (`docs/maxusai/`). The | | upstream ollama | this fork | record | |---|---|---|---| | **`think` + `format` in one request** | since v0.34.4, one pass: the grammar applies from the first token, after free thinking that ends at the parser's closing strings | the same by default. `OLLAMA_FORMAT_TWO_PASS=1` keeps the fork's two-pass flow: pass one thinks without the grammar and stops at the think-close marker, and pass two answers under it. Production runs it, because on MLX it lets the thinking draft | ADR 0045 (0002/0004/0010) | -| **whitespace in structured output (MLX)** | the JSON grammar allows an unlimited whitespace run between tokens, so a stalled decode can fill `num_predict` with indentation and never close the answer | runs bounded at 32 characters per separator (`max_whitespace_cnt` on the `json_schema` element) | ADR 0035 (proposed) | +| **whitespace in structured output (MLX)** | the JSON grammar allows an unlimited whitespace run between tokens, so a stalled decode can fill `num_predict` with indentation and never close the answer | runs bounded at 22 characters per separator (`max_whitespace_cnt` on the `json_schema` element), the length of llama.cpp's longest run for a JSON schema but not its shape (any 22 whitespace characters); GGUF's `format: "json"` grammar uses llama.cpp's own rule, one per gap | ADR 0035 (proposed; amended 2026-10-12) | | **drafting under a grammar (MLX)** | always on | on by default to match upstream; `OLLAMA_MLX_DRAFT_UNDER_GRAMMAR=0` restores the gate, and production sets it | ADR 0033 | | **stop sequences (MLX)** | not honoured by the MLX runner | honoured, with a possible stop prefix held back until it matches or the stream ends | `mlxrunner/stopper.go` | | **KV cache type** | one global `OLLAMA_KV_CACHE_TYPE` | per model, with K/V pair syntax; production runs f16 on every platform, and a quantized cache is per model or per request | ADR 0005, ADR 0043 | diff --git a/docs/maxusai/adr/0035-structured-output-bounds-whitespace-between-tokens.md b/docs/maxusai/adr/0035-structured-output-bounds-whitespace-between-tokens.md index b1df9dfa0a9..ec61457d181 100644 --- a/docs/maxusai/adr/0035-structured-output-bounds-whitespace-between-tokens.md +++ b/docs/maxusai/adr/0035-structured-output-bounds-whitespace-between-tokens.md @@ -68,3 +68,25 @@ per document, and does not touch the schema it wraps. documents that were lost is the confirmation, and it has not been done. - Nothing changes for the GGUF path, which passes the schema to llama-server and is bounded by llama.cpp's own converter. + +## Amendment 2026-10-12: 22, the GGUF run, and both bounds checked by a matcher + +**What moved it.** muse-glimmer on MLX in single pass (`OLLAMA_FORMAT_TWO_PASS=0`), answering after its message header, padded the first and last slot of every object to the bound. `format_schema` on `30b-nvfp4` returned `{` + 33 spaces, the value, about 23 spaces and `}` (#463, #464). The bound held: stock 0.40.2 has no bound and filled all 16384 tokens with whitespace (FAIL). The cost of the padding was the bound itself. + +**Decision, amended:** +- **MLX:** `max_whitespace_cnt` is **22**. That is the length of the longest run llama.cpp allows for a JSON schema: `SPACE_RULE`, two newlines and 20 spaces or tabs. The runners agree on that length only, not on the shape: `max_whitespace_cnt` counts any whitespace, so MLX accepts 22 spaces in one gap, where `SPACE_RULE` allows at most one space. A newline plus ten levels of two-space indent (21 characters) still fits. +- **GGUF, `format: "json"`:** the GGUF path is not wholly bounded. A schema goes to llama-server's converter, which is bounded. But `format: "json"` uses the fork's own grammar (`grammarJSON`, `llm/llama_server.go`), and its `ws ::= ([ \t\n] ws)?` was unbounded. It is now `ws ::= (" " | "\n"{1,2} [ \t]{0,20})?`: llama.cpp's rule, written as an optional group, because xgrammar's parser rejects a leading empty alternative. + - **Placement (2026-10-13, after review on the PR):** the rule's bound held only per rule, and two rules sat side by side in three places: `"{" ws (...)? ws "}"`, the same for arrays, and a literal's trailing `ws` before `ws "}"` or `ws "]"`. A gap could run to 44 characters, and `{ }` was accepted. The grammar now has the shape of llama.cpp's `grammars/json.gbnf`: whitespace follows `{`, `[`, `:`, `,` and every value, and nothing puts it before `}` or `]`. So each gap is one `ws`, and no run is longer than 22. `{ }` is refused, as json.gbnf refuses it; `{ }` and `{` + newline + `}` are accepted. + - Two differences from json.gbnf were kept from the old grammar. The root object ends at its `}`, with no trailing whitespace, so the grammar completes there. And there is no whitespace between a key and its `:`. + - One placement is new: whitespace after a string, number, object or array value, before a `,`. The old grammar allowed it only after `true`, `false` and `null`; json.gbnf allows it after every value. +- **Not taken: patching xgrammar to llama.cpp's exact rule.** On MLX that would cut glimmer's padding to one space per slot. It costs a native patch on every xgrammar bump, and it moves every MLX format request's mask. xgrammar v0.2.7 hardcodes `any_whitespace=true` for structural-tag `json_schema` elements (`cpp/structural_tag.cc`), so only the count can be set. Revisit if xgrammar adds a whitespace-rule option. + +**Verification, amended:** +- `TestRequestGrammarBoundsWhitespaceRunLength` drives `requestGrammar`'s tag through xgrammar one byte at a time, for a plain format, a format after thinking, and `"json"` after a header. It checks that 22 spaces are accepted, 23 refused, and two newlines plus 20 spaces accepted. It proves the run's length, not llama.cpp's shape. +- `TestGrammarJSONBoundsWhitespace` checks `grammarJSON` with a reference GBNF recognizer written in Go (`llm/gbnf_reference_test.go`), so it runs everywhere. The cases include two gaps at the bound side by side in an empty object, in an empty array and after a literal, and `{ }`. `TestGrammarJSONBoundsWhitespaceXGrammar` drives the same cases through xgrammar's grammar parser one byte at a time. llama.cpp's engine has no Go entry point. +- Both fail on the old grammars. The xgrammar tests run wherever the MLX payload is built and skip elsewhere. The "not yet verified end to end" above still holds for MLX documents; the real-model check of this amendment is the GGUF and MLX A/B on the branch's PR. + +**Upstream:** +- ollama/ollama has no bound on MLX. ollama/ollama#18567 (a whitespace runaway) was closed as standard behaviour. +- ollama/ollama#18569 (a cap of 32) was closed untested. +- The fork's bound has not been offered upstream. diff --git a/llm/gbnf_reference_test.go b/llm/gbnf_reference_test.go new file mode 100644 index 00000000000..2ba7fe89338 --- /dev/null +++ b/llm/gbnf_reference_test.go @@ -0,0 +1,376 @@ +package llm + +import ( + "fmt" + "maps" + "slices" + "strconv" + "strings" + "unicode/utf8" +) + +// gbnfRecognizer is a reference recognizer for the GBNF subset the fork's +// own grammars use (rules, literals, character classes with ranges and +// negation, groups, alternation, ?, *, + and {m,n}), so grammarJSON's +// language can be checked in plain Go wherever the tests run. It parses as +// llama.cpp's grammar parser does (a rule ends at a newline outside +// parentheses; # comments run to the end of the line) and recognises by +// carrying the set of positions every expression can end at, which is exact +// for any grammar without left recursion. It is not a decoder: it answers only +// whether a whole input is in the language. +type gbnfRecognizer struct { + rules map[string]*gbnfNode +} + +type gbnfNode struct { + kind byte // 'l' literal, 'c' class, 'r' rule reference, 's' sequence, 'a' alternation, 'x' repetition + lit string + ranges [][2]rune + negate bool + ref string + kids []*gbnfNode + min, max int // repetition bounds; max < 0 is unbounded +} + +func parseGBNF(src string) (*gbnfRecognizer, error) { + p := &gbnfParser{src: src} + rules := map[string]*gbnfNode{} + for { + p.space(true) + if p.pos >= len(p.src) { + break + } + name := p.name() + if name == "" { + return nil, p.errorf("expected a rule name") + } + p.space(false) + if !strings.HasPrefix(p.src[p.pos:], "::=") { + return nil, p.errorf("expected ::= after %s", name) + } + p.pos += 3 + p.space(true) + n, err := p.alternation(false) + if err != nil { + return nil, err + } + rules[name] = n + } + for _, n := range rules { + if err := checkRefs(n, rules); err != nil { + return nil, err + } + } + if rules["root"] == nil { + return nil, fmt.Errorf("gbnf: no root rule") + } + return &gbnfRecognizer{rules: rules}, nil +} + +func checkRefs(n *gbnfNode, rules map[string]*gbnfNode) error { + if n.kind == 'r' && rules[n.ref] == nil { + return fmt.Errorf("gbnf: undefined rule %s", n.ref) + } + for _, k := range n.kids { + if err := checkRefs(k, rules); err != nil { + return err + } + } + return nil +} + +type gbnfParser struct { + src string + pos int +} + +func (p *gbnfParser) errorf(format string, args ...any) error { + return fmt.Errorf("gbnf at %d: %s", p.pos, fmt.Sprintf(format, args...)) +} + +// space skips blanks and comments, and newlines when newlines are allowed. +func (p *gbnfParser) space(newlines bool) { + for p.pos < len(p.src) { + switch c := p.src[p.pos]; { + case c == ' ' || c == '\t': + p.pos++ + case c == '#': + for p.pos < len(p.src) && p.src[p.pos] != '\n' { + p.pos++ + } + case (c == '\n' || c == '\r') && newlines: + p.pos++ + default: + return + } + } +} + +func (p *gbnfParser) name() string { + start := p.pos + for p.pos < len(p.src) { + c := p.src[p.pos] + if c != '-' && c != '_' && (c < '0' || c > '9') && (c < 'a' || c > 'z') && (c < 'A' || c > 'Z') { + break + } + p.pos++ + } + return p.src[start:p.pos] +} + +func (p *gbnfParser) alternation(nested bool) (*gbnfNode, error) { + alt := &gbnfNode{kind: 'a'} + for { + seq, err := p.sequence(nested) + if err != nil { + return nil, err + } + alt.kids = append(alt.kids, seq) + if p.pos >= len(p.src) || p.src[p.pos] != '|' { + return alt, nil + } + p.pos++ + p.space(true) + } +} + +func (p *gbnfParser) sequence(nested bool) (*gbnfNode, error) { + seq := &gbnfNode{kind: 's'} + for p.pos < len(p.src) { + var n *gbnfNode + switch c := p.src[p.pos]; { + case c == '|' || c == ')' || c == '\n' || c == '\r': + return seq, nil + case c == '"': + p.pos++ + var b strings.Builder + for p.pos < len(p.src) && p.src[p.pos] != '"' { + r, err := p.char() + if err != nil { + return nil, err + } + b.WriteRune(r) + } + if p.pos >= len(p.src) { + return nil, p.errorf("unterminated literal") + } + p.pos++ + n = &gbnfNode{kind: 'l', lit: b.String()} + case c == '[': + p.pos++ + n = &gbnfNode{kind: 'c'} + if p.pos < len(p.src) && p.src[p.pos] == '^' { + n.negate = true + p.pos++ + } + for p.pos < len(p.src) && p.src[p.pos] != ']' { + lo, err := p.char() + if err != nil { + return nil, err + } + hi := lo + if p.pos+1 < len(p.src) && p.src[p.pos] == '-' && p.src[p.pos+1] != ']' { + p.pos++ + if hi, err = p.char(); err != nil { + return nil, err + } + } + n.ranges = append(n.ranges, [2]rune{lo, hi}) + } + if p.pos >= len(p.src) { + return nil, p.errorf("unterminated character class") + } + p.pos++ + case c == '(': + p.pos++ + p.space(true) + var err error + if n, err = p.alternation(true); err != nil { + return nil, err + } + if p.pos >= len(p.src) || p.src[p.pos] != ')' { + return nil, p.errorf("expected )") + } + p.pos++ + default: + name := p.name() + if name == "" { + return nil, p.errorf("unexpected %q", c) + } + n = &gbnfNode{kind: 'r', ref: name} + } + for p.pos < len(p.src) { + lo, hi := 0, 0 + switch p.src[p.pos] { + case '?': + lo, hi = 0, 1 + case '*': + lo, hi = 0, -1 + case '+': + lo, hi = 1, -1 + case '{': + end := strings.IndexByte(p.src[p.pos:], '}') + if end < 0 { + return nil, p.errorf("unterminated repetition") + } + body := p.src[p.pos+1 : p.pos+end] + var err error + if m, n, ok := strings.Cut(body, ","); ok { + if lo, err = strconv.Atoi(m); err == nil { + hi = -1 + if n != "" { + hi, err = strconv.Atoi(n) + } + } + } else if lo, err = strconv.Atoi(body); err == nil { + hi = lo + } + if err != nil { + return nil, p.errorf("bad repetition {%s}", body) + } + p.pos += end + default: + lo = -1 + } + if lo < 0 { + break + } + p.pos++ + n = &gbnfNode{kind: 'x', kids: []*gbnfNode{n}, min: lo, max: hi} + } + seq.kids = append(seq.kids, n) + p.space(nested) + } + return seq, nil +} + +// char reads one character of a literal or class, decoding escapes. +func (p *gbnfParser) char() (rune, error) { + if p.src[p.pos] != '\\' { + r, size := utf8.DecodeRuneInString(p.src[p.pos:]) + p.pos += size + return r, nil + } + if p.pos+1 >= len(p.src) { + return 0, p.errorf("dangling escape") + } + c := p.src[p.pos+1] + p.pos += 2 + switch c { + case 'n': + return '\n', nil + case 't': + return '\t', nil + case 'r': + return '\r', nil + case 'x': + if p.pos+2 > len(p.src) { + return 0, p.errorf("short \\x escape") + } + v, err := strconv.ParseUint(p.src[p.pos:p.pos+2], 16, 8) + if err != nil { + return 0, p.errorf("bad \\x escape") + } + p.pos += 2 + return rune(v), nil + default: + return rune(c), nil + } +} + +// accepts reports whether the whole of in is in the grammar's language. +func (g *gbnfRecognizer) accepts(in string) bool { + m := &gbnfMatch{g: g, in: in, memo: map[gbnfMemoKey][]int{}} + return slices.Contains(m.match(&gbnfNode{kind: 'r', ref: "root"}, 0), len(in)) +} + +type gbnfMemoKey struct { + rule string + pos int +} + +type gbnfMatch struct { + g *gbnfRecognizer + in string + memo map[gbnfMemoKey][]int +} + +// match returns every position n can end at when it starts at pos. +func (m *gbnfMatch) match(n *gbnfNode, pos int) []int { + switch n.kind { + case 'l': + if strings.HasPrefix(m.in[pos:], n.lit) { + return []int{pos + len(n.lit)} + } + case 'c': + if pos < len(m.in) { + r, size := utf8.DecodeRuneInString(m.in[pos:]) + in := slices.ContainsFunc(n.ranges, func(rg [2]rune) bool { return rg[0] <= r && r <= rg[1] }) + if in != n.negate { + return []int{pos + size} + } + } + case 'r': + key := gbnfMemoKey{n.ref, pos} + if ends, ok := m.memo[key]; ok { + return ends + } + m.memo[key] = nil // a left-recursive reference matches nothing + ends := m.match(m.g.rules[n.ref], pos) + m.memo[key] = ends + return ends + case 's': + cur := []int{pos} + for _, k := range n.kids { + cur = m.each(k, cur) + if len(cur) == 0 { + break + } + } + return cur + case 'a': + set := map[int]bool{} + for _, k := range n.kids { + for _, e := range m.match(k, pos) { + set[e] = true + } + } + return slices.Sorted(maps.Keys(set)) + case 'x': + out := map[int]bool{} + seen := map[int]bool{pos: true} + if n.min == 0 { + out[pos] = true + } + cur := []int{pos} + for i := 1; n.max < 0 || i <= n.max; i++ { + cur = m.each(n.kids[0], cur) + if n.max < 0 { + // Unbounded: a position reached before adds nothing new. + cur = slices.DeleteFunc(cur, func(e int) bool { return seen[e] && i > n.min }) + } + if len(cur) == 0 { + break + } + for _, e := range cur { + seen[e] = true + if i >= n.min { + out[e] = true + } + } + } + return slices.Sorted(maps.Keys(out)) + } + return nil +} + +// each matches n from every start position and returns the union of the ends. +func (m *gbnfMatch) each(n *gbnfNode, starts []int) []int { + set := map[int]bool{} + for _, s := range starts { + for _, e := range m.match(n, s) { + set[e] = true + } + } + return slices.Sorted(maps.Keys(set)) +} diff --git a/llm/grammar_json_ws_test.go b/llm/grammar_json_ws_test.go new file mode 100644 index 00000000000..063a81aa2ba --- /dev/null +++ b/llm/grammar_json_ws_test.go @@ -0,0 +1,158 @@ +package llm + +import ( + "encoding/json" + "errors" + "io/fs" + "path/filepath" + "strings" + "testing" + + "github.com/ollama/ollama/mlx" + "github.com/ollama/ollama/mlxrunner/xgrammar" +) + +// grammarJSONWhitespaceCases are what format "json"'s GBNF must accept and +// refuse. llama-server's own JSON-schema converter allows, in one gap between +// tokens, nothing, one space, or one or two newlines and up to 20 spaces or +// tabs (common/json-schema-to-grammar.cpp, SPACE_RULE); format "json" is the +// fork's own grammar and must allow no more, or a model that degenerates into +// whitespace runs to num_predict, or to llama-server's 100-repeat abort (an +// HTTP 500 with no output). +// +// Bounding the rule is not enough if two of them can sit side by side: until +// 2026-10-13 "{" ws ... ws "}" and a literal's trailing ws before ws "}" each +// let two gaps run together, 44 characters in a row. The grammar now has +// llama.cpp's grammars/json.gbnf shape (a value ends in ws, and nothing puts +// ws before "}" or "]"), so an empty "{ }" is refused there as it is here. +// gapAtBound is the longest run one gap allows: two newlines and 20 spaces. +var gapAtBound = "\n\n" + strings.Repeat(" ", 20) + +var grammarJSONWhitespaceCases = []struct { + name, in string + accept bool +}{ + {"compact", `{"a":1,"b":[true,null]}`, true}, + {"one space after each literal", `{ "a": 1, "b": [ true ] }`, true}, + {"newline and ten levels of two-space indent", "{\n \"a\": 1\n}", true}, + {"two newlines and a tab", "{\n\n\t\"a\": 1\n}", true}, + {"pretty-printed nesting", "{\n \"a\": {\n \"b\": [\n 1,\n true\n ]\n },\n \"c\": null\n}", true}, + {"one gap at the bound after a literal", `{"a":true` + gapAtBound + `}`, true}, + {"one gap at the bound in an empty object", `{` + gapAtBound + `}`, true}, + {"empty object with one space", `{ }`, true}, + {"empty object with a newline", "{\n}", true}, + {"empty array with one space", `{"a":[ ]}`, true}, + {"a space after a value before a comma", `{"a":1 ,"b":"x" ,"c":{} ,"d":[] }`, true}, + {"two spaces", `{ "a":1}`, false}, + {"a run of spaces", `{` + strings.Repeat(" ", 32) + `"a":1}`, false}, + {"three newlines", "{\n\n\n\"a\":1}", false}, + {"a newline and 21 spaces", "{\n" + strings.Repeat(" ", 21) + `"a":1}`, false}, + {"empty object with two spaces", `{ }`, false}, + {"empty array with two spaces", `{"a":[ ]}`, false}, + // Two gaps at the bound run together: 44 characters in a row. + {"empty object with two gaps at the bound", `{` + gapAtBound + gapAtBound + `}`, false}, + {"empty array with two gaps at the bound", `{"a":[` + gapAtBound + gapAtBound + `]}`, false}, + {"two gaps at the bound after a literal in an object", `{"a":true` + gapAtBound + gapAtBound + `}`, false}, + {"two gaps at the bound after a literal in an array", `{"a":[null` + gapAtBound + gapAtBound + `]}`, false}, + {"two spaces after a literal", `{"a":false }`, false}, + {"a space between a key and its colon", `{"a" :1}`, false}, + {"whitespace after the closing brace", `{"a":1} `, false}, +} + +// TestGrammarJSONBoundsWhitespace checks grammarJSON against +// grammarJSONWhitespaceCases with the reference GBNF recognizer, in plain Go, +// so it runs wherever the tests do. +func TestGrammarJSONBoundsWhitespace(t *testing.T) { + g, err := parseGBNF(grammarJSON) + if err != nil { + t.Fatal(err) + } + for _, c := range grammarJSONWhitespaceCases { + t.Run(c.name, func(t *testing.T) { + if got := g.accepts(c.in); got != c.accept { + t.Errorf("accepted = %v, want %v for %q", got, c.accept, c.in) + } + }) + } +} + +// TestGrammarJSONBoundsWhitespaceXGrammar drives the same cases through a +// real grammar engine one byte at a time, as a decoder would be masked. +// llama.cpp's grammar engine has no Go entry point, so this is xgrammar's, +// which reads the same GBNF; it runs wherever the MLX payload is built and +// skips elsewhere. +func TestGrammarJSONBoundsWhitespaceXGrammar(t *testing.T) { + compiler := byteGrammarCompiler(t) + for _, c := range grammarJSONWhitespaceCases { + t.Run(c.name, func(t *testing.T) { + got := acceptsBytes(t, compiler, grammarJSON, c.in) + if got != c.accept { + t.Errorf("accepted = %v, want %v for %q", got, c.accept, c.in) + } + }) + } +} + +const ( + grammarTestVocab = 257 + grammarTestEOS int32 = 256 +) + +func byteGrammarCompiler(t *testing.T) *xgrammar.Compiler { + t.Helper() + path, err := mlx.LoadedLibraryPath() + if err != nil { + t.Skipf("native MLX payload is not built: %v", err) + } + pieces := make([]string, grammarTestVocab) + for i := range 256 { + pieces[i] = string([]byte{byte(i)}) + } + pieces[grammarTestEOS] = "" + compiler, err := xgrammar.New(filepath.Dir(path), pieces, grammarTestVocab, []int32{grammarTestEOS}, 8, 128<<20) + if err != nil { + if errors.Is(err, fs.ErrNotExist) { + t.Skipf("native xgrammar payload is not built: %v", err) + } + t.Fatal(err) + } + t.Cleanup(compiler.Close) + return compiler +} + +// acceptsBytes reports whether grammar accepts in, fed one byte per token +// under the mask as a decoder would be, and can end there. +func acceptsBytes(t *testing.T, compiler *xgrammar.Compiler, grammar, in string) bool { + t.Helper() + g, err := json.Marshal(grammar) + if err != nil { + t.Fatal(err) + } + matcher, err := compiler.Compile(`{"type":"structural_tag","format":{"type":"grammar","grammar":` + string(g) + `}}`) + if err != nil { + t.Fatalf("compile: %v", err) + } + defer matcher.Close() + row := make([]int32, (grammarTestVocab+31)/32) + allowed := func(id int32) bool { return uint32(row[id/32])&(uint32(1)<"}}, want: `{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"any_text","excludes":[""]},` + `{"type":"optional","content":{"type":"sequence","elements":[{"type":"const_string","value":""},` + - `{"type":"json_schema","max_whitespace_cnt":32,"json_schema":` + schema + `}]}}]}}`, + `{"type":"json_schema","max_whitespace_cnt":22,"json_schema":` + schema + `}]}}]}}`, }, { name: "schema after thinking with two closings", req: llm.CompletionRequest{Format: json.RawMessage(schema), ThinkingClose: []string{"<|final|>", "<|final|>json"}}, want: `{"type":"structural_tag","format":{"type":"sequence","elements":[{"type":"any_text","excludes":["<|final|>","<|final|>json"]},` + `{"type":"optional","content":{"type":"sequence","elements":[{"type":"or","elements":[{"type":"const_string","value":"<|final|>"},{"type":"const_string","value":"<|final|>json"}]},` + - `{"type":"json_schema","max_whitespace_cnt":32,"json_schema":` + schema + `}]}}]}}`, + `{"type":"json_schema","max_whitespace_cnt":22,"json_schema":` + schema + `}]}}]}}`, }, {name: "thinking without a format", req: llm.CompletionRequest{ThinkingClose: []string{""}}}, } { diff --git a/mlxrunner/client_whitespace_test.go b/mlxrunner/client_whitespace_test.go new file mode 100644 index 00000000000..3974f64ef8f --- /dev/null +++ b/mlxrunner/client_whitespace_test.go @@ -0,0 +1,116 @@ +package mlxrunner + +import ( + "encoding/json" + "errors" + "io/fs" + "path/filepath" + "strings" + "testing" + + "github.com/ollama/ollama/llm" + "github.com/ollama/ollama/mlx" + "github.com/ollama/ollama/mlxrunner/xgrammar" +) + +// TestRequestGrammarBoundsWhitespaceRunLength drives the structural tag the +// runner decodes under through xgrammar one byte at a time (ADR 0035 checked the +// bound only on the tag's JSON). It proves the length of the longest whitespace +// run between tokens: 22 characters, the longest run llama.cpp's SPACE_RULE +// allows for a JSON schema on the GGUF path (two newlines and 20 spaces or tabs), +// which still fits a newline and ten levels of two-space indent. The runners +// agree on that length only, not on the shape: max_whitespace_cnt counts any +// whitespace, so MLX accepts 22 spaces in one gap where SPACE_RULE allows at most +// one space. muse-glimmer answering after its message header pads the first and +// last slot to the bound in single pass (2026-10-12), so the bound is also what +// that padding costs. +func TestRequestGrammarBoundsWhitespaceRunLength(t *testing.T) { + compiler := byteXGrammar(t) + schema := json.RawMessage(`{"type":"object","properties":{"a":{"type":"integer"}},"required":["a"]}`) + for _, tt := range []struct { + name, prefix string + req llm.CompletionRequest + }{ + {name: "format", req: llm.CompletionRequest{Format: schema}}, + {name: "format after thinking", prefix: "plan", req: llm.CompletionRequest{Format: schema, ThinkingClose: []string{""}}}, + {name: "json after a header", prefix: " to=user<|message|>", req: llm.CompletionRequest{Format: json.RawMessage(`"json"`), ThinkingClose: []string{" to=user<|message|>"}}}, + } { + t.Run(tt.name, func(t *testing.T) { + tag := string(requestGrammar(tt.req)) + for _, c := range []struct { + name, in string + accept bool + }{ + {"compact", `{"a":1}`, true}, + {"22 spaces", `{` + strings.Repeat(" ", 22) + `"a":1}`, true}, + {"23 spaces", `{` + strings.Repeat(" ", 23) + `"a":1}`, false}, + {"two newlines and 20 spaces", "{\n\n" + strings.Repeat(" ", 20) + `"a":1}`, true}, + {"a newline and ten levels of indent", "{\n" + strings.Repeat(" ", 10) + "\"a\": 1\n}", true}, + } { + if got := xgrammarAccepts(t, compiler, tag, tt.prefix+c.in); got != c.accept { + t.Errorf("%s: accepted = %v, want %v", c.name, got, c.accept) + } + } + }) + } +} + +const ( + byteVocab = 257 + byteEOS int32 = 256 +) + +func byteXGrammar(t *testing.T) *xgrammar.Compiler { + t.Helper() + path, err := mlx.LoadedLibraryPath() + if err != nil { + t.Skipf("native MLX payload is not built: %v", err) + } + pieces := make([]string, byteVocab) + for i := range 256 { + pieces[i] = string([]byte{byte(i)}) + } + pieces[byteEOS] = "" + compiler, err := xgrammar.New(filepath.Dir(path), pieces, byteVocab, []int32{byteEOS}, 8, 128<<20) + if err != nil { + if errors.Is(err, fs.ErrNotExist) { + t.Skipf("native xgrammar payload is not built: %v", err) + } + t.Fatal(err) + } + t.Cleanup(compiler.Close) + return compiler +} + +// xgrammarAccepts reports whether tag accepts in, fed one byte per token under +// the mask as a decoder would be, and can end there. +func xgrammarAccepts(t *testing.T, compiler *xgrammar.Compiler, tag, in string) bool { + t.Helper() + matcher, err := compiler.Compile(tag) + if err != nil { + t.Fatalf("compile %s: %v", tag, err) + } + defer matcher.Close() + row := make([]int32, (byteVocab+31)/32) + allowed := func(id int32) bool { return uint32(row[id/32])&(uint32(1)<