lmstudio-ai / lmstudio-ai/mlx-engine
Backend scheduler crashes with broadcast_shapes ValueError in BatchRotatingKVCache.merge under concurrent requests with shared prefix
- Dominant language
- Python
- Stars
- 1.2k
- Forks
- 133
- Avg merge
- 21h 6m
- Merged PRs (30d)
- 1
Description
**Environment**
LM Studio backend: app-mlx-generate-mac14-arm64@34 (mlx_engine, Python 3.11)
Model: qwen/qwen3-4b-2507 (MLX, 4-bit)
Hardware: Mac mini, Apple Silicon, 16 GB
Server settings at crash: Max Concurrent Predictions 16, context 32768, KV cache quantization off. Also reproduces at the default 4 slots / 8192 ctx.
**Reproduction**
8 concurrent OpenAI-format streaming clients, each in a tight loop: identical ~2000-token prompt prefix + short unique suffix, max_tokens: 64, temperature: 0, no delay between turns. Crashes within ~60 seconds. Reproduced in two independent runs (23 and 10 failed requests respectively; 4 model crashes total). At concurrency ≤4 the same workload runs clean (hundreds of requests, zero errors).
**Observed**
Client side: streams open (HTTP 200) then deliver zero content chunks; some requests receive HTTP 500; surviving requests see TTFT p95 ~18s (reload gap). The model then auto-reloads — with default settings, discarding the user-configured slot count and context length (see "Impact").
Server log (repeated 23×, followed by 4× "The model has crashed without additional information. (Exit code: null)"):
[batched_model_kit][ERROR]: Encountered fatal exception in the backend scheduler:
Traceback (most recent call last):
... mlx_engine/model_kit/batched_model_kit.py ... self._generate()
... mlx_lm/generate.py ... prompt_responses, generation_responses = batch_generator.next()
... mlx_lm/generate.py ... self._prompt_batch.extend(self._make_batch(n))
... mlx_lm/generate.py ... return PromptProcessingBatch(
... mlx_lm/generate.py ... self.prompt_cache = _merge_caches(caches)
... mlx_lm/generate.py ... batch_cache.append(caches[0][i].merge([c[i] for c in caches]))
... mlx_lm/models/cache.py ... return BatchRotatingKVCache.merge(caches)
... mlx_lm/models/cache.py ...
keys[i : i + 1, :, p : p + l] = c._temporal_order(c.keys)[..., -l:, :]
ValueError: [broadcast_shapes] Shapes (1,8,348,128) and (1,8,907,128) cannot be broadcast.
**Analysis (best effort)**
The scheduler batches prompt-processing for concurrent requests and merges their per-request prompt caches. When the cached prefixes have different retained lengths (here 348 vs 907 tokens — plausibly different truncation / rotation states of the same shared prefix), BatchRotatingKVCache.merge writes one cache's keys into a slot sized for another and the shape mismatch raises, killing the whole scheduler rather than the one request.
**Impact**
All in-flight requests fail (empty streams / 500s).
The model auto-reloads with default load settings, silently discarding the configured Max Concurrent Predictions and context length — so a server that was configured for 16 slots continues running at 4 after recovery. For unattended server deployments this is a second bug: crash recovery causes silent configuration drift.
**Suggested severity**: high for headless/server use — the trigger (several concurrent clients sharing a prompt prefix) is the normal shape of agentic traffic behind a gateway.
Contributor guide
Research direction
Start with the concurrent streaming reproduction described in the issue, then inspect mlx_engine/model_kit/batched_model_kit.py and mlx_lm/models/cache.py, especially BatchRotatingKVCache.merge. Use the reported 348-versus-907 token shape mismatch as the initial failure to investigate. Done means the workload no longer crashes the scheduler or fails in-flight requests, and recovery does not silently lose the configured slot count and context length.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- backend, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100