The Crash

Fatal error: [broadcast_shapes] Shapes (1,2,335,256) and (1,2,292,256) cannot be broadcast.

Two MLXArrays with incompatible sequence dimensions (335 vs 292) combined in the computation graph. The 292 is a stale pre-trim array that should have been replaced.

ScenarioRequests Before CrashTimeCrash Shapes
Stress test (7 users)59~10 min(1,2,695,256) vs (1,2,524,256)
Reproducer run 1107~2 min(1,2,43,256) vs (1,2,35,256)
Reproducer run 212~15 sec(1,2,332,256) vs (1,2,325,256)
Reproducer run 37~9 sec(1,2,335,256) vs (1,2,292,256)

Root Cause

The KV Cache Trim Bug

KVCacheSimple.trim() only decrements offset — it does NOT physically slice self.keys or self.values:

// KVCache.swift lines 429-433
public override func trim(_ n: Int) -> Int {
    let trimmed = min(offset, n)
    offset -= trimmed       // Only this changes. self.keys untouched.
    return trimmed
}

After restoring a cached state with 292 tokens and trimming to effectivePrefix=97:

self.keys shape: [1, 2, 292, 256] ← full pre-trim array, unchanged self.offset: 97 ← decremented by trim() Array layout: position: 0 .............. 96 97 ................. 291 [ valid cached data ][ stale data from prior gen ] ↑ offset = 97 ↑ still in self.keys, shape 292

During the forward pass, update() writes new suffix tokens starting at position 97. Since 97 + suffix_len ≤ 292 (capacity sufficient), no reallocation happens. The in-place subscript assignment creates a new lazy MLX graph node that depends on the original shape-292 array. When MLX's scheduler runs operations that reference this stale shape alongside the expected shape (97 + suffix), the shapes don't match → broadcast_shapes crash.

The Fix

Round-trip through the state getter (slices to offset) and setter (replaces self.keys/self.values with the sliced result):

// After trim, physically truncate arrays
for i in 0..<generationCache.count {
    if generationCache[i].isTrimmable && generationCache[i].offset > 0 {
        generationCache[i].state = generationCache[i].state
    }
}

This replaces self.keys (shape [1,2,292,256]) with self.keys[..., 0..<97, ...] (shape [1,2,97,256]). The stale pre-trim array is fully dereferenced from the computation graph.

Why It Didn't Crash Every Time

The crash required a specific chain of conditions to all be true simultaneously. Any condition failing meant the request succeeded normally:

Condition 1 — Cache Hit
The request must match a prefix stored in the radix tree. First requests, requests with unique prefixes, and multimodal inputs always miss.
With 4 system prompts x 30 questions, the reproducer achieved high hit rates quickly.
Condition 2 — Trim Needed
The stored state must be longer than effectivePrefix. The radix cache stores state after full generation (prompt + output tokens), so stored_len = prompt + generated (e.g., 97 + 195 = 292). This was almost always true.
Condition 3 — No Reallocation (the key variable)
KVCacheSimple.update() uses a 2x growth strategy. When effectivePrefix + suffix_len > stored_len, it allocates a fresh array and copies only valid data — the stale pre-trim array is dereferenced and the bug doesn't manifest.
Since stored_len includes generated tokens (100-500+) and suffix_len is typically short (20-80 tokens), reallocation rarely happened — the stale array almost always persisted.
Condition 4 — MLX Lazy Graph Evaluation Order
MLX defers all computation. The stale shape-N array and the correct shape-(P+S) slice are both lazy graph nodes. The crash only manifests when MLX's scheduler runs operations in an order that exposes the shape mismatch between the two.
This is the non-deterministic factor — even with all other conditions met, the crash depended on graph scheduling.

Why the reproducer crashed fast

Why normal use crashed slower

Fix Attempts

1
isTrimmable guard — skip MambaCache restore
Crashed at 28 requests. Wrong theory: MambaCache state is fixed-size conv state and doesn't cause shape mismatches. The bug was in KVCacheSimple's lazy trim, not MambaCache. Also wasteful — skipping MambaCache restore throws away caching benefit for 30/40 layers.
2
arr.eval() after trim — materialize lazy slices
Crashed at 7 requests. Calling arr.eval() on the state getter's lazy slices materializes those slices, but doesn't replace the underlying self.keys/self.values that update() uses directly. The stale full-size arrays persist in the graph.
3
layer.state = layer.state — round-trip getter → setter
700 requests, 0 crashes. Getter slices self.keys to [..., 0..<offset, ...]. Setter assigns the sliced copy back as self.keys and derives offset from its shape. Stale pre-trim array fully dereferenced from computation graph.

Verification

700
Requests (0 crashes)
0
Errors
1.72
Requests/sec
312
Max cached tokens
Test RunRequestsErrorsDurationRPSConfig
Fix 3 — run 12000116s1.72isTrimmable guard + state round-trip
Fix 3 — run 25000291s1.72isTrimmable guard + state round-trip
Fix 3 — run 3 (guard removed)5000299s1.67state round-trip only (final fix)

Run 3 confirmed the isTrimmable guard (fix 1) was unnecessary — the state round-trip alone is sufficient. Removing the guard restores MambaCache prefix caching for 30/40 layers.

Architecture Context

Qwen3.5-35B-A3B: 40 layers, fullAttentionInterval=4 Layer: 0 1 2 3 4 5 6 7 8 9 10 11 ... 36 37 38 39 Type: Mam Mam Mam KV Mam Mam Mam KV Mam Mam Mam KV Mam Mam Mam KV Mam = MambaCache (30 layers) Fixed-size conv state. isTrimmable=false. KV = KVCacheSimple (10 layers) Sequence-indexed K/V. isTrimmable=true. ← BUG HERE

Cache Save/Restore Flow

SAVE (after generation completes): layerStates = generationCache.map { $0.state } radixCache.insert(inputTokens, layerStates) // KVCacheSimple.state getter returns keys[..., 0..<offset, ...] // Stored state shape = [1, heads, prompt+generated, dim] RESTORE (on cache hit): (prefixLen, states) = radix.findPrefix(inputTokens) effectivePrefix = min(prefixLen, inputTokens.count - 16) generationCache[i].state = states[i] // KVCacheSimple.state setter: self.keys = states[i][0] // Now: self.keys shape = [1, heads, stored_len, dim], offset = stored_len trim(stored_len - effectivePrefix) // offset = effectivePrefix (correct) // self.keys shape = [1, heads, stored_len, dim] (STALE — not sliced!) generationCache[i].state = generationCache[i].state // THE FIX // getter: returns keys[..., 0..<offset, ...] = [1, heads, effectivePrefix, dim] // setter: self.keys = sliced array, offset = effectivePrefix // Now: self.keys shape = [1, heads, effectivePrefix, dim] (CLEAN)

Performance Impact

None measurable. The state round-trip only runs on cache restore (not every request). It replaces one set of arrays with sliced copies — a lightweight operation (<1ms) compared to the model forward pass (500-700ms). No graph recompilation is triggered because the operation uses the same getter/setter code path that MLX already optimizes for.

Files Changed

FileChange
Sources/MacLocalAPI/Models/MLXModelService.swift Add layer.state = layer.state after trim in both streaming (~line 841) and non-streaming (~line 482) restore paths

Generated 2026-03-11 | maclocal-api | Reproducer: Scripts/repro-prefix-crash.py | Stress test: Scripts/stress-test-7users.py