The MLX runner accepted the API's format field but did not enforce it:
requests asking for JSON or a JSON Schema got unconstrained text, and
clients had no way to tell.
Enforce format with xgrammar: each sampling step masks the logits to
the tokens the grammar allows, so every emitted token and the end of
generation are valid under the constraint. Sampling, penalties, and
logprobs see the constrained distribution, and "json" yields a JSON
object, as the API documents and the llama-server path already
enforces. Only sampling waits on the mask; the forward pass is
dispatched before it, so constrained decoding stays pipelined.
The grammar engine is a dynamic library alongside MLX; when it is
missing, plain inference is unaffected and structured requests fail
with an explicit error. Constrained requests decode without
speculative decoding for now.
Decoding 256 tokens of a book-list schema on qwen3.8:27b-mlx (M5 Max,
seed 42, thinking off); pre-decode is the request time spent before
the first token:
unconstrained ~65 tok/s pre-decode ~70 ms
unconstrained, no draft ~32 tok/s pre-decode ~70 ms
JSON schema ~32 tok/s pre-decode ~70 ms
Schema and draft-less decoding are equal to within 0.1 tok/s in
paired adjacent requests, and a cold grammar compile adds nothing
measurable to pre-decode. The gap to unconstrained decoding is the
disabled draft model.
Fixes#16563
Co-authored-by: Daniel Hiltgen <daniel@ollama.com>
Token ids are int32 throughout the runner, so every caller reading ids
out of an int32 array narrowed the widened value right back. Make Int
and Ints return int32 and Float return float32, matching Floats, and
require the exact dtype instead of accepting and widening every
integer and float width: no caller read anything through those paths
but int32 tokens.
Ints and Floats also copied out of the array's buffer without
evaluating it first, so reading an array still in flight after an
async dispatch could return unwritten data, and correctness depended
on every call site remembering an explicit Eval. Evaluate in every
reader, matching the scalar readers, which already wait through item.
An available array costs a status check and an in-flight one waits
for its event; only a never-dispatched array evaluates a graph.
Nothing has set Grammar since the CGO engine removal took its writers
out; it survived as a read-only pass-through on the llama-server path
and a comment claiming it is set before dispatch. Remove the field and
the dead pass-through. llama-server keeps its wire-level grammar field,
which the "json" format conversion still uses.
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
A prefill that resumes partway into cached history — routine once
client timeouts interrupt long prompts — used to attach its captures
onto a node extended in place, so the stored snapshot spanned only the
tokens the prefill evaluated while the node's edge reached further
back. Restores walk node by node and trust each snapshot to cover its
node's edge; the short snapshot stranded the caches at mismatched
offsets and, on models with recurrent layers, ended up freeing all
cache state — a request matching 46k of a 47k-token prompt reprocessed
from zero.
Growth now never extends a node underneath its snapshots. New tokens
become a child node that carries exactly its own captures, and the
path stays compressed because non-user segments merge back into their
parent through the caches' snapshot Merge. Close already pages out
what it records, so every merge combines adjacent covered snapshots
and every stored snapshot spans exactly its node's edge.
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.
Page-in restores a path node by node and trusts each stored snapshot to
cover its node's whole edge. A capture taken during prefill spans from
the previous capture or the prefill base, which need not line up with
the node it lands on: when a prefill resumes partway into cached
history, a capture can reach back before its node's start, and a
capture landing on a node that already has snapshots replaced them
with a shorter span that page-in then could not serve.
Clip each capture to its node's edge on attach, and keep the snapshots
the node already has instead of replacing them.
A prefill settles the drafter with the seed token after its last chunk,
leveling the draft caches with the targets; a cancelled prefill
returned before that, leaving the targets one token past the draft
caches and the recorded keys. The next request then had to move every
cache, and models with recurrent layers, which cannot rewind, fell back
to the last snapshot: a retry after a client timeout lost up to a full
snapshot interval of the prompt it had just evaluated.
Settle with the next prompt token on the cancelled path too. The caches
then rest level with the recorded keys, and a retry resumes exactly
where the prefill stopped.
ModelOpt checkpoints apply a float32 global scale to every projection
output on top of the per-group quantization scales. Running the
multiply and the cast back to the activation dtype as separate eager
ops costs an extra kernel launch and a materialized intermediate per
projection.
Compile the multiply and cast into one kernel. On an M5 Max (medians
of order-swapped A/B runs against main; greedy outputs byte-identical):
qwen3.6:27b prefill 703 -> 769 t/s +7.9%
muse-glimmer:30b prefill 790 -> 843 t/s +6.7%
Speculative decode is unchanged within noise on both models. Only
checkpoints with a global scale are affected; single-scale nvfp4,
mxfp8, and affine checkpoints take the unchanged path.
Request options are the model's published parameters and the request's
own options layered over the server defaults, so the default
repeat_penalty of 1.1 reaches every model whose parameters leave it
unset. No maker of the library's current models recommends 1.1: their
generation configs either omit the penalty, meaning 1.0, or pin 1.05.
llama.cpp dropped the same 1.1 default in 2024; vLLM, SGLang, and
transformers apply no penalty. An always-on penalty also distorts
output that legitimately repeats tokens, such as code, JSON, and long
reasoning traces.
The penalty is especially costly for speculative decoding, where
drafts are proposed without it: the penalized target rejects drafted
tokens and the depth controller backs off. On muse-glimmer 30B (DFlash
on M5 Max, HumanEval) the 1.1 default costs 13-16% of end-to-end
throughput at greedy and temperature 1 alike, and drops prose
acceptance at temperature 0.8 from 0.44 to 0.30. On qwen3.6-35B it
cuts the mean accepted draft length from 4.3 to 3.5 tokens and makes
the controller stop speculating on prose.
Defaulting to 1.0 disables the penalty unless a model's parameters or
the request set one. Across the library:
- qwen3, qwen3.6, and qwen3-coder pin their own values (1.0, 1.0, and
Qwen's recommended 1.05) and are unchanged.
- Everything else local now matches its maker's no-penalty
recommendation, including gemma2 through gemma4, muse-glimmer, both
laguna 2.1 models, qwen3.5 (previously 1.1 stacked on its
presence_penalty of 1.5), gpt-oss, deepseek-r1 and v3.1, the
nemotron family, granite4, the mistral and llama3/llama4 families,
phi4, glm4, llava, and devstral.
- qwen2.5 recommends 1.05 but ships no parameters, so it moves from
1.1 to 1.0 and still needs a parameters layer to conform.
- Cloud models (kimi-k3, deepseek-v4-flash) never receive these
defaults.
Small older models may repeat themselves more without the penalty
masking it; the remedy is a per-model parameter, not a penalty applied
to every model.
One vision path serves every qwen3.5/qwen3.6 registration, dense and
MoE. Rope positions are precomputed at prepare time as the request's
layout — the family uses interleaved M-RoPE — while text-only requests
keep the fused 1D rope path, which is numerically identical for
uniform channels. Image expansions are causal for this family, so
prefill chunks split them. The MTP head embeds prompt tokens, so it
scatters the delivered image features and applies the same position
tables, keeping speculative decoding working on image prompts. The
merger's exact erf GELU adds an Erf op to the MLX bindings.
A checkpoint whose config declares vision must ship its tower: missing
vision weights or a deepstack_visual_indexes request fail the load
rather than silently serving text-only or skipping the injections.
Text-only checkpoints, which carry no vision_config, load as before.
Verified tensor-by-tensor against HF transformers for all eight family
members and live on every published -mlx tag; published towers are
already bf16, so no re-import is needed.
Each media item's features are encoded lazily when a prefill chunk
first overlaps its expansion and stay pinned until the expansion is
fully evaluated. A chunk never ends strictly inside an atomic
expansion: a bidirectional run's early rows attend its later keys, so
its first evaluation must cover the whole run in one forward. Items
marked Causal are exempt and split at any boundary.
Draft models need the same request state — reference MTP drafters
embed prompt tokens with the image features merged in, and an M-RoPE
drafter cannot compute positions without the request's layout — so the
layout is stamped on every forward, target and draft alike, and the
MTP session holds feature rows across its deferred flush. The dflash
drafter ignores media: its context rows are target hiddens.
A prompt that references media arrives as text containing [img-N] tags
plus the media bytes. Prepare now splits on the tags, tokenizes the text
between them, and hands the model the resulting segments — text runs and
media in stream order — in a single PrepareMedia call. The model returns
the expanded stream with each media segment's placeholder expansion
spliced in place, described per item so the runner can key identity and
schedule encoding, along with any opaque request-scoped layout state it
derives while building the stream. Building the whole stream in one call
is what lets a model derive values that span items, and lets it choose
item granularity (one per image, or one per independently evaluable
tile).
The runner validates the model-authored items before trusting them —
ranges ordered, non-overlapping, in bounds, and covering every media
segment, since prefix-cache identity is keyed on them.
Unknown tag IDs fail the request, media the prompt never references is
ignored with a warning, and duplicate references are allowed, matching
the previous engine. A media request still produces no image output:
nothing feeds the features to the model yet, and no model implements
the media interface.
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.
MLX checkpoints that include a vision tower are already tagged with the
vision capability at import, so the server accepts image chats and ships
the image bytes with the completion request. The MLX client dropped the
bytes, and the prompt's image tags were answered as literal text.
Carry the media through to the runner and fail the request with a clear
error when the loaded model has no media support. Nothing implements the
new media interface yet, so every media request now returns the error
rather than a silently wrong answer; later changes build the image path
on top of the same interface.
Vision towers are much more sensitive to weight quantization than
language layers: measured against the reference encoder on a real
image, 4-bit types and scale-only mxfp8 distort the projected image
features by 26-34% mean relative error (worst tokens near-orthogonal),
which shows up as degraded image recognition — down to complete
blindness for the small e-series towers under nvfp4. Affine 8-bit was
the only quantized format that matched the bf16 tower.
Keep vision tower tensors at source precision instead, matching the
audio tower's treatment and every vision component Ollama publishes in
GGUF form, including gemma4's own GGUF tags, which ship f16/f32 vision
beside 4-bit language weights. Towers are small and run once per image,
so neither size nor decode bandwidth argues for quantizing. Existing
MLX imports keep their quantized towers until re-imported.
Implements the DFlash draft checkpoint format: a few decoder layers
over fused target-layer outputs, which enter every layer as key/value
context while the block being drafted supplies the queries. The draft
has no embedding table or output head of its own; it borrows the
target's.
One model covers the known checkpoints. Attention weights normalize at
load to a q projection plus a fused k|v, stacking split checkpoints and
slicing fused ones, exact for quantized tensors; gate and up fuse the
same way. Optional tensors decide the output gate and per-tap norms,
config decides attention shape, and the architecture name decides only
laguna's context-norm convention.
A manifest can pair any draft with any target, so construction checks
the fit: tap ids inside the target's layers, matching hidden width, and
the target vocabulary covering the mask token. A bad pairing fails at
load.
A DFlash draft proposes a whole block per forward, which doesn't fit
the MTP session's one-token-per-call chain. Add a second drafting
session for block drafts: committed target features write straight
into the draft's context caches, and each round drafts a block in one
forward and samples it in one batched call, rolling the block's cache
entries back with the same mechanism speculative rounds use on the
target caches. The depth controller's search is capped at the deepest
draft the drafter can produce, since a depth it can never measure
would otherwise always look best.
Distribution aligns its rows with the end of the draft chain, so when
the caller passes no chain, every row sees the slot history unchanged.
That case already worked; only the row-count guard rejected it. The
guard now applies only when a chain is present, which is where more
rows than chain positions would silently drop history. A block drafter
needs the chainless case to sample its whole proposal batch in one
call.
The runner used to build caches by probing the model for an optional
NewCaches method, with one KV cache per layer as the fallback. A model
with a draft head appended the draft's cache slots to its own list, and
the speculative engine later recovered the two groups by comparing slot
identities, panicking when the lists didn't line up.
NewCaches is now a required method on both the model and the draft, and
each returns only the slots it writes. The runner concatenates the two
lists for the prefix cache and passes them to the speculative engine
separately, so snapshots and rollback apply to the target's slots and
the draft forward receives both groups as arguments. The identity
comparison, its panics, and the per-request rebinding are gone; the two
groups are fixed at load time.
A draft model conditions on state that the target produces during its
own forward pass. For an MTP head or an assistant model that state is
the final hidden state; for a block draft it is the concatenated
outputs of several layers. The choice belongs to the model, so Forward
now returns the conditioning state along with the hidden state to
unembed. Models without a special conditioning state return the final
hidden state for both, and the decode paths hand the value to the
drafter without looking at it.
A lazy KV snapshot indexes into the cache's live buffer instead of owning
a copy, so it must be copied out before an append overwrites the slots it
names. appendKV checked for that only on the first append after a rewind,
and only against that append's own range: a still-lazy snapshot further
ahead in the buffer was overwritten without a copy when a later append
reached it. This happens when a request reuses a short prefix of a longer
cached conversation and prefills past one of the old conversation's
snapshots; restoring that snapshot later silently serves the new request's
KV in place of the old conversation's.
Scan every append instead. The overlap test already limits copies to
snapshots the current write clobbers, and appends outside a rewind refill
sit above every snapshot, so the steady-state scan walks a short list and
finds nothing. This restores the invariant Restore's lazy fast path relies
on: a snapshot still in its lazy state has never been overwritten.
Load the MTP head from the mtp.* tensors instead of freeing them and implement
Draft to propose one token per step, gated solely on the tensors being
present; a model whose head ships inline is its own draft via base.SelfDraft.
The runtime keeps sole ownership of the +1 RMSNorm shift (conversion passes
tensors through verbatim), and the head's norms shift under the same
original-format detection as the main stack, so nothing shifts twice.
Decode-length scans spend more time in launch gaps than math: the q/k
norms, decay gate, and recurrence each dispatched separately per layer.
Fuse the step into one Metal kernel over the activated conv output,
with per-token boundary states available from the same pass. The graph
implementation remains as the fallback and contract-miss path, and pins
the kernel bit-for-bit in the parity test.
The activation belongs to the conv stage: downstream consumers see
activated values however the conv is computed. WithConvSiLU routes to a
fused depthwise conv+SiLU kernel when the conv fits its contract and
the same computation as graph ops otherwise; cached conv state is the
raw input tail, unaffected by activation placement.
Each custom kernel repeated the same host-side creation and launch
boilerplate plus a CUDA-then-Metal-then-graph dispatch at every call
site. gpuKernel declares the sources (either backend may be absent) and
a graph fallback; run executes the first that works.
Split checkpoints ran four input projections per recurrent layer, and
native combined checkpoints paid a per-forward slice-and-concat to
rebuild the contiguous qkv rows the causal conv consumes. Normalize
both at load to packed [q|k|v|z] and [beta|alpha] rows: split tensors
concatenate, native interleaved tensors permute once. The forward keeps
a single projection path, and the packed rows are the layout a fused
scan can consume directly.
Pairs with mismatched quantization dequantize before packing rather
than keeping a split fallback path alive.
The C-level dequantize accepts a global_scale argument but rejects it
on the Metal backend, so dequantize-fallback sites hand-rolled the
same post-multiply. Take the scale in the Go wrapper and apply it on
top of the op, cast back to the output dtype. The quantized embedding
passes its scale; laguna's expert paths keep their own multiplies,
which shape per-expert scales and differ on result dtype.
mlx item<T> reinterprets without checking the dtype, so Array.Int's
8-byte read of int32 scalars took in neighboring pool bytes — masked by
Metal's zeroed allocations, corrupting token IDs on CUDA's warm pool.
Read at the element's width.
The per-request stats are the main diagnostic for speculative
throughput, so log them at info; the controller line stays debug.
Recording chosen depths at the next beginRound dropped rounds with no
successor, so record at endRound and count resume as a depth-0 round.
Draft token embeddings were kept at source precision. A draft that
reuses its embedding as the output projection (the gemma4 assistant)
then reads the whole 537MB bf16 tensor on every draft step — about half
the step's cost. Draft quality only affects how many drafts are
accepted, so the output head now takes the requested type instead of the
8-bit type that protects a target's output quality.
gemma4:26b-mlx, M5 Max: MTP code decode 148 -> 157 tok/s (+26% -> +37%
over plain); prose goes from roughly zero to +2-5%; acceptance unchanged.
Gathering gate and up separately cost a third expert gather per MoE
layer. Keep gate_up packed as one tensor, joining it at load when the
checkpoint ships the halves separately, and split the gather's output
instead.
Output is byte-identical; decode is 4% faster (7.89 -> 7.58 ms/token on
M5 Max) and prefill 9% faster.
The lm_head rule was asymmetric: the fp modes kept an untied head at
source precision (even under mxfp8, leaving it the only bf16 matmul in
the model), while int4 quantized it at 4 bits with no promotion. The
tied-embedding overrides (gemma4, cohere2moe) already resolve the head
to the 8-bit family type and hold quality close to bf16.
Apply the same decision to untied heads: the 8-bit type in the
requested family when it fits the shape, source precision otherwise.
int4 now promotes the head to int8, and the fp modes quantize it to
mxfp8 instead of keeping bf16.
Per-token cost of the batched head forward keeps falling until the
flush is large enough to reach the fastest kernels: NAX matmul tiles
for dense heads, and the segmented gather path for MoE heads, which
needs tokens*topK/experts >= 4. Measured across the qwen3.6 heads,
256 is the smallest cap past every threshold and within a few percent
of each head's per-token floor. The cost is bounded: up to 2.5 MiB of
pinned hiddens per request and a flush stall under one decode step.
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.
Keeping the recurrent conv state small was handled unevenly: the committed
live state was recopied on every commit — wasted work on single-token
decode, where the window is already tiny — while boundary states captured as
snapshots could still be plain slices of the forward-sized convolution
buffer. A cached slice pins that whole buffer even though the trie's eviction
accounting only counts the slice's bytes, so recurrent cache memory piled up
across requests and eviction could never reclaim it.
Compact each boundary state to its real size once, where it is produced in
the conv wrapper, so live state and snapshots own only their own bytes and
eviction sees the true cost. Single-token decode leaves the already-tiny
window as a slice.
Fixes#16698
The MTP validation forward schedules a snapshot at every drafted token, which
made CausalConv1D re-run the depthwise conv once per segment to recover each
boundary's conv tail. A conv boundary state is just the trailing convTail
input positions, so run the conv once over the whole window and slice each
boundary tail from the shared buffer, removing the per-token conv launches.
Qwen3.5/Qwen3-Next architecture strings contain the substring "qwen3", so the
broad qwen3 match claimed them for the generic parser and qwen3-coder
renderer, whose template doesn't frame the thinking block — an empty
<think></think> leaked into content and think=false was ignored. Match the
family first via isQwen35Family so the parser, renderer, and
thinking-capability checks share one variant list.
The heuristic schedule grew the draft toward a fixed cap on acceptance alone,
maximizing accepted-tokens-per-step rather than throughput, and on a
steep-forward target it regressed below no speculation. Replace it with an
engine-level controller that drafts the depth maximizing
committed-tokens-per-wallclock from live per-position acceptance and persisted
per-width forward cost, with no draft-length cap; the heuristic schedule and
the OLLAMA_MLX_MTP_* env vars go with it.
Acceptance took two blocking evals per round: one to read the accepted mask,
then a second for the bonus or residual token whose graph needed the
host-known rejection point. Sample the residual at every rejection point in
one batched draw alongside the bonus row, so a single eval covers acceptance
and the next token.
Each speculative round ran the target stack twice — once for the current
token's hidden and base logits, once to validate the drafts — capping
throughput below plain decode. Fuse them into one forward over [current,
draft_0..draft_{N-1}], whose hidden rows already line up with the acceptance
math, so the separate base-logits unembed disappears from the drafted path.
Sampler.Distribution built row i as if draftTokens[:i] were appended, leaving
a single-row proposal call with no draft history, so a drafter skipped the
repeat/presence penalties the target's validation applies and re-proposed
penalized tokens. Align rows with the end of the draft chain instead: the
final row sees every draft token, each earlier row one fewer.
Generalize the draft path so a head that maintains a KV cache (EAGLE-style)
and Gemma's read-only single-position assistant both fit one drafter
interface with no per-model branches, and make the committed stream the
drafter's maintenance mechanism — every committed run is reported, the
drafter pairs each draft slot with its look-ahead token and flushes completed
pairs to the draft caches. The draft KV thus stays prefix-cached alongside
the target in every session, drafting or not.
The pipeline and the MTP decoder each owned a decode loop with duplicated
prefill, budget, and emission handling. Split the pipeline into prefill and
decode phases behind a decoder interface, with the decode loop the sole
emitter enforcing the NumPredict budget, and split speculation into a generic
engine that returns the accepted run and a drafter interface that owns only
how proposals are made.
Greedy is a special case of sampled decoding — at temperature 0 the sampler
yields a point mass, so rejection-sampling acceptance reduces to argmax-match
— so collapse the separate greedy, sampled, and serial paths into one. MTP
now honors any temperature, penalty, and top-k/p/min-p setting; logprobs
remain the only gated feature.
The batched MTP accept paths advance the cache by the whole accepted run
before streaming it to the client. If the stream was cancelled partway
(e.g. the caller disconnects), the loop returned before recording the
remaining accepted tokens, leaving the cache offset ahead of
session.outputs. close() then indexed the token log past its end and
panicked with a slice-bounds error.
Record the whole run to session.outputs before streaming any of it, so a
cancelled stream can no longer desync the cache from the token log.
The same bug is present on main, with identical mechanics: the accept
paths there commit the cache to before+accepted and then stream in a loop
that returns on cancellation before recording the rest.
Prefill no longer splits its batch at each requested snapshot offset. The
session schedules the pending offsets on every cache before prefill, runs the
forward in full-size chunks, and attaches the captured snapshots to the trie
afterward. Offsets the prefill never crosses (it leaves one token for decode
seeding) are dropped instead of materializing a node for tokens never written,
and snapshots from an abandoned prefill are released on session close.
Speculation used a parallel hierarchy of wrapper cache types that shadowed
the live caches and reconciled against them on commit. Replace it with
snapshot/restore on the live caches themselves: a cache snapshots itself as
a write crosses each offset, and the runner commits a batched draft by
restoring to the accepted count. The wrappers and the comparison plumbing
around them are gone.
Snapshots are lazy. A KV or rotating capture indexes into the live buffer and
owns no memory until a destructive write forces a copy-out, so rejecting a
draft is free.
Recurrent layers now validate in the same batched pass rather than falling
back to serial. A gated-delta layer reports its interior split offsets and
hands back the recurrent state at each one, which the cache records as a
snapshot.
CausalConv1D and GatedDelta now run their scan in segments cut at optional
WithSnapshotSplits offsets and return the recurrent state at each boundary
instead of just the final state. The output is identical to the unsegmented
scan; segmenting only adds a few kernel launches, not extra recurrence compute.
This lets a batched forward capture interior recurrent state without re-running
the scan, which the cache will use for speculative validation rollback points.
RecurrentCache.Put and the Qwen3.5 layer now thread the boundary-state slices,
committing the final entry as the live state.
cache.go had grown to hold every cache kind. Move KVCache (and its
speculative wrappers) to kvcache.go and RotatingKVCache (and its
sliding-window mask applier) to rotating.go, leaving cache.go with the
shared interfaces and the Speculation transaction. Pure relocation;
no behavior change.
Work that panics on the locked MLX worker goroutine was recovered and
re-raised on the caller, so the printed trace pointed at the re-panic
site in this package rather than the code that actually panicked.
Capture the worker stack at recovery and carry it through a value that
implements error, so the runtime prints the original location in the
fatal trace.
Split the gated-delta Metal/CUDA kernels' dtype template into separate
input (InT) and state (StT) types so activations can stay in bf16/fp16
while the accumulated delta state stays in float32. Allocate the delta
state and qwen3_5's no-cache zero state in float32 to match.
Previously the draft architecture was hardcoded to
Gemma4AssistantForCausalLM. Read it from the draft model's config so
any draft architecture can be packaged.
This reverts commit 98e26b8c37.
The DFlash integration is too invasive to keep at this stage: it
threads DFlash-specific logic through the pipeline, base model
interfaces, and the cache layer. The recurrent cache also now
has qwen3.5 model-specific code. Revert it now and reintroduce
the self-contained, generally-useful pieces (YaRN RoPE DRY-out, draft
architecture autodetection, gated-delta fp32 state) as separate
follow-up commits.
Models build their own attention masks and read K/V directly from
the cache's buffers, which ties them to the cache's storage layout.
That blocks multi-sequence batching — right-padded rows need a
query-padding mask composed onto every model — and rules out
variants like paged attention where K/V isn't one contiguous tensor.
Caches now hand back a per-layer KVHistory holding post-update K, V,
and a MaskApplier that merges the cache's storage restrictions into
the model's logical mask. Models describe their mask in logical
terms; SDPA composes model, padding, and applier contributions and
dispatches to the kernel's causal or no-mask fast path when it can.
KVHistory still exposes K, V, and the composed mask for manual
attention paths (e.g. CUDA prefill at head_dim > 128).
Performance for single-sequence inference is unchanged.
Switch RoPE from the scalar-offset kernel (mlx_fast_rope) to the
array-offset one (mlx_fast_rope_dynamic) so each batch row can start
at its own position. The pipeline tracks the current position locally
and passes it to the model through Batch.SeqOffsets; each model
materializes that slice into an int32 array for the RoPE call.
Single-sequence behavior is unchanged; this is the wiring needed
before the runner can batch independent sequences.
Gives a single extension point for per-call context (positions,
sequence IDs, masks) as multi-sequence batching grows, without having
to churn every model's Forward signature again.
Register sequences with Add/Remove; each Sample call takes any subset of
registered slots and samples one token per row, appending to each slot's
ring-buffer history. When all slots share Options and penalty rings are
full, one fused transform pass runs over the whole batch via a persistent
pooled history tensor; otherwise calls fall back to per-slot serial
processing indexed against the same pool.
Performance is unchanged for a single sequence, which is all that is
exposed for now.
AppendToken used to concatenate the new token onto the history tensor
and slice it back to RepeatLastN every decode step, churning the graph
shape and reallocating a fresh tensor each call. The stateful penalties
don't care about order within the window, so a fixed-capacity ring with
one SliceUpdate per append keeps the tensor shape constant across
steps.
Move tokenization out of the single GPU processing goroutine and
into each request's HTTP handler goroutine. This allows the next
request's prompt to be tokenized on the CPU while the current
request is executing on the GPU.
Use atomic.Int32 for Array.pinned and a sync.Mutex for the global
arrays slice so MLX arrays can be created and pinned from multiple
goroutines without racing on those structures. Convert Array value
receivers to pointer receivers and struct fields from Array to
*Array to avoid copying the atomic.
This does not fully achieve thread safety even when building
completely independent graphs. The tracing flag and traceScratch
slice in compile.go are unprotected, so concurrent Compile calls
will race. MLX itself is not fully thread-safe either although
it is working to improve.
When both filters are active, avoid paying for a full sort in top-P
and a partial sort in top-K. Single-filter paths are unchanged.
Improves generation throughput on gemma4:e4b by 1.5%.
Match the ollamarunner and OpenAI semantics: raw, full-vocab log-softmax
with the top-K ranked by probability. Skipped on the GPU when the request
doesn't ask for logprobs so decode doesn't pay for it otherwise.
DeepSeek-V2-style aux-loss-free routing computes sigmoid(gates) once but
needs it twice: the raw sigmoid output is gathered after top-k, while the
post-bias negation is the argpartition key. Fuse into a single multi-output
Compiled kernel returning both, saving two launches on the routing path
per token. Exposed as a general SigmoidRouter since the same pattern is
shared across DeepSeek-V2 descendants.
Improves glm4.7 generation performance by approximately 1%.
Converts SiLU/GELUApprox to compiled kernels and adds SwiGLU,
matching upstream mlx/mlx_lm's activations pattern. Routes llama,
qwen3, qwen3_5 (dense + MoE), and glm4_moe_lite MLP paths through
mlx.SwiGLU so each MLP invocation runs as one fused Metal/CUDA
kernel rather than a chain of per-op launches.
Wraps MLX's mlx_compile API so Go functions can be traced into fused
kernels. Contiguous elementwise chains collapse into a single
Metal/CUDA kernel instead of launching one per op.
Exposes Compile plus arity helpers (Compile1/2/3) that mirror Python's
@mx.compile decorator shape, lazily building the closure on first call
so package-level declarations work before the MLX dylib loads.
cublasGemmBatchedEx fails during graph capture when pool allocations
return fake pointers. This is triggered when NUM_PARALLEL is greater
than 1 for models like gemma4 that use batched matmuls. Skip it
during reservation since the memory tracking is already handled by
the pool allocations.
Fixes#15249
Add periodic snapshots every 8k tokens and near the end of the prompt
so that long prompts can be partially restored and thinking/generation
can be retried without full reprocessing.
Update LRU last used time just on the nodes that actually used
during processing rather than all snapshots along the path. This
allows eviction to remove nodes more accurately so we can avoid
other heuristics to auto-merge nodes.
mlx.Copy shares the backing buffer with its source (via
copy_shared_buffer) rather than allocating independent storage.
When used to snapshot a slice of the KV cache, the snapshot array
holds the entire original cache buffer alive through the shared
data pointer — even after eval detaches the computation graph.
Replace Copy with Contiguous in Snapshot and Split. Contiguous
allocates a compact buffer when the source buffer is significantly
larger than the logical slice (Contiguous::eval checks
buffer_size > nbytes + 16384), which is always the case for KV
cache slices.
Previously, a partial match within a node's edge would truncate the path
to the parent snapshot - effectively making all cache types behave as
recurrent caches. Caches with only transformer layers can rewind to
arbitrary boundary so this restores this capability to improve cache
hits
Receiving from a buffered chan error consumes the value, so only the
first caller (WaitUntilRunning, HasExited, or Close) sees the signal.
Subsequent receivers block or take the wrong branch. Replace with a
closed chan struct{} which can be received from any number of times,
and store the error in a separate field.
The stderr reader used bufio.Scanner which has a 64KB max line size.
If the subprocess wrote a line exceeding this limit, the scanner would
stop reading, the OS pipe buffer would fill, and the subprocess would
deadlock.
Replace the scanner with a statusWriter that wraps io.Copy. The writer
forwards all stderr to os.Stderr while capturing the last short line
(≤256 bytes) for error reporting, avoiding both the deadlock and the
need to buffer arbitrarily long lines.
Enable multiple conversations to reuse cached computations when they
share token prefixes (e.g. the same system prompt). A prefix trie
tracks shared regions so switching between conversations only
recomputes tokens that diverge. Inactive conversation state is paged
from active GPU memory to other memory and restored on demand, with LRU
eviction to keep memory usage bounded.
Slice used cmp.Or to resolve a zero stop value to the dimension size,
intended to support open-ended slices like a[i:]. This made Slice(0, 0)
indistinguishable from Slice(), so any slice with a zero stop would
silently include the entire dimension instead of being empty.
Replace cmp.Or with an explicit End sentinel and resolve negative
indices against the dimension size, matching Python/PyTorch semantics.
Claude Code sends an x-anthropic-billing-header that changes on every
request. This is embedded in the system prompt and consequently
breaks the KV cache for every request. Given the size of the prompts
that Claude Code usees, this has significant performance impact.
MLX runners (image generation and LLM) previously bypassed the
scheduler's standard load path via a separate loadMLX method. This meant
they skipped VRAM fitting checks and couldn't participate in model
eviction.
Now all model types flow through the same load function. Model eviction
for MLX is based on weights as KV cache and compute graph are dynamic.
This means that eviction does not take into account the worst case
memory and models can still compete for memory but it is a significant
improvement.
Only the last token's processing time is included in prompt processing,
giving an artificially high rate. In addition, the number of tokens
only included the tokens that miss the cache, instead of our historic
total tokens.
Currently, context length is unbounded - the cache will keep
growing forever independent of the model's trained context
length. This caps it and enforces semantics similar to most
cloud services:
- Long prompts will result in an error, not truncation.
- Generation that exceeds the context will be stopped
Errors that occur during pipeline processing are currently only
logged but not sent back to the client. Rather than using HTTP
status codes as we have historically done, this serializes errors
as messages to allow sending them at any time during the stream.
The MLX runner previously reported a static VRAM estimate that was
computed at load time and consisted only of the weights. This is
strictly less than the actual memory usage, as it does not include
the KV cache or compute graph.
When the entire prompt was already cached (e.g. repeated prompt),
findRemaining returned an empty slice, causing FromValues to panic
on an index-out-of-range accessing a zero-length byte slice.
Fix by always keeping at least one token to re-evaluate so the
pipeline can seed token generation. Also reject empty prompts
early rather than panicking.
Currently, a canceled request can result in computation continuing
in the background to completion. It can also trigger a deadlock
when there is nobody to read the output tokens and the pipeline
cannot continue to the next request.
Particularly in error cases, it can be difficult to ensure that
all pinned memory is unpinned, MLX buffers are released and cache
state is consistent. This encapsulates those pieces and sets up
proper deferrals so that this happens automatically on exit.
Pass subprocess stdout/stderr through to the parent's stderr directly
instead of re-wrapping each line with slog. The subprocess already
writes structured slog output, so the re-wrapping produced nested
timestamps, levels, and message fields that were hard to read.
Also downgrade verbose KV cache debug logs to trace level.
The KV cache previously used a tree structure which could
store multiple divergent sequences, which is good for cache
reuse. However, this is typically used in conjunction with
paged attention so each node in the tree can store just a
chunk of the KV cache and they can be stitched together later.
We don't currently do this, so the cache was storing copies of
the full cache for each past sequence.
This redundancy plus the lack of resource limits, caused significant
memory use as a conversation grew. Instead, this changes to store
a single entry for the cache, which can be prefix matched. Although
it is less ideal for multiple users, it largely matches Ollama's
current behavior. It can be improved as additional pieces are fleshed
out.
Jesse Gross