fix(attention): preserve per-group KV cache formats

[voipmonitor]

Resulting behavior

Status: implemented.

Each attention group resolves its KV-cache dtype from the layers assigned to that group. Model-owned attention implementations receive the engine-selected physical cache layout without replacing their group-specific logical dtype.

A speculative draft can therefore use a cache format supported by its attention backend while the target uses a different format. Target-only serving and attention groups whose layers share one global cache format retain their established behavior.

Source contract

• Base: local-inference-lab/vllm:dev/jovian-judgement at c79f35ca00e8e93e0943a0d79b85b22b18aac939.
• Head: eca8a6bca19f210a30fbb0bb6942556f8243cb14.
• All attention layers within one cache group must resolve to one cache dtype. A mixed dtype inside one group fails validation.
• The physical layout remains selected by the engine cache configuration and is propagated to every model-owned implementation.

Validation

• tests/v1/attention/test_group_head_counts.py and tests/v1/worker/test_attn_utils.py: 23 passed.
• Repository pre-commit hooks, including Ruff, mypy, SPDX, forbidden-import, and configuration checks: passed.
• git diff --check: passed.
• The composed TP4 GLM-5.3 target-plus-DFlash2 runtime initialized separate target and draft attention groups, captured all configured graph shapes, and served an OpenAI-compatible completion.

Duplicate-work check

Related open pull requests implement different contracts:

• vllm-project/vllm#53979 enables non-causal FlashAttention 2 reads from NVFP4 pages; it does not resolve cache dtype per attention group.
• vllm-project/vllm#48392 replicates dense draft KV heads for decode-context parallelism; it does not preserve independent target and draft cache formats under decode-context parallel size 1.
• local-inference-lab/vllm#488 implements GLM CKV gathering and retention under decode-context parallelism; it does not provide this per-group dtype and physical-layout propagation.

Review disclosure

OpenAI Codex assisted with implementation, tests, runtime qualification, and pull-request preparation. Human review of every changed line and the attention-group cache contract is required before merge.

[coderabbitai]

Warning

Review limit reached

Next included review available in 59 minutes.

View limit details

Limit details: You’ve used the included review currently available.

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

Learn how review limits work.

Review configuration:

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: b5550539-d63e-4f7e-85a7-f61198716956

📥 Commits

Reviewing files that changed from the base of the PR and between c79f35c and eca8a6b.

📒 Files selected for processing (5)

• tests/v1/attention/test_group_head_counts.py
• tests/v1/worker/test_attn_utils.py
• vllm/v1/attention/backends/flashinfer.py
• vllm/v1/attention/backends/utils.py
• vllm/v1/worker/gpu/attn_utils.py

---

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