[pull] main from ollama:main - #828
Merged
Merged
Conversation
KV snapshots must cover a node's edge exactly. Recurrent and sliding-window state is only useful at a node's end, and a node may have none: a request resuming there lands on the previous checkpoint and begin schedules a capture at the match. The header claimed every node carries its snapshots from creation, which a node split out of an existing edge at close cannot. That hid a gap: when a response is a prefix of a stored one, close lands on the split-off head with the caches resting at its end, and pageOut skipped the capture because the node already had a KV snapshot. Restate the header as the rules that hold, and make pageOut capture whatever layers a node is missing. The scheduling comment also said eviction preserves user nodes; it only resists compaction.
The tokens of a non-causal media item attend to each other in both directions, so the item has to be evaluated in one forward. Prefill honors that when it picks chunk boundaries, but the prefix cache did not: a snapshot could be taken partway through an item, and a request that resumed there would evaluate the rest of the item alone and compute different attention for it. Snapshots scheduled inside a non-causal item now land at its end, and a match that ends inside one resumes at its start.
When a request resumes partway through a cached edge, the node holding that edge was dropped from the active path, because the path has to end at the live offset for close and the prefill captures to extend the trie from its last node. Off the path, the node was an ordinary leaf with a stale last-used time, so eviction removed it first. The captures taken during that request start at the resume offset, but attach rebuilds the missing node from the path's last node, so the new node's edge begins earlier than its KV snapshot. A later request resuming there was refused by the KV cache and re-prefilled from scratch. If eviction first merged the node into its parent, the two snapshots were concatenated as if adjacent, and the restore reported a hit while the buffer held tokens from other positions. Split the node at the live offset instead. The head stays on the path and gets the last-used update. Only the unused tail can be evicted, and losing it costs nothing. The split only happens when every layer can rewind into the edge, so it never involves a recurrent layer, and the head gets the same KV-only snapshots a close-time split already produces. When the request follows the edge, compaction merges the halves back.
…budget
Eviction skipped every node on the active path, so a conversation's own
turn checkpoints were never reclaimed no matter how far over budget the
trie was. On models with sliding-window or recurrent layers each turn's
checkpoint is a full copy of that state, 800 MiB per turn on
gemma4:31b-mlx, and a long chat grows without bound. The scheduler then
counts that memory as in use and evicts the model to load anything
else.
Only the frontier and branch points are protected now. Any other node,
active or not, is evicted least recently used first. On the active path
that merges a turn into the next one: the merged node keeps the newer
whole-state, and the KV snapshots there are lazy views of the live
buffer, so nothing is copied. Rewinding to an evicted turn resumes at
the newest surviving checkpoint before it.
qwen3.8:27b-mlx on an M5 Max, the same short question every turn with
24 tokens generated per reply, 8 GiB budget, 17.2 GiB of weights:
turn | before: paged out nodes reported | after: paged out nodes reported
11 | 4.61 GiB 33 21.6 GiB | 4.61 GiB 33 21.6 GiB
21 | 7.91 GiB 56 24.9 GiB | 7.92 GiB 56 24.9 GiB
31 | 8.46 GiB 60 25.5 GiB | 7.94 GiB 56 25.0 GiB
41 | 9.90 GiB 70 26.9 GiB | 7.96 GiB 56 25.0 GiB
50 | 11.19 GiB 79 28.2 GiB | 7.98 GiB 56 25.0 GiB
Fixes #17783
…re loaded On Apple silicon the scheduler's free-memory figure for the GPU is the Metal working set minus what Ollama's own runners report. It does not see memory held by other applications, so a second MLX model can pass the fit check on a machine that is already short of memory, and the load pushes the system into swap and compression. While other models are loaded, the MLX fit check now also bounds the available memory by the system's free memory on shared-memory GPUs, the same rule llama-server loads already apply. A miss evicts an idle model and retries instead of starting the load. First loads are unchanged: with nothing else loaded, the model loads against the working-set figure alone, as both engines do today. The check also does not cover memory that grows after load, such as KV caches and prefix-cache snapshots.
…s the next model The scheduler starts the next load as soon as Close returns. The MLX client sent SIGINT, gave the process five seconds, then sent SIGKILL and returned without waiting, so a runner that could not take the signal was still exiting, with its memory still held, when the next load began. The runner has no signal handler, so SIGINT was already a kill. Load also started the process and recorded it without the client's mutex, so a Close racing with a load at server shutdown could find nothing to stop and leave the runner it missed running. Close now kills the process and waits for it to be reaped, as the llama-server client does. Load starts and records the process under the mutex and refuses to start once Close has run.
The bindings freed arrays by sweeping everything not pinned, so freeing anything required knowing what every other caller still held, and code that never swept accumulated until memory ran out. The prefix cache's eviction of a long stored path did exactly that: each merge copied the KV snapshots and nothing freed the consumed copies until the request ended, which drove a second long request past physical memory. Every array now belongs to a scope. A function scope, entered with Scoped or one of the ScopedEval forms, frees what was created in it when the function returns; results leave only by being returned. A held scope is closed by its holder and frees what was attached to it. A graph is built in a function scope and evaluated after it, so the eval frees each intermediate as it consumes it. Pin, Unpin, Sweep, and the array list's mutex are gone. On an M5 Max with qwen3.8:27b-mlx, the second 84k-token request after a stored one peaks at 35 GB instead of 57 GB; the cold path is unchanged. The copies themselves are untouched, so restoring an owned path can still exceed memory.
Models and layers checked optional weights for nil and also for a handle that no longer refers to an array, and evaluation and weight collection skipped such handles. No path produces one: a missing tensor is nil, and a handle only loses its array when its scope frees it, after which using it is a bug. The nil checks stay; the validity check is internal to the bindings now.
Loading a model can transform tensors after reading them: qwen3.5 models pack their linear-attention projections into one layout, and MoE models fuse the gate and up expert stacks. The buffers those transforms consume go back to MLX's allocator pool rather than to the system, and nothing releases the pool until the first request finishes. On qwen3.8:27b-mlx that is 2.15 GiB held idle on top of 16.9 GiB of weights, counted in the runner's reported memory the whole time. Clear the pool once the weights are evaluated. Models whose tensors load unchanged, such as gemma4, leave nothing in the pool and are unaffected.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )