[KV Connector] Canonical CPU layout for parallelism-agnostic KV offload - #48414
Conversation
Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
|
This pull request has merge conflicts that must be resolved before it can be |
All parallelism reasoning moves to vllm/v1/kv_offload/sharding.py, covering TP/DCP/PCP, packed and split KV layouts, with single-writer election. Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
CPU pages follow the refs' canonical mappings behind an opt-in flag; the persisted format id joins the file-tier namespace. Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
93d8816 to
f5966e8
Compare
Etelis
changed the title
# Conflicts: # tests/v1/kv_offload/test_file_mapper.py # vllm/v1/kv_offload/cpu/gpu_worker.py # vllm/v1/kv_offload/file_mapper.py Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
This pull request has merge conflicts that must be resolved before it can be |
|
Hi @Etelis — while preparing the gate relaxation @orozery suggested in #48906 (extending replicated_layout to UniformTypeKVCacheSpecs / all-MLA multi-group), my duplicate check landed back on this PR — the canonical layout covers a superset of that goal, so I don't want to open a competing path if this is landing. Two questions:
|
I will continue tomorrow :) rebase and continue |
| kv_bytes_per_block=self.kv_bytes_per_chunk, | ||
| cpu_page_size=self.cpu_page_size_per_worker, | ||
| ) | ||
| canonical_layout = bool(self.extra_config.get("canonical_layout", False)) |
There was a problem hiding this comment.
Can we parse this at build_offloading_config?
| canonical_layout=canonical_layout, | ||
| ) | ||
|
|
||
| def _validate_canonical_refs(self, kv_caches: CanonicalKVCaches) -> None: |
There was a problem hiding this comment.
same, can we move this when parsing OffloadingConfig?
| """Allocate a strided int8 view shared by all workers for one | ||
| canonical tensor (canonical layout). Canonical views are carved from | ||
| the start of each block row; workers write disjoint bytes within them | ||
| as described by their mappings. Must be called once per canonical | ||
| tensor, instead of create_next_worker_view.""" |
There was a problem hiding this comment.
Can we add a diagram like in create_next_worker_view?
create_next_worker_view
|<----- cpu_page_size ----->|<----- cpu_page_size ----->|
| worker 0 slot | worker 1 slot |
| | |
block 0: | [ tensor0 | tensor1 ] | [ tensor0 | tensor1 ] |
block 1: | [ tensor0 | tensor1 ] | [ tensor0 | tensor1 ] |
block 2: | [ tensor0 | tensor1 ] | [ tensor0 | tensor1 ] |
^ ^
worker0 offset=0 worker1 offset=cpu_page_size
create_next_canonical_view:
|<------- canonical area ------->|<--- unused padding --->|
| all workers share this area | |
| | |
block 0: | [ canonical_t0 | canonical_t1 ]| (waste) |
block 1: | [ canonical_t0 | canonical_t1 ]| (waste) |
block 2: | [ canonical_t0 | canonical_t1 ]| (waste) |
^ ^
_canonical_offset=0 then advances by each tensor's size
| is_parallelism_agnostic = ( | ||
| not vllm_config.use_v2_model_runner | ||
| and single_group_spec is not None | ||
| and isinstance(single_group_spec, FullAttentionSpec) | ||
| and not isinstance(single_group_spec, MLAAttentionSpec) | ||
| and single_group_spec.num_kv_heads * parallel_config.tensor_parallel_size | ||
| == vllm_config.model_config.get_total_num_kv_heads() | ||
| and not single_group_spec.kv_quant_mode.is_per_token_head | ||
| and parallel_config.decode_context_parallel_size == 1 | ||
| and parallel_config.prefill_context_parallel_size == 1 |
There was a problem hiding this comment.
It seems like we are just adding conditions.
With this PR we should have more cases where is_parallelism_agnostic is True, no?
| fill_group_ops = ( | ||
| self._fill_legacy_ops | ||
| if self._canonical_copy_plans is None | ||
| else self._fill_canonical_ops | ||
| ) | ||
|
|
There was a problem hiding this comment.
Let's bind _fill_group_ops at init instead of branching on self._canonical_copy_plans is None in every transfer_async call.
| block_bases_dst[:, None] + plan.frag_offsets_dst[None, :] | ||
| ).ravel() | ||
| all_sizes[op_idx:end_idx] = np.tile(plan.frag_sizes, num_active_blocks) | ||
| num_bytes += num_active_blocks * int(plan.frag_sizes.sum()) |
There was a problem hiding this comment.
We can avoid plan.frag_sizes.sum() by introducing total_bytes=int(frag_sizes_arr.sum()) to CopyPlan.
| all_dst[op_idx:end_idx] = ( | ||
| block_bases_dst[:, None] + plan.frag_offsets_dst[None, :] | ||
| ).ravel() | ||
| all_sizes[op_idx:end_idx] = np.tile(plan.frag_sizes, num_active_blocks) |
There was a problem hiding this comment.
Claude suggests this avoids allocation:
| all_sizes[op_idx:end_idx] = np.tile(plan.frag_sizes, num_active_blocks) | |
| all_sizes[op_idx:end_idx].reshape(num_active_blocks, plan.num_frags)[:] = ( | |
| plan.frag_sizes | |
| ) |
It claims:
This works because all_sizes[op_idx:end_idx] is a contiguous slice of a 1D array, so .reshape() returns a view (no copy).
The [:] = broadcasts frag_sizes (shape (num_frags,)) across rows.
But need to verify.
| block_bases_src = np.empty(group_size, dtype=np.uint64) | ||
| block_bases_dst = np.empty(group_size, dtype=np.uint64) |
There was a problem hiding this comment.
Claude's suggestion (have not verified):
Currently allocates two np.empty(group_size, uint64) arrays per (group, ref) per transfer. Pre-allocate once at init, sized to the max possible group.
In init, after computing blocks_per_chunk etc.:
+ max_group_size = max(t.tensor.shape[0] for t in kv_caches.tensors)
+ self._scratch_bases_src = np.empty(max_group_size, dtype=np.uint64)
+ self._scratch_bases_dst = np.empty(max_group_size, dtype=np.uint64)
Then in _fill_canonical_ops:
- block_bases_src = np.empty(group_size, dtype=np.uint64)
- block_bases_dst = np.empty(group_size, dtype=np.uint64)
+ block_bases_src = self._scratch_bases_src[:group_size]
+ block_bases_dst = self._scratch_bases_dst[:group_size]
Note: after _filter_writer_blocks applies a boolean mask, the result is a new array anyway (fancy indexing copies). So the scratch is only reused for the compute_sub_block_ptrs output — the filtered result is still a fresh allocation. This is fine; the savings are on the common path (num_writers == 1) where no filtering happens and the scratch is used directly for the broadcast.
| all_src[op_idx:end_idx] = ( | ||
| block_bases_src[:, None] + plan.frag_offsets_src[None, :] | ||
| ).ravel() | ||
| all_dst[op_idx:end_idx] = ( | ||
| block_bases_dst[:, None] + plan.frag_offsets_dst[None, :] | ||
| ).ravel() |
There was a problem hiding this comment.
Claude suggestion (have not verified, and if it's big we can do as follow-up):
Avoid 2D broadcast intermediate in _fill_canonical_ops
(block_bases_src[:, None] + frag_offsets_src[None, :]).ravel() allocates an (N, M) temporary array every transfer just to immediately flatten it into the output.
Replace with np.add(..., out=...) writing directly into a reshaped view of the destination buffer — zero allocation.
The blocker is a dtype mismatch: the output buffer is backed by a torch.int64 tensor (kernel requirement) but the operands are np.uint64 (natural for addresses). Clean fix: at init, create a uint64 numpy view over the same memory:
self._np_src = self._all_src.numpy().view(np.uint64)
This is zero-copy — int64 and uint64 are bit-equivalent for pointer values.
All numpy arithmetic then works in uint64 natively, and the CUDA kernel sees the same bytes as int64. No per-call casts.
| @@ -0,0 +1,171 @@ | |||
| # SPDX-License-Identifier: Apache-2.0 | |||
There was a problem hiding this comment.
Claude:
Missing test: cross-topology roundtrip for canonical layout
The canonical layout's core value proposition — KV written with one parallelism config can be read with a different one — has no integration test.
The existing tests in test_canonical_layout.py verify the building blocks (_build_copy_plan, _canonical_page_ids, _canonical_block_sizes) and a single-topology GPU roundtrip, but nobody tests that a TP=2 writer's data is actually usable by a TP=4 reader.
Suggested test (in test_canonical_layout.py):
- Create a shared
SharedOffloadRegion(simulates the canonical mmap file) - N writer workers (
tp_size=writer_tp), each scatters its KV fragment into canonical positions (GPU→CPU) - M reader workers (
tp_size=reader_tp), each gathers its fragment from the same canonical region (CPU→GPU) - Assert: reassembled reader data == original writer data (same logical KV, different physical sharding)
Parametrize over (writer_tp, reader_tp) pairs like (2, 4), (4, 2), (2, 1).
This is the one test that proves the feature works end-to-end, not just that the pieces are individually correct.
…, direct-layout naming Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
|
✅ @Etelis, CI is now available for this PR.
|
Canonical certification is stride-verified per layer at registration and fails closed, so the static gate needs only the derivation preconditions; this admits the v2 model runner and replicated GQA heads. Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
|
/ci run |
|
✅ Triggered Buildkite CI #83056 for commit |
Seems like it, on it. |
The Triton load kernel dereferences CPU pointers on the GPU, which is only legal on pinned memory; production pins via CPUOffloadingWorker. Signed-off-by: Itay Etelis <itay.etelis@ibm.com>
|
/ci run |
|
✅ Triggered Buildkite CI #83063 for commit |
|
/ci retry |
|
✅ No failed, timed-out, or expired jobs need retrying: https://buildkite.com/vllm/ci/builds/83063 |
|
/ci retry |
|
✅ Queued 1 failed job(s) for retry in Buildkite CI #83063. |
|
@orozery green. |
…ad (vllm-project#48414) Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com>
Map replicated MLA host KV into one pinned region for eligible single-node MP tensor-parallel workers. Each rank retains its own DMA submission so its CUDA stream continues to order host writes before subsequent reads. This is a HiSparse integration of the shared offload substrate from vllm-project#48414, not a duplicate of that general layout work. Tests: 141 CPU and multiprocess tests passed; 32 CUDA offload tests passed. No model evaluation was required because storage bytes and attention results are unchanged. AI assistance was used for implementation and review. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Map replicated MLA host KV into one pinned region for eligible single-node MP tensor-parallel workers. Each rank retains its own DMA submission so its CUDA stream continues to order host writes before subsequent reads. This is a HiSparse integration of the shared offload substrate from vllm-project#48414, not a duplicate of that general layout work. Tests: 141 CPU and multiprocess tests passed; 32 CUDA offload tests passed. No model evaluation was required because storage bytes and attention results are unchanged. AI assistance was used for implementation and review. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Map replicated MLA host KV into one pinned region for eligible single-node MP tensor-parallel workers. Each rank retains its own DMA submission so its CUDA stream continues to order host writes before subsequent reads. This is a HiSparse integration of the shared offload substrate from vllm-project#48414, not a duplicate of that general layout work. Tests: 141 CPU and multiprocess tests passed; 32 CUDA offload tests passed. No model evaluation was required because storage bytes and attention results are unchanged. AI assistance was used for implementation and review. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Map replicated MLA host KV into one pinned region for eligible single-node MP tensor-parallel workers. Each rank retains its own DMA submission so its CUDA stream continues to order host writes before subsequent reads. This is a HiSparse integration of the shared offload substrate from vllm-project#48414, not a duplicate of that general layout work. Tests: 141 CPU and multiprocess tests passed; 32 CUDA offload tests passed. No model evaluation was required because storage bytes and attention results are unchanged. AI assistance was used for implementation and review. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Map replicated MLA host KV into one pinned region for eligible single-node MP tensor-parallel workers. Each rank retains its own DMA submission so its CUDA stream continues to order host writes before subsequent reads. This is a HiSparse integration of the shared offload substrate from vllm-project#48414, not a duplicate of that general layout work. Tests: 141 CPU and multiprocess tests passed; 32 CUDA offload tests passed. No model evaluation was required because storage bytes and attention results are unchanged. AI assistance was used for implementation and review. Co-authored-by: Codex <codex@openai.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
4 participants
Stacked on #48408. Stores offloaded KV in the canonical (parallelism-free) layout
described by the refs' page mappings: each worker scatters its page fragments to
their canonical positions in a CPU area shared by the whole worker group.
MLA latent and replicated GQA heads are stored once instead of once per rank
(empty store runs on non-writers). Copy expansions are precomputed per ref at
init; per-block placement is one vectorized base + offsets sum, with no
parallelism inputs anywhere in the transfer path.
Configuration
kv_connector_extra_config: {"canonical_layout": true}onTieringOffloadingSpec(requires the shared mmap). Requesting it on an uncertifiable config fails at
startup instead of silently downgrading. The persisted format identity
(
v1-nhd/v1-hnd) joins the FileMapper namespace so canonical, legacy, andcross-family bytes can never resolve to the same files; the
parallel_agnosticgate now also excludes replicated GQA heads, per-token-head scales, and CP.
Test plan
pytest tests/v1/kv_offload/cpu/test_canonical_layout.py tests/v1/kv_offload/test_file_mapper.py(no GPU): expansion offsets, empty store runs, per-tensor areas, namespace
separation, parallel_agnostic gates. Cross-topology byte semantics are covered by
the schema tests in #48408.
E2E on 4x H100 (eviction-forced CPU reloads, block-aligned prompts, greedy
token comparison against a GPU-cached baseline, lookup hits verified in logs):
Note for reviewers:
reset_prefix_cachealso resets the offload manager, soreload tests must evict via cache pressure instead.