BatchRotatingKVCache.merge() crashes on concurrent requests with different prompt lengths

[juliensimon]

Description

mlx_lm.server crashes when handling concurrent requests to models that use RotatingKVCache (e.g. models with sliding window attention). The crash occurs in BatchRotatingKVCache.merge() when trying to merge caches with different _idx values, producing a shape mismatch.

This is reproducible with any client that sends multiple requests concurrently (e.g. OpenCode, which fires a short system probe and a full prompt simultaneously).

Environment

• mlx-lm: 0.31.1
• mlx: 0.24.2
• macOS: 15.5 (Apple M3 Max, 48 GB)
• Python: 3.12
• Model: mlx-community/Trinity-Nano-Preview-8bit (afmoe architecture, uses RotatingKVCache)

Also reproduced with mlx-community/Trinity-Mini-8bit.

Steps to Reproduce

1. Start the server:

uvx --from "mlx-lm>=0.28.4" mlx_lm.server \
    --model mlx-community/Trinity-Nano-Preview-8bit \
    --port 8080

2. Send two concurrent requests with different prompt lengths:

# In parallel (e.g. using & or a client that sends concurrent requests)
curl -s http://localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages": [{"role": "user", "content": "Hi"}], "max_tokens": 10}' &

curl -s http://localhost:8080/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{"messages": [{"role": "system", "content": "You are a helpful assistant. '"$(python3 -c "print('x ' * 5000)")"'"}, {"role": "user", "content": "Hello"}], "max_tokens": 10}' &

wait

The first request succeeds, but the second crashes the server's generate thread.

Error

Exception in thread Thread-2 (_generate):
Traceback (most recent call last):
  File ".../mlx_lm/server.py", line 948, in _generate
    responses = batch_generator.next()
  File ".../mlx_lm/generate.py", line 1329, in next
    return self._next()
  File ".../mlx_lm/generate.py", line 1257, in _next
    batch = self._process_prompts(prompts)
  File ".../mlx_lm/generate.py", line 1126, in _process_prompts
    prompt_cache = _merge_caches(caches)
  File ".../mlx_lm/generate.py", line 922, in _merge_caches
    batch_cache.append(caches[0][i].merge([c[i] for c in caches]))
  File ".../mlx_lm/models/cache.py", line 580, in merge
    return BatchRotatingKVCache.merge(caches)
  File ".../mlx_lm/models/cache.py", line 1364, in merge
    keys[i : i + 1, :, p : p + c._idx] = c._temporal_order(c.keys)
ValueError: [broadcast_shapes] Shapes (1,2,3935,128) and (1,2,2048,128) cannot be broadcast.

Analysis

In cache.py:1364, BatchRotatingKVCache.merge() assumes all caches being merged have compatible _idx values. When two requests have very different prompt lengths (e.g. 519 tokens vs 10,081 tokens), their RotatingKVCache entries end up with different _idx values (e.g. 3935 vs 2048), and the slice assignment fails because the shapes don't match.

The issue is in the merge logic — it allocates a target tensor based on one cache's dimensions but then tries to copy data from another cache with a different _idx.

Workaround

Disabling batching in server.py avoids the crash:

# In ModelProvider.load(), after the is_batchable check:
self.is_batchable = False

This forces sequential request processing. The --prompt-cache-size 1 and --decode-concurrency 1 flags do not prevent the crash because the server still attempts to batch concurrent requests.

Expected Behavior

Concurrent requests with different prompt lengths should either:

• Be merged correctly (padding/truncating caches to compatible shapes), or
• Fall back to sequential processing when cache shapes are incompatible

[Giustino98]

Hi @juliensimon, I just ran into this crash and spent a bit of time tracing it down in cache.py. The root cause is actually quite straightforward.

The shape mismatch in BatchRotatingKVCache.merge() happens because the code is using c._idx (the circular buffer's internal write pointer) to determine the size of the slice receiving the temporal-ordered tokens.

When a request has wrapped around the buffer (offset > max_size), _idx resets and points somewhere in the middle of the buffer but c._temporal_order() correctly unrolls and returns the full max_size valid tokens.
Trying to fit a [..., 2048, ...] tensor into a slice of size _idx ([..., 904, ...]) triggers the broadcast ValueError.

The fix is simply to use the actual number of valid tokens (which is c.size(), i.e., min(offset, max_size)) instead of the internal _idx pointer when defining the target slice:

cache.py ~L1364

valid_length = c.size()
keys[i : i + 1, :, p : p + valid_length] = c._temporal_order(c.keys)
values[i : i + 1, :, p : p + valid_length] = c._temporal_order(c.values)

I've tested this locally and it seems to resolve the crash for models using RotatingKVCache without needing to disable batching but I will do further tests.
I'll put together a quick PR with this fix shortly