[MLX] Support radix cache

[yeahdongcn]

Motivation

This PR adds a radix cache to the MLX backend, bringing it closer to feature parity with the PyTorch-based implementation and improving prefill throughput for workloads with shared prefixes.

Modifications

Added a kv_cache/ sub-package under the MLX backend with the core radix-cache components: a flat KV pool (MlxKVPool) for MLX-native data storage (auto-sized from available Metal memory), two cache flavors for prefill (PoolBackedCache, zero-copy gather) and decode (ContiguousKVCache, slice-assignment), and a batched attention wrapper (MLXAttentionWrapper) for BS>1 decode with vectorized RoPE.

Prefix-trie lookup, eviction, and slot allocation are handled by the scheduler's existing RadixCache and TokenToKVPoolAllocator — the MLX runner receives prefix_slot_ids and new_slot_ids from the scheduler and uses MlxKVPool purely as a data store.

Updated model_runner.py to wire scheduler-provided prefix/slot info, pool management, and deferred sync into the prefill/decode lifecycle. Updated tp_worker.py to extract prefix_indices, out_cache_loc, and req_pool_idx from ModelWorkerBatch and pass them to the MLX runner. Updated bench_one_batch.py to pass disable_radix_cache, mem_fraction_static, and pool_size for benchmarking.

Logs

# First request
[2026-03-27 11:40:03] INFO:     127.0.0.1:56520 - "POST /v1/chat/completions HTTP/1.1" 200 OK
[2026-03-27 11:40:08] Prefill batch, #new-seq: 1, #new-token: 3662, #cached-token: 0, token usage: 0.09, #running-req: 0, #queue-req: 0, cuda graph: False, input throughput (token/s): 0.01

# Second request with the same prompt
[2026-03-27 11:40:42] INFO:     127.0.0.1:56520 - "POST /v1/chat/completions HTTP/1.1" 200 OK
[2026-03-27 11:40:44] Prefill batch, #new-seq: 1, #new-token: 1, #cached-token: 3661, token usage: 0.09, #running-req: 0, #queue-req: 0, cuda graph: False, input throughput (token/s): 102.45

Accuracy Tests

Benchmarking and Profiling

Model: Qwen/Qwen3-0.6B

bs=1 input-len=60 output-len=100

Warmup ...
Prefill. latency: 0.12373 s, throughput:    484.91 token/s
Decode 0. Batch size: 1, latency: 0.02902 s, throughput:     34.46 token/s
Decode 1. Batch size: 1, latency: 0.02970 s, throughput:     33.67 token/s
Decode 2. Batch size: 1, latency: 0.02997 s, throughput:     33.37 token/s
Decode 3. Batch size: 1, latency: 0.02590 s, throughput:     38.61 token/s
Decode 4. Batch size: 1, latency: 0.02596 s, throughput:     38.53 token/s
Decode.  median latency: 0.02463 s, median throughput:     40.60 token/s
Total. latency:  0.907 s, throughput:    101.43 token/s
Benchmark ...
Prefill. latency: 0.05981 s, throughput:   1003.24 token/s
Decode 0. Batch size: 1, latency: 0.02457 s, throughput:     40.70 token/s
Decode 1. Batch size: 1, latency: 0.02656 s, throughput:     37.66 token/s
Decode 2. Batch size: 1, latency: 0.02859 s, throughput:     34.97 token/s
Decode 3. Batch size: 1, latency: 0.03229 s, throughput:     30.97 token/s
Decode 4. Batch size: 1, latency: 0.02955 s, throughput:     33.84 token/s
Decode.  median latency: 0.02510 s, median throughput:     39.85 token/s
Total. latency:  2.547 s, throughput:     62.83 token/s

bs=4 input-len=60 output-len=100

Warmup ...
Prefill. latency: 0.35028 s, throughput:    685.16 token/s
Decode 0. Batch size: 4, latency: 0.05203 s, throughput:     76.88 token/s
Decode 1. Batch size: 4, latency: 0.04963 s, throughput:     80.60 token/s
Decode 2. Batch size: 4, latency: 0.05592 s, throughput:     71.54 token/s
Decode 3. Batch size: 4, latency: 0.05310 s, throughput:     75.34 token/s
Decode 4. Batch size: 4, latency: 0.05279 s, throughput:     75.77 token/s
Decode.  median latency: 0.05212 s, median throughput:     76.75 token/s
Total. latency:  1.951 s, throughput:    188.63 token/s
Benchmark ...
Prefill. latency: 0.23874 s, throughput:   1005.29 token/s
Decode 0. Batch size: 4, latency: 0.05667 s, throughput:     70.59 token/s
Decode 1. Batch size: 4, latency: 0.05219 s, throughput:     76.64 token/s
Decode 2. Batch size: 4, latency: 0.05203 s, throughput:     76.88 token/s
Decode 3. Batch size: 4, latency: 0.05087 s, throughput:     78.63 token/s
Decode 4. Batch size: 4, latency: 0.05011 s, throughput:     79.83 token/s
Decode.  median latency: 0.05360 s, median throughput:     74.63 token/s
Total. latency:  5.544 s, throughput:    115.44 token/s

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[jonahbernard]

@yeahdongcn good work on this!

Can you please make the PR description more detailed and concrete to help reviewers and future contributors better understand this PR?

For example:

1. What exact changes have been made to both the prefill and decode forward passes individually to make them faster?
2. What existing SGLang code is reused in this PR, and what code should be reused indefinitely, and what code should be rewritten for the MLX backend in future PRs?
3. What does "auto-sized" mean, and what is the purpose of that?
4. What does this mean concretely: "two cache flavors for prefill (pool-backed, zero-copy) and decode (contiguous, slice-assignment)"?
5. Can you elaborate more on what this means: "Updated model_runner.py to wire prefix matching, pool management, and deferred sync into the prefill/decode lifecycle."
6. What flag did you add to bench_one_batch.py and why?

...and anything else that you think can be explained more clearly.

Thank you for addressing these! It will help document the developed of the MLX backend much more clearly for future contributors to build off of your great code!

[yeahdongcn]

Can you please make the PR description more detailed and concrete to help reviewers and future contributors better understand this PR?

Thanks for your questions. Please give me some time to prepare a write-up (AI-assisted).

1. What exact changes have been made to both the prefill and decode forward passes individually to make them faster?

Prefill

	Before	After
Prefix reuse	None — full sequence computed every time	Scheduler's RadixCache finds longest cached prefix; runner receives prefix_slot_ids
KV source	Computed from scratch	PoolBackedCache gathers cached KV from MlxKVPool by index (zero-copy in MLX lazy graph)
Tokens computed	All	Only uncached suffix tokens (new_token_ids)

Decode

	Before	After
KV management	BatchKVCache.merge() copies/stacks all per-request KV each step, then _extract_kv_cache() splits back	Pre-allocated ContiguousKVCache buffer per request
Token write	Concatenation + reallocation	Slice-assignment: keys[:,:,offset:offset+1,:] = k
BS>1 batching	Merge all caches → forward → extract	MLXAttentionWrapper reads per-request caches from thread-local context, handles batched SDPA directly

2. What existing SGLang code is reused in this PR, and what code should be reused indefinitely, and what code should be rewritten for the MLX backend in future PRs?

Component	Status	Notes
Scheduler, ForwardMode, request lifecycle	✅ Reused permanently	MLX hooks in via GenerationBatchResult
TpModelWorker	✅ Reused (subclassed)	MlxTpModelWorker overrides only forward_batch_generation
RadixCache / TokenToKVPoolAllocator	✅ Reused permanently	Scheduler handles all trie/slot management; MLX runner is a pure consumer
MLXAttentionWrapper	🔄 Rewrite later	Per-request loop with padding; replace with native batched/ragged attention when mx.fast.scaled_dot_product_attention supports variable-length sequences

3. What does "auto-sized" mean, and what is the purpose of that?

When pool_size isn't explicitly set, _compute_pool_size() checks total RAM, Metal working set limit, and current MLX usage (model weights), then divides the remaining budget by bytes_per_slot = 2 × layers × kv_heads × head_dim × dtype.size. Users don't need to manually tune --max-total-tokens.

4. What does this mean concretely: "two cache flavors for prefill (pool-backed, zero-copy) and decode (contiguous, slice-assignment)"?

	PoolBackedCache (prefill)	ContiguousKVCache (decode)
KV source	Gathers from shared MlxKVPool by slot index	Pre-allocated per-request buffer
Write strategy	Zero-copy gather in MLX lazy graph	Slice-assignment: keys[:,:,offset:offset+1,:] = k — no concat/realloc
Lifecycle	Created at prefill start	Used for all decode steps

After prefill, PoolBackedCache KV is bulk-copied into a ContiguousKVCache for fast decode.

5. Can you elaborate more on what this means: "Updated model_runner.py to wire prefix matching, pool management, and deferred sync into the prefill/decode lifecycle."

Mechanism	Where	What it does
Prefix matching	prefill()	Receives prefix_slot_ids and new_slot_ids from the scheduler (via req.prefix_indices and out_cache_loc). Creates PoolBackedCache for cached prefix, runs forward on new tokens only.
Pool sync (prefill)	prefill() / extend()	Immediately writes new KV to MlxKVPool via _sync_new_kv_to_pool() so the scheduler's trie can expose those slots for future prefix sharing.
Pool sync (decode)	remove_request()	Decode KV stays in ContiguousKVCache during generation. On request removal, _sync_decode_kv_to_pool() flushes un-synced positions to the pool using slot IDs from req_to_token_pool.

6. What flag did you add to bench_one_batch.py and why?

_MlxBenchRunner now passes disable_radix_cache, mem_fraction_static, and pool_size (from max_total_tokens) to MlxModelRunner. Radix cache is always disabled in bench mode because the benchmark runs without a scheduler — there is no RadixCache or TokenToKVPoolAllocator to manage trie/slot state.

[hnyls2002]

Question about the deferred decode KV sync in model_runner.py:

Prefill KV is synced to MlxKVPool eagerly, but decode KV is deferred until remove_request(). In a multi-turn conversation, the scheduler calls cache_finished_req() (inserting the full token sequence including generated tokens into the radix trie) before the worker calls remove_request(). If a new request matches a prefix that includes decode-generated tokens, PoolBackedCache would gather stale/zero KV from the pool.

Timeline:

1. Request A finishes decode → scheduler calls cache_finished_req(A) → trie exposes full prefix
2. Request B arrives, prefix matches A's sequence (including generated tokens)
3. Next forward pass (extend mode) → tp_worker does NOT cleanup stale requests (only decode mode does)
4. B's PoolBackedCache.update_and_fetch() gathers from pool → reads zeros for A's decode tokens

Would it be safer to sync each new decode token to the pool eagerly (one scatter per step), or do you see a reason the current deferral is needed?

[yeahdongcn]

Would it be safer to sync each new decode token to the pool eagerly (one scatter per step), or do you see a reason the current deferral is needed?

Thanks for taking such a close look at the code!

Yes, this is a real bug. Originally both prefill and decode sync were deferred to remove_request() since mx.eval() is expensive. But during E2E testing, prefill reads the pool immediately via PoolBackedCache, so prefill sync had to move to eager (right after _eval_with_cache). Decode sync stayed deferred because decode steps never read from the pool themselves.

The fix follows the same principle as prefill (sync before the data is read). But after some experimentation, eager-per-step turned out to be wasteful. Pool reads only happen during extend/prefill (PoolBackedCache), never during decode. So instead of paying a per-step sync cost, flush_all_decode_kv() runs once at the start of each extend batch, right before any pool read. Decode steps pay zero overhead. remove_request() still syncs as a defensive fallback.

@hnyls2002 Could you please check the latest commit and see if it makes sense? Thanks!

[yeahdongcn]

@hnyls2002 @alexnails Ping 😄

[alexnails]

does this impl break the chained decode stuff @changminbark is working on ?

[yeahdongcn]

I'm not entirely sure. There are a few new changes built on top of my branch, so we may need to do some rebase/merge work.

[changminbark]

There's a good chance it may break it, so I will need to do some merge work and apply fixes.

[alexnails]

as this is self contained to MLX, LGTM. Just make sure CI passes and there are no build errors, etc

[yeahdongcn]

/tag-and-rerun-ci

[yeahdongcn]

/rerun-failed-ci

[yeahdongcn]

/rerun-failed-ci

[yeahdongcn]

/rerun-failed-ci

[yeahdongcn]

/rerun-failed-ci