[Bugfix][Core] Isolate hidden-state cache from DeepSeek-V4 MLA groups

[my0901]

Purpose

Prevent extract_hidden_states from inflating DeepSeek-V4 KV-cache page sizes and blocking engine startup.

Details

Problem

HiddenStateCacheSpec inherits from MLAAttentionSpec, so DeepSeek-V4's model-specific grouping path treats the hidden-state cache as a normal MLA layer. Its large page is then packed with MLA/SWA tuples and pads their pages.

With four bf16 hidden states of size 4096 and block size 256, the hidden-state page is 8 MiB. On DeepSeek-V4-Flash with TP=4 on 4x H100 80GB, this caused the startup check for an 8,192-token request to report 581.48 GiB of required KV cache with only 21.61 GiB available per worker.

What this PR changes

• Exclude HiddenStateCacheSpec from DeepSeek-V4 MLA/SWA tuple grouping.
• Add hidden-state layers back exactly once as an independent group.
• Let ExampleHiddenStatesConnector find the wrapped hidden-state group and
  use its block size.
• Add focused CPU regression tests for grouping and connector lookup.

Related PRs

• #39949 isolates the hidden-state cache for hybrid-attention models, but its generic path is unreachable after DeepSeek-V4's model-specific early return.
• #50894 fixes a separate TP page-size alignment issue in hidden-state extraction; it does not change DeepSeek-V4 MLA grouping.
• #49811 adds Model Runner V2 support for extract_hidden_states; it does not change DSV4 KV grouping.
• zupengwang/vllm#1 is a PP follow-up; its generic hybrid block-size fix runs after the DSV4-specific early return and does not fix this grouping bug.

Out of scope

• Model Runner V2 or pipeline-parallel hidden-state extraction support.
• Changes to MLA/SWA cache layouts outside hidden-state extraction.

Test Plan

Add directed regression coverage for:

• Excluding hidden-state cache specs from DeepSeek-V4 MLA/SWA tuple grouping.
• Restoring hidden-state cache specs exactly once as an independent uniform group without changing attention-group packing.
• Resolving a wrapped hidden-state group and using its independent block size in ExampleHiddenStatesConnector.

Run the existing extraction integration test and reproduce DeepSeek-V4-Flash startup with TP=4, max_model_len=8192, and hidden-state extraction enabled. The Docker wrapper accepts VLLM_WHEEL so this can be reproduced with the exact build under test.

Test Result

• Directed core regression cases:
  • test_group_and_unify_kv_cache_specs_excludes_hidden_state: hidden-state specs do not participate in DSV4 tuple grouping or alter the resulting attention-group page sizes.
  • test_hidden_state_page_does_not_force_dsv4_packing: a large hidden-state page does not force an otherwise uniform MLA/SWA layout into tuple packing.
  • test_get_kv_cache_groups_isolates_hidden_state: the hidden-state layer is restored exactly once in a standalone UniformTypeKVCacheSpecs group.
• Connector regression coverage verifies that the connector finds a wrapped hidden-only group, uses that group's block size, and reports the new isolation error when zero or multiple hidden groups are present.
• Validation runs with the patch overlaid on the exact 793ca6998 wheel:
  • Three directed node IDs from tests/v1/core/test_kv_cache_utils.py — 3 passed.
  • python -m pytest tests/v1/kv_connector/unit/test_hidden_states_connector.py — 10 passed.
• Four-GPU DeepSeek-V4-Flash smoke test with TP=4, max_model_len=8192, and hidden-state extraction enabled:
  • Available KV cache memory: 21.61 GiB.
  • GPU KV cache size: 6,199 tokens.
  • Maximum concurrency for 8,192 tokens per request: 0.76x.
  • KV-cache initialization and connector creation completed successfully; the engine proceeded to DeepGEMM warmup. The unpatched baseline failed at the KV-cache check with a 581.48 GiB estimate.

AI assistance was used to investigate and implement this change. Every changed line has been reviewed by the human submitter. The targeted behavior was validated with focused unit tests and the four-GPU startup test above.

Where should the reviewer start?

Start with group_and_unify_kv_cache_specs() and get_kv_cache_groups() invllm/v1/core/kv_cache_utils.py, then ExampleHiddenStatesConnector._find_cache_kv_group_id().

[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.

🚀

[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.