repeat_penalty causes fatal backend-scheduler crashes in multi-turn conversations (32-bit slice error / KV cache merge broadcast error)

[thefrenchbaron]

Environment

• LM Studio: 0.4.20 (macOS)
• Runtime: mlx-llm-mac-arm64-apple-metal-nax-advsimd@1.11.0 (crashes also observed as silent model death on mlx-llm-mac-arm64-apple-metal-advsimd@1.10.1)
• OS: macOS 26 (Darwin 25.6.0), Apple Silicon M5, 24 GB unified memory
• Model: Daizee/Impish_Bloodmoon_12B-mlx-4bit (Mistral-Nemo-architecture 12B, MLX 4-bit), loaded via lms load --context-length 16384
• API: /v1/chat/completions, non-streaming

Summary

Setting repeat_penalty in the request payload reliably crashes the MLX backend scheduler during multi-turn conversations with a long shared system prompt. The same workload with repeat_penalty omitted completes 8/8 turns without issues. After the fatal exception, the model process dies and every subsequent request fails with "The model has crashed without additional information. (Exit code: null)" until the model is reloaded.

Reproduction (self-contained)

import json, urllib.request

URL = "http://localhost:1234/v1/chat/completions"
PARA = (
    "You are a creative writing assistant. Maintain a consistent narrative voice, "
    "keep continuity with earlier scenes, avoid repeating earlier sentences, and "
    "answer in the same language as the user. Prefer concrete sensory detail over "
    "abstraction, vary sentence length, and never break character. "
)
SYSTEM = PARA * 100  # ~5-6k tokens
STOPS = ["[", "**", "```", "###", "\n# ", "\nUser", "User:"]

history = []
for turn in range(8):
    msgs = ([{"role": "system", "content": SYSTEM}] + history
            + [{"role": "user", "content": f"Continue the story, part {turn + 1}."}])
    payload = {
        "model": "impish_bloodmoon_12b-mlx",
        "messages": msgs,
        "stream": False,
        "max_tokens": 650,
        "temperature": 0.78,
        "top_p": 0.9,
        "repeat_penalty": 1.15,   # <-- remove this line and 8/8 turns succeed
        "stop": STOPS,
    }
    req = urllib.request.Request(URL, data=json.dumps(payload).encode(),
                                 headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=300) as r:
        rep = json.load(r)["choices"][0]["message"]["content"]
    print(f"turn {turn + 1}: ok ({len(rep)} chars)")
    history += [{"role": "user", "content": f"Continue the story, part {turn + 1}."},
                {"role": "assistant", "content": rep}]

Observed: with repeat_penalty: 1.15, generation degrades on turn 7 (20-char reply) and turn 8 fails with HTTP 400 carrying a fatal scheduler exception. Without the repeat_penalty line (all else identical, fresh model load), all 8 turns succeed.

I also reproduced a crash with repeat_penalty on single-shot requests (3/3) when the server had previously processed several multi-turn requests sharing a long prompt prefix — so accumulated prompt-cache state appears to be part of the trigger.

frequency_penalty / presence_penalty do not trigger the crash (probe: 5/5 OK with frequency_penalty=0.4, presence_penalty=0.2 after a clean reload).

Fatal tracebacks observed (two distinct failure points, same trigger)

1. Logprob indexing (batched_model_kit.py:410):

File ".../mlx_engine/model_kit/batched_model_kit.py", line 410, in _generate
    token_logprob = r.logprobs[r.token].item()
ValueError: Slice indices must be 32-bit integers.

(Possibly an int64 token index hitting MLX's 32-bit slice requirement — same class as ml-explore/mlx#2710.)

2. Prompt-cache merge (mlx_lm/models/cache.py:1462 via _merge_caches):

File ".../mlx_lm/generate.py", line 878, in _merge_caches
    batch_cache.append(caches[0][i].merge([c[i] for c in caches]))
File ".../mlx_lm/models/cache.py", line 1462, in merge
    keys[i : i + 1, :, p : p + l] = c._temporal_order(c.keys)[..., -l:, :]
ValueError: [broadcast_shapes] Shapes (1,8,8,128) and (1,8,8777,128) cannot be broadcast.

Aftermath: once either exception fires, the model process is dead; all following requests return
{"error":"The model has crashed without additional information. (Exit code: null)"} until manual reload.
Occasionally the same workload instead ends in
RuntimeError: [METAL] Command buffer execution failed: Insufficient Memory at the same batched_model_kit.py site.

Expected

Either repeat_penalty works with the batched scheduler's prompt cache, or the parameter is rejected/ignored gracefully — a request parameter should not be able to kill the model process.