fix(mlxrunner): grammar-aware speculation — stages 2b/3, and the fix that makes them correct

[glennneuber]

First of four stages toward speculative decoding under a grammar. Nothing calls this yet — the serial constrained path is untouched and remains the only behaviour, so this is inert by construction.

Why

Constrained decoding runs one token per forward pass (constrainedDecoder.next) while unconstrained decoding is speculative. That is the whole of the cost. Measured on CUDA, gemma4:31b-nvfp4, one 1920×1080 image, n=3:

engine	format:"json"	unconstrained	penalty
MLX	21.13	36.53	−42%
GGUF	57.23	58.17	−1.6%

llama-server pays almost nothing because it verifies drafts against the grammar instead of giving up on drafting. Background: mlx-constrained-decode-disables-speculation.md (#187).

What this stage does

draftPrefix answers one question: given a draft, how much of it could the grammar possibly accept? Verification then never has to consider a token the format forbids.

It does not decide acceptance — the target model's masked distribution does, and that's the device half, not in this commit.

The load-bearing property

The caller's matcher is not mutated. Grammar state is a stack machine and a draft is a guess; advancing the live matcher over a guess that verification then rejects would leave the request generating against a state its own output never reached. So it walks a Clone() and returns it, for the caller to adopt over exactly the tokens that commit — which is why the count and the matcher come back together.

EOS ends the walk rather than extending it: Vocab.Mask admits EOS only where the grammar CanComplete, so one reaching this point is legal by construction, but nothing follows it and advancing past it would ask the matcher to consume sentinel text.

Tests

Nine, including the two that actually matter:

• the caller's matcher is unchanged after both an accepted and a rejected draft
• the returned matcher is advanced over exactly the accepted tokens (verified against a hand-walk)

plus truncation at the first illegal token, truncation mid-draft, immediate rejection, EOS stopping the walk, empty draft, and an out-of-range id ending the walk rather than panicking.

One test was wrong before the code was

It drafted { " } a expecting truncation at 2, reasoning that } can't close an object while a string is open. draftPrefix accepted all four — and was right: after {" the parser is inside a string, where } and a are ordinary content.

The premise is now : after {, which no open string can absorb. The reasoning error is recorded in the test comment, because the next reader will make it too.

Remaining stages

2. Per-row masked verification in speculationSession.accept, behind a flag. Each drafted position sits in a different grammar state, so this needs a [k+1, V] bias — ~4 MiB/step at k=3 on a 262,144 vocabulary. Affordable against a 42% penalty, but to be measured, not assumed.
3. Measure on the GPU host: acceptance under grammar, tok/s against the 21.13 baseline, and that output still parses as JSON every time.
4. Flip the default only if 2 and 3 hold.

Stage 2 touches accept, whose own comments warn that every exit between schedule and commit must drain the snapshot schedule — a bug there is silent KV corruption, not a crash. It wants its own review.

gofmt clean · go vet clean · go test ./x/mlxrunner ./x/structured: ok

🤖 Generated with Claude Code

[glennneuber]

Reviewing as consolidator. The staging and the inertness are right, and the load-bearing property is correctly identified.

Not mutating the caller's matcher is exactly the thing that would have been easy to get wrong and expensive to debug: a draft is a guess, grammar state is a stack machine, and advancing the live matcher over a rejected guess leaves the request generating against a state its own output never reached. Returning the clone alongside the count — so the caller adopts it over precisely the committed tokens — is the correct shape. Those two tests (caller's matcher unchanged after accept and after reject; returned matcher advanced over exactly the accepted tokens, hand-walked) are the ones I would have asked for.

The EOS reasoning is right for a subtle reason worth keeping in the comment: Vocab.Mask only admits EOS where the grammar CanComplete, so one arriving here is legal by construction — the walk stops because nothing follows it, not because it might be illegal.

And the corrected test is the best thing in the PR. { " } a truncating at 2 is what a reader reasons their way to; that after {" the parser is inside a string where } is ordinary content is what the grammar actually does. Recording the reasoning error in the comment rather than silently fixing the expectation is what stops the next person re-deriving it — and it is a genuine near-miss, since a test asserting an over-strict prefix would have made a correct implementation look broken.

One thing I would ask before stage 4 claims the win

The 42% is measured on scene_hd.png, one image. The suite's heaviest cells are three-image (multi_3img, and six of the eight bbox_contract_* probes), and there is direct evidence that cell shape changes this kind of cost sharply: #185's allocator sweep moved scene_single by 1% and bbox_contract_multi by 54%.

So when the device half lands, I would want the before/after on at least one three-image cell as well as scene_single. Two reasons it might differ: a longer prompt changes the draft-acceptance rate (unconstrained measured 0.76 with avg_draft 2–3, on a single-image prompt), and the JSON these probes emit is more deeply nested, which changes how often the grammar truncates a draft early. A 42% headline that turns out to be 15% on the cells people actually run would be worth knowing before it is quoted.

Happy to run that comparison here once the stage that changes behaviour exists — the harness, the fixtures and a clean CUDA host are all in place.

No objection to this stage. Inert, well-tested, and the one property that must hold is the one the tests pin.

[glennneuber]

Stage 3: it runs, it is 40% faster, and it breaks the format

Do not enable. OLLAMA_MLX_GRAMMAR_SPECULATION stays off.

gate-off  22.41 / 21.98 / 22.37 tok/s   eval=83   valid-json
gate-ON   29.28 / 31.41 / 32.83 tok/s   eval=88   INVALID JSON x3
          acceptance=0.00  avg_draft=0.49  max_draft=1

It is faster because it stopped doing the job. eval=88 is exactly what the unconstrained arm produced earlier, against the constrained path's 83 — the gated path is generating unconstrained output. The ~40% is the grammar penalty being avoided by skipping the grammar, not by verifying drafts against it. Throughput alone would have called this a win.

The remaining defect is the math, not the plumbing

acceptance=0.00 with drafts being produced means every step falls through to the residual or bonus token. targetDist is masked, but the residual comes from

targetDist.SliceRows(0, draftCount).ResidualAgainst(draftDist)

a masked target minus an unmasked draft distribution. That difference isn't a distribution over the masked support, and a token drawn from it need not respect the mask. Masked speculative decoding needs the draft side masked too, or the residual re-masked before sampling. I masked one side of a two-sided comparison — a design error, not a wiring one.

Three plumbing defects were found and fixed first

#	defect	symptom
1	truncation after scheduleSpeculation	cache: offset 34 exceeds 33 stored keys
2	reassigned candidates.tokens to an unpinned slice	Ints requires DTypeInt32, got BOOL — use-after-free
3	same reassignment broke the caller's defer Unpin	leaks the original, unpins the wrong array

2 and 3 were one decision: mutating a field whose pinning the caller owns. Verification now runs through local verifyTokens/verifyDist views. A draftCount of 0 returns errNoLegalDraft and falls back to a serial step.

What this settles about reviewing device code

Stage 2b shipped with the claim that reviewing it is not the same as running it. Four defects later that is a measurement, not a caution — every one compiled, passed vet, passed the host tests, and read correctly. The second was invisible at the Go level entirely: a pointer swap that only misbehaves once Sweep() runs.

The JSON parse in the harness is the only reason this isn't being reported as a 40% win.

Stages 1 and 2a stand on their own — pure host logic, 15 tests, no device dependency. If this PR is too large to carry, they are worth splitting out and landing; stages 2b/3 need the residual masking before they mean anything.

[glennneuber]

Stages 1 and 2a are split out into #193 as requested — pure host logic, 15 tests, no device dependency, green.

This branch is rebased onto feat/grammar-draft-host-logic and its base is retargeted to #193, so it now carries only the two commits that do not work: stage 2b (device bias, gated off) and stage 3 (the wiring that produces invalid JSON). The diff shown here is the untrusted half alone.

Blocked on the same thing as before: targetDist.SliceRows(0, draftCount).ResidualAgainst(draftDist) subtracts an unmasked draft distribution from a masked target, so the residual token need not respect the mask. Until that is fixed, the gate-ON speedup is just the grammar being skipped.

[glennneuber]

Root cause found, and it was not the residual

The note I left on 6fa913d blamed ResidualAgainst subtracting an unmasked draft from a masked target. That claim is false and no code changed for it. The -Inf bias makes the masked target's probability for a forbidden token exactly zero, so max(target − draft, 0) is zero there too, in both the sparse and dense layouts — and min(1, p_masked/q_unmasked) is the correct acceptance rule for sampling from the masked target with an arbitrary proposal. Masking the draft side would have been a pure throughput regression bought for nothing.

The actual defect was that the gated path never applied the grammar to a single emitted token. eval_count=88 matching the no-format arm exactly was the clue: not slightly-wrong output, entirely unconstrained output.

Three emission paths were unmasked or unadopted, and they compound:

1. park() built a plain pipelinedDecoder — raw logits, no mask, no matcher advance. The depth controller starts at zero, so the first round of every request took this path.
2. resume() drains a token the inner decoder sampled but never fed through its own next, so nothing advanced the matcher over it either.
3. accept() adopted commitIDs[:keep] but built that slice before appending nextID — the residual-or-bonus token every round returns. Even a fully working round left the matcher one token behind.

Any one desynchronises the matcher; then mask row 0 is the state after some earlier prefix, legal drafts look illegal, the round parks, and the parked token is emitted unadopted too. It does not converge, it compounds — which is why acceptance=0.00 and avg_draft=0.49: propose → errNoLegalDraft → park → resume, forever.

The fix

Every token the request emits is drawn from a masked distribution and advances the one matcher the session owns.

• park() under a grammar builds a constrained step over the session's matcher (constrainedStep), so a parked round is masked and stays level with the drafting rounds it alternates with. Sharing rather than cloning is load-bearing — park and draft alternate within one request.
• resume() adopts the token it drains, before close() unpins it.
• accept() adopts nextID as well as the committed drafts.
• adoptGrammar returns an error instead of discarding Advance's bool. The serial decoder has treated this as fatal since it was written; the gated path swallowing it is why the drift survived a whole campaign.

Also fixed, found while tracing: an admitted drafted EOS ends the mask sequence a row early, which indexed sampleTokenAt past the distribution and handed Sampler.Distribution as many rows as draft tokens — which it reads as a proposal-shaped call and silently shifts every row's repeat-penalty history by one. verifyPlan drops the drafted EOS instead, restoring the bonus row universally at a cost of one unaccepted draft per request.

maskedTargetLogits was dead code that disagreed with the block accept() actually ran, and pipeline.go's comment pointed readers at it. Replaced with grammarRows/verifyPlan/maskRows, all live.

Verification

isEOS is carried on the session rather than reached for through spec.r.Tokenizer, which makes the grammar bookkeeping constructible without a model. 15 host-side tests, no MLX — the bonus-row contract over every shape draftMasks can return, adoption, the illegal-token guard, and the invariant that the matcher tracks every emitted token. Full ./x/... suite green, gofumpt and vet clean.

Still gated off. The on-GPU arm has not been re-run yet — measurement to follow.

Not changed, deliberately

Sampler.Distribution's guard is one-directional (it panics only when rows exceed len(draftTokens)+1, and is silent when short because the proposal path relies on that shape). Tightening it to == would need a validation-only variant in shared sampler code; accept() now asserts the row contract on its own side instead.

[glennneuber]

Measurement status: not run — build environment, not the code

Six attempts to build a CUDA image on the 10.8.0.6 workstation, none completed. Recording this so nobody assumes the fix was measured.

• / is at 99% (26 GB free). The image build needs more transient space than that: two attempts were killed at the 3 GB mark to protect the other containers on the host.
• Routing buildkit to /mnt/8TB_SN850X_RAID1_BTRFS via a bind-mounted buildx volume works — root held at 26–28 GB for the whole build, and 74 GB of compile cache accumulated on the array. The pre-existing bigdisk builder does not help; its state volume is on / despite the name.
• The two runs that got furthest both died when the buildkit daemon container stopped mid-build — once SIGKILL (137), once a graceful goaway. No memory limit, 632 GB RAM free, no docker daemon restarts, no prune timer. Cause unidentified; something external to this session is stopping that container.
• Also worth knowing: with a genuinely cold cache the NVIDIA dnf install step failed once (it succeeded on retry). The layer cache may be load-bearing for building this repo from scratch.

A retry is cheap now — the builder gramfix and its 74 GB cache are still in place:

docker start buildx_buildkit_gramfix0 && docker buildx build --builder gramfix --load -t maxusai/ollama:gramfix /opt/github/MaxusAI/ollama

The measurement harness is written and ready: gate-off vs gate-ON, gemma4:31b-nvfp4, one 1920×1080 image, format:"json", n=3, reporting eval_count / tok/s / JSON validity / acceptance.

Acceptance criterion when it does run: eval_count must diverge from the no-format arm's 88 (it should land near the gate-off arm's 83), JSON must parse in all three repeats, and grammar_no_legal_draft must be well below the round count — an acceptance of 0.00 alongside a high grammar_no_legal_draft would mean the matcher is still drifting. Those two counters were added precisely so that reading is unambiguous.

Until then this stays draft.

[glennneuber]

Measured on CUDA. The format bug is fixed; the feature does not speed anything up.

RTX PRO 6000 Blackwell, gemma4:31b-nvfp4, visimgs/scene_hd.png (1920×1080), same
prompt in both arms, think:false, num_ctx:8192, warmed, n=3.

binary	format:"json"	unconstrained
published baseline (eb0ad43, ADR note)	21.13	36.53
pristine eb0ad43, re-probed (n=1)	22.35	—
main @ 640a349c (n=1)	22.98	—
this branch, gate OFF	23.63 (23.81 / 23.09 / 23.98)	37.47 (36.92 / 38.22 / 37.27)
this branch, gate ON	22.25 (21.96 / 22.48 / 22.32)	34.00 (34.64 / 33.08 / 34.29)

What is fixed

Every constrained response on every arm came back as bare, parseable JSON —
bare_json=True parses=True, no markdown fence. Stage 3 shipped output that ignored
the grammar entirely; f1d3d223 and 99dc5c88 close that. The correctness claim in
the writeup above holds on device, not just in the host tests.

What does not work

Gate ON is not faster — it is marginally slower, and the reason is that it never
drafts. The stats, taken from the same requests as the table:

iterations=126 drafted=0 accepted=0 acceptance=0.00 avg_draft=0.00 max_draft=0 \
  grammar_truncated=0 grammar_no_legal_draft=0 depth_over_time="0.0/0 0.0/0 ... 0.0/0"
iterations=121 drafted=1 accepted=0 acceptance=0.00 avg_draft=0.01 max_draft=1 ...

drafted=0 on nearly every round, and not because the grammar rejected the
proposals — grammar_no_legal_draft=0 and grammar_truncated=0, so the
errNoLegalDraft fallback never fired either. Nothing was ever proposed.
depth_over_time is flat zero for the whole request.

The reading the code supports: speculativeDecoder only proposes when
s.limit > 0 (speculate.go:232), and s.limit is refreshed solely in
endRound from depth.next(). The acceptance model only learns when
observed > 0 (speculate.go:167), and a parked round reports observed=0.
Depth starts at zero, so round 1 parks, which teaches the controller nothing,
so depth stays zero — a cold-start deadlock. The one round that did draft
(max_draft=1) was a cadence probe, and its single token was rejected, which
is evidence against raising depth.

For contrast, the same binary unconstrained climbs normally —
depth_over_time="3.6/5 5.1/6 5.0/5 ... 5.9/7", acceptance 0.72–0.93.

So the feature as it stands is correct and inert. The gap it was built to close
(−37% here, −42% published) is untouched.

Caveat on the unconstrained column

Gate ON reads 34.00 vs 37.47 unconstrained. The gate should not touch that path
at all. This host is shared and its load moved a great deal across the session,
so I would not call that a real regression on n=3 without re-running both arms
back to back on a quiet box. Flagging it rather than explaining it away.

Recommendation

Keep the gate off and keep the branch — the correctness work is worth landing
on its own, and the harness now exists. The next change is the depth controller:
it needs a way to enter a grammar-constrained request at depth ≥ 1, or to learn
from parked rounds, or speculation under a grammar can never start.

Method note: measured with a Go-only binary swap over
maxusai/ollama:0.32.14-rc0-dynres-mlxfix. The branch changes only x/mlxrunner/*.go;
git diff eb0ad43...HEAD over x/mlxrunner/mlx, MLX_VERSION, MLX_C_VERSION,
CMakeLists.txt, cmake/, llama/ and ml/ is empty apart from two Go files and a
README, so the native payload is identical and the swap is sound. Verified by sha256
inside the image and by main reproducing the published baseline through the same path.