[6/N][KV-Cache Layout Refactor] Standardize KV cache layout by LucasWilkinson · Pull Request #51718 · vllm-project/vllm

LucasWilkinson · 2026-08-10T18:30:58Z

Purpose

The core of the KV-cache layout standardization series (RFC #42082). Stacked on #51704 → #51612 — the diff shown includes them until they land and this retargets.

Standardizes every KV cache allocation on the logical [L, B, H, N, C] vocabulary:

KVCacheLayout enumerates the physical stride permutations (LBHNC, LBNHC, LHBNC, BLHNC, BLNHC, BHLNC); reshape_kv_cache views each flat allocation as dense [num_blocks, block_bytes] tiles; 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.
The AITER backends opt into separate K/V head groups ([B, 2, N, H*hs], LBHNC required) through customize_spec, keeping the fused QK-norm+RoPE+cache kernel dispatched with the exact block-interior contract aiter v0.1.19 asserts — fusion never regresses at any commit in the series.
The KV block zeroer matches main's per-segment form (block dim is always 0 in standardized views); block copies are a plain unflatten + index copy; connectors (NIXL, mooncake, moriio, hf3fs, offloading, example) and the attention benchmark runner consume the standardized views and published layout.
indexes_kv_by_block_stride and per-backend cache-shape/stride-order hooks are removed — the standardized views make them constant.

Replaces #44458 (same content, restacked).

Test Plan / Result

Validated on B200 (full history in #44458): backend-correctness suites (FlashAttention/FlashInfer/Triton/Flex 127P), full MLA suite and test_prefix_prefill at exact main failure-set parity, connector unit battery at parity, zeroer/attn-utils/model-runner suites green, e2e smokes across FLASH_ATTN / TRITON(+int8 per-token-head) / FLASHINFER / FLEX / fp8 / CUTLASS_MLA (DeepSeek-V2-Lite). ROCm/AITER validation in progress on AMD.

AI-assisted (Claude); under review by submitter

Validation

CPU suites (kv-cache-utils, packing, zeroer, attn-utils, gpu-model-runner,
nixl / mooncake / moriio / offloading / simple-offload): green. The only
failures on this host are environmental and reproduce identically on the base
branch (same 5 test ids): GPU 0 is held by another process, and NIXL is not
installed.

tests/v1/attention/test_mla_backends.py: 1729 passed. 20 more OOM'd against
a GPU another process had filled; rerun on a reserved GPU they pass (290
passed, 0 failed).

E2E on this branch, all deterministic on re-run:

model	exercises	result
`Qwen/Qwen3-0.6B`	dense attention	pass
`deepseek-ai/DeepSeek-V2-Lite-Chat`	MLA specs, compressed KV	pass under LBNHC, LBHNC and BLHNC
`ibm-granite/granite-4.0-h-tiny`	hybrid attention + Mamba	pass

Not run here (needs a bigger box): DeepSeek-V4 and MiniMax-M3 themselves, so
the DSV4 multi-group overlay, the fp8_ds_mla 584/656-byte paths, the compact
MXFP4 indexer and M3's sparse/indexer backends have unit coverage but no e2e.
Those paths are covered by tests/v1/core/test_contiguous_kv_packing.py and
tests/v1/worker/test_dsv4_packed_zeroer_geometry.py using the real DSV4 page
sizes. To reproduce e2e:

vllm serve deepseek-ai/DeepSeek-V4 --kv-cache-dtype fp8_ds_mla
vllm serve MiniMaxAI/MiniMax-M3
# and repeat each with VLLM_KV_CACHE_LAYOUT=BLHNC

AI-assisted (Claude); every changed line reviewed and tests run by the
submitter.

mergify · 2026-08-10T18:31:40Z

Documentation preview: https://vllm--51718.org.readthedocs.build/en/51718/

mergify · 2026-08-10T18:39:02Z

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

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

Rework of the extensible-kv-cache V2 integration on the #51718 single-allocation model, substantially simplified by it: - The shared allocate_and_reshape_kv_cache gains a reserve(size, num_segments) hook that stands in for the plain torch.zeros backing allocation; one VMM buffer backs the whole KV cache. - Buffer segmentation falls out of the KVCacheTensor stride model: num_segments = size / (num_blocks * block_stride), i.e. one segment per contiguous run of blocks (1 for block-outermost layouts, one per layer region under layer-compact layouts), replacing per-backend shape/stride/block-dim probing. - Connector registration views are a plain narrow(0) of the logical [B, H, N, C] per-layer views. Ported unchanged: measured post-warmup sizing, deferred KV-transfer init, sleep release/recommit, VMM probe/fallback, elastic-EP skip, and encoder-cache profiling reservation. Co-authored-by: Zhuohan Li <zhuohan123@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ho2gVXA5r7PhrM2sk8oJUn Signed-off-by: Nick Hill <nickhill123@gmail.com>

mergify · 2026-08-11T15:20:20Z

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

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

okorzh-amd · 2026-08-12T05:28:49Z

+                else:
+                    block_stride = cache.stride(0) * cache.element_size()
+                storage_is_block_major = num_blocks * block_stride == storage.nbytes()
+                hnc_contiguous = (


For an MLA layer under BLHNC, hnc_contiguous is trivially true (H=1), so every layer takes the per-layer region branch, replacing base's _register_packed_kv_cache (num_regions=1, num_descs=num_blocks, one contiguous packed-row chunk per block) with one region per layer. For a ~100-layer DSV4 and a ~400k-block pool that is 1 region/400k descs → ~100 regions/~40M descs: prep_xfer_dlist init time and memory ~100×, and each block transfer becomes ~100 scatter/gather entries of ~42 KB instead of a single ~4.2 MB copy. Scoped to the DSv4-style packed path ordinary per-layer models already had num_regions == num_layers.

mergify · 2026-08-12T06:43:48Z

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

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

LucasWilkinson · 2026-08-13T18:50:11Z

/ci run

github-actions · 2026-08-13T18:50:28Z

✅ Triggered Buildkite CI #83782 for commit 9d61ce258b5c.

Rework of the extensible-kv-cache V2 integration on the #51718 single-allocation model, substantially simplified by it: - The shared allocate_and_reshape_kv_cache gains a reserve(size, num_segments) hook that stands in for the plain torch.zeros backing allocation; one VMM buffer backs the whole KV cache. - Buffer segmentation falls out of the KVCacheTensor stride model: num_segments = size / (num_blocks * block_stride), i.e. one segment per contiguous run of blocks (1 for block-outermost layouts, one per layer region under layer-compact layouts), replacing per-backend shape/stride/block-dim probing. - Connector registration views are a plain narrow(0) of the logical [B, H, N, C] per-layer views. Ported unchanged: measured post-warmup sizing, deferred KV-transfer init, sleep release/recommit, VMM probe/fallback, elastic-EP skip, and encoder-cache profiling reservation. Co-authored-by: Zhuohan Li <zhuohan123@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ho2gVXA5r7PhrM2sk8oJUn Signed-off-by: Nick Hill <nickhill123@gmail.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>

@okorzh-amd

- benchmarks/attention_benchmarks/runner.py: import get_kv_cache_layout (resolve_kv_cache_layout was renamed) and pass reshape_kv_cache's offset/layer_stride/block_stride. Fixes the non-optional MI355 "Attention Benchmarks Smoke Test" job. Reported by @okorzh-amd. - test_rocm_aiter_fa: replace the deleted get_kv_cache_shape assertions with the customize_spec packing contract. - test_sparse_mla_backends: squeeze the head axis before calling forward_mha directly, mirroring bind_kv_cache. Without it the chunked-context prefill read the cache at the wrong stride. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

Rework of the extensible-kv-cache V2 integration on the #51718 single-allocation model, substantially simplified by it: - The shared allocate_and_reshape_kv_cache gains a reserve(size, num_segments) hook that stands in for the plain torch.zeros backing allocation; one VMM buffer backs the whole KV cache. - Buffer segmentation falls out of the KVCacheTensor stride model: num_segments = size / (num_blocks * block_stride), i.e. one segment per contiguous run of blocks (1 for block-outermost layouts, one per layer region under layer-compact layouts), replacing per-backend shape/stride/block-dim probing. - Connector registration views are a plain narrow(0) of the logical [B, H, N, C] per-layer views. Ported unchanged: measured post-warmup sizing, deferred KV-transfer init, sleep release/recommit, VMM probe/fallback, elastic-EP skip, and encoder-cache profiling reservation. Co-authored-by: Zhuohan Li <zhuohan123@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Ho2gVXA5r7PhrM2sk8oJUn Signed-off-by: Nick Hill <nickhill123@gmail.com>

mergify Bot added documentation Improvements or additions to documentation deepseek Related to DeepSeek models performance Performance-related issues kimi k3 nvidia labels Aug 10, 2026

github-project-automation Bot added this to NVIDIA Aug 10, 2026

mergify Bot added the rocm Related to AMD ROCm label Aug 10, 2026

github-project-automation Bot added this to AMD Aug 10, 2026

github-project-automation Bot moved this to Todo in AMD Aug 10, 2026

mergify Bot added intel-gpu Related to Intel GPU cpu Related to CPU backends mrv2 Model Runner V2 specific kv-connector labels Aug 10, 2026

LucasWilkinson changed the title ~~[4/N][KV-Cache Layout Refactor] Standardize KV cache layout~~ [6/N][KV-Cache Layout Refactor] Standardize KV cache layout Aug 10, 2026

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

LucasWilkinson force-pushed the lwilkinson/kv-layout/core-standardize-stacked branch from 34241fe to e04d271 Compare August 10, 2026 18:45

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

njhill mentioned this pull request Aug 10, 2026

[Core] Extensible (growable) KV cache #50779

Draft

LucasWilkinson force-pushed the lwilkinson/kv-layout/core-standardize-stacked branch from 5555e98 to e5b3e58 Compare August 11, 2026 01:03

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

njhill mentioned this pull request Aug 11, 2026

[Bugfix][ROCm] Give KV-first attention blocks their own page in hybrid models #51837

Merged

4 tasks

LucasWilkinson force-pushed the lwilkinson/kv-layout/core-standardize-stacked branch from 2518149 to fb79a88 Compare August 11, 2026 21:31

okorzh-amd reviewed Aug 12, 2026

View reviewed changes

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

LucasWilkinson force-pushed the lwilkinson/kv-layout/core-standardize-stacked branch from f29cb31 to 9689aba Compare August 12, 2026 19:15

mergify Bot added the ci/build label Aug 12, 2026

LucasWilkinson force-pushed the lwilkinson/kv-layout/core-standardize-stacked branch from 9689aba to cf5e6e5 Compare August 12, 2026 21:50

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

LucasWilkinson force-pushed the lwilkinson/kv-layout/core-standardize-stacked branch from cf5e6e5 to 6397d65 Compare August 13, 2026 06:26

LucasWilkinson force-pushed the lwilkinson/kv-layout/core-standardize-stacked branch from 6397d65 to 9d61ce2 Compare August 13, 2026 18:10

LucasWilkinson added the ready-run-all-tests Trigger CI with all tests for wide-ranging PRs label Aug 13, 2026

LucasWilkinson force-pushed the lwilkinson/kv-layout/core-standardize-stacked branch from 9d61ce2 to d77c00a Compare August 14, 2026 03:01

Base automatically changed from lwilkinson/kv-packing-customize-spec to main August 14, 2026 13:47

LucasWilkinson force-pushed the lwilkinson/kv-layout/core-standardize-stacked branch from d77c00a to b4abb9e Compare August 14, 2026 13:54

LucasWilkinson added the ready ONLY add when PR is ready to merge/full CI is needed label Aug 14, 2026

LucasWilkinson and others added 3 commits August 14, 2026 20:42

Simplify KV cache layout stride handling

28da35b

Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>

LucasWilkinson force-pushed the lwilkinson/kv-layout/core-standardize-stacked branch from b4abb9e to 28da35b Compare August 14, 2026 20:43

Navigation Menu

Uh oh!

Uh oh!

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

[6/N][KV-Cache Layout Refactor] Standardize KV cache layout#51718
LucasWilkinson wants to merge 3 commits into
mainfrom
lwilkinson/kv-layout/core-standardize-stacked

LucasWilkinson commented Aug 10, 2026 •

edited

Loading

Uh oh!

mergify Bot commented Aug 10, 2026

Uh oh!

mergify Bot commented Aug 10, 2026

Uh oh!

mergify Bot commented Aug 11, 2026

Uh oh!

okorzh-amd Aug 12, 2026

Uh oh!

mergify Bot commented Aug 12, 2026

Uh oh!

LucasWilkinson commented Aug 13, 2026

Uh oh!

github-actions Bot commented Aug 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Uh oh!

Conversation

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

Uh oh!

Purpose

Test Plan / Result

Validation

Uh oh!

mergify Bot commented Aug 10, 2026

Uh oh!

mergify Bot commented Aug 10, 2026

Uh oh!

mergify Bot commented Aug 11, 2026

Uh oh!

okorzh-amd Aug 12, 2026

Choose a reason for hiding this comment

Uh oh!

mergify Bot commented Aug 12, 2026

Uh oh!

LucasWilkinson commented Aug 13, 2026

Uh oh!

github-actions Bot commented Aug 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

LucasWilkinson commented Aug 10, 2026 •

edited

Loading