[Core][DSV4] Compact MXFP4 indexer KV cache and packed group overlays by GirasoleY · Pull Request #48993 · vllm-project/vllm

GirasoleY · 2026-07-17T22:33:37Z

Context

When DeepSeek V4 uses MXFP4 indexer K values, engine reserves the larger FP8 row for them. Its packed KV planner also buckets layers by page size, leaving avoidable holes when different cache groups have different mixtures of page sizes.

This PR:

sizes MXFP4 indexer K rows from their packed values and UE8M0 scales (head_dim / 2 + head_dim / 32 bytes), reducing the 128-dimensional row from 132 to 68 bytes;
lays out every packed cache group densely in one per-block slab and overlays the group layouts, while keeping layers within each group disjoint; and
removes cross-group SWA page padding that the dense planner no longer needs.

The packing change is generic to the existing packed-cache path. DeepSeek V4 uses that path by default; other hybrid models retain the existing opt-in via enable_cross_layers_blocks.

Memory results

Matched DeepSeek-V4-Pro DEP8 startup measurements on GB200:

Backend / KV budget	Baseline blocks	This change	Increase
FlashInfer sparse / 26.46 GiB	21,709	23,120	+6.50%
FlashMLA sparse / 26.09 GiB	19,015	21,306	+12.05%

For the FlashInfer layout, physical block stride decreases from 1,308,672 to 1,228,800 bytes (-6.10%). For the FlashMLA fp8_ds_mla layout, it decreases from 1,473,408 to 1,315,008 bytes (-10.75%).

Allocate each indexer K row from the packed MXFP4 value and UE8M0 scale layout instead of reusing the larger FP8 row. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: girasoley <girasoley@inferact.ai>

Lay out each cache group densely within a shared per-block slab. Since cache groups draw block IDs from one pool, group layouts may overlap while layers within each group remain disjoint. Remove cross-group SWA page padding and cover the DSV4 FP32 state overlay, production stride, offsets, and group-owned block reuse. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: girasoley <girasoley@inferact.ai>

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.

ivanium

LGTM! cc @tlrmchlsmth @LucasWilkinson for a second look too as this is related (actually enabled) by the DSV4 packed kv cache layout, and I feel the idea is smart but not very straightforward.

ivanium · 2026-07-18T00:04:10Z

@claude review

claude · 2026-07-18T00:15:46Z

        for spec in group.kv_cache_specs.values()
    )

-    # Split each SWA UniformKV group into smaller groups to align their #(layer tuples)
-    # Possibly padding layer tuples for this.
-    # Additionally, we also pad KV blocks in each SWA layer, to align the page size
-    # with the corresponding layer in the full-MLA group.
-    all_page_sizes = full_mla_spec.get_page_sizes()
+    # Split each SWA UniformKV group into smaller groups to align their
+    # numbers of layer tuples. The packed block planner overlays groups, so
+    # their page sizes do not need to match.
    swa_mla_groups = []
    for sm_spec in swa_mla_specs:
-        sm_page_sizes = sm_spec.get_page_sizes()
        layers_per_size: dict[int, list[str]] = defaultdict(list)
-        assert max(sm_page_sizes) <= max(all_page_sizes)
-
-        # Unify page size by padding layers' page_size to the nearest larger page_size.
-        # Compute candidate (nearest larger page_size) for each unique page size.
-        size_to_candidate: dict[int, int] = {}
-        for ps in sm_page_sizes:
-            size_to_candidate[ps] = min(x for x in all_page_sizes if x >= ps)
-        # Pad and collect layer names per page size.
+
        for layer_name, layer_spec in sm_spec.kv_cache_specs.items():
-            current_size = layer_spec.page_size_bytes
-            candidate = size_to_candidate[current_size]
-            if current_size < candidate:
-                object.__setattr__(layer_spec, "page_size_padded", candidate)
-            layers_per_size[candidate].append(layer_name)
+            layers_per_size[layer_spec.page_size_bytes].append(layer_name)
        # NOTE(yifan): for now, inside a UniformKV group, each page_size should
        # have the same number of layers. This also means we don't need to pad layers
        # inside a partial-full layer tuple.


🔴 This PR removes the padding in _get_kv_cache_groups_uniform_groups that used to force SWA/state group page sizes to align with the full-MLA group, but the unmodified _max_memory_usage_bytes_from_groups (~line 1834) still derives layer_tuple_bytes from only the full-MLA group's page sizes and reuses it for every group — exactly the invariant its own comment says is required and this PR removes.

Extended reasoning...

The bug: _get_kv_cache_groups_uniform_groups (vllm/v1/core/kv_cache_utils.py:1640-1654) used to pad each SWA/state UniformTypeKVCacheSpecs group's layer page sizes up to the nearest page size in the full-MLA group, guaranteeing that all groups shared the same "layer-tuple" byte width. This PR removes that padding — the new comment explicitly says "The packed block planner overlays groups, so their page sizes do not need to match." That statement is true for the new _get_packed_kv_cache_layout/_get_kv_cache_config_packed overlay path (which now computes block_stride = max over groups of each group's own summed page sizes), but it is not true for _max_memory_usage_bytes_from_groups, an unmodified function a few hundred lines below that still assumes the old invariant.

Where it breaks: In _max_memory_usage_bytes_from_groups's DeepseekV4 branch (the "all groups are UniformTypeKVCacheSpecs" case), it computes:

full_mla_spec = kv_cache_groups[0].kv_cache_spec layer_tuple_bytes = sum(full_mla_spec.get_page_sizes()) # ONLY group[0], the full-MLA group ... g_max_mem_usage_page_bytes = num_layer_tuples * g_max_mem_usage_pages * layer_tuple_bytes

This layer_tuple_bytes value — derived solely from the full-MLA group — is then applied uniformly to every group's page count, including SWA/state groups whose own per-tuple byte width may now be larger (since it's no longer padded to match the full group). The function's own comment ("They must already be page_size aligned and share a common padded layer-tuple layout") documents the precondition that this PR's diff removes without updating this function.

Concrete proof, using the PR's own new test data (test_deepseek_v4_pro_stride):

full-MLA group: [32768, 4608]*30 + [1024]*31 → per-tuple bytes = 32768+4608 = 37376, group total = 3037376 + 311024 = 1,153,024

c4_state group: [32768, 8192]*30 → per-tuple bytes = 32768+8192 = 40960, group total = 30*40960 = 1,228,800

After this PR, _get_packed_kv_cache_layout correctly sets block_stride = 1,228,800 (driven by the larger c4_state group), matching the PR's own documented FlashInfer stride reduction. But _max_memory_usage_bytes_from_groups computes layer_tuple_bytes from only the full-MLA group: num_layer_tuples * layer_tuple_bytes = 31 * 37376 ≈ 1,158,656-ish (using the group[0] per-tuple width), not the 1,228,800 that the c4_state group actually consumes per physical block in the overlay layout. This is a systematic underestimate of required KV cache memory (~3% in this realistic DSv4-Pro configuration, ~9% in the PR's test_compact_cache_overlays_fp32_state_group fp32-state scenario: 74752 vs actual 81920 bytes/tuple).

Why nothing else catches this: _max_memory_usage_bytes_from_groups is not touched by this PR's diff and has no test coverage tying it to the new overlay layout — the new tests (test_compact_cache_overlays_fp32_state_group, test_deepseek_v4_pro_stride) only exercise _get_kv_cache_config_packed/_get_packed_kv_cache_layout, not the memory-estimate path.

Impact: _max_memory_usage_bytes_from_groups feeds check_enough_kv_cache_memory (the "at least one max-length request fits" admission check, via get_kv_cache_configs) and _estimate_max_model_len_from_groups (the max_model_len=-1 auto-fit binary search). Underestimating required memory in the unsafe direction means: (1) the admission check can pass even when the real packed layout (whose actual per-block cost is block_stride, driven by the largest group) cannot serve one max-length request, and (2) auto-fit can select a max_model_len a few percent larger than what physically fits, since DSv4 uses this packed path by default.

Fix direction: Compute layer_tuple_bytes per-group — sum(cast(UniformTypeKVCacheSpecs, group.kv_cache_spec).get_page_sizes()) inside the loop instead of hoisting it from kv_cache_groups[0] — or, more precisely, key the per-block cost off the actual block_stride returned by _get_packed_kv_cache_layout so the estimate matches the real overlay allocation exactly.

LucasWilkinson

Makes sense to me; thanks for the contribution!

Replace the packed-vs-unpacked allocator split with a single rule: every cache group packs its layers densely (running sum of page sizes) into a per-block window W = max group packing, and groups overlay each other (a block ID is owned by one group at a time) — #48993's dense placement, generalized to every model. Bucketing layers by page size is gone. KVCacheTensor now describes where a set of same-shaped layers lives in the one allocation: layer l of block b starts at `offset + l * layer_stride + b * block_stride`. The layout picks which stride is large — layer-outermost layouts give each layer a contiguous region (main's per-layer arrangement), block-outermost layouts make each block a window of all layers' pages — so the allocation is identical regardless of layout. Models whose groups overlay with different page sizes need the layer dim inside the block dim, so they publish a block-outermost layout requirement (DSV4 packs by default, as on main); plane layouts are rejected for mixed page sizes. Binding collapses to one strided view per layer in both runners. `shared_by` is gone: cross-group aliasing is expressed by overlapping address ranges, which also describes DSV4's partial overlaps. Also: - Fix a duplicated `state_content_bytes` kwarg in sparse_swa (syntax error). - Fix `compress_ratio` -> `tokens_per_state` in the offloading config. - Fix KV byte accounting in the offloading config and CPU-offload manager: every tensor's size is the whole backing, so summing double-counted. - Re-plan (instead of rescaling) when unifying num_blocks across workers, so strides and offsets stay consistent on asymmetric TP/PP ranks. - Assert kernel block splitting only where blocks are dense and unpadded. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Standardize every KV cache allocation on the logical [L, B, H, N, C] vocabulary (RFC #42082). - `KVCacheLayout` enumerates the physical stride permutations (LBHNC, LBNHC, LHBNC, BLHNC, BLNHC, BHLNC); per-layer caches are always logical [B, H, N, C] views with the physical layout expressed via strides (MLA binds the squeezed [B, N, C]). - Layout resolution has a single writer: attention-backend selection publishes the layout into CacheConfig (test override > backend-required > VLLM_KV_CACHE_LAYOUT > connector preference > LBNHC); every consumer reads it from there, and conflicting backend requirements fail loudly. - One allocator, no packed-vs-unpacked split: every cache group packs its layers densely (running sum of page sizes) into a per-block window W = max group packing, and groups overlay each other (a block ID is owned by one group at a time). This is #48993's dense placement, generalized to every model; bucketing layers by page size is gone. - `KVCacheTensor` describes where a set of same-shaped layers lives in the one allocation: layer l of block b starts at `offset + l * layer_stride + b * block_stride`. The layout decides which stride is large — layer-outermost layouts give each layer a contiguous region, block-outermost layouts make each block a window of all layers' pages — so the allocation is identical regardless of layout. Models whose groups overlay with different page sizes need the layer dim inside the block dim and publish a block-outermost requirement (DeepSeek V4 packs by default); plane layouts are rejected for mixed page sizes. `shared_by` is gone: cross-group aliasing is expressed by overlapping address ranges. - Backends publish their KV packing through `customize_spec`: AITER opts into separate K/V head groups ([B, 2, N, H*hs], LBHNC) so the fused QK-norm+RoPE+cache kernel keeps the block-interior contract aiter v0.1.19 asserts, and ROCM_ATTN keeps main's K/V-plane geometry (LHBNC) with the native HIP write/decode kernels and x-packed block interior. - `indexes_kv_by_block_stride` and the per-backend cache-shape/stride-order hooks are removed; the standardized views make them constant. The KV block zeroer, block copies, and the connectors (NIXL, mooncake, moriio, hf3fs, offloading, example) consume the standardized views. Replaces #44458. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Standardize every KV cache allocation on the logical [L, B, H, N, C] vocabulary (RFC #42082). - KVCacheLayout enumerates the physical stride permutations; per-layer caches are always logical [B, H, N, C] views with the layout expressed via strides (MLA binds the squeezed [B, N, C]). - Layout resolution has a single writer: backend selection publishes it into CacheConfig and every consumer reads it from there; conflicting backend requirements fail loudly. - One allocator: each cache group packs its layers densely into a block and groups overlay, so bytes per block is the largest group's packing (#48993's placement, generalized). KVCacheTensor gives the placement of a set of same-shaped layers: layer l of block b starts at offset + l * layer_stride + b * block_stride, so the allocation is the same under every layout. Models whose groups overlay with different page sizes publish a block-outermost layout requirement. - Backends publish their packing through customize_spec: AITER opts into separate K/V head groups for its fused QK-norm+RoPE+cache kernel, and ROCM_ATTN keeps its K/V-outermost geometry and native HIP kernels. - The old per-backend get_kv_cache_shape / get_kv_cache_stride_order / get_kv_cache_block_dim / indexes_kv_by_block_stride hooks and the cross-layer KV machinery are removed; compress_ratio becomes tokens_per_state. Replaces #44458. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Port upstream vLLM vllm-project#48993, retaining distinct tensors for lockstep DCP/PCP groups and updating the fork's capacity accounting for dense group overlays. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Yoshikazu Aoyama <175681+blauerberg@users.noreply.github.com>

Standardize every KV cache allocation on the logical [L, B, H, N, C] vocabulary (RFC #42082). - KVCacheLayout enumerates the physical stride permutations; per-layer caches are always logical [B, H, N, C] views with the layout expressed via strides (MLA binds the squeezed [B, N, C]). - Layout resolution has a single writer: backend selection publishes it into CacheConfig and every consumer reads it from there; conflicting backend requirements fail loudly. - One allocator: each cache group packs its layers densely into a block and groups overlay, so bytes per block is the largest group's packing (#48993's placement, generalized). KVCacheTensor gives the placement of a set of same-shaped layers: layer l of block b starts at offset + l * layer_stride + b * block_stride, so the allocation is the same under every layout. Models whose groups overlay with different page sizes publish a block-outermost layout requirement. - Backends publish their packing through customize_spec: AITER opts into separate K/V head groups for its fused QK-norm+RoPE+cache kernel, and ROCM_ATTN keeps its K/V-outermost geometry and native HIP kernels. - The old per-backend get_kv_cache_shape / get_kv_cache_stride_order / get_kv_cache_block_dim / indexes_kv_by_block_stride hooks and the cross-layer KV machinery are removed; compress_ratio becomes tokens_per_state. Replaces #44458. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

GirasoleY and others added 2 commits July 17, 2026 22:25

[DSV4] Right-size MXFP4 indexer K cache

6a4aad8

Allocate each indexer K row from the packed MXFP4 value and UE8M0 scale layout instead of reusing the larger FP8 row. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: girasoley <girasoley@inferact.ai>

mergify Bot added the v1 label Jul 17, 2026

GirasoleY force-pushed the codex/dsv4-compact-packed-kv-public branch from c9540c3 to 8ad3d3f Compare July 17, 2026 22:34

GirasoleY marked this pull request as ready for review July 17, 2026 23:04

GirasoleY requested review from ApostaC, WoosukKwon, alexm-redhat, heheda12345, ivanium, njhill, orozery, robertgshaw2-redhat, ywang96 and zyongye as code owners July 17, 2026 23:04

claude Bot reviewed Jul 17, 2026

View reviewed changes

ivanium approved these changes Jul 18, 2026

View reviewed changes

claude Bot reviewed Jul 18, 2026

View reviewed changes

LucasWilkinson approved these changes Jul 20, 2026

View reviewed changes

ivanium added the ready ONLY add when PR is ready to merge/full CI is needed label Jul 20, 2026

ivanium enabled auto-merge (squash) July 20, 2026 23:48

mergify Bot and others added 2 commits July 20, 2026 23:48

Merge branch 'main' into codex/dsv4-compact-packed-kv-public

e483bf2

Merge branch 'main' into codex/dsv4-compact-packed-kv-public

a451d79

mgoin mentioned this pull request Jul 22, 2026

[Bugfix][KV cache] Support sparse-MLA targets with SWA drafts #48776

Merged

Merge branch 'main' into codex/dsv4-compact-packed-kv-public

7f5cdd5

WoosukKwon merged commit f3a920a into vllm-project:main Jul 22, 2026
59 of 62 checks passed

QiliangCui2023 mentioned this pull request Jul 26, 2026

[DSv4] Adapt KV cache creation to vLLM packed byte-offset overlay groups vllm-project/tpu-inference#3267

Merged

fank mentioned this pull request Jul 31, 2026

Bump base to nightly-0f173945 for the DSV4 KV-capacity fixes (vllm#48993, vllm#48317); drop vendored vllm#48748 enthus-appdev/vllm-pcai#36

Merged

LucasWilkinson mentioned this pull request Aug 10, 2026

[6/N][KV-Cache Layout Refactor] Standardize KV cache layout #51718

Merged

yewentao256 mentioned this pull request Aug 13, 2026

[Feature]: Performance Optimization for Deepseek V4 #45861

Open

21 tasks

Navigation Menu

Uh oh!

Uh oh!

[Core][DSV4] Compact MXFP4 indexer KV cache and packed group overlays - #48993

[Core][DSV4] Compact MXFP4 indexer KV cache and packed group overlays#48993
WoosukKwon merged 5 commits into
vllm-project:mainfrom
GirasoleY:codex/dsv4-compact-packed-kv-public

GirasoleY commented Jul 17, 2026 •

edited

Loading

Uh oh!

claude Bot left a comment

Uh oh!

ivanium left a comment •

edited

Loading

Uh oh!

ivanium commented Jul 18, 2026

Uh oh!

claude Bot Jul 18, 2026

Uh oh!

LucasWilkinson left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Uh oh!

Conversation

GirasoleY commented Jul 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Context

Memory results

Uh oh!

claude Bot left a comment

Choose a reason for hiding this comment

Claude Code Review

Uh oh!

ivanium left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ivanium commented Jul 18, 2026

Uh oh!

claude Bot Jul 18, 2026

Choose a reason for hiding this comment

Uh oh!

LucasWilkinson left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

GirasoleY commented Jul 17, 2026 •

edited

Loading

ivanium left a comment •

edited

Loading