From ec3cc23076c489c988e4e8a85659e48e92dbb734 Mon Sep 17 00:00:00 2001 From: Jesse Gross Date: Wed, 16 Sep 2026 20:05:34 -0700 Subject: [PATCH] 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. --- mlxrunner/pipeline.go | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/mlxrunner/pipeline.go b/mlxrunner/pipeline.go index eabe2e0d401..827dc0bc2f3 100644 --- a/mlxrunner/pipeline.go +++ b/mlxrunner/pipeline.go @@ -278,7 +278,9 @@ func (r *Runner) decode(ctx context.Context, request Request, session *cacheSess now := time.Now() // Release MLX's cached free buffers every clearCacheInterval tokens so the - // allocator's pool does not grow unbounded over a long generation. + // allocator's pool does not grow unbounded over a long generation. A + // speculative round emits several tokens at once, so the clear fires on + // crossing a multiple of the interval, not on landing exactly on one. const clearCacheInterval = 256 generated := 0 @@ -289,6 +291,7 @@ func (r *Runner) decode(ctx context.Context, request Request, session *cacheSess var done bool var err error + before := generated mlx.Scoped(func() { var results []sampler.Result results, err = d.next(request.Options.NumPredict - generated) @@ -348,7 +351,7 @@ func (r *Runner) decode(ctx context.Context, request Request, session *cacheSess break } - if generated%clearCacheInterval == 0 { + if generated/clearCacheInterval != before/clearCacheInterval { mlx.ClearCache() } }