Prefix cache reuse is broken for all hybrid-architecture models (sliding window, SSM/Mamba)

[skcadri]

Prefix cache reuse is broken for all hybrid-architecture models (sliding window, SSM/Mamba)

Summary

Prompt prefix caching — the mechanism that reuses computed KV states across requests sharing a common prefix — only works for pure full-attention models. Any model using sliding window attention, Mamba/SSM layers, or mixed attention types silently falls back to full prompt recomputation on every request. This makes multi-turn conversations unusably slow for the majority of modern open-weight models.

This is not a single bug but a systemic gap in how mlx-lm handles non-standard cache types. I'm filing this as a unifying issue because the symptoms are spread across many separate reports that all trace back to the same root causes.

Empirical Evidence

Tested on Mac Studio M3 Ultra (512GB unified memory), LM Studio 0.4.6, mlx-engine. Three sequential requests with identical system prompts, different user messages, max_tokens=30:

MiniMax M2.5 (pure attention, MoE) — caching works

Request	Time	Speedup
1 (cold)	29.33s	—
2 (warm)	6.15s	4.8x
3 (warm)	2.79s	10.5x

GPT-OSS 120B (sliding_attention + full_attention hybrid) — no caching

Request	Time	Speedup
1 (cold)	1.54s	—
2 (warm)	1.77s	none
3 (warm)	1.67s	none

Qwen 3.5 9B (attention + Mamba/SSM hybrid) — no caching

Request	Time	Speedup
1 (cold)	5.02s	—
2 (warm)	7.76s	none (slower)
3 (warm)	8.00s	none (slower)

MiniMax shows clear prefix reuse (29s → 3s). GPT-OSS and Qwen 3.5 show zero improvement — the full prompt is recomputed every turn.

Root Cause Analysis

There are two distinct failure modes, both stemming from the assumption that all layers use identical, trimmable KV caches:

1. Sliding window models → RotatingKVCache can't be trimmed

Models like GPT-OSS 120B and Gemma 3 27B alternate between sliding window and full attention layers:

# GPT-OSS 120B config.json
"layer_types": ["sliding_attention", "full_attention", "sliding_attention", "full_attention", ...]
"sliding_window": 128

# Gemma 3 27B — 5:1 pattern
5x local attention (window=1024) + 1x global attention, repeated

Sliding window layers use RotatingKVCache (circular buffer). When the cache wrapper attempts to trim to a common prefix for reuse, the circular buffer state can't be meaningfully trimmed — so the entire cache is erased and recomputed from scratch. See lmstudio-ai/mlx-engine#177.

2. SSM/Mamba hybrid models → non-trimmable state

Models like Qwen 3.5 (all sizes) use attention + Mamba layers:

# Qwen 3.5 — hybrid attention + SSM
Attention layers: standard KVCache (trimmable)
Mamba/SSM layers: recurrent state (NOT trimmable)

The Mamba state is fundamentally different from a KV cache — it's a compressed recurrent state that can't be split at an arbitrary token boundary. Additionally, the KVCache.make_mask() interface requires window_size and return_array arguments that don't apply to SSM state, causing TypeError on multi-turn prefill (see QwenLM/Qwen3.6#37).

Affected Models

Every popular hybrid-architecture model is affected. This covers the majority of modern open-weight models:

Model	Architecture	Cache behavior
Qwen 3.5 (all sizes)	Attention + Mamba/SSM	Broken — crashes or no reuse
GPT-OSS 120B / 20B	Sliding + full attention	Broken — full recompute
Gemma 3 (all sizes)	5:1 sliding + global	Broken — full recompute
Llama 4 Scout/Maverick	iRoPE chunked (8K) + NoPE	Likely broken
Qwen2.5-VL	Partial sliding window	Likely broken
MiniMax M2.5	Pure full attention (MoE)	Works

As of March 2026, MiniMax M2.5 appears to be the only major model where prefix caching works correctly on MLX.

Impact

This is particularly painful for agentic workloads where:

• System prompts are large (tool definitions, personas, instructions)
• Conversations are multi-turn (each turn should only process new tokens)
• Multiple agents share the same model (each request recomputes from scratch)

Without prefix caching, a 40K-token context takes ~200 seconds to process vs ~5 seconds with cache reuse. For agentic frameworks running on local MLX models, this is the difference between usable and unusable.

Proposed Solution

Implement per-layer cache logic instead of assuming uniform cache types:

1. For sliding window layers: Either make RotatingKVCache trimmable to a prefix boundary, or maintain a parallel standard cache for the prefix portion that gets replayed into the rotating buffer on reuse.

2. For SSM/Mamba layers: Use make_prompt_cache(model) (which correctly creates ArrayCache for linear attention layers) instead of uniform KVCache() allocation. The workaround in [Bug] Prefill Failure of Qwen3.5 Model Using KV Cache in the mlx‑lm Framework QwenLM/Qwen3.6#37 demonstrates this works for Qwen 3.5 at the code level.

3. Cache type introspection: The cache wrapper should inspect what type of cache each layer requires and handle trim/reuse differently per type.

PR #923 proposes a partial fix for Qwen 3.5 specifically. The RotatingKVCache trim issue has PRs at lmstudio-ai/mlx-engine#188 and #192. These should be unified into a comprehensive solution.

Related Issues

• Caching doesn't seem to be working for Qwen3.5 #903 — Caching doesn't work for Qwen3.5
• KV cache cross-contamination between concurrent requests in mlx_lm.server #965 — KV cache cross-contamination between concurrent requests
• Strange cache behavior with 0.31.0 in server mode #975 — Strange cache behavior in server mode (cache corruption)
• Prompt caching returns different logits for repeat prompts #259 — Prompt caching returns different logits for repeat prompts
• Hybrid cache for Qwen3.5 #923 — PR: Qwen3.5 cache fix (not merged)
• MLX RotatingKVCache trim behavior causes context overflow policies to always erase the whole cache lmstudio-ai/mlx-engine#177 — RotatingKVCache trim erases entire cache
• Cache reuse and cache fixes lmstudio-ai/mlx-engine#188, Is there a way to use params JSON? #192 — PRs for RotatingKVCache fix (not merged)
• Qwen3.5-35B-A3B: KV cache reuse not supported — full prompt recompute on every request lmstudio-ai/lmstudio-bug-tracker#1563 — Qwen3.5 full recompute in LM Studio
• KV Caching broken for MLX on Mac LM Studio 0.3.35(Build 1) lmstudio-ai/lmstudio-bug-tracker#1319 — KV caching broken in LM Studio 0.3.35
• [Bug] Prefill Failure of Qwen3.5 Model Using KV Cache in the mlx‑lm Framework QwenLM/Qwen3.6#37 — Prefill failure with KV cache (workaround found)
• Eval bug: Qwen3.5 always re-processes the full prompt ggml-org/llama.cpp#19858 — Qwen3.5 always re-processes full prompt (fixed in b8212)

Environment

• macOS 26.2.0 (Tahoe), Mac Studio M3 Ultra, 512GB unified memory
• LM Studio 0.4.6 (build 1)
• Models tested: MiniMax-M2.5-MLX-8bit, gpt-oss-120b-mlx-8bit, Qwen3.5-9B (MLX)

[Giustino98]

Thanks for the detailed summary. Is it ok if I take this on?

I'll review the existing partial PRs (#188, #192, #923) and work on a unified fix. The goal is to refactor the cache manager to properly handle the different cache types (StandardKVCache, RotatingKVCache, ArrayCache) during the trim/reuse phase, preventing the current TypeErrors while keeping the speedup for standard models.

[angeloskath]

Hi @skcadri, could you provide a reproduction with some code?

I would say that the non-trimmable cache issue has been fixed. I regularly use the latest Qwen models and there is the expected reuse.

It reprocesses part of the latest model response but that has to do with the chat template which removes the thinking traces thus changing the history of the conversation.

When the interaction is long, in agentic interactions it can easily be 10s of thousands of tokens per user message, then we have to process the last interaction again. That is not something that can be fixed without changing the chat template.

[Giustino98]

Hi @angeloskath, I spent some time going through the codebase and running a few tests to understand the current state.

Normal multi-turn chat works fine for hybrid models. The trie-based LRUPromptCache has a "shorter cache" path in fetch_nearest_cache that doesn't require trimming since it just deepcopies the shorter cached state and processes the remaining tokens. That matches your experience with Qwen models.

The thinking trace issue is also as you described. The batch codepath mitigates this somewhat via _compute_prompt_checkpoint (saving a checkpoint before the thinking tokens), but it's fundamentally a chat template problem as you said.

Parallel/concurrent requests sharing a system prompt but with different user messages represent a problematic scenario instead. For example, request A finishes and its cache is stored at [sys_prompt, user_A, gen_A]. Request B comes in with [sys_prompt, user_B, ...]: the trie finds request A's cache as a "longer" match with common_prefix = len(sys_prompt), but reusing it requires trimming. For pure KVCache this works, but for hybrid caches (ArraysCache.is_trimmable() -> False, RotatingKVCache.is_trimmable() -> False when full), it falls through to full recompute. This matters mostly for agentic setups where multiple requests share long system prompts but diverge at the user message.

[skcadri]

Hi @angeloskath, thanks for looking into this — and great analysis @Giustino98, you nailed the exact scenario.

You're right that normal sequential multi-turn chat works fine for hybrid models. The trie's "shorter cache" path in fetch_nearest_cache doesn't require trimming — it deepcopies the shorter prefix and processes the remaining tokens. That path is solid.

The problem is specifically the divergent-suffix scenario that @Giustino98 described, which is the dominant pattern in agentic workloads:

1. Request A completes with [system_prompt + user_A + gen_A] → cache stored at that full sequence
2. Request B arrives with [system_prompt + user_B] → trie finds A's cache as longer with common_prefix = len(system_prompt)
3. can_trim_prompt_cache() returns False because any ArraysCache layer (Mamba/SSM) inherits _BaseCache.is_trimmable() → False (cache.py L144)
4. The "longer" path is skipped (server.py L308), there's no standalone system-prompt "shorter" entry, so it falls through to full recompute

For RotatingKVCache (sliding window models like GPT-OSS, Gemma 3), the same issue occurs once the buffer fills: is_trimmable() returns self.offset < self.max_size → False (cache.py L540-541).

Reproduction

Tested on Mac Studio M3 Ultra (512GB), LM Studio 0.4.6. Three sequential requests with the same ~1200-token agentic system prompt, different user messages, max_tokens=30:

Model	Architecture	Req 1 (cold)	Req 2 (warm)	Req 3 (warm)	Prefix reuse?
Devstral 24B	Pure attention	3.21s	0.78s (4.1×)	0.79s (4.1×)	✓ Yes
Magistral 24B	Pure attention	3.59s	1.25s (2.9×)	1.41s (2.5×)	✓ Yes
Qwen3.5 9B	Attention + Mamba/SSM	3.12s	2.46s (1.3×)	2.45s (1.3×)	✗ No
GPT-OSS 120B	Sliding + full attention	1.56s	1.56s (1.0×)	1.55s (1.0×)	✗ No

Pure-attention models show clear prefix reuse (3–4× speedup on warm requests). Hybrid models show no improvement — both the SSM hybrid (Qwen3.5) and sliding-window hybrid (GPT-OSS) recompute the full prompt every time.

Here's a self-contained reproduction script using the Python API directly:

#!/usr/bin/env python3
"""
Demonstrates that can_trim_prompt_cache() returns False for hybrid models,
preventing prefix reuse when the cached entry is longer than the new request.

Usage:
  python repro.py --model mlx-community/Qwen3.5-9B-bf16       # hybrid → False
  python repro.py --model mlx-community/MiniMax-M2.5-MLX-8bit  # pure attn → True
"""

import argparse, copy, time
import mlx.core as mx
from mlx_lm import load
from mlx_lm.models.cache import can_trim_prompt_cache, make_prompt_cache, trim_prompt_cache
from mlx_lm.generate import generate_step

SYSTEM = (
    "You are an AI assistant with tool access. Available tools: "
    "read_file, write_file, search_code, run_command, git_diff. "
    "Always explain your reasoning before calling a tool. "
    "Follow best practices for code quality and testing. "
) * 5  # ~500+ tokens

def run(model_path):
    model, tokenizer = load(model_path)

    # Show per-layer cache types
    cache = make_prompt_cache(model)
    types = {}
    for c in cache:
        t = type(c).__name__
        types[t] = types.get(t, 0) + 1
    print(f"Cache types: {types}")
    print(f"can_trim_prompt_cache: {can_trim_prompt_cache(cache)}")

    # Build two requests sharing a system prompt but with different user messages
    msgs_a = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": "Read src/main.py"}]
    msgs_b = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": "List the tests/ dir"}]
    toks_a = tokenizer.apply_chat_template(msgs_a, add_generation_prompt=True)
    toks_b = tokenizer.apply_chat_template(msgs_b, add_generation_prompt=True)
    common = sum(1 for a, b in zip(toks_a, toks_b) if a == b)

    # Warmup
    wc = make_prompt_cache(model)
    for _ in generate_step(mx.array(tokenizer.encode("Hi")), model, max_tokens=3, prompt_cache=wc):
        pass

    # Request A: prefill + generate
    cache_a = make_prompt_cache(model)
    t0 = time.perf_counter()
    for tok, _ in generate_step(mx.array(toks_a), model, max_tokens=20, prompt_cache=cache_a):
        pass
    t_a = time.perf_counter() - t0

    can_trim = can_trim_prompt_cache(cache_a)
    print(f"\nRequest A: {len(toks_a)} tokens, {t_a:.2f}s")
    print(f"Common prefix: {common} tokens")
    print(f"can_trim for reuse: {can_trim}")

    if can_trim:
        cache_b = copy.deepcopy(cache_a)
        trim_prompt_cache(cache_b, len(toks_a) + 20 - common)
        t0 = time.perf_counter()
        for tok, _ in generate_step(mx.array(toks_b[common:]), model, max_tokens=20, prompt_cache=cache_b):
            pass
        t_b = time.perf_counter() - t0
        print(f"Request B (reused prefix): {len(toks_b)-common} new tokens, {t_b:.2f}s — {t_a/t_b:.1f}x faster")
    else:
        cache_b = make_prompt_cache(model)
        t0 = time.perf_counter()
        for tok, _ in generate_step(mx.array(toks_b), model, max_tokens=20, prompt_cache=cache_b):
            pass
        t_b = time.perf_counter() - t0
        print(f"Request B (full recompute): {len(toks_b)} tokens, {t_b:.2f}s — no speedup")

if __name__ == "__main__":
    p = argparse.ArgumentParser()
    p.add_argument("--model", required=True)
    run(p.parse_args().model)

Output for Qwen3.5-9B (hybrid):

Cache types: {'ArraysCache': 24, 'KVCache': 8}
can_trim_prompt_cache: False

Request A: 642 tokens, 1.40s
Common prefix: 607 tokens
can_trim for reuse: False
  Cache is NOT trimmable — server falls through to full recompute
  Request B (full recompute): 1.29s — no speedup

Why this matters for agentic use

In agentic frameworks, the pattern is:

• Large system prompt (tool definitions, persona, instructions) — often 2K–10K+ tokens
• Each tool call is a new request with the same system prompt but a different user message (the tool result)
• Requests don't extend each other sequentially — they diverge at the user message boundary

This is exactly the divergent-suffix pattern @Giustino98 identified. Without the trim path working, every tool call reprocesses the entire system prompt from scratch. As context grows (multi-turn agentic conversations easily reach 40K+ tokens), this becomes the dominant bottleneck.

The specific code involved

• ArraysCache inherits _BaseCache.is_trimmable() → False with no override (cache.py L592-679)
• RotatingKVCache.is_trimmable() returns self.offset < self.max_size — False once full (cache.py L540)
• can_trim_prompt_cache() requires all(c.is_trimmable() for c in cache) — one non-trimmable layer poisons the whole cache (cache.py L86-90)
• fetch_nearest_cache() gates the "longer" reuse path on can_trim_prompt_cache() (server.py L308)
• Qwen3.5 make_cache(): ArraysCache for Mamba layers, KVCache for attention (qwen3_5.py L297-298)

Happy to run additional tests or help with a fix. @Giustino98 if you're still interested in taking this on, the fix would need to either make ArraysCache trimmable (by replaying the SSM state from a prefix boundary) or change fetch_nearest_cache to store a separate prefix-only cache entry at the system-prompt boundary for non-trimmable models.

[Giustino98]

@skcadri great analysis! If it's okay, I'm going to open a PR by Sunday

[Giustino98]

Hi @skcadri and @angeloskath,

I've just opened PR #999 as a Draft. It implements the fix for the sliding window models (RotatingKVCache and BatchRotatingKVCache) by linearizing the ring buffer via _temporal_order() during a trim(). Tested locally with gpt-oss and prefix caching now triggers correctly on divergent-suffix tool calls.

I intentionally left ArraysCache untouched here to avoid silent hallucinations with Mamba models, which likely requires server-side checkpointing in LRUPromptCache instead.

PR is still in draft since I want to do a final sanity check on the pointer math. I'll run some more tests but if anyone has time to look at the trimming logic it would be really helpful. I'll mark it as ready for review once I'm 100% confident.

[UPDATE]
I've completed the sanity checks and moved the PR out of Draft.
During testing, I noticed that the single RotatingKVCache class was using the same offset variable for both the physical buffer size and the absolute sequence position (RoPE). Trimming the buffer physically would therefore break the absolute RoPE embeddings for subsequent tokens. I refactored it to strictly separate _offset (physical tensor size) from offset (logical position), aligning its architecture with BatchRotatingKVCache.

[skcadri]

Test results for PR #999 (RotatingKVCache trim fix)

Tested PR #999 on GPT-OSS 120B (8-bit, sliding_window=128) — Mac Studio M3 Ultra.

Test scenario

Simulated an agentic workload: 5 independent requests sharing a large system prompt + context document (~2700 tokens of shared prefix) but with different user messages. This is the exact pattern where RotatingKVCache.is_trimmable() returning False prevented the server's "longer cache" trim path from activating — each request was a full recompute.

Results: PR #999 (patched mlx-lm server) vs stock LM Studio (unpatched)

                       Cache Hit    Avg Warm   Speedup   Time Saved
                        --------    --------   -------   ----------
    PR#999 (patched)         98%       0.79s     3.37x        7.44s
    LM Studio (stock)         0%       2.83s     0.98x       -0.27s

Per-request breakdown (patched):

Req  Type   Prompt   Cached   Hit%  Uncached     Time    Saved
---  ----   ------   ------   ----  --------     ----    -----
  1  COLD     2704       69     2%      2635    2.65s         
  2  warm     2692     2660    98%        32    0.79s   +1.85s
  3  warm     2692     2661    98%        31    0.78s   +1.87s
  4  warm     2698     2660    98%        38    0.81s   +1.84s
  5  warm     2690     2660    98%        30    0.77s   +1.88s

Without the fix (stock), all 5 requests recompute the full ~2700 tokens every time (0% cache hit, ~2.8s each).

What's happening

After request 1 completes, its KV cache is stored at the full sequence length ([system + context + user_1 + generation_1]). When request 2 arrives with [system + context + user_2], the trie finds request 1's cache as a "longer" entry with common_prefix ≈ 2660 tokens. Reusing it requires trimming the cache back to the common prefix length.

Before this PR, RotatingKVCache.is_trimmable() returned False once offset >= max_size (i.e., after the cache wraps), so the server skipped the trim path entirely — falling through to a full recompute. With the fix, is_trimmable() always returns True and trim() properly handles the rotating buffer (including reordering via _temporal_order and resizing).

Unit tests

All 20 tests in test_prompt_cache.py pass, including the existing trim tests.

Non-sliding-window models unaffected

Verified Qwen3.5-9B (which uses ArraysCache for linear attention layers) — still shows 0% cache hit on this divergent-suffix pattern, as expected. ArraysCache.is_trimmable() is unchanged and correctly returns False.

[skcadri]

Following up on #999 (which fixed cache reuse for sliding-window models like GPT-OSS), I've opened #1006 to address the remaining non-sliding-window hybrid models — Qwen3.5, Nemotron-H, Jamba, and other architectures that use ArraysCache for their recurrent/SSM layers.

The core issue: ArraysCache stores compressed recurrent state (conv + RNN hidden state) that fundamentally cannot be trimmed — you can't remove N tokens from a hidden state vector. So can_trim_prompt_cache() returns False, blocking the longer-cache trim path entirely → 0% cache reuse.

The fix extends the existing prompt checkpoint mechanism to save the cache at the last-message boundary during prefill. Subsequent requests with the same prefix find it via the "shorter cache" path in fetch_nearest_cache — no trimming needed.

Tested on two hybrid architectures:

• Qwen3.5-9B (attention + linear recurrence): 98% cache hit, 2.86x speedup on 2.7K-token agentic prompts
• NVIDIA Nemotron-H 30B (attention + Mamba SSM, 6/29 layers trimmable): 98% cache hit, 3.1x speedup

Together with #999, this should fully resolve #980 for both sliding-window and non-sliding-window hybrid models.

[allenwlee]

Confirming this at larger scale — Qwen3.5-397B-A17B MLX 8-bit (mlx-lm 0.31.1) on Mac Studio M3 Ultra 512GB with a ~21,000 token codebase prefix (TypeScript project). This is an agent workload pattern: same large context, different questions per request.

Results (stock mlx-lm, no fix applied)

Call	TTFT	vs Call 1
1 (cold)	169.9s	baseline
2 (warm)	169.2s	-0% ❌
3 (warm)	170.4s	-0% ❌

Zero cache reuse — every call pays the full 169s prefill cost. At 21K tokens the economics are brutal: an agent making 10 sequential calls against the same codebase spends ~28 minutes on prefill alone.

Comparison: same model in GGUF Q8_0 via llama-server

Running the identical prompt against llama-server (llama.cpp) with --swa-full to get correct checkpoint behavior:

Call	TTFT	vs Call 1
1 (cold)	230.7s	baseline
2 (warm)	6.9s	-97% ✅
3 (warm)	6.6s	-97% ✅

This validates that the hardware and model are fully capable of the caching behavior — it's purely a stack difference.

Notes on scale

The issue scales badly with context length. At the ~2.7K token prompts in the tests above, cold prefill is ~2.6s — painful but survivable. At 21K tokens it's 170s per call. Any agent loop with a large system prompt or codebase context hits this hard.

Looking forward to testing #1006 on this setup once it's ready — if Qwen3.5-9B at 2.7K tokens achieves 98% cache hit, the same fix should theoretically scale to 397B at 21K tokens since the architectural issue (ArraysCache non-trimmability) is the same regardless of model size or context length.