Fix batched decode crash for hybrid cache models (Qwen3.5)

[laudney]

Summary

Interim fix for the mlx-native (non-paged) path so hybrid models like Qwen3.5 can serve today.

• Hybrid models (e.g. Qwen3.5-35B-A3B) use mixed cache types (ArraysCache for linear/SSM layers + KVCache for attention layers). BatchKVCache.offset returns mx.array but hybrid attention code uses cache.offset as a Python int for mask slicing, causing ValueError: Slice indices must be integers or None.
• Detects hybrid caches at model load time via make_prompt_cache() and falls back to sequential decode for incompatible models.
• Core detection logic in vllm_metal/v1/cache_utils.py (standalone pure function), keeping model_runner.py changes minimal per [Refactor] Refactor model runner to keep it minimal and easy to read #122.

NOTE: This is a band-aid until per-layer attention dispatching (#201) and a paged linear attention kernel (roadmap #148) land, at which point hybrid models will be handled properly at the dispatch level.

Test plan

• Verified fix resolves ValueError when serving Qwen3.5 via vllm-metal
• Verified standard (non-hybrid) models still use batched decode path (Qwen3-0.6B: all KVCache → True)
• Verified Qwen3.5-35B-A3B hybrid model falls back gracefully ({ArraysCache, KVCache} → False)

[WindChimeRan]

Thanks for the PR @laudney

Is it possible to also add support to paged kv cache?

VLLM_METAL_USE_PAGED_ATTENTION=1 VLLM_METAL_MEMORY_FRACTION=0.7

[ericcurtin]

@laudney please sign your commit to pass the DCO build

[ricky-chaoju]

Hi @laudney, please complete the DCO requirement:
https://github.com/vllm-project/vllm-metal/pull/121/checks?check_run_id=66030669295
Thanks!

[laudney]

Hey @ericcurtin @ricky-chaoju — thanks for the heads up! Rebased onto latest main and added the DCO sign-off. Should be all green now. Let me know if anything else needs updating!

[ericcurtin]

@laudney sadly there's conflicts now

[laudney]

@ericcurtin Merge conflict is resolved — rebased onto latest main and force-pushed. DCO check is passing.

To get this merged, I believe I still need:

1. Code review approval from a maintainer
2. Paged KV cache support — @WindChimeRan asked about compatibility with VLLM_METAL_USE_PAGED_ATTENTION=1. I'll look into this as a follow-up unless it's a blocker for this PR.

Let me know if there's anything else needed!

[WindChimeRan]

unreachable code?

[laudney]

Fixed — the init was placed after a return in the is_stt property. Now correctly in __init__.

[WindChimeRan]

Thanks for the PR @laudney ! It's been a few weeks and main has changed significantly. Could you check if the original crash still reproduces on current main? If it does, I think the fix should be more surgical and ideally live outside model_runner.py in a utils file.

Clarification: Paged KV cache support is non-blocking.

Also noticed the two test plan items are still unchecked. How's that going?

[laudney]

@WindChimeRan Thanks for the thorough review. Addressed all feedback:

Crash still reproduces on current main

Confirmed with mlx-lm 0.31.1:

• KVCache.offset → int (e.g. 4)
• BatchKVCache.offset → mx.array (e.g. array([4, 4]))
• Hybrid model attention code does mask[..., :cache.offset] → ValueError: Slice indices must be integers or None.

Tested with mlx-community/Qwen3.5-35B-A3B-4bit: cache is {ArraysCache, KVCache} → detection correctly returns False.

Refactored: detection logic moved out of model_runner.py

Per your suggestion and #122's direction, the core logic now lives in vllm_metal/v1/cache_utils.py as a standalone pure function. model_runner.py only has a thin 7-line wrapper. Net change to model_runner.py is minimal (+22/-1).

Fixed unreachable code

Good catch on line 673 — the _supports_batched_decode init was after a return in the is_stt property. Now correctly in __init__.

Scoping: this is an interim fix

I understand this PR sits within a broader effort:

• Refactor paged attention dispatch to support multiple attention types #201 — per-layer attention dispatching (SDPA vs linear)
• [Kernel] Fused GDN linear attention kernel for Qwen3.5 #186 — fused GDN linear attention kernel
• [RoadMap] [Paged KV] Continuous Batching + Chunked Prefilling + paged varlen flash att #148 roadmap — paged linear attention kernel for Qwen3.5

This PR is specifically an interim workaround for the mlx-native (non-paged) path so Qwen3.5 can serve today via sequential decode. Once #201 lands and a paged linear attention kernel is available, hybrid models will be handled properly at the dispatch level and this fallback becomes unnecessary.

Test plan items

• Verified fix resolves ValueError when serving Qwen3.5 via vllm-metal
• Verified standard (non-hybrid) models still use batched decode path — tested Qwen3-0.6B (all KVCache → supports_batched_decode = True)
• Verified Qwen3.5-35B-A3B hybrid model falls back gracefully ({ArraysCache, KVCache} → supports_batched_decode = False)

[laudney]

@WindChimeRan Re: paged KV cache support — this PR intentionally targets only the mlx-native (non-paged) path as an interim fix.

Proper paged attention support for hybrid models like Qwen3.5 requires the per-layer attention dispatching you're building in #201 (routing SDPA layers vs GatedDeltaNet linear attention layers separately) plus a paged linear attention kernel (roadmap #148). That's the right long-term approach and this PR doesn't try to duplicate that effort.

Happy to help with #201 or the linear attention kernel if useful.

[WindChimeRan]

I couldn't reproduce the crash Qwen3.5-35B-A3B from the main branch. (M2 max, 64GB)

Could you please take a look? @LxYuan0420. It seems like the problem has already been fixed by #110 , but I'm not very sure.

[LxYuan0420]

Could you provide exact repro commands on current main (including model, mlx-lm version)?

Also, the current change looks like a temp workaround rather than a proper fix. Ideally we want the fix to live in the attention dispatch layer for hybrid models, rather than globally gating batching?

[ericcurtin]

Closing this, conflicts and inactivity. Feel free to open a new PR again in future