[Bugfix][KV Offload] Scale UniformTypeKVCacheSpecs groups by DCP by drakosha · Pull Request #50883 · vllm-project/vllm

drakosha · 2026-08-03T15:41:58Z

Purpose

#49964 gated the DCP scaling of a group's token span on isinstance(group.kv_cache_spec, AttentionSpec) so that Mamba groups stay unscaled. UniformTypeKVCacheSpecs is a container rather than an AttentionSpec subclass, so DSA models — whose indexer and MLA layers merge into a single uniform group — lost the DCP factor too:

build_offloading_config reports tokens_per_block = block_size (64),
resolve_kv_cache_block_sizes still returns block_size * dcp (128) for a single group,

so enabling CPU offload under DCP dies at startup:

AssertionError: tokens_per_block=64 not divisible by tokens_per_hash=128.
Hybrid models (e.g. Mamba+Attention) need --enable-prefix-caching to align block sizes.

Reported by @Leoyzen on GLM-5.2, TP8/DCP2, fp8_ds_mla, MTP, OffloadingConnector: #46514 (comment)

The scaling is now keyed off a helper that unwraps the container, so Mamba groups stay unscaled and every attention group keeps its DCP span. Same call site is fixed in resolve_kv_cache_block_sizes for multi-group configs.

Related: #50823 fixes the other half of the same container-not-unwrapped family (block table width).

Test

Three tests, one per call site plus the hybrid case, next to the Mamba tests from #49964:

offloading_connector/test_config.py::test_dcp_scales_uniform_type_attention_group_blocks
offloading_connector/test_config.py::test_dcp_scales_uniform_type_group_alongside_mamba
core/test_kv_cache_utils.py::test_resolve_block_sizes_scales_uniform_type_group_by_dcp

Verified fail→pass in a container built from the current nightly (0.26.1rc1.dev255+g5e35a6f4f); each test fails on unpatched main with the reported assert (tokens_per_block=16 not divisible by tokens_per_hash=32) or a wrong span, and the two suites pass with the fix:

python3 -m pytest tests/v1/core/test_kv_cache_utils.py \
  tests/v1/kv_connector/unit/offloading_connector/test_config.py -q
# 115 passed

Config-level change, no effect on model output; no eval run.

AI assistance (Claude) was used for this work; all changes reviewed.

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.

mergify · 2026-08-11T21:46:29Z

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

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

mergify · 2026-08-13T17:13:22Z

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

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

mergify · 2026-08-22T01:58:57Z

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

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

mergify · 2026-08-29T03:13:24Z

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

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

cjackal · 2026-08-29T04:55:15Z

Seems ~~superceded by~~we can reuse the resolve_dcp_kv_block_size from the commit 43bb773 in #53324 (which has already been merged)

cjackal · 2026-08-29T05:12:31Z

                * (
                    parallel_config.decode_context_parallel_size
-                    if isinstance(group.kv_cache_spec, AttentionSpec)
+                    if is_kv_cache_spec_dcp_sharded(group.kv_cache_spec)


the whole tokens_per_block calculation is now doable with tokens_per_block=resolve_dcp_kv_block_size(group.kv_cache_spec, parallel_config.decode_context_parallel_size).

isinstance(spec, AttentionSpec) so that Mamba groups stay unscaled. UniformTypeKVCacheSpecs is a container rather than an AttentionSpec subclass, so DSA models, whose indexer and MLA layers merge into one uniform group, stopped being scaled as well: build_offloading_config reports tokens_per_block=block_size while resolve_kv_cache_block_sizes still returns block_size * dcp for a single group, and enabling CPU offload under DCP dies on the divisibility assert. Key the scaling off a helper that unwraps the container, so Mamba stays unscaled and every attention group keeps its DCP span. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Mikhail Kostryukov <mike@triptrack.net>

…cache spec Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

isinstance(spec, AttentionSpec) so that Mamba groups stay unscaled. UniformTypeKVCacheSpecs is a container rather than an AttentionSpec subclass, so DSA models, whose indexer and MLA layers merge into one uniform group, stopped being scaled as well: build_offloading_config reports tokens_per_block=block_size while resolve_kv_cache_block_sizes still returns block_size * dcp for a single group, and enabling CPU offload under DCP dies on the divisibility assert. Key the scaling off a helper that unwraps the container, so Mamba stays unscaled and every attention group keeps its DCP span. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Mikhail Kostryukov <mike@triptrack.net>

…cache spec Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

build_offloading_config() decides DCP scaling with isinstance(spec, AttentionSpec) so that Mamba groups stay unscaled. UniformTypeKVCacheSpecs is a container rather than an AttentionSpec subclass, so DSA models, whose indexer and MLA layers merge into one uniform group, are not scaled either: the connector reports tokens_per_block=block_size while resolve_kv_cache_block_sizes() returns block_size * dcp for the same group, and enabling CPU offload under DCP dies on the divisibility assert. resolve_kv_cache_block_sizes() already answers this question with resolve_dcp_kv_block_size(), which unwraps the container. Use that helper on the offloading side too, so both sides agree by construction and Mamba groups keep their unscaled per-rank span. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Mikhail Kostryukov <mike@triptrack.net>

drakosha · 2026-09-02T07:15:45Z

@cjackal You are right, and half of this PR is now redundant: resolve_kv_cache_block_sizes
goes through resolve_dcp_kv_block_size since #53324, so my predicate duplicated yours. Dropped
it, and my test_kv_cache_utils case with it.

What is left is the connector side: build_offloading_config still scales by
isinstance(spec, AttentionSpec), and UniformTypeKVCacheSpecs is a container, not a subclass.
On DSA models, where indexer and MLA merge into one uniform group, the connector reports
tokens_per_block=block_size while the scheduler reports block_size * dcp, and CPU offload
under DCP dies on the divisibility assert. Rebased on c00091e026: nine lines in that file
calling your helper, plus the connector test. On nightly a9a17e709 with current-main files
overlaid, test_config.py is 47 passed with the fix, 2 failed (both dcp=2) without it.

Aside: resolve_dcp_kv_block_size treats a group as sharded when every layer spec is an
AttentionSpec, so sliding window scales, while dcp_world_size_for_kv_cache_spec answers the
same question with FullAttentionSpec only and unwraps containers by next(iter(...)). Not
this PR's business, but they disagree.

AI assistance was used for this work.

drakosha · 2026-09-02T07:28:52Z

/ci run

github-actions · 2026-09-02T07:29:09Z

❌ @drakosha, A reviewer with write access must run /ci run, approve the PR, or add the ready label first.

orozery · 2026-09-02T11:36:45Z

/ci run

github-actions · 2026-09-02T11:36:51Z

✅ @drakosha, CI is now available for this PR.

/ci run starts upstream CI; /amd-ci run starts AMD CI only.
/ci retry retries failed jobs in the CI build for the current PR head. If the current head has no CI build, it starts a new CI build for the current head containing only jobs that failed in the latest earlier CI build for this PR.
/amd-ci retry retries failed jobs in AMD CI for the current PR head. Use /amd-ci run when the current head has no AMD CI build.
/ci cancel cancels scheduled or running CI builds for this PR branch; /amd-ci cancel does the same for AMD CI only.

github-actions · 2026-09-02T11:37:53Z

✅ Triggered Buildkite CI #86859 for commit 6567ad1d8ecd.

…m-project#50883) Signed-off-by: Mikhail Kostryukov <mike@triptrack.net> Co-authored-by: Claude <noreply@anthropic.com>

drakosha requested review from ApostaC, NickLucche, WoosukKwon, alexm-redhat, heheda12345, ivanium, njhill, orozery, robertgshaw2-redhat, xuechendi and ywang96 as code owners August 3, 2026 15:41

claude Bot reviewed Aug 3, 2026

View reviewed changes

mergify Bot added bug Something isn't working kv-connector labels Aug 3, 2026

drakosha mentioned this pull request Aug 3, 2026

[Bugfix] Shard UniformTypeKVCacheSpecs block table width under DCP #50823

Merged

drakosha force-pushed the fix-offloading-dcp-uniform-spec branch from 5353ac9 to 08f985f Compare August 3, 2026 16:50

drakosha mentioned this pull request Aug 10, 2026

[Attention][MLA] FlashMLA sparse: DCP on the fp8_ds_mla mixed-batch path + MTP #46514

Merged

mergify Bot added needs-rebase and removed needs-rebase labels Aug 11, 2026

mergify Bot added needs-rebase kv-cache-manager and removed needs-rebase labels Aug 13, 2026

mergify Bot added the needs-rebase label Aug 22, 2026

p5hgryktp4 mentioned this pull request Aug 24, 2026

[Bug][KV Offload] OffloadingConnector fs tier: multi-group MLA+DSA (DeepSeek-V4) TP=2 lookup fully misses across restart; single-group MHA hits 99.6% #53569

Open

drakosha force-pushed the fix-offloading-dcp-uniform-spec branch from d572f3e to 02b320d Compare August 25, 2026 09:24

drakosha mentioned this pull request Aug 25, 2026

[Bug]: DC unavailable for GlmMoeDsa (GLM-5.2 FP8) on SM90 — sparse MLA backend lacks decode-LSE support #53134

Open

mergify Bot removed the needs-rebase label Aug 25, 2026

mergify Bot added the needs-rebase label Aug 29, 2026

cjackal reviewed Aug 29, 2026

View reviewed changes

LucasWilkinson added a commit to neuralmagic/vllm that referenced this pull request Aug 31, 2026

[PR vllm-project#50883] Resolve single-group DCP block size from its …

c6f9b29

…cache spec Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

LucasWilkinson added a commit to neuralmagic/vllm that referenced this pull request Aug 31, 2026

[PR vllm-project#50883] Resolve single-group DCP block size from its …

e8b3687

…cache spec Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>

nood-co1 mentioned this pull request Sep 1, 2026

[KV Offload] Scope offload group configs to prefix-cacheable KV cache groups #54743

Open

drakosha force-pushed the fix-offloading-dcp-uniform-spec branch from 02b320d to 6567ad1 Compare September 2, 2026 07:10

mergify Bot removed the needs-rebase label Sep 2, 2026

orozery added the ready ONLY add when PR is ready to merge/full CI is needed label Sep 2, 2026

orozery approved these changes Sep 2, 2026

View reviewed changes

orozery merged commit ba6c60e into vllm-project:main Sep 2, 2026
31 checks passed

mylibrar pushed a commit to tanyuqian/vllm that referenced this pull request Sep 3, 2026

[Bugfix][KV Offload] Scale UniformTypeKVCacheSpecs groups by DCP (vll…

08a28ba

…m-project#50883) Signed-off-by: Mikhail Kostryukov <mike@triptrack.net> Co-authored-by: Claude <noreply@anthropic.com>

Uh oh!

Uh oh!

Conversation

drakosha commented Aug 3, 2026

Purpose

Test

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

mergify Bot commented Aug 11, 2026

Uh oh!

mergify Bot commented Aug 13, 2026

Uh oh!

mergify Bot commented Aug 22, 2026

Uh oh!

mergify Bot commented Aug 29, 2026

Uh oh!

cjackal commented Aug 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cjackal Aug 29, 2026

Choose a reason for hiding this comment

Uh oh!

drakosha commented Sep 2, 2026

Uh oh!

drakosha commented Sep 2, 2026

Uh oh!

github-actions Bot commented Sep 2, 2026

Uh oh!

orozery commented Sep 2, 2026

Uh oh!

github-actions Bot commented Sep 2, 2026

Uh oh!

github-actions Bot commented Sep 2, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

cjackal commented Aug 29, 2026 •

edited

Loading