[Core] Prefer a block-outermost KV cache layout when layers can be packed

[LucasWilkinson]

Stacked on #58207. Only the top commit ([Core] Prefer a block-outermost KV cache layout when layers can be packed) is new here; review #58207 first.

Purpose

When KV cache specs have mixed page sizes, e.g. the DFlash drafter for Qwen3.6-35B-A3B (4 KiB/token/layer) next to the target (2 KiB/token/layer), the default layer-outermost layout unifies page sizes by shrinking the drafter's block size, which keeps target and drafter layers in separate buckets and forces padding and extra KV cache groups (#58207 picks 17 groups with 12 KiB/token of padding at the default --min-kv-cache-group-layers 3).

A block-outermost layout (BLHNC / BLNHC) lets _get_packed_kv_cache_groups pack layers with different page sizes into one block instead: the drafter's full attention layer occupies 2 page slots in the same group as the target's 10 full attention layers. For Qwen3.6-35B-A3B + DFlash this gives 5 groups and no padding (3 mamba, 1 full attention = 10 target + 1 drafter, 1 drafter sliding window), with slightly more KV cache capacity than main's 46 groups.

Legend:  F = target full attn (2 KiB/token)   f = drafter full attn (4 KiB/token)
         M = mamba/GDN   s = drafter sliding window (4 KiB/token)
         ▓ = full attn padding (grows with context)   ░ = idle/padding, bounded (window / state)

#58207 uniform-page layout (LBHNC), gs=3: 17 groups

M  [MMM] ×10
F  [FFF][FFF][FF▓][FF▓]
s  [sss][ss░]
f  [f▓▓]                  ← 4 padded full attn slots: 12 KiB/token (+50%)

this PR, packed layout (BLNHC): 5 groups. Each row is a separate group; every group's block is the same size (12 pages of 640 tokens × 2 KiB = 15 MiB)

full-attn ×1   │ F │ F │ F │ F │ F │ F │ F │ F │ F │ F │   f   │   10 target + drafter (2 pages): no padding
mamba ×3       │ M │ M │ M │ M │ M │ M │ M │ M │ M │ M │ ░ │ ░ │   10 GDN states each, 2 idle pages
drafter sw ×1  │   s   │   s   │   s   │   s   │   s   │ ░ │ ░ │   5 layers × 2 pages, 2 idle pages

Group counts, block sizes and KV capacity are from the served runs below. The per-group layer split is from _get_packed_kv_cache_groups (not logged).

Block-outermost layouts cannot split a manager block into smaller kernel blocks (create_kv_cache_views requires dense pages for that). So this PR only changes the worker's reported layout preference, moving block-outermost layouts first, when:

1. the worker's KV cache specs have mixed page sizes (when the packed path applies), and

2. every attention backend accepts each attention layer's manager block size as a kernel block (no splitting).

3. the BLHNC cost check passes: both groupings are computed from the specs, and block-outermost is preferred only if the layer-outermost grouping would pad full attention, or packing needs fewer KV cache groups. BLHNC is not always better: for MiMo-V2.5 + DFlash packing gives 12 groups vs 7, with no full attention padding to remove.

Otherwise the preference order is unchanged. VLLM_KV_CACHE_LAYOUT and KV connector layout preferences still take precedence.

On B300, FlashInfer (16/32/64-token kernel blocks) and FA4's hd256 kernel (128-token pages) need the 640/1280-token hybrid manager blocks split, so with default backends nothing changes; TRITON_ATTN (any multiple of 16) gets the packed layout automatically. Making FlashInfer/FA work with packed blocks needs in-kernel re-paging, similar to what #55219 / #57169 do for sparse MLA.

Why not duplicate: #55219 / #57169 route GLM-5.3-Flash through the packed layout and re-page for sparse MLA kernels; #54795 / #55340 change how per-worker / draft layout sets are intersected. None prefers block-outermost layouts automatically for mixed-page-size models.

Test Plan

pytest tests/v1/worker/test_utils.py -q

End to end, Qwen3.6-35B-A3B-FP8 + z-lab/Qwen3.6-35B-A3B-DFlash, 1× B300, BF16 mamba states, GSM8K, c=1..32, no VLLM_KV_CACHE_LAYOUT:

vllm serve Qwen/Qwen3.6-35B-A3B-FP8 --max-num-seqs 32 --no-enable-prefix-caching --block-size 128 \
  --attention-backend TRITON_ATTN --mamba-cache-dtype bfloat16 --mamba-ssm-cache-dtype bfloat16 \
  --speculative-config '{"model":"z-lab/Qwen3.6-35B-A3B-DFlash","method":"dflash","num_speculative_tokens":7,"attention_backend":"TRITON_ATTN"}'

Test Result

test_prefers_packed_kv_cache covers mixed / uniform page sizes and a backend that would need kernel-block splitting. All of tests/v1/worker/test_utils.py passes.

config	layout	groups	KV cache capacity
main grouping, FlashInfer/FA	LBHNC	46	7.42M tokens
#58207, FlashInfer/FA (unchanged by this PR)	LBHNC	17	5.19M tokens
#58207, Triton	LBHNC	17	5.20M tokens
this PR, Triton	BLNHC (auto)	5	7.53M tokens

With the same Triton backend, the packed layout improves output throughput by ~4% at c=1 and ~12% at c=32 over #58207's grouping, and recovers the KV cache capacity #58207 spends on padding. Acceptance length is 5.1–5.2 in all runs (single runs).

Cost check, served with default backends (no VLLM_KV_CACHE_LAYOUT):

model	layout picked	groups	KV cache capacity
Qwen3.6-35B-A3B-FP8 + DFlash, TRITON_ATTN	BLNHC	5	7.53M
Qwen3.6-35B-A3B-FP8 + DFlash, FlashInfer/FA	LBHNC (can't pack: kernel-block split)	17	5.19M
MiMo-V2.5 + DFlash (TP2)	LBNHC (cost check fails: 12 vs 7 groups, no full attn padding)	7	5.96M

See the RFC for the full design: #58638.

---

AI assistance (Claude Code) was used; I have reviewed the change.

🤖 Generated with Claude Code

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

[tomasruizt]

What should happen in pipeline parallelism (assume two stages are assigned different layers), and they compute different preferences? e.g. one stage prefers packed layout and the other doesn't. Atm the function resolve_kv_cache_layout() raises an AssertionError when the supported layouts of each rank are ordered differently.

vllm/vllm/v1/attention/backends/utils.py

Lines 285 to 288 in 741edee

	assert all(names == supported_layouts[0] for names in supported_layouts[1:]), (
	f"Workers disagree on supported KV cache layouts: {supported_layouts}."
	)
	candidates = [_layout_from_name(name) for name in supported_layouts[0]]

We could compute the intersection of the layouts across all ranks, and then e.g. use the ordering of rank0 as an ordering heuristic.

# intersection
common = reduce(set.intersection, [set(l) for l in supported_layouts])
if len(common) == 0:
    raise ValueError("No KV cache layout is supported by every worker.")

# ordering heuristic: use rank0 preferences
candidates = [
    _layout_from_name(name)
    for name in supported_layouts[0]
    if name in common
]

Code to reproduce the issue:

• Refactor to pure functions: tomasruizt@c3a53da
• Add failing test case: tomasruizt@b1c5f28
• Suggested fix: tomasruizt@adb5936

[mergify]

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

[tomasruizt]

Another case: Currently, the packed layout is preferred when uniform grouping pads full-attentions layers, even when packing more more memory. We could instead compare worst-case allocation for both layouts.

• Test to surface the problem: tomasruizt@102302b
• Proposed fix: tomasruizt@7702575