[KV Connector] Support NIXL heterogeneous P/D block sizes for hybrid models - #49612
Merged
Merged
Conversation
njhill
requested review from
ApostaC,
NickLucche,
ivanium,
orozery and
xuechendi
as code owners
July 23, 2026 16:06
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
…models Hybrid (mamba) models previously asserted out heterogeneous block sizes entirely — yet they are the models where P/D block sizes most readily diverge, since the mamba-padded attention block size varies with TP sharding (vllm-project#41037). Lift the restriction: - Mamba state blocks are indivisible, so their descriptors always use local page geometry and their desc ids are never ratio-expanded; attention descriptors remain expanded to remote-block granularity. - A shared per-group mapping (pull READ and push WRITE) expands attention groups to remote granularity and clips to the transferred coverage, while mamba groups map 1:1. - The receive post-process permutes only attention-layer caches (a lazily cached list) and zeroes the untransferred token tail of the last local block, whose zeroing the scheduler skipped for the load. - The multi-read hetero-TP path (P_TP > D_TP with sharded SSM state, vllm-project#49297) builds its per-source split handles from the remote- granularity descriptors, keyed by (tp_ratio, remote_block_size); replicated and single-source attention descriptors pass through whole, and SSM chunking is unaffected by the ratio. Still unsupported, now rejected loudly at handshake time: head-sharded attention reads combined with a block-size mismatch, and host-buffer mode with a block-size ratio. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
…-strict count tolerance Adds an end-to-end descriptor-geometry harness for hybrid MLA+SSM models under heterogeneous P/D block geometry (TP-sharded KDA-style state making the mamba-aligned logical block size differ between P and D while kernel pages stay equal): a recording NIXL wrapper resolves every prepared transfer's local and remote descriptor ids back to byte ranges, and a sweep over four geometries and 38 prompt lengths asserts three invariants per read: every local write lands within the requesting request's own blocks (a violation would corrupt a co-resident request mid-decode), every local/remote descriptor pair is token-aligned, and the matched tokens are fully covered. The sweep found one real defect: _apply_prefix_caching's count-mismatch tolerance used max(local_ppl, remote_ppl) as a strict bound, but allocation rounding legitimately leaves up to ppl-1 trailing dead kernel blocks per side (e.g. 12-token local vs 8-token remote blocks at N=25 gives |9-6| == max_padding), crashing legitimate transfers with an AssertionError. Widen the bound to the sum of the ratios, inclusive. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_013TQHxUrXWp7k5kk1oz95Cy Signed-off-by: Nick Hill <nickhill123@gmail.com>
With equal kernel pages but differing logical block sizes (hybrid heterogeneous TP, e.g. TP8 prefill block 768 -> TP1 decode block 5760), the pull transfer is front-trimmed to min(local, remote) kernel blocks. The scheduler excluded the matched-range blocks from alloc-time KV zeroing (it would race the RDMA write), and the existing receive tail-zeroing only runs for block_size_ratio > 1 and non-MLA models, so the untransferred tail of the last local logical block kept stale recycled bytes. Once decode grew into that tail, requests emitted garbage mid-response; only the prefill-block < decode-block direction is affected, matching observed failures in that direction only. Adds a receive-time zeroing of the clipped attention kernel blocks and extends the descriptor-geometry suite with a KimiLinear-scale geometry (5760/768, kernel 64, tp_ratio=-8) plus a sentinel-based transferred-or-zeroed invariant that reproduced the bug. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QVJLHN2a1DP6iBJjayotrs Signed-off-by: Nick Hill <nickhill123@gmail.com>
Receive-side zeroing of blocks the transfer didn't write lived in two places with two different notions of granularity: the block-size-ratio tail inside post_process_device_kv_on_receive (sub-block units, non-MLA only, last covered block only) and _zero_untransferred_hetero_ppl_tail (whole kernel blocks, mamba only). Neither knew about the other, so a hybrid running both regimes at once - differing kernel block sizes and differing physical_blocks_per_logical, reachable with prefix caching off - compared local kernel blocks against remote sub-blocks and left whole untransferred blocks unzeroed. Fold both into post_process_device_kv_on_receive, which now takes the covered sub-block count in the finest (remote-block) granularity and zeroes everything past it: the token tail of the last partially covered block, then whole blocks beyond. MLA enters the path too, with the block-size/layout conversion gated separately, so its clipped blocks are also zeroed. The invariant is now uniform: every byte of the request's local attention blocks is either transferred or zeroed. Also switch the block-id H2D copies to async_tensor_h2d - a pageable torch.tensor(..., device=cuda) copy is host-blocking and stream-ordered, so it can stall the engine loop behind queued forward kernels - and build the tensor only when a consumer needs it. tests/v1/kv_connector/unit/test_nixl_desc_geometry.py and test_nixl_connector_hma.py: 236 passed on GB200 (1 unrelated failure, test_fewer_blocks_with_hma, needs HF auth for a gated repo). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
…ry suite The desc-geometry suite only ever ran the two heterogeneous regimes in isolation: differing physical_blocks_per_logical with equal kernel block sizes, or (elsewhere) a block-size ratio alone. The combination - local kernel 8 / remote kernel 4 (ratio 2) with ppl 3 vs 2 - is what the two former receive-side zeroing paths each handled only half of. Generalize the harness to take a separate remote kernel block size: the remote's advertised kernel page scales down by the ratio, and descriptor resolution now works in absolute token offsets with an explicit desc page, so local sub-block descs pair correctly against whole remote pages. Coverage is checked at the finest transfer granularity (remote kernel). Against the previous two-path implementation the new case fails 11 of 16 prompt lengths on the transferred-or-zeroed invariant, e.g. N=29 leaves 192 stale bytes per region in local block 3; all 16 pass with the unified path. tests/v1/kv_connector/unit/test_nixl_desc_geometry.py and test_nixl_connector_hma.py: 252 passed on GB200 (1 unrelated failure, test_fewer_blocks_with_hma, needs HF auth for a gated repo). Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
It holds the by-remote-block-size map of src_blocks_data and sits beside src_xfer_handles_by_block_size; both siblings are public-styled. No references outside base_worker.py. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Nick Hill <nickhill123@gmail.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
async_tensor_h2d's result was discarded, leaving indices None: the has_stale zeroing path then failed its assertion and the convert path would pass None to the postprocess kernels. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QVJLHN2a1DP6iBJjayotrs Signed-off-by: Nick Hill <nickhill123@gmail.com>
The geometry suite fakes a CUDA platform for the worker but builds its
VllmConfig with the real one, and KVTransferConfig.kv_buffer_device
defaults to current_platform.device_type. On the cpu-small CI queue
(image-build-cpu) that resolves to "cpu" while the worker's faked
device_type is "cuda", so use_host_buffer flips on. Host xfer buffers
are allocated per layer, so the HMA shared-tensor pooling this suite
relies on stops deduplicating: the 6 layers register as 6 regions
instead of 2, tripping the region-count check in add_remote_agent
("assert len(kv_caches_base_addr) == len(block_len_per_layer)"), and the
new block_size_ratio > 1 cases additionally hit "Heterogeneous block
sizes are not supported with host buffer".
Pin kv_buffer_device to the faked device type. The other nixl cpu_test
files are unaffected because they don't fake the platform, so their
device_type is "cpu" and use_host_buffer is forced off.
Verified against a forced CpuPlatform (which reproduces the CI failures
exactly - 17 failed - without this change): 196 passed.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Nick Hill <nickhill123@gmail.com>
NickLucche
approved these changes
Jul 27, 2026
Contributor
|
This pull request has merge conflicts that must be resolved before it can be |
Conflict in nixl base_worker._build_mamba_local: vllm-project#49988 added a descriptor-layout diagram to its docstring and an `assert block_size_ratio == 1`, while this branch had dropped the function's block_size_ratio parameter entirely (mamba state pages are indivisible, so their descriptors are ratio-independent). Kept the ratio-free signature and body, kept main's diagram, and dropped the assert -- lifting that restriction is what this branch is for. The rest of vllm-project#49988 (K/V packed into one region per layer, get_backend_aware_kv_block_len removed) merged cleanly and subsumes this branch's _build_fa_local edits; no callers of the removed helper remain.
This was referenced Aug 19, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Hybrid (mamba) models previously asserted out heterogeneous block sizes entirely — yet they are the models where P/D block sizes most readily diverge, since the mamba-padded attention block size varies with TP sharding (#41037). Lift the restriction:
Still unsupported, now rejected loudly at handshake time: head-sharded attention reads combined with a block-size mismatch, and host-buffer mode with a block-size ratio.
Note this is an alternative to #45575, which inflates the block-size on one side to make it common between P and D.
Fixes #41037.