ollama/mlxrunner
Jesse Gross ec3cc23076 mlxrunner: Release freed KV buffers during speculative decode
The decode loop releases MLX's pool of freed buffers every 256 generated
tokens, which is also how often the KV cache grows and drops its previous,
smaller buffers. The check fires only when the token count lands exactly on
a multiple of 256. Speculative decoding emits several tokens per round, so
most rounds step over the boundary and the pool is never released. Each
growth at a long context leaves several GB of buffers that no later
allocation can reuse, so the runner's footprint keeps climbing over a long
generation until the system runs out of memory.

We now release the pool whenever a round crosses a multiple of 256 tokens,
which is what a single-token round already did. With qwen3.8:27b-mlx at a
98k-token context on a 128 GB machine, a long speculative generation
previously grew the runner past 90 GB and panicked the kernel; it now stays
flat at 30 GB.
2026-09-17 09:45:01 -07:00
..
batch mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
cache mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
model mlxrunner: lay out model by contract, checkpoint and construction 2026-09-16 14:06:08 -07:00
nn mlxrunner: lay out nn one layer kind per file 2026-09-16 14:06:08 -07:00
sample mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
tokenizer mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
xgrammar mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
cache_trie.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
cache_trie_test.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
client.go mlxrunner: read model manifests through the manifest package 2026-09-16 14:06:08 -07:00
client_test.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
dflash.go mlxrunner, cmd, server: fold three small packages into their users 2026-09-16 14:06:08 -07:00
dflash_test.go mlxrunner, cmd, server: fold three small packages into their users 2026-09-16 14:06:08 -07:00
grammar.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
grammar_mask_test.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
grammar_test.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
media.go mlxrunner, cmd, server: fold three small packages into their users 2026-09-16 14:06:08 -07:00
media_test.go mlxrunner, cmd, server: fold three small packages into their users 2026-09-16 14:06:08 -07:00
mtp.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
mtp_test.go mlxrunner, cmd, server: fold three small packages into their users 2026-09-16 14:06:08 -07:00
pipeline.go mlxrunner: Release freed KV buffers during speculative decode 2026-09-17 09:45:01 -07:00
pipeline_test.go mlxrunner, cmd, server: fold three small packages into their users 2026-09-16 14:06:08 -07:00
prefix_cache.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
prefix_cache_scenario_test.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
prefix_cache_test.go mlxrunner, cmd, server: fold three small packages into their users 2026-09-16 14:06:08 -07:00
runner.go mlxrunner: lay out model by contract, checkpoint and construction 2026-09-16 14:06:08 -07:00
server.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
speculate.go mlxrunner, cmd, server: fold three small packages into their users 2026-09-16 14:06:08 -07:00
speculate_depth.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
speculate_depth_test.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
speculate_stats.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
status_memory.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
status_memory_test.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
utf8_buffer.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
utf8_buffer_test.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00
verify.go mlxrunner, cmd, server: fold three small packages into their users 2026-09-16 14:06:08 -07:00
verify_test.go mlx, mlxrunner: move the MLX engine out of x/ 2026-09-16 14:06:08 -07:00