[Bugfix][Attention] Use packed MLA sizes for hybrid block alignment

[lucamotz]

Purpose

Hybrid block alignment uses the semantic head size to estimate MLA pages, but packed FP8/NVFP4 caches have different physical byte widths. Use the same width lookup as MLA attention so padded Mamba pages match the final MLA pages. This fixes inconsistent cache sizing; the packed layouts themselves are unchanged.

Duplicate check: #55219 is a separate GLM layout refactor; no other open PR was found for this sizing calculation when this PR was opened. AI-assisted contribution.

Test Plan

HF_HUB_OFFLINE=1 .venv/bin/python -m pytest tests/v1/core/test_kv_cache_utils.py \
  -k 'hybrid_mla_block_alignment or glm5 or mamba_hybrid' -q --tb=short
HF_HUB_OFFLINE=1 .venv/bin/python -m pytest tests/v1/attention/test_mla_backends.py \
  -k mla_kv_cache_spec_uses_layer_cache_dtype -q --tb=short
HF_HUB_OFFLINE=1 .venv/bin/python -m pytest tests/v1/core/test_kv_cache_utils.py -q --tb=short
.venv/bin/pre-commit run --files vllm/platforms/interface.py \
  vllm/v1/kv_cache_interface.py \
  vllm/model_executor/layers/attention/mla_attention.py \
  tests/v1/core/test_kv_cache_utils.py

Test Result

13 alignment/GLM/Mamba tests and 2 existing MLA-spec tests passed. The alignment regression also checks that Mamba padding matches the spec built by MLA attention. Both packed-format cases fail on base e473e90; ordinary FP8 passes. Full cache test file: 105 passed

Checked the real GLM-5.3-Flash TP2/MTP5 fp8_ds_mla configuration through the cache-spec and allocator APIs:

	Block tokens	Padded Mamba page	Final MLA page	Physical manager block
Base e473e90	4,608	2,359,296 B	3,022,848 B	38,098,944 B
PR a120883	3,584	2,351,104 B	2,351,104 B	29,632,512 B

The PR makes early Mamba padding agree with the final MLA page.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Team

Run ID: 51b3f354-5d84-42f1-94ad-c0ea5bdddb32

📥 Commits

Reviewing files that changed from the base of the PR and between f229587 and a120883.

📒 Files selected for processing (4)

• tests/v1/core/test_kv_cache_utils.py
• vllm/model_executor/layers/attention/mla_attention.py
• vllm/platforms/interface.py
• vllm/v1/kv_cache_interface.py

Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review.

---

Walkthrough

Hybrid MLA alignment now uses a shared helper for cache-dtype-specific packed Mamba state sizes. MLA cache specification and platform alignment use the helper. A parametrized test covers fp8_ds_mla, nvfp4_ds_mla, and fp8.

Changes

Hybrid MLA alignment

Layer / File(s)	Summary
Packed state size alignment and coverage vllm/v1/kv_cache_interface.py, vllm/model_executor/layers/attention/mla_attention.py, vllm/platforms/interface.py, tests/v1/core/test_kv_cache_utils.py	get_mla_state_content_bytes returns packed state sizes for supported MLA cache dtypes. MLA cache specification and platform alignment use the helper. The parametrized test verifies block sizes, Mamba state sizing, padded page sizes, and cache specification consistency.

Estimated code review effort: 2 (Simple) | ~10 minutes

Merge Risk: ⚪ Minimal · up to a1208

Hybrid MLA alignment now uses the packed physical cache width for supported FP8 and NVFP4 formats, with matching cache-spec and alignment behavior covered by tests. No current merge-blocking risk remains.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 62.50% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 8 functions across 4 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Title check	✅ Passed	The title clearly and concisely describes the main change: using packed MLA sizes for hybrid block alignment.
Description check	✅ Passed	The description directly explains the sizing bug, the implementation change, and the test results. It is relevant to the changeset.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment /ci run for upstream CI or /amd-ci run for AMD CI only whenever CI signals are needed.

Once the PR is approved or has the ready label, the PR author can also use the corresponding /ci run, /ci retry, and /ci cancel commands, or their /amd-ci variants. New commits do not start upstream CI automatically.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀