A long prompt records restore points during prefill, but they only
reached the prefix trie when the prefill completed; a cancelled request
closed and released everything it had captured. Agent clients routinely
cancel long prefills — their timeouts are shorter than the minutes a
40k-token prompt takes — so every retry started the whole prompt over
and never got further than the timeout allowed, which presents as the
model hanging forever.
Closing a session now attaches every snapshot the prefill crossed, so a
retry resumes from the last one and makes progress across timeouts.
Scenario tests cover retries resuming exactly where a cancelled attempt
stopped and cancellations on divergent conversation variants.
Fixes#17839
When a session closes, every cache rests exactly at the end of the
segment the trie is about to record. That is the one moment the
segment's state can be captured for every layer, so close now pages
the new segment out itself instead of recording it without snapshots
and leaving the capture to a later path switch.
Path switching then has nothing left to capture and only rewinds and
pages in. The whole-state entry taken at close is released when the
next request grows past the segment; sliding-window layers pay the
same window copy a scheduled capture already costs.
Add MLX support for Nemotron 3 Nano Omni, including the model implementation, Mamba2/recurrent pieces, MoE routing, and quantized NVFP4/MXFP8 expert paths.
Use a shared mapped MoE GatherQMM fast path under the generic moe_gather_qmm_mapped naming, with Metal-optimized NVFP4/MXFP8 block-mapped kernels and generic fallbacks for unsupported backends.
Serve the model's multi-token prediction head as a self-draft speculator, so speculative decoding needs no separate draft model.
Render the Nemotron prompt from the published chat template. The template the renderer was based on had drifted from the current reference; refreshing it surfaced five mismatches: stray leading newlines, the wrong turn separator and a trailing newline before the generation prompt; /think and /no_think toggles left in user turns; a trimmed system message the template leaves intact; a user block opened by a leading tool message; and Go scalar syntax for schema extras where the template applies Python str(), sending true/false/<nil> in place of True/False/None. Reference tests now render every case through the template itself.
Also harden the Nemotron parser path shared by both backends: while collecting thinking, preserve whitespace before partial </think>, <think>, and <tool_call> fakeouts, with streaming tests covering those cases.
Media placeholders repeat one token ID, so two prompts with different
images would produce identical trie keys and falsely share cached state.
Substitute a per-item hash of the media bytes and preprocessing shape
across each item's expansion range at the key layer; the model still
sees real token IDs. Fold values carry a bit no token ID has, so a media
stream can never alias text, and the bigram packing for draft caches
composes unchanged, so draft restore points inherit the same identity.
Text-only prompts key exactly as before. Nothing records media items yet;
the change is inert until the prompt preparation wires them.
A draft cache pairs each slot with the token that follows it, so the
deepest stored pair always names one token past what a prefix match
can verify - at generation end, the sampled-but-never-committed
final token. Restoring at the match point reuses that pair blind: a
stop token stripped from the next prompt, or any divergence at the
boundary, leaves it stale, and pairing never rewrites below the
resume position, quietly lowering draft acceptance.
Key the trie by token pairs instead: the key for offset i packs
(token i, token i+1), so matching k keys verifies k+1 tokens and
every match is a valid restore point. A pair is reused only if the
token it names matched, and prefill re-evaluates the boundary token,
rebuilding its pair with the token that actually follows. A token
gets a key only once its successor is recorded, so endings record the
final sampled token - never forwarded - and the trie stays level with
the caches. Without a look-ahead the keys are the tokens and behavior
is unchanged.
The recorded tokens' slice bounds used to reject state past them for
free; close now checks the invariant against the stored keys
directly. The test harness rests requests the way the pipeline does -
the deepest recorded token never enters the caches.
The caches, the speculation binding, and the drafter were each built
lazily inside the first request: begin constructed the caches, and open
bound the cache partition and made a fresh drafter every time. All of
it is a property of the loaded model, so build it once at load.
newPrefixCache replaces the lazy construction in begin, speculation
binds when it is created, and the drafter splits the way speculation
does: a persistent mtpDrafter constructed at load opens each request's
mtpDraftSession, whose constructor syncs the pairing cursor to the
draft caches' restored offset.
