fix: pass size to ArraysCache in BatchMambaCache for Qwen3.5 hybrid models

[neomody77]

Summary

Fixes --continuous-batching mode crashing with Qwen3.5 models (hybrid Attention + Mamba/SSM architecture).

Root cause: BatchMambaCache.__init__() conditionally skips passing size to ArraysCache.__init__() when HAS_MAMBA_CACHE=True. Since MambaCache was removed in mlx-lm >= 0.30.6, the fallback assigns ArraysCache as MambaCache, but the HAS_MAMBA_CACHE flag becomes unreliable. ArraysCache.__init__() requires size as a positional argument, causing:

Engine loop error: ArraysCache.__init__() missing 1 required positional argument: 'size'

This error loops infinitely, making the server unresponsive.

The same issue exists in BatchMambaCache.extract(), which calls MambaCache() without size.

Fix

Always pass size to super().__init__() — this is safe for both ArraysCache (requires it) and legacy MambaCache (accepts it). Remove the HAS_MAMBA_CACHE conditional branches that are no longer reliable.

Impact

• Before: --continuous-batching with Qwen3.5 → infinite error loop, server hangs
• After: continuous batching + prefix cache works correctly

Tested on M1 Pro 32GB with mlx-community/Qwen3.5-4B-4bit:

Request	Before (SimpleEngine, no cache)	After (BatchedEngine + prefix cache)
Cold	4.9s	1.17s
Warm	4.9s (no cache)	1.08s (cache hit)

Context

Qwen3.5 uses a hybrid architecture where model.make_cache() returns a mix of KVCache (attention layers) and ArraysCache (Mamba/SSM layers). This is a known challenge:

• Prefix cache reuse is broken for all hybrid-architecture models (sliding window, SSM/Mamba) ml-explore/mlx-lm#980 — prefix cache fails for hybrid models
• [Bug] Prefill Failure of Qwen3.5 Model Using KV Cache in the mlx‑lm Framework QwenLM/Qwen3.5#37 — ArraysCache vs KVCache in hybrid architecture
• Related to continuous batching incompatible with MLLM (ArraysCache vs KVCache) #159 (continuous batching incompatible with ArraysCache)

Test plan

• BatchGenerator with Qwen3.5-4B-4bit no longer crashes
• --continuous-batching --enable-prefix-cache serves requests correctly
• Prefix cache provides ~2x speedup on repeated prompts
• Tool calling (qwen3_coder parser) works with batched engine

[Thump604]

Nice catch. The HAS_MAMBA_CACHE conditional branches are indeed unreliable since mlx-lm >= 0.30.6 made MambaCache just an alias for ArraysCache. The fix is clean — always passing size is correct for both code paths since MambaCache (when it existed separately) also accepted size, and ArraysCache requires it.

I've been running Qwen3.5-122B-A10B on M2 Ultra 128GB and can confirm this exact error class. Our PR #165 fixes hybrid cache handling in the MLLM batch generator (_make_batch_cache, MLLMBatch.extend(), etc.) but doesn't touch mamba_cache.py — so these two PRs are complementary, not conflicting. Both are needed for full continuous batching support on Qwen3.5 hybrid models.

One small note: the extract() fix is especially important. Without size, the extracted single-sequence MambaCache would also fail during batch operations that need to reconstitute individual caches.

+1 for merge. This is a straightforward, well-tested bugfix that unblocks continuous batching on hybrid architectures.

[waybarrios]

Reviewed the change. The fix is correct, ArraysCache requires size as a positional argument and the old HAS_MAMBA_CACHE conditional was skipping it on the True branch. Since mlx-lm >= 0.30.6 removed MambaCache entirely, that flag is always False anyway, so the conditional was dead code hiding a real bug.

Always passing size is safe for both ArraysCache (requires it) and the legacy MambaCache (inherited it from ArraysCache).

One minor note: after this PR, the HAS_MAMBA_CACHE variable is still assigned in the try/except import block but never read anywhere. Could be cleaned up in a follow up.