Fix/qwen3.5 init failed

[Alex-ai-future]

Summary

This PR enables Qwen3.5 hybrid models (SDPA + GDN layers) to run on Metal by implementing update_block_size_for_backend() to unify KV cache page sizes, adding MLA (Multi-Token Latent Attention) support, and improving memory reporting for the MLX path.

Problem

When trying to run Qwen3.5 (a hybrid model with SDPA + GDN layers) on Metal, the following issues were encountered:

1. Hybrid Model Page Size Alignment Failure

vLLM's KV cache validation fails because SDPA page size and Mamba page size are not divisible:

NotImplementedError: The page size of the layer is not divisible by the maximum page size.
Cannot unify by adjusting block_size.

Test failure on main branch:

FAILED tests/test_attention_dispatch.py::test_qwen35_paged_attention_hybrid - 
NotImplementedError: The page size of the layer is not divisible by the maximum page size.

2. Forced Paged Attention Causes OOM

Previously, paged attention was auto-enabled for hybrid models, which allocates a large contiguous memory buffer that exceeds the capacity of smaller Metal devices.

3. Inaccurate Memory Reporting

The MLX path reported memory based on max_model_len, which gave misleadingly small values for the scheduler.

Solution

1. Add update_block_size_for_backend() to MetalPlatform

Implements a 5-step process to unify page sizes for hybrid models:

1. Compute attention page size per token - Uses MLAAttentionSpec for MLA models or FullAttentionSpec otherwise
2. Get Mamba page size - Queries model class for mamba state shape and dtype
3. Calculate block_size - Ensures SDPA page_size >= Mamba page_size using kernel_block_alignment_size=32
4. Sync mamba_block_size - If using align mode
5. Pad mamba_page_size - Matches SDPA page_size exactly

Key insight: This is a "logical" fix for vLLM's scheduler validation only. The Metal plugin manages KV cache internally via MLX's make_prompt_cache(), independent of vLLM's calculations.

2. Add MLA Support

Hybrid models with MLA (e.g., DeepSeek variants) now use MLAAttentionSpec for correct page size calculation:

if getattr(model_config, "use_mla", False):
    attn_page_size_1_token = MLAAttentionSpec(...).page_size_bytes
else:
    attn_page_size_1_token = FullAttentionSpec(...).page_size_bytes

3. Add Early Error for Hybrid + Paged Attention

Raises a clear ValueError when users attempt to use paged attention with hybrid models:

ValueError: Hybrid models (e.g., Qwen3.5) are not supported with paged attention on Metal. 
The Metal paged attention kernel only supports block_size in {8, 16, 32}, but hybrid models 
require block_size=160. Please remove VLLM_METAL_USE_PAGED_ATTENTION=1.

Root cause: Metal paged attention kernels only support block_size ∈ {8, 16, 32}, but Qwen3.5 requires block_size=160.

4. Improve Memory Reporting

Changed MLX path to report 80% of remaining Metal memory instead of one max-length sequence:

_MLX_MEMORY_BUDGET_FRACTION = 0.8
metal_limit = mx.device_info()["max_recommended_working_set_size"]
model_memory = self._get_model_memory_usage()
available = int((metal_limit - model_memory) * _MLX_MEMORY_BUDGET_FRACTION)

Before: reporting 4.29GB for scheduler admission control (one max-length sequence, max_model_len=2048)

After: reporting 11.20 GB for scheduler (Metal limit: 16.00 GB, Model: 2.00 GB, Remaining: 14.00 GB, KV budget: 11.20 GB)

5. Remove Auto-Enable Paged Attention for Hybrid Models

MLX's make_prompt_cache() handles hybrid KV cache natively. Paged attention is now opt-in rather than forced.

Changes

New Files

File	Lines	Description
tests/test_platform_update_block_size.py	576	Comprehensive unit tests (16 test cases)

Modified Files

File	Changes	Description
vllm_metal/platform.py	+172	Add update_block_size_for_backend() with MLA support
vllm_metal/v1/worker.py	+41, -23	Improve memory reporting, remove auto-enable paged
vllm_metal/v1/model_runner.py	+1, -1	Fix MLX API: mx.device_info()
tests/test_v1_worker.py	+22	Update memory reporting test

Test Results

Unit Tests (16/16 Pass)

$ source .venv-vllm-metal/bin/activate && python -m pytest tests/test_platform_update_block_size.py -v

tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_hybrid_model_success PASSED
tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_hybrid_model_block_size_already_sufficient PASSED
tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_non_hybrid_model_skipped PASSED
tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_model_config_none PASSED
tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_model_resolution_failure PASSED
tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_get_mamba_state_shape_failure PASSED
tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_get_mamba_state_dtype_failure PASSED
tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_mamba_page_size_zero PASSED
tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_invalid_architecture PASSED
tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_block_size_increased_to_minimum PASSED
tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_mamba_cache_mode_align PASSED
tests/test_platform_update_block_size.py::TestUpdateBlockSizeForBackend::test_hybrid_with_paged_attention_raises_error PASSED
tests/test_platform_update_block_size.py::TestMLAModels::test_mla_hybrid_model_uses_mla_spec PASSED
tests/test_platform_update_block_size.py::TestMLAModels::test_mla_non_hybrid_skipped PASSED
tests/test_platform_update_block_size.py::TestMLAModels::test_mla_with_cache_dtype[bfloat16] PASSED
tests/test_platform_update_block_size.py::TestMLAModels::test_mla_with_cache_dtype[float16] PASSED

16 passed, 2 warnings in 2.69s

Lint Checks

$ ruff check vllm_metal/ tests/
All checks passed!

$ ruff format --check vllm_metal/ tests/
5 files already formatted

Usage

Default Path (Recommended for Hybrid Models)

# No env var needed - uses MLX's native KV cache via make_prompt_cache()
vllm serve Qwen/Qwen3.5-0.8B
vllm serve Qwen/Qwen3.5-4B
vllm serve Qwen/Qwen3.5-14B
vllm serve Qwen/Qwen3.5-32B

Paged Attention Path (Non-Hybrid Models Only)

# Only for non-hybrid models
VLLM_METAL_USE_PAGED_ATTENTION=1 vllm serve HuggingFaceTB/SmolLM2-135M-Instruct

Unsupported Configuration (Will Raise Clear Error)

# This will now raise a helpful ValueError instead of cryptic kernel failure
VLLM_METAL_USE_PAGED_ATTENTION=1 vllm serve Qwen/Qwen3.5-0.8B
# ValueError: Hybrid models (e.g., Qwen3.5) are not supported with paged attention on Metal...

Known Limitations

Hybrid + Paged Attention is unsupported on Metal due to kernel limitations:

• Metal paged attention kernels only instantiate block_size ∈ {8, 16, 32}
• Hybrid models require block_size=160 to satisfy vLLM's page size divisibility validation
• Users should use the native MLX KV cache path (default) for hybrid models

Related Issues

• Fixes Engine crashes when prefix caching hits new complete prefill in Metal paged unified path #184 (hybrid model initialization failure)
• Addresses reviewer feedback from Fix/qwen3.5 init failed #230 (@ricky-chaoju, @ericcurtin)
• Related to [Bugfix] Derive head_dim for Qwen3 paged attention path  #232 (head_dim AttributeError on Qwen3)
• Related to [Qwen3.5] C++ Metal kernel for hybrid SDPA + GDN paged inference(Stage C) #226 (regression fix)

Commits

This PR includes the following commits:

964e9df [Metal] Fix inaccurate docstring and test comments
c87263b [Metal] Fix import order and add test for hybrid + paged error case
f7161cc [Metal] Address reviewer feedback on hybrid + paged attention
774c8bd  use new device
fb5e2d5 [Metal] Fix ruff format issues
5c7379b [Metal] Fix lint issues in MLA support changes
80b8d19 [Metal] Add MLA support to update_block_size_for_backend
733612b [Metal] Improve error handling in update_block_size_for_backend
36d8cb7 [Metal] Add unit tests for update_block_size_for_backend and improve error handling
b2973f0 Fix test for determine_available_memory single_sequence mode
a292980 Restore model_runner.py to upstream version
3f22b84 [Metal] Fix Qwen3.5 hybrid model initialization
d49d0fb [Metal] Fix hybrid model KV cache page size alignment

[ericcurtin]

update_block_size_for_backend has no callers, no tests for new platform method, mamba_page_size_padded may be unset if update_block_size_for_backend is never invoked, magic 0.8 factor not configurable

[Alex-ai-future]

update_block_size_for_backend has no callers, no tests for new platform method, mamba_page_size_padded may be unset if update_block_size_for_backend is never invoked, magic 0.8 factor not configurable

Thank you for your review! I would like to discuss the hardcoded coefficient of 0.8 within the determine_available_memory() function.
Issue: The GPU path measures memory usage via actual performance profiling, whereas the MLX path relies on a hardcoded estimate. This approach introduces inconsistency and prevents users from fine-tuning the settings to suit their specific application scenarios.
Question: Should this function report the actual memory capacity (similar to the GPU path), or should it establish a configurable memory budget? Is it necessary to make this function user-adjustable, or should it simply return an actual, measured value?

[Alex-ai-future]

1. "update_block_size_for_backend has no callers"

The method is called by vLLM core in:

• File: vllm/v1/executor/uniproc_executor.py
• Function: _UniProcExecutor.init() (line 53)
• Context: During executor initialization, after load_model() completes

1 # vllm/v1/executor/uniproc_executor.py:53
2 self.driver_worker.load_model()
3 current_platform.update_block_size_for_backend(self.vllm_config) # ← Called here

This is the standard integration point for all platform plugins.

2. "no tests for new platform method"

Agreed. I've added 11 unit tests in tests/test_platform_update_block_size.py covering success, failure, and edge cases.

3. "mamba_page_size_padded may be unset"

Good catch. I've improved error handling to raise exceptions for hybrid models instead of silently returning.

[ricky-chaoju]

Thanks for working on this. A few issues found during review:

block_size will crash the Metal kernel. Verified with all Qwen3.5 variants (0.8B/4B/14B/32B) the formula at platform.py:413-417 computes attn_block_size=160, but Metal paged attention only has kernel instantiations for {8, 16, 32}. Users who set VLLM_METAL_USE_PAGED_ATTENTION=1 will hit RuntimeError: Unable to load function ...bs160.... The docstring claims this is "scheduler validation only" but cache_config.block_size is consumed by the kernel in the paged path.

Magic constant. The 0.8 factor in worker.py:420 should be a named constant. Agree with ericcurtin's earlier feedback.

Stale test comments. test_mla_hybrid_model_uses_mla_spec and test_mla_with_cache_dtype say "will FAIL until MLA support is added" but the implementation landed in the same commit — all 15 tests pass. Comments should be cleaned up.

Related: #232 fixes another #226 regression (head_dim AttributeError on Qwen3).

[ricky-chaoju]

Verified with actual model parameters — all Qwen3.5 variants (0.8B/4B/14B/32B) compute attn_block_size=160 here.

Model	mamba_page_size	attn_page_size_1_token	block_size
Qwen3.5-0.8B	548,864	4,096	160
Qwen3.5-4B	573,440	4,096	160
Qwen3.5-14B	585,728	4,096	160
Qwen3.5-32B	585,728	4,096	160

Metal paged attention kernel only has template instantiations for block_size ∈ {8, 16, 32} (pagedattention.metal:1452-1457). When a user sets VLLM_METAL_USE_PAGED_ATTENTION=1, this crashes:

RuntimeError: Unable to load function paged_attention_..._bs160_...

The docstring says "scheduler validation only" — inaccurate because cache_config.block_size is used by the Metal kernel in the paged path.

Related: #232 fixes another #226 regression (head_dim AttributeError on Qwen3).

[Alex-ai-future]

My understanding is as follows:

1. The latest version of vLLM now validates the required block size; if it remains unmodified, it will fail this check.
2. In paged mode, if the block size is modified, there is no corresponding kernel available to match it.

However, how were you able to successfully run the application while utilizing the paged kernel?

[Alex-ai-future]

And, do we have a plan to support dynamic kernel or to skip validation? Otherwise it is not likely to work on the latest vllm. Tell me is I was misunderstood!

[Alex-ai-future]

Is it possible to detect whether Paged Mode is enabled?

If enabled, skip modifying the block size.

If disabled, modify the block size.

As part of a broader plan, at the very least, an option could be provided to allow running hybrid models without Paged Mode on latest vllm.

[ricky-chaoju]

Confirmed — none of the Metal-supported block sizes {8, 16, 32} pass vLLM's page size divisibility check for Qwen3.5 (mamba_page_size=548864). The minimum required block_size is 134, which has no kernel instantiation. So hybrid + paged is genuinely unsupported right now.

Your approach is correct: update_block_size_for_backend solves the non-paged validation, and removing
auto-enable paged for hybrid is the right guard. My earlier review about the block_size crash still holds, but only when a user manually forces VLLM_METAL_USE_PAGED_ATTENTION=1 on a hybrid model — which is an unsupported configuration.

Suggestion: add an early error in _setup_paged_attention() (or update_block_size_for_backend) when both hybrid and paged attention are active, so users get a clear message instead of a cryptic kernel load failure.

[Alex-ai-future]

thank you for your advice

[ricky-chaoju]

Agree with ericcurtin's feedback. Should at least be extracted to a named constant (e.g. _MLX_MEMORY_BUDGET_FRACTION = 0.8).

[ricky-chaoju]

This docstring says "will FAIL until MLA support is added" but the implementation was added in the same commit — all 15 tests pass. Same issue at line 479 and 528. Comments should be cleaned up.

[ericcurtin]

Need to sign to pass DCO

[Alex-ai-future]

signed, thank you！

[LxYuan0420]

This replaces the one-sequence admission logic (30eb33b, refined in #229) with a heuristic '80% of remaining memory.' That changes scheduler semantics and makes the block-alignment fix irrelevant on the MLX path. Keep the one-sequence estimate unless this is an explicit policy change.

[Alex-ai-future]

Regarding determine_available_memory: Does the result calculated using the one-sequence method actually represent the available memory? During actual execution, I observed that the resulting value was extremely low, which prevented the program from running.

[LxYuan0420]

This hard ValueError blocks hybrid + paged attention unconditionally, which conflicts with PR #235’s block-size translation. Gate it on actual kernel support or downgrade to a warning.

[Alex-ai-future]

Even if execution doesn't halt here, the program would still fail during actual runtime in this version due to the lack of a matching kernel.
Are you sure you want to allow this branch to execute within this PR? If so, I can change it to a warning for now.

[LxYuan0420]

_MLX_MEMORY_BUDGET_FRACTION = 0.8 ignores cache_config.gpu_memory_utilization, removing user control over memory limits on the MLX path. Use the config value instead of a hard constant.

[LxYuan0420]

Duplicated MLAAttentionSpec/FullAttentionSpec instantiation. Use a SpecClass selection to remove repetition.

[Alex-ai-future]

@LxYuan0420 Thank you for the thoughtful review! You raised valid concerns, and I've addressed them:

1. Hard ValueError → Warning ✅ Since PR #235 is now merged, hybrid + paged attention works via block-size translation. Changed the hard error to an informative warning explaining the mechanism.

2. Memory reporting strategy ✅ You were right — my change from one-sequence estimate to gpu_memory_utilization altered scheduler semantics without proper discussion. I've reverted to PR #229's design (one max-length sequence) to maintain conservative admission control for MLX's dynamic allocation.

3. Code duplication ✅ Refactored to use spec_class selection pattern instead of duplicating MLAAttentionSpec/FullAttentionSpec instantiation.

4. Tests & Lint ✅ All tests pass (16/16 + 11/11), and ruff check/format are clean.

I appreciate you catching these issues — especially the memory reporting strategy change, which I should have discussed more carefully before implementing. 🙏

[LxYuan0420]

I noticed a separate issue in linear_cache_bytes_per_slot() where recurrent GDN memory is undercounted because GDNPagedStateCache forces mx.float32. That came from PR #226. It’s out of scope for this PR, so I’ll follow up with a small fix.