Skip to content

[KV Connector] Fragment-major canonical CPU layout for KV offload#48414

Draft
Etelis wants to merge 2 commits into
vllm-project:mainfrom
Etelis:kv-offload-fragment-layout
Draft

[KV Connector] Fragment-major canonical CPU layout for KV offload#48414
Etelis wants to merge 2 commits into
vllm-project:mainfrom
Etelis:kv-offload-fragment-layout

Conversation

@Etelis

@Etelis Etelis commented Jul 12, 2026

Copy link
Copy Markdown
Contributor

Uses the per-layer head-region schema (#48408) to store offloaded KV in a
TP-invariant layout: behind an opt-in flag, each worker scatters its page
fragments to rank-interleaved positions in the shared CPU row.

Groundwork for cross-TP cache consumption (PD-disagg, parallel-
agnostic file sharing); stacked on #48408.

legacy row:     [worker0: T0|T1][worker1: T0|T1]...     axis: [worker][tensor][region]
canonical row:  [T0: K interleaved | V interleaved][T1: ...]   axis: [tensor][region][heads]

Configuration

kv_connector_extra_config: {"canonical_layout": true} on
TieringOffloadingSpec only (requires the shared mmap); auto-disabled unless
world_size == tensor_parallel_size. Flag off (default): the legacy transfer
path is byte-for-byte unchanged. HND KV-cache layout recommended when
enabling (NHD fragments per token: more, smaller copy ops).

Test plan

Unit (tests/v1/kv_offload/cpu/test_fragment_layout.py, no GPU): NHD/HND
expansions against hand-computed offsets, opaque single-slot, store/load
offset mirror, and all-slots tiling (destinations partition the region with
no gaps/overlaps).

E2E on 2x H100 (Qwen2.5-0.5B, generate -> reset_prefix_cache -> regenerate,
compare token ids):

case TP tokens_equal
legacy (flag off) 1 true
canonical 1 true
legacy (flag off) 2 true
canonical (2 workers interleaving one row) 2 true

Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants