[Bugfix] Fix KV cache undercount in MLX path for large block sizes

[samwarren]

Summary

• Fix _one_sequence_kv_bytes in MetalWorker to use block-aligned token counts, matching the upstream scheduler's cdiv(max_model_len, block_size) * page_size_bytes accounting
• Prevents server startup failure on Mamba-hybrid models (e.g. Granite 4.0-H) where block_size is padded to 400 to match the mamba page size

Problem

_one_sequence_kv_bytes computes KV cache bytes using max_model_len directly (e.g. 2048 tokens). But the upstream _check_enough_kv_cache_memory in vLLM core uses block-aligned sizes: cdiv(2048, 400) = 6 blocks = 2400 tokens. This causes "needed > available" even though the intent is to report exactly enough for one sequence:

ValueError: To serve at least one request with the models's max seq len (2048),
(0.18 GiB KV cache is needed, which is larger than the available KV cache memory
(0.16 GiB).

For the default block_size=16, cdiv(2048, 16) * 16 = 2048 — no padding, so this never triggers. It only manifests with large block sizes like 400, which occurs on Mamba-hybrid models (GraniteMoeHybridForCausalLM) where the attention block size is padded to match the mamba page size.

Fix

Round max_model_len up to the nearest block_size boundary in _one_sequence_kv_bytes:

block_size = self.vllm_config.cache_config.block_size
max_tokens = -(-self.model_config.max_model_len // block_size) * block_size

Reproduction

vllm serve mlx-community/granite-4.0-h-tiny-3bit-MLX --max-model-len 2048 --enforce-eager
# Fails with KV cache memory error

# After fix:
# Server starts successfully

Test plan

• Added test_block_alignment_rounds_up_token_count — verifies block-aligned calculation with block_size=400
• Updated existing test_non_hybrid_counts_all_layers and test_hybrid_adds_linear_state to include vllm_config.cache_config.block_size in mocks
• All 10 tests in test_v1_worker.py pass
• Verified vllm serve mlx-community/granite-4.0-h-tiny-3bit-MLX --max-model-len 4096 --enforce-eager starts and serves requests on M4 Pro 48GB