[Spec Decode] Support hybrid attention models in extract_hidden_states

[mgoin]

Summary

Hidden-state extraction now works on hybrid-attention models (e.g. Qwen3.5). The kv-transfer config no longer force-disables HMA, it stays on for connectors that declare SupportsHMA.

The KV-cache grouping handles the cache-only hidden-state layer alongside a hybrid attention/Mamba layout by filtering it out before unification and adding it back as its own page-aligned 1-layer group, with a strided reshape in the model runner to span the padded page.

ExampleHiddenStatesConnector is updated accordingly, and a new CI job runs the integration tests.

Test plan

• tests/v1/kv_connector/extract_hidden_states_integration/test_extraction.py — Llama end-to-end (GPU)
• Qwen3.5-4B + extract_hidden_states — hybrid model end-to-end (GPU), hidden states shape [N, 3, 2560] with non-zero values
• pre-commit run ruff-check / ruff-format / mypy-3.10 — all passing
• CI

🤖 AI-assisted (Claude)

[gemini-code-assist]

Code Review

This pull request introduces the HiddenStateCacheSpec to support hidden-state extraction within the vLLM V1 engine. Key changes include updating the KV cache grouping heuristics to prevent singleton cache-only layers from collapsing group sizes, refactoring the ExampleHiddenStatesConnector to utilize attn_metadata.slot_mapping directly, and implementing dynamic HMA (Hybrid Memory Architecture) support checks for connectors. Feedback is provided regarding the max_memory_usage_bytes implementation in HiddenStateCacheSpec, which currently fails to account for context parallelism, potentially leading to memory over-estimation during initialization.

[benchislett]

Besides the other included comments, my main concern here is that we're adding a very manual specialization for this extractor. This creates deviation from existing standard practices (native padding e.g. in mamba) and requires us to maintain the codepath as the HMA changes. This seems quite fragile, especially since bugs in the hidden extractor would not manifest during routine debugging / correctness testing.

I'm not convinced that it isn't feasible to simply apply padding to the hiddenstatecachespec (or to the other blocks, whichever is smaller, either case should be fine) and handle that natively: either in the kv connector itself or in the client code.

Specifically, I have problems with the code in gpu_model_runner.py which manually strides the view to avoid the padding. This seems to break some conventions about how we handle the KV connectors (which, AFAICT, seem to keep the padding?). Also, the special handling of the hidden state group in vllm/v1/core/kv_cache_utils.py seems fragile and prone to issues.

[johnnynunez]

this is good @mgoin... seems to enable nemotron too?

[benchislett]

I'm pretty sure the manual striding was merged into gpu model runner for DSV4. Assuming we can piggy-back on that, I think we can get this in without much change. Can you merge main?

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @mgoin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[DarkLight1337]

I'm trying to replicate the setup for training a model like https://huggingface.co/z-lab/Qwen3.5-27B-DFlash and the memory usage is quite unreasonable on AMD, it tells me I need over 200 GB VRAM for a 27B model...

Docker:

services:
  vllm:
    image: vllm/vllm-openai-rocm:nightly-bf0d2dc6d764f7ab1a69504f60a55883ec6d9b39
    container_name: vllm
    command: |
      --model Qwen/Qwen3.5-27B --enable-prefix-caching --limit-mm-per-prompt '{"image": 1, "video": 0}' --reasoning-parser qwen3 --default-chat-template-kwargs '{"enable_thinking": false}' --allowed-local-media-path / --log-error-stack --speculative_config '{"method": "extract_hidden_states", "num_speculative_tokens": 1, "draft_model_config": {"hf_config": {"eagle_aux_hidden_state_layer_ids": [1, 16, 31, 46, 61]}}}' --kv_transfer_config '{"kv_connector": "ExampleHiddenStatesConnector", "kv_role": "kv_producer", "kv_connector_extra_config": {"shared_storage_path": "/tmp/hidden_states"}}'
    environment:
      - OPENAI_API_KEY=${OPENAI_API_KEY}
      - HF_TOKEN=${HF_TOKEN}
    network_mode: host
    group_add:
      - video
    ipc: host
    cap_add:
      - SYS_PTRACE
    security_opt:
      - seccomp=unconfined
    devices:
      - /dev/kfd
      - /dev/dri
    ports:
      - 8000:8000
    restart: unless-stopped
    volumes:
      - type: bind
        source: ~/.cache/huggingface/hub/
        target: /root/.cache/huggingface/hub/

Error message:

ValueError: To serve at least one request with the models's max seq len (262144), (204.06 GiB KV cache is needed, which is larger than the available KV cache memory (158.0 GiB). Based on the available memory, the estimated maximum model length is 202952. Try increasing `gpu_memory_utilization` or decreasing `max_model_len` when initializing the engine. See https://docs.vllm.ai/en/latest/configuration/conserving_memory/ for more details.

Things I tried but didn't help:

• Setting --max-num-seqs 1
• Disabling prefix caching (this is required regardless of OOM because of block size assertion in HybridKVCacheCoordinator)
• Eager mode

[NickLucche]

Changes here went a bit too deep with the MultiConnector bit @mgoin :(

[NickLucche]

we were tracking this change here #41847 and #42024 !

[DarkLight1337]

I worked around the OOM by setting a lower value of --max-model-len. But now I'm getting some hidden states mismatch issues when running offline data generation in vllm-project/speculators#495 (I set --validate-outputs so the input tokens should be correct), e.g.:

WARNING  Skipping sample 17 due to error: Sequence length of hidden states 390 in output/speculators/hidden_states/hs_17.safetensors doesn't match num tokens 496
WARNING  Skipping sample 458 due to error: Sequence length of hidden states 625 in output/speculators/hidden_states/hs_458.safetensors doesn't match num tokens 3888
WARNING  Skipping sample 3460 due to error: Sequence length of hidden states 133 in output/speculators/hidden_states/hs_3460.safetensors doesn't match num tokens 449
WARNING  Skipping sample 3555 due to error: Sequence length of hidden states 401 in output/speculators/hidden_states/hs_3555.safetensors doesn't match num tokens 433
WARNING  Skipping sample 3968 due to error: Sequence length of hidden states 35 in output/speculators/hidden_states/hs_3968.safetensors doesn't match num tokens 397
WARNING  Skipping sample 4252 due to error: Sequence length of hidden states 48 in output/speculators/hidden_states/hs_4252.safetensors doesn't match num tokens 271
WARNING  Skipping sample 4333 due to error: Sequence length of hidden states 184 in output/speculators/hidden_states/hs_4333.safetensors doesn't match num tokens 488
WARNING  Skipping sample 4399 due to error: Sequence length of hidden states 110 in output/speculators/hidden_states/hs_4399.safetensors doesn't match num tokens 798
WARNING  Skipping sample 5365 due to error: Sequence length of hidden states 127 in output/speculators/hidden_states/hs_5365.safetensors doesn't match num tokens 343

This only happens roughly once every 500 samples, and the sample numbers that are failing are not deterministic. This leads me to think that it's a race condition of some sort.

[DarkLight1337]

Ok this happens with Qwen/Qwen3-VL-30B-A3B-Instruct as well, so maybe it is a problem for multimodal models in general rather than hybrid models. Turning off CUDA graph and async scheduling didn't help either.

[DarkLight1337]

cc @fynnsu

[fynnsu]

Hmm strange that the hidden state shape is less than the num tokens. I will try to repro this.

[DarkLight1337]

If I set --concurrency 1 in the offline generation script then the problem goes away.

[fynnsu]

@DarkLight1337 Okay, I was able to reproduce the issue with Qwen/Qwen3-VL-30B-A3B-Instruct. I actually still see the problem when using --concurrency 1.

Claude seems to think its an issue with chunked prefill, and I'm trying to develop a fix. Unfortunately this model doesn't seem to work if you just --no-enable-chunked-prefill. I'll keep working on this, and let you know when i have an update.

[fynnsu]

This branch has a WIP fix: main...fynnsu:vllm:fix_chunked_prefill_hs_connector

That seems to solve the problem but it also makes the save happen after the request finishes so we need to add handling in speculators for this. (To test this I just temporarily hardcoded in a 2-second wait before reading the file in data_generation_offline.py).

#37374 also enforces a single save (by requiring no-chunked-prefill) and adds async lock file logic so that the downstream client know when it's safe to read the file. I think we can build a proper fix off of that pr. This should also get around the no-chunked-prefill requirement because instead we just save all the blocks when the request finishes which allows chunked prefill to run as normal.

[DarkLight1337]

This branch has a WIP fix: main...fynnsu:vllm:fix_chunked_prefill_hs_connector

This also works on my end, thanks for the quick fix! Let's also put this in #37374.

To test this I just temporarily hardcoded in a 2-second wait before reading the file in data_generation_offline.py

We could add a check in generate_hidden_states_* on speculators side to validate that vLLM has actually finished saving the file before reading it. To avoid TOCTOU issues we could use a sentinel file approach:

# vllm
safetensors.torch.save_file(tensors, pending.filename)
finished_sending.add(pending.req_id)
Path(f"{pending.filename}.done").touch()

# speculators
hs_filepath = extract_output(res, token_ids)
done_path = Path(f"{hs_filepath}.done")

while not done_path.exists():
    await asyncio.sleep(1)

done_path.unlink()
return hs_filepath