lmstudio-ai / lmstudio-ai/mlx-engine

Backend scheduler crashes with broadcast_shapes ValueError in BatchRotatingKVCache.merge under concurrent requests with shared prefix

Open
#363 0 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.