From 84d3062d446bcec631bcd098dd4576c414793d5e Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Wed, 23 Sep 2026 14:09:28 -0700 Subject: [PATCH 01/13] [KV Connector][NIXL] Shard a GQA draft's KV by head under an MLA target NIXL builds one TPMapping per remote engine from the model-wide topology. For an MLA target that mapping treats every region as replicated (rank_offset_factor=0, total_num_kv_heads=1), including the KV regions of a GQA draft such as DFlash, which are split by head across TP. With heterogeneous TP (e.g. prefill TP1 -> decode TP8) this: - rejects the correct draft geometry at handshake ("KV cache sizes must match between P and D when replicated"), since model_replicated is forced by use_mla, while accepting an unscaled draft block; - reads draft head 0 on every decode rank once past the check, which collapses draft acceptance (output stays correct, the target verifies); - skips the hetero-TP layout guards for the draft because use_mla is set. Resolve the draft's KV head count from speculative_config when an MLA target runs a non-MLA draft, and build a second TPMapping with compute_tp_mapping(head_sharded_kv_heads=...) for the draft's SPLIT regions. Reads, handshake size checks and layout guards use it for those regions only; MLA regions are unchanged. Prefill TP > decode TP and DCP are rejected with NotImplementedError for such drafts. CSA-linear and mamba models, and MLA drafts such as MTP, keep today's behaviour. Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Stu Cao --- .../kv_connector/unit/test_nixl_connector.py | 122 ++++++++++++++++++ .../unit/test_nixl_push_connector.py | 1 + tests/v1/kv_connector/unit/test_tp_mapping.py | 36 ++++++ .../kv_connector/v1/nixl/base_worker.py | 112 ++++++++++++++-- .../kv_connector/v1/nixl/tp_mapping.py | 20 ++- 5 files changed, 278 insertions(+), 13 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 52ac746391ff..6513b7ea74b4 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1383,6 +1383,128 @@ def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init): with pytest.raises(AssertionError): worker2.add_remote_agent(bad_meta, remote_tp_size=1) + @staticmethod + def _mla_target_with_gqa_draft(tp_size: int, tp_rank: int, remote_tp_size: int): + """Worker for an MLA target (REPLICATE region) plus a GQA draft with + 8 KV heads (SPLIT region), and the prefill metadata it pairs with.""" + base_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker" + with ( + patch( + f"{base_worker}.get_tensor_model_parallel_world_size", + return_value=tp_size, + ), + patch( + f"{base_worker}.get_tensor_model_parallel_rank", return_value=tp_rank + ), + ): + worker = FakeNixlConnectorWorker( + create_vllm_config(), "engine", hand_shake_latency=0 + ) + worker.use_mla = True + worker._head_sharded_draft_kv_heads = 8 + worker.transfer_topo = TransferTopology( + tp_rank=tp_rank, + tp_size=tp_size, + block_size=worker.block_size, + engine_id=worker.engine_id, + is_mla=True, + is_mamba=False, + total_num_kv_heads=1, + attn_backends=worker.attn_backends, + tensor_shape=None, + ) + mla_len = 576 * worker.block_size + draft_len = 128 * worker.block_size + worker.slot_size_per_layer = [576, 128] + worker.block_len_per_layer = [mla_len, draft_len] + worker.block_stride_per_layer = [mla_len, draft_len] + worker._region_is_mla = [True, False] + worker.num_blocks = 1 + worker.dst_num_blocks[worker.engine_id] = worker.num_blocks + worker.src_blocks_data = np.array( + [(0, mla_len, worker.tp_rank), (0, draft_len, worker.tp_rank)], + dtype=np.uint64, + ) + worker.num_descs = len(worker.src_blocks_data) + + remote_draft_len = ( + draft_len * max(1, 8 // remote_tp_size) // max(1, 8 // tp_size) + ) + meta = NixlAgentMetadata( + engine_id=FakeNixlConnectorWorker.REMOTE_ENGINE_ID, + agent_metadata=FakeNixlWrapper.AGENT_METADATA, + kv_caches_base_addr=[0x10000, 0x20000], + device_id=0, + num_blocks=1, + block_lens=[mla_len, remote_draft_len], + block_strides=[mla_len, remote_draft_len], + kv_cache_layout=worker.kv_cache_layout, + block_size=worker.block_size, + ssm_sizes=(0, 0), + attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=1, + ) + return worker, meta, mla_len, draft_len + + @patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, + ) + @pytest.mark.parametrize( + "tp_rank,remote_tp_size,draft_slice", [(3, 1, 3), (5, 2, 1)] + ) + def test_handshake_mla_target_reads_draft_head_slice( + self, default_vllm_config, dist_init, tp_rank, remote_tp_size, draft_slice + ): + """A GQA draft under an MLA target is head-sharded even though the + target is replicated: at D_TP=8 each rank owns one draft head, so the + handshake must accept the larger prefill draft block and each rank + must read its own slice of it, while the MLA region stays at offset 0. + """ + worker, meta, mla_len, draft_len = self._mla_target_with_gqa_draft( + tp_size=8, tp_rank=tp_rank, remote_tp_size=remote_tp_size + ) + worker.add_remote_agent(meta, remote_tp_size=remote_tp_size) + + plan = worker.tp_mappings[FakeNixlConnectorWorker.REMOTE_ENGINE_ID] + assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [ + [0x10000, mla_len, 0], + [0x20000 + draft_slice * draft_len, draft_len, 0], + ] + + @patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, + ) + def test_handshake_mla_target_rejects_unsharded_draft_block( + self, default_vllm_config, dist_init + ): + """The draft region is validated against its own head ratio, so a + prefill draft block that was not scaled by it is still rejected.""" + worker, meta, mla_len, draft_len = self._mla_target_with_gqa_draft( + tp_size=8, tp_rank=3, remote_tp_size=1 + ) + meta.block_lens = [mla_len, draft_len] + meta.block_strides = [mla_len, draft_len] + with pytest.raises(AssertionError, match="SPLIT region 1"): + worker.add_remote_agent(meta, remote_tp_size=1) + + @patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, + ) + def test_handshake_mla_target_rejects_draft_with_prefill_tp_above_decode( + self, default_vllm_config, dist_init + ): + """With P_TP > D_TP the draft needs several prefill ranks while the + MLA mapping reads one, so the pairing is rejected rather than reading + a fraction of the draft heads.""" + worker, meta, _, _ = self._mla_target_with_gqa_draft( + tp_size=1, tp_rank=0, remote_tp_size=8 + ) + with pytest.raises(NotImplementedError, match="head-sharded draft"): + worker.add_remote_agent(meta, remote_tp_size=8) + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index a36c7e5b2a71..4f4cac6aaa6e 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -371,6 +371,7 @@ def fresh(cls) -> _StubWriterWorker: w._has_packed_cache = False w._reqs_to_process = set() w._reqs_to_send = {} + w.draft_tp_mappings = {} w.consumer_notification_counts_by_req = defaultdict(int) w.tp_rank = 0 w.pcp_rank = 0 diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 3ddc568a524d..5bb0805c5efe 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -38,6 +38,7 @@ def _compute_mapping( group_spec_types: tuple[type, ...] = (FullAttentionSpec,), dcp_size: int = 1, remote_dcp_size: int = 1, + **overrides, ) -> TPMapping: transfer_topology = object.__new__(TransferTopology) transfer_topology.tp_rank = tp_rank @@ -50,6 +51,7 @@ def _compute_mapping( remote_tp_size=remote_tp_size, group_spec_types=group_spec_types, remote_dcp_size=remote_dcp_size, + **overrides, ) @@ -72,6 +74,40 @@ def test_source_ranks_p_gt_d(self): assert m.all_source_ranks == (0, 1) +@pytest.mark.parametrize( + "tp_rank,tp_size,remote_tp_size,draft_kv_heads,expected_offset", + [ + (3, 8, 1, 8, 3), # one draft head per decode rank + (5, 8, 2, 8, 1), # rank 5 reads P rank 1, second head slot + (5, 16, 1, 8, 2), # D_TP > heads: two ranks share each head + ], +) +def test_head_sharded_draft_under_mla_target_gets_gqa_offset( + tp_rank, tp_size, remote_tp_size, draft_kv_heads, expected_offset +): + """An MLA target's mapping reads every head at offset 0, which is wrong + for a GQA draft's head-sharded KV. Mapping the draft as GQA with its own + head count recovers each rank's slice from the same source rank.""" + mla = _compute_mapping( + tp_rank=tp_rank, + tp_size=tp_size, + remote_tp_size=remote_tp_size, + is_mla=True, + num_kv_heads=1, + ) + draft = _compute_mapping( + tp_rank=tp_rank, + tp_size=tp_size, + remote_tp_size=remote_tp_size, + is_mla=True, + num_kv_heads=1, + head_sharded_kv_heads=draft_kv_heads, + ) + assert mla.rank_offset_factor == 0 + assert draft.rank_offset_factor == expected_offset + assert draft.all_source_ranks == mla.all_source_ranks + + @pytest.mark.parametrize( "tp_rank,tp_size,remote_tp_size,dcp_size,remote_dcp_size,expected_ranks", [ diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 7ce6e93ebb46..6e331711046e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -419,6 +419,36 @@ def _is_region_replicated(self, region_idx: int) -> bool: """ return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx] + def _resolve_head_sharded_draft_kv_heads(self) -> int | None: + """Total KV heads of a non-MLA speculative draft under an MLA target. + + The TP mapping is built from the target, so an MLA target yields a + replicated mapping (offset 0) that is wrong for a GQA draft's sharded + heads. Returns None when the draft, if any, shares the target's layout. + """ + if not self.use_mla or self._has_mamba or self._is_csa_linear: + return None + spec_config = self.vllm_config.speculative_config + draft_config = spec_config.draft_model_config if spec_config else None + if draft_config is None or draft_config.use_mla: + return None + return draft_config.get_total_num_kv_heads() + + def _is_head_sharded_draft_region(self, region_idx: int) -> bool: + return self._head_sharded_draft_kv_heads is not None and ( + not self._is_region_replicated(region_idx) + ) + + def _has_head_sharded_draft_regions(self) -> bool: + return any( + self._is_head_sharded_draft_region(i) + for i in range(len(self.block_len_per_layer)) + ) + + def _draft_kv_heads_per_rank(self, tp_size: int) -> int: + assert self._head_sharded_draft_kv_heads is not None + return max(1, self._head_sharded_draft_kv_heads // tp_size) + def _set_region_layers(self, region_layers: list[list[str]]) -> None: layer_names = [layer_name for region in region_layers for layer_name in region] assert len(layer_names) == len(set(layer_names)), ( @@ -960,6 +990,11 @@ def __init__( self._transfer_layer_group_ids = tuple[int, ...]() # Per-engine TP mappings. Generated during handshake. self.tp_mappings: dict[EngineId, TPMapping] = {} + # A non-MLA draft under an MLA target (e.g. a GQA DFlash or EAGLE-3 + # draft) keeps head-sharded KV while the target's is replicated, so its + # SPLIT regions are mapped with the draft's own KV heads. + self._head_sharded_draft_kv_heads = self._resolve_head_sharded_draft_kv_heads() + self.draft_tp_mappings: dict[EngineId, TPMapping] = {} self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config( "enforce_handshake_compat", True @@ -2116,9 +2151,7 @@ def _build_fa_remote( fa_group_idx = next( i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t) ) - # SPLIT regions read their head slice from this many remote ranks at a - # per-rank offset; REPLICATE regions read the whole block once. - split_reads = len(plan.source_ranks_per_group[fa_group_idx]) + draft_plan = self.draft_tp_mappings.get(nixl_agent_meta.engine_id) region_num_blocks = nixl_agent_meta.region_num_blocks or [ nixl_agent_meta.num_blocks ] * len(nixl_agent_meta.kv_caches_base_addr) @@ -2139,10 +2172,23 @@ def _build_fa_remote( local_block_len = remote_kv_block_len # REPLICATE reads the whole block once at offset 0; SPLIT gathers - # its head slice from `split_reads` remote ranks at a per-rank offset. - num_reads = 1 if replicated else split_reads + # its head slice from each source rank at a per-rank offset. A + # draft's SPLIT regions under an MLA target use the draft mapping. + region_plan = ( + draft_plan + if draft_plan is not None + and self._is_head_sharded_draft_region(local_region) + else plan + ) + num_reads = ( + 1 + if replicated + else len(region_plan.source_ranks_per_group[fa_group_idx]) + ) rank_offset = ( - 0 if replicated else plan.rank_offset_factor * remote_kv_block_len + 0 + if replicated + else region_plan.rank_offset_factor * remote_kv_block_len ) local_block_len = local_block_len // num_reads @@ -2360,6 +2406,14 @@ def add_remote_agent( group_spec_types=self._group_spec_types, remote_dcp_size=remote_dcp_size, ) + if self._has_head_sharded_draft_regions(): + self.draft_tp_mappings[engine_id] = compute_tp_mapping( + transfer_topology=transfer_topo, + remote_tp_size=remote_tp_size, + group_spec_types=self._group_spec_types, + remote_dcp_size=remote_dcp_size, + head_sharded_kv_heads=self._head_sharded_draft_kv_heads, + ) remote_agent_name = self.nixl_wrapper.add_remote_agent( nixl_agent_meta.agent_metadata @@ -2549,6 +2603,19 @@ def _validate_remote_agent_handshake( assert not ( tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id) ) + # A head-sharded draft under an MLA target reads from the target's + # single MLA source rank, so it only transfers when that rank also + # holds every draft head the local rank needs. + has_draft_regions = self._has_head_sharded_draft_regions() + if has_draft_regions and ( + tp_ratio < 0 or self.dcp_size > 1 or remote_dcp_size > 1 + ): + raise NotImplementedError( + "NIXL cannot transfer a head-sharded draft's KV under an MLA " + "target when prefill TP exceeds decode TP or with DCP: " + f"local TP={self.transfer_topo.tp_size}, remote TP={remote_tp_size}, " + f"local DCP={self.dcp_size}, remote DCP={remote_dcp_size}." + ) remote_physical_per_logical = ( nixl_agent_meta.physical_blocks_per_logical_kv_block @@ -2641,6 +2708,23 @@ def _validate_remote_agent_handshake( "Heterogeneous TP head-dimension splitting requires contiguous heads. " "Use a block-contiguous layout (e.g. LBHNC) on the prefill side." ) + # The MLA exemption above does not cover a head-sharded draft, which + # slices the prefill draft block by head. + if ( + has_draft_regions + and self._draft_kv_heads_per_rank(remote_tp_size) + != self._draft_kv_heads_per_rank(self.transfer_topo.tp_size) + and ( + nixl_agent_meta.kv_cache_layout != kv_cache_layout + or not KVCacheLayout[kv_cache_layout].is_block_contiguous + ) + ): + raise RuntimeError( + "Head-sharded draft KV under an MLA target needs the same " + "block-contiguous kv_cache_layout (e.g. LBHNC) on prefill and " + f"decode, got local={kv_cache_layout}, " + f"remote={nixl_agent_meta.kv_cache_layout}." + ) # Per-region block_len validation enforcing the P/D invariant. # REPLICATE regions (MLA, or a whole-model MLA / replicated-KV transfer) @@ -2674,12 +2758,21 @@ def _validate_remote_agent_handshake( remote_engine_id ) total_kv_heads = self.transfer_topo.total_num_kv_heads - local_heads = self.transfer_topo.local_physical_heads - remote_heads = max(1, total_kv_heads // remote_tp_size) for i, region_idx in enumerate(local_regions): local_len = self.block_len_per_layer[region_idx] - replicated = model_replicated or self._is_region_replicated(region_idx) remote_len = nixl_agent_meta.block_lens[i] + if self._is_head_sharded_draft_region(region_idx): + replicated = False + local_heads = self._draft_kv_heads_per_rank( + self.transfer_topo.tp_size + ) + remote_heads = self._draft_kv_heads_per_rank(remote_tp_size) + else: + replicated = model_replicated or self._is_region_replicated( + region_idx + ) + local_heads = self.transfer_topo.local_physical_heads + remote_heads = max(1, total_kv_heads // remote_tp_size) if replicated: assert local_len // block_size_ratio == remote_len, ( "KV cache sizes must match between P and D when " @@ -3824,6 +3917,7 @@ def _cleanup_remote_engine( self.dst_uses_region_group_mapping.pop(engine_id, None) self.dst_region_mem_types.pop(engine_id, None) self.tp_mappings.pop(engine_id, None) + self.draft_tp_mappings.pop(engine_id, None) if self.transfer_topo is not None: self.transfer_topo.unregister_remote_engine(engine_id) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py index 150fdabe354f..5a022f32b8f2 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py @@ -72,6 +72,7 @@ def compute_tp_mapping( remote_tp_size: int, group_spec_types: tuple[type[KVCacheSpec], ...], remote_dcp_size: int = 1, + head_sharded_kv_heads: int | None = None, ) -> TPMapping: """Build the complete local-to-remote TP mapping. @@ -80,13 +81,24 @@ def compute_tp_mapping( DCP support is scoped to MLA only, with a side is either fully replicated or fully sharded. DCP-branch reuses the same rank set used at handshake selection. + + Args: + transfer_topology: The local rank's view of the transfer. + remote_tp_size: TP size of the remote engine. + group_spec_types: KV cache spec type of each KV cache group. + remote_dcp_size: DCP size of the remote engine. + head_sharded_kv_heads: If set, map head-sharded KV with this many total + KV heads instead of the topology's, e.g. a GQA draft whose regions + are sharded under an MLA target that is replicated. + """ tp_rank = transfer_topology.tp_rank tp_size = transfer_topology.tp_size - total_num_kv_heads = transfer_topology.total_num_kv_heads + is_mla = transfer_topology.is_mla and head_sharded_kv_heads is None + total_num_kv_heads = head_sharded_kv_heads or transfer_topology.total_num_kv_heads # --- Attention source ranks --- - if transfer_topology.is_mla or tp_size >= remote_tp_size: - if transfer_topology.is_mla and remote_dcp_size > 1: + if is_mla or tp_size >= remote_tp_size: + if is_mla and remote_dcp_size > 1: attn_ranks = transfer_topology.dcp_source_ranks( remote_tp_size, remote_dcp_size ) @@ -137,7 +149,7 @@ def compute_tp_mapping( } # --- Rank offset factor --- - if transfer_topology.is_mla or tp_size <= remote_tp_size: + if is_mla or tp_size <= remote_tp_size: # We don't index into remote for reading, no offset needed. rank_offset_factor = 0 elif tp_size > total_num_kv_heads: From 20878d1d164f49b3660aeb7304a550d01435b07d Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Wed, 23 Sep 2026 20:48:25 -0700 Subject: [PATCH 02/13] [KV Connector][NIXL] Test the draft head slice under a block-outermost layout With BLHNC one page holds every layer of a block, so a region's stride is the page rather than its block_len. Cover that the draft head offset still lands inside the draft layer of each page across several blocks. Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Stu Cao --- .../kv_connector/unit/test_nixl_connector.py | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 6513b7ea74b4..d12e17ce68e1 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1472,6 +1472,43 @@ def test_handshake_mla_target_reads_draft_head_slice( [0x20000 + draft_slice * draft_len, draft_len, 0], ] + @patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, + ) + def test_handshake_mla_target_reads_draft_head_slice_block_outermost( + self, default_vllm_config, dist_init + ): + """With a block-outermost layout (BLHNC) one page holds every layer of a + block, so the region stride is the page, not the region's block_len. The + draft head slice must still land inside the draft layer of each page.""" + tp_rank = 3 + worker, meta, mla_len, draft_len = self._mla_target_with_gqa_draft( + tp_size=8, tp_rank=tp_rank, remote_tp_size=1 + ) + remote_draft_len = meta.block_lens[1] + local_page = mla_len + draft_len + remote_page = mla_len + remote_draft_len + remote_page_base = 0x100000 + num_blocks = 3 + + worker.kv_cache_layout = "BLHNC" + worker.block_stride_per_layer = [local_page, local_page] + worker.num_blocks = num_blocks + worker.dst_num_blocks[worker.engine_id] = num_blocks + meta.kv_cache_layout = "BLHNC" + meta.num_blocks = num_blocks + meta.kv_caches_base_addr = [remote_page_base, remote_page_base + mla_len] + meta.block_strides = [remote_page, remote_page] + worker.add_remote_agent(meta, remote_tp_size=1) + + plan = worker.tp_mappings[FakeNixlConnectorWorker.REMOTE_ENGINE_ID] + pages = [remote_page_base + b * remote_page for b in range(num_blocks)] + assert worker._build_fa_remote(plan, meta, block_size_ratio=1).tolist() == [ + *([page, mla_len, 0] for page in pages), + *([page + mla_len + tp_rank * draft_len, draft_len, 0] for page in pages), + ] + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, From c25bf02074cc3fd896e44cf2dd45f80d962378a2 Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Wed, 23 Sep 2026 21:02:36 -0700 Subject: [PATCH 03/13] [KV Connector][NIXL] Pass head_sharded_kv_heads explicitly in the test helper Replace the **overrides passthrough in _compute_mapping with an explicit keyword, matching the helper's other parameters. Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Stu Cao --- tests/v1/kv_connector/unit/test_tp_mapping.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 5bb0805c5efe..34321b0731c5 100644 --- a/tests/v1/kv_connector/unit/test_tp_mapping.py +++ b/tests/v1/kv_connector/unit/test_tp_mapping.py @@ -38,7 +38,7 @@ def _compute_mapping( group_spec_types: tuple[type, ...] = (FullAttentionSpec,), dcp_size: int = 1, remote_dcp_size: int = 1, - **overrides, + head_sharded_kv_heads: int | None = None, ) -> TPMapping: transfer_topology = object.__new__(TransferTopology) transfer_topology.tp_rank = tp_rank @@ -51,7 +51,7 @@ def _compute_mapping( remote_tp_size=remote_tp_size, group_spec_types=group_spec_types, remote_dcp_size=remote_dcp_size, - **overrides, + head_sharded_kv_heads=head_sharded_kv_heads, ) From 54db0e4c36ef2c52e16cdb860f66aa12b71c51ee Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Wed, 23 Sep 2026 21:25:41 -0700 Subject: [PATCH 04/13] [KV Connector][NIXL] Simplify the draft read path and test the resolver - Branch once on replicated in _build_fa_remote instead of two ternaries, and index the draft mapping directly: it always exists for draft regions. - Replace the compute_tp_mapping Args block with one sentence. - Trim comments that restated the code or the error message. - Test _resolve_head_sharded_draft_kv_heads, which gates the fix. Co-authored-by: Claude Opus 5.5 (1M context) Signed-off-by: Stu Cao --- .../kv_connector/unit/test_nixl_connector.py | 26 +++++++++++ .../unit/test_nixl_push_connector.py | 1 + .../kv_connector/v1/nixl/base_worker.py | 44 ++++++------------- .../kv_connector/v1/nixl/tp_mapping.py | 11 +---- 4 files changed, 43 insertions(+), 39 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index d12e17ce68e1..99f27cf59e51 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1383,6 +1383,32 @@ def test_handshake_mixed_fa_mla_hetero_tp(self, default_vllm_config, dist_init): with pytest.raises(AssertionError): worker2.add_remote_agent(bad_meta, remote_tp_size=1) + @pytest.mark.parametrize( + "target_mla,draft_mla,has_draft,expected", + [ + (True, False, True, 8), # GQA draft under an MLA target + (True, True, True, None), # MLA draft (e.g. MTP) is replicated too + (True, False, False, None), # no speculative decoding + (False, False, True, None), # GQA target maps the draft already + ], + ) + def test_resolve_head_sharded_draft_kv_heads( + self, target_mla, draft_mla, has_draft, expected + ): + draft = SimpleNamespace(use_mla=draft_mla, get_total_num_kv_heads=lambda: 8) + worker = SimpleNamespace( + use_mla=target_mla, + _has_mamba=False, + _is_csa_linear=False, + vllm_config=SimpleNamespace( + speculative_config=( + SimpleNamespace(draft_model_config=draft) if has_draft else None + ) + ), + ) + resolve = NixlConnectorWorker._resolve_head_sharded_draft_kv_heads + assert resolve(worker) == expected + @staticmethod def _mla_target_with_gqa_draft(tp_size: int, tp_rank: int, remote_tp_size: int): """Worker for an MLA target (REPLICATE region) plus a GQA draft with diff --git a/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index 4f4cac6aaa6e..0ea58892f1d2 100644 --- a/tests/v1/kv_connector/unit/test_nixl_push_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_push_connector.py @@ -372,6 +372,7 @@ def fresh(cls) -> _StubWriterWorker: w._reqs_to_process = set() w._reqs_to_send = {} w.draft_tp_mappings = {} + w._head_sharded_draft_kv_heads = None w.consumer_notification_counts_by_req = defaultdict(int) w.tp_rank = 0 w.pcp_rank = 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 6e331711046e..26d09ded1c36 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -420,12 +420,9 @@ def _is_region_replicated(self, region_idx: int) -> bool: return region_idx < len(self._region_is_mla) and self._region_is_mla[region_idx] def _resolve_head_sharded_draft_kv_heads(self) -> int | None: - """Total KV heads of a non-MLA speculative draft under an MLA target. + """Total KV heads of a non-MLA draft under an MLA target, else None. - The TP mapping is built from the target, so an MLA target yields a - replicated mapping (offset 0) that is wrong for a GQA draft's sharded - heads. Returns None when the draft, if any, shares the target's layout. - """ + Such a draft's KV is head-sharded while the target's is replicated.""" if not self.use_mla or self._has_mamba or self._is_csa_linear: return None spec_config = self.vllm_config.speculative_config @@ -990,9 +987,7 @@ def __init__( self._transfer_layer_group_ids = tuple[int, ...]() # Per-engine TP mappings. Generated during handshake. self.tp_mappings: dict[EngineId, TPMapping] = {} - # A non-MLA draft under an MLA target (e.g. a GQA DFlash or EAGLE-3 - # draft) keeps head-sharded KV while the target's is replicated, so its - # SPLIT regions are mapped with the draft's own KV heads. + # GQA draft under an MLA target: its SPLIT regions use a draft TP mapping. self._head_sharded_draft_kv_heads = self._resolve_head_sharded_draft_kv_heads() self.draft_tp_mappings: dict[EngineId, TPMapping] = {} @@ -2151,7 +2146,6 @@ def _build_fa_remote( fa_group_idx = next( i for i, t in enumerate(self._group_spec_types) if _is_attention_spec(t) ) - draft_plan = self.draft_tp_mappings.get(nixl_agent_meta.engine_id) region_num_blocks = nixl_agent_meta.region_num_blocks or [ nixl_agent_meta.num_blocks ] * len(nixl_agent_meta.kv_caches_base_addr) @@ -2172,24 +2166,17 @@ def _build_fa_remote( local_block_len = remote_kv_block_len # REPLICATE reads the whole block once at offset 0; SPLIT gathers - # its head slice from each source rank at a per-rank offset. A - # draft's SPLIT regions under an MLA target use the draft mapping. - region_plan = ( - draft_plan - if draft_plan is not None - and self._is_head_sharded_draft_region(local_region) - else plan - ) - num_reads = ( - 1 - if replicated - else len(region_plan.source_ranks_per_group[fa_group_idx]) - ) - rank_offset = ( - 0 - if replicated - else region_plan.rank_offset_factor * remote_kv_block_len - ) + # its head slice from each source rank at a per-rank offset. + if replicated: + num_reads, rank_offset = 1, 0 + else: + region_plan = ( + self.draft_tp_mappings[nixl_agent_meta.engine_id] + if self._is_head_sharded_draft_region(local_region) + else plan + ) + num_reads = len(region_plan.source_ranks_per_group[fa_group_idx]) + rank_offset = region_plan.rank_offset_factor * remote_kv_block_len local_block_len = local_block_len // num_reads block_arange = np.arange(region_num_blocks[i], dtype=np.uint64) @@ -2603,9 +2590,6 @@ def _validate_remote_agent_handshake( assert not ( tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id) ) - # A head-sharded draft under an MLA target reads from the target's - # single MLA source rank, so it only transfers when that rank also - # holds every draft head the local rank needs. has_draft_regions = self._has_head_sharded_draft_regions() if has_draft_regions and ( tp_ratio < 0 or self.dcp_size > 1 or remote_dcp_size > 1 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py index 5a022f32b8f2..a0bd8c97cb15 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/tp_mapping.py @@ -82,15 +82,8 @@ def compute_tp_mapping( DCP support is scoped to MLA only, with a side is either fully replicated or fully sharded. DCP-branch reuses the same rank set used at handshake selection. - Args: - transfer_topology: The local rank's view of the transfer. - remote_tp_size: TP size of the remote engine. - group_spec_types: KV cache spec type of each KV cache group. - remote_dcp_size: DCP size of the remote engine. - head_sharded_kv_heads: If set, map head-sharded KV with this many total - KV heads instead of the topology's, e.g. a GQA draft whose regions - are sharded under an MLA target that is replicated. - + ``head_sharded_kv_heads`` maps head-sharded KV with that many total KV heads + instead of the topology's, e.g. a GQA draft under an MLA target. """ tp_rank = transfer_topology.tp_rank tp_size = transfer_topology.tp_size From 4e14299b72c53d63c824757fc7e6e6c99f1e5f71 Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Tue, 29 Sep 2026 16:08:56 -0700 Subject: [PATCH 05/13] [Bugfix] Address review: role-correct hetero-TP guard, region identity check Two fixes for review comments on #58471 (not yet posted): 1. The new "head-sharded draft under prefill TP > decode TP" guard used tp_ratio's sign directly, but tp_ratio is purely local-vs-remote and which side is local flips between pull (local=decode) and push (local=prefill, since prefill is the side that needs decode's write-target addresses). Reusing the pull-mode sign check in push mode would reject the common P_TP < D_TP topology instead of the one it's meant to catch. Derive the condition from _TRANSFER_MODE instead. 2. _is_head_sharded_draft_region inferred "SPLIT under a fully-MLA target" implies draft purely by exclusion. Track each SPLIT region's own KV head count at registration (mirroring _region_is_mla) and assert it's consistent with the resolved draft's head count, tying the classification to an actual per-layer model property instead of the absence of one. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Stu Cao --- .../kv_connector/unit/test_nixl_connector.py | 1 + .../kv_connector/v1/nixl/base_worker.py | 68 ++++++++++++++++--- 2 files changed, 58 insertions(+), 11 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 99f27cf59e51..88271796ec99 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1445,6 +1445,7 @@ def _mla_target_with_gqa_draft(tp_size: int, tp_rank: int, remote_tp_size: int): worker.block_len_per_layer = [mla_len, draft_len] worker.block_stride_per_layer = [mla_len, draft_len] worker._region_is_mla = [True, False] + worker._region_num_kv_heads = [None, max(1, 8 // tp_size)] worker.num_blocks = 1 worker.dst_num_blocks[worker.engine_id] = worker.num_blocks worker.src_blocks_data = np.array( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 26d09ded1c36..a5cb3e1c3d3d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -432,9 +432,29 @@ def _resolve_head_sharded_draft_kv_heads(self) -> int | None: return draft_config.get_total_num_kv_heads() def _is_head_sharded_draft_region(self, region_idx: int) -> bool: - return self._head_sharded_draft_kv_heads is not None and ( - not self._is_region_replicated(region_idx) + if self._head_sharded_draft_kv_heads is None or self._is_region_replicated( + region_idx + ): + return False + # Verify against the region's own recorded KV head count instead of + # inferring "SPLIT under a fully-MLA target" implies draft purely by + # exclusion. `_resolve_head_sharded_draft_kv_heads` already returns + # None for the other constructs that can produce a SPLIT region here + # (mamba, CSA-linear), so every SPLIT region left should genuinely be + # this draft -- this assertion is what actually confirms that, tied to + # a real per-layer model property rather than the absence of one. + region_kv_heads = self._region_num_kv_heads[region_idx] + assert ( + region_kv_heads is not None + and region_kv_heads <= self._head_sharded_draft_kv_heads + ), ( + f"Region {region_idx} is SPLIT under a fully-MLA target but its " + f"recorded KV head count ({region_kv_heads}) is not consistent " + "with the resolved head-sharded draft " + f"({self._head_sharded_draft_kv_heads} total KV heads); expected " + "every SPLIT region here to belong to that draft." ) + return True def _has_head_sharded_draft_regions(self) -> bool: return any( @@ -956,6 +976,12 @@ def __init__( # (MLA), False -> SPLIT (head-sharded full-attn). Mixed only for models # combining both (e.g. GQA main + MLA Eagle-3 draft). self._region_is_mla = list[bool]() + # Per-region KV head count for non-MLA (SPLIT) regions, 1:1 with + # block_len_per_layer; None where not tracked (MLA regions, or a + # layer_spec without num_kv_heads). Lets a head-sharded draft region be + # verified by an actual model property instead of inferred purely from + # "not MLA" -- see `_is_head_sharded_draft_region`. + self._region_num_kv_heads = list[int | None]() self._ssm_region_indices = list[int]() # Regions holding a scratch cache (the CSA compressor circular buffer). # Tracked explicitly because a scratch page shares its address with the @@ -1592,6 +1618,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): is_mla_region = isinstance( layer_spec, (MLAAttentionSpec, SlidingWindowMLASpec) ) + region_num_kv_heads = ( + None if is_mla_region else getattr(layer_spec, "num_kv_heads", None) + ) logger.debug( "Registering layer %s with cache shape: %s", layer_name, cache.shape ) @@ -1757,6 +1786,8 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): region_index = seen_base_addresses.index(base_addr) assert region_mem_types[region_index] == mem_type self._region_is_mla[region_index] |= is_mla_region + if not is_mla_region: + self._region_num_kv_heads[region_index] = region_num_kv_heads if is_mla_region: self.block_len_per_layer[region_index] = block_len self.block_stride_per_layer[region_index] = block_stride @@ -1772,6 +1803,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): self.region_names.append(layer_name) self.region_num_blocks.append(num_blocks) self._region_is_mla.append(is_mla_region) + self._region_num_kv_heads.append(region_num_kv_heads) region_mem_types.append(mem_type) if track_region_layers: @@ -1823,6 +1855,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): == len(self.block_stride_per_layer) == len(self.region_group_ids) == len(self.region_names) + == len(self._region_num_kv_heads) == len(self.region_num_blocks) ) # Descriptor ids must be region-ordered, matching the remote side. @@ -2591,15 +2624,28 @@ def _validate_remote_agent_handshake( tp_ratio < 0 and self.transfer_topo.is_kv_replicated(remote_engine_id) ) has_draft_regions = self._has_head_sharded_draft_regions() - if has_draft_regions and ( - tp_ratio < 0 or self.dcp_size > 1 or remote_dcp_size > 1 - ): - raise NotImplementedError( - "NIXL cannot transfer a head-sharded draft's KV under an MLA " - "target when prefill TP exceeds decode TP or with DCP: " - f"local TP={self.transfer_topo.tp_size}, remote TP={remote_tp_size}, " - f"local DCP={self.dcp_size}, remote DCP={remote_dcp_size}." - ) + if has_draft_regions: + # tp_ratio is local-vs-remote, and which side is local flips between + # pull (local=decode, since decode calls the handshake to read from + # prefill) and push (local=prefill, since prefill calls it to learn + # where to write into decode). The condition this guards against is + # always "prefill TP > decode TP", not "tp_ratio < 0" itself, so it + # has to be derived per direction rather than read off the sign + # directly -- reusing the pull-mode sign check here would reject + # ordinary push topologies (P_TP < D_TP is the common case) instead + # of the one this is meant to catch. + if self._TRANSFER_MODE == "push": + prefill_tp_above_decode = tp_ratio > 1 + else: + prefill_tp_above_decode = tp_ratio < 0 + if prefill_tp_above_decode or self.dcp_size > 1 or remote_dcp_size > 1: + raise NotImplementedError( + "NIXL cannot transfer a head-sharded draft's KV under an MLA " + "target when prefill TP exceeds decode TP or with DCP: " + f"local TP={self.transfer_topo.tp_size}, remote TP={remote_tp_size}, " + f"local DCP={self.dcp_size}, remote DCP={remote_dcp_size}, " + f"transfer_mode={self._TRANSFER_MODE}." + ) remote_physical_per_logical = ( nixl_agent_meta.physical_blocks_per_logical_kv_block From 02cc6e7596a6073a818fce0e612e401dc24aad63 Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Tue, 29 Sep 2026 17:11:21 -0700 Subject: [PATCH 06/13] Fix test fixture missing _region_num_kv_heads after review-fix commit _is_head_sharded_draft_region now reads worker._region_num_kv_heads, so the object.__new__-based worker fixture in test_overlaid_transfer_groups_share_region_geometry needs the attribute initialized alongside _region_is_mla or registration raises AttributeError. Co-Authored-By: Claude Sonnet 5 Signed-off-by: Stu Cao --- tests/v1/kv_connector/unit/test_nixl_desc_geometry.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 9f6bc56deeb5..341d470a6e1e 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -413,6 +413,7 @@ def _register_overlaid_mla_worker( worker._transfer_layer_region_indices = () worker._transfer_layer_group_ids = () worker._region_is_mla = [] + worker._region_num_kv_heads = [] worker.block_len_per_layer = [] worker.block_stride_per_layer = [] worker.device_id = 0 From 6baba61519808cd8b3c6eb090f37d2cd205843fa Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Thu, 1 Oct 2026 22:58:23 +0000 Subject: [PATCH 07/13] [KV Connector][NIXL] Verify draft layer identity and reject unsupported push Signed-off-by: Stu Cao --- .../kv_connector/unit/test_nixl_connector.py | 29 +++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 31 +++++++------------ 2 files changed, 41 insertions(+), 19 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 88271796ec99..2e678f24ae28 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1446,6 +1446,10 @@ def _mla_target_with_gqa_draft(tp_size: int, tp_rank: int, remote_tp_size: int): worker.block_stride_per_layer = [mla_len, draft_len] worker._region_is_mla = [True, False] worker._region_num_kv_heads = [None, max(1, 8 // tp_size)] + worker.region_names = [ + "model.layers.0.self_attn", + "draft_model.model.layers.0.self_attn", + ] worker.num_blocks = 1 worker.dst_num_blocks[worker.engine_id] = worker.num_blocks worker.src_blocks_data = np.array( @@ -1473,6 +1477,31 @@ def _mla_target_with_gqa_draft(tp_size: int, tp_rank: int, remote_tp_size: int): ) return worker, meta, mla_len, draft_len + @patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, + ) + def test_head_sharded_draft_region_checks_layer_identity( + self, default_vllm_config, dist_init + ): + worker, _, _, _ = self._mla_target_with_gqa_draft(4, 0, 1) + assert worker._is_head_sharded_draft_region(1) + worker.region_names[1] = "model.layers.1.self_attn" + worker._region_num_kv_heads[1] = 1 + assert not worker._is_head_sharded_draft_region(1) + + @patch( + "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", + FakeNixlWrapper, + ) + def test_handshake_mla_target_rejects_push_draft( + self, default_vllm_config, dist_init + ): + worker, meta, _, _ = self._mla_target_with_gqa_draft(1, 0, 4) + worker._TRANSFER_MODE = "push" + with pytest.raises(NotImplementedError, match="push mode"): + worker.add_remote_agent(meta, remote_tp_size=4) + @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", FakeNixlWrapper, diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index a5cb3e1c3d3d..22ec466783ef 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -436,13 +436,12 @@ def _is_head_sharded_draft_region(self, region_idx: int) -> bool: region_idx ): return False - # Verify against the region's own recorded KV head count instead of - # inferring "SPLIT under a fully-MLA target" implies draft purely by - # exclusion. `_resolve_head_sharded_draft_kv_heads` already returns - # None for the other constructs that can produce a SPLIT region here - # (mamba, CSA-linear), so every SPLIT region left should genuinely be - # this draft -- this assertion is what actually confirms that, tied to - # a real per-layer model property rather than the absence of one. + # Draft models are loaded with the "draft_model" layer prefix. + # An MLA target may have its own SPLIT attention regions; those retain + # the target mapping rather than using the draft's KV head count. + region_name = self.region_names[region_idx] + if not region_name.startswith("draft_model."): + return False region_kv_heads = self._region_num_kv_heads[region_idx] assert ( region_kv_heads is not None @@ -2625,19 +2624,13 @@ def _validate_remote_agent_handshake( ) has_draft_regions = self._has_head_sharded_draft_regions() if has_draft_regions: - # tp_ratio is local-vs-remote, and which side is local flips between - # pull (local=decode, since decode calls the handshake to read from - # prefill) and push (local=prefill, since prefill calls it to learn - # where to write into decode). The condition this guards against is - # always "prefill TP > decode TP", not "tp_ratio < 0" itself, so it - # has to be derived per direction rather than read off the sign - # directly -- reusing the pull-mode sign check here would reject - # ordinary push topologies (P_TP < D_TP is the common case) instead - # of the one this is meant to catch. if self._TRANSFER_MODE == "push": - prefill_tp_above_decode = tp_ratio > 1 - else: - prefill_tp_above_decode = tp_ratio < 0 + raise NotImplementedError( + "NIXL push mode does not support a head-sharded draft's KV " + "under an MLA target. Use pull mode." + ) + # In pull mode local is decode and remote is prefill. + prefill_tp_above_decode = tp_ratio < 0 if prefill_tp_above_decode or self.dcp_size > 1 or remote_dcp_size > 1: raise NotImplementedError( "NIXL cannot transfer a head-sharded draft's KV under an MLA " From 92f41429a8fe074d44fd1f1e7e35d80627033ef4 Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Fri, 2 Oct 2026 00:53:45 +0000 Subject: [PATCH 08/13] [KV Connector][NIXL] Recognize numbered DFlash draft layers Signed-off-by: Stu Cao --- .../kv_connector/unit/test_nixl_connector.py | 12 ++++++++- .../kv_connector/v1/nixl/base_worker.py | 27 ++++++++++++------- 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 2e678f24ae28..1cbacc2b3f12 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -1486,9 +1486,19 @@ def test_head_sharded_draft_region_checks_layer_identity( ): worker, _, _, _ = self._mla_target_with_gqa_draft(4, 0, 1) assert worker._is_head_sharded_draft_region(1) - worker.region_names[1] = "model.layers.1.self_attn" + worker.region_names[1] = "model.layers.77.self_attn.attn" worker._region_num_kv_heads[1] = 1 + worker.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(get_total_num_hidden_layers=lambda: 78), + speculative_config=SimpleNamespace( + draft_model_config=SimpleNamespace( + get_total_num_hidden_layers=lambda: 6 + ) + ), + ) assert not worker._is_head_sharded_draft_region(1) + worker.region_names[1] = "model.layers.78.self_attn.attn" + assert worker._is_head_sharded_draft_region(1) @patch( "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker.NixlWrapper", diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 22ec466783ef..438d6819dfe1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -436,22 +436,31 @@ def _is_head_sharded_draft_region(self, region_idx: int) -> bool: region_idx ): return False - # Draft models are loaded with the "draft_model" layer prefix. - # An MLA target may have its own SPLIT attention regions; those retain - # the target mapping rather than using the draft's KV head count. + # Draft layers may use a draft_model prefix or be numbered after the + # target's layers (e.g. DFlash's model.layers.78 after 78 target layers). + # A target SPLIT region must keep the target mapping. region_name = self.region_names[region_idx] if not region_name.startswith("draft_model."): - return False + spec_config = self.vllm_config.speculative_config + draft_config = spec_config.draft_model_config if spec_config else None + if draft_config is None: + return False + match = re.search(r"(?:^|\.)layers\.(\d+)(?:\.|$)", region_name) + if match is None: + return False + layer_idx = int(match.group(1)) + target_layers = self.vllm_config.model_config.get_total_num_hidden_layers() + draft_layers = draft_config.get_total_num_hidden_layers() + if not target_layers <= layer_idx < target_layers + draft_layers: + return False region_kv_heads = self._region_num_kv_heads[region_idx] assert ( region_kv_heads is not None and region_kv_heads <= self._head_sharded_draft_kv_heads ), ( - f"Region {region_idx} is SPLIT under a fully-MLA target but its " - f"recorded KV head count ({region_kv_heads}) is not consistent " - "with the resolved head-sharded draft " - f"({self._head_sharded_draft_kv_heads} total KV heads); expected " - "every SPLIT region here to belong to that draft." + f"Draft region {region_idx} ({region_name!r}) has " + f"{region_kv_heads} local KV heads, inconsistent with the " + f"draft's {self._head_sharded_draft_kv_heads} total KV heads." ) return True From 415c9704c88f29db8b3714b5621edff3fac94732 Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Fri, 9 Oct 2026 09:38:46 -0700 Subject: [PATCH 09/13] [KV Connector][NIXL] Wrap a long error message line ruff E501. Signed-off-by: Stu Cao --- .../kv_transfer/kv_connector/v1/nixl/base_worker.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 438d6819dfe1..280088e886e5 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -2644,7 +2644,8 @@ def _validate_remote_agent_handshake( raise NotImplementedError( "NIXL cannot transfer a head-sharded draft's KV under an MLA " "target when prefill TP exceeds decode TP or with DCP: " - f"local TP={self.transfer_topo.tp_size}, remote TP={remote_tp_size}, " + f"local TP={self.transfer_topo.tp_size}, " + f"remote TP={remote_tp_size}, " f"local DCP={self.dcp_size}, remote DCP={remote_dcp_size}, " f"transfer_mode={self._TRANSFER_MODE}." ) From 21b768ab681e3368069cdd94488f493fde98ce54 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=91=9B=E6=98=8A=E5=AE=87?= Date: Sat, 19 Sep 2026 13:00:28 +0800 Subject: [PATCH 10/13] [Bugfix][NIXL] Dedup transfer regions by (base_addr, block_len) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Region registration deduplicated NIXL transfer regions by base address alone, and the collapsed branch only refreshed block geometry for MLA regions. When a hybrid-attention HMA layer with a different block length aliases another layer's base address (DeepSeek-V4.1-Flash compressor state_cache over the swa_cache it serves, and its indexer k_cache), the collapsed layer silently inherited the first region's block length, so its transferred state was truncated and misaligned — producing garbled output in P/D disaggregated serving while identical-geometry models ran fine. Deduplicate by (base_addr, block_len) so regions with distinct transfer geometry stay separate; identical-geometry aliases still fold as before. Co-authored-by: ZCode Co-authored-by: chaterm Signed-off-by: 葛昊宇 --- .../unit/test_nixl_desc_geometry.py | 136 ++++++++++++++++-- .../kv_connector/v1/nixl/base_worker.py | 17 ++- 2 files changed, 140 insertions(+), 13 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 341d470a6e1e..91f1a2f3b775 100644 --- a/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py +++ b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py @@ -1561,17 +1561,19 @@ def test_csa_linear_remote_ple_is_copied_whole(): @pytest.mark.cpu_test def test_csa_linear_ple_page_wider_than_the_shared_region_page(): - """The PLE page is discovered in the region of the main KV page registered - first (both start at byte 0 of the block); its descriptors must still cover - the whole PLE page, not the region's block_len.""" + """The PLE page shares its base address with the main KV page (both start + at byte 0 of the block) but has a wider page, so it registers as its own + region; its descriptors must still cover the whole PLE page.""" worker = _make_csa_linear_ple_worker(ple_page_size=384) - assert worker._ple_region_index == 0 + assert worker._ple_region_index == 1 assert worker.block_len_per_layer[0] == 256 assert worker.block_stride_per_layer[0] == 384 + assert worker.block_len_per_layer[1] == 384 assert worker._ple_block_len == 384 - bases = [0x10000, 0x20000, 0x30000, 0x40000] + # Region 1 (the PLE page) aliases region 0's base with a wider block_len. + bases = [0x10000, 0x10000, 0x20000, 0x30000, 0x40000] mamba = worker._build_mamba_local(bases) assert mamba[-2:, 0].tolist() == [0x10000, 0x10000 + 384] assert mamba[-2:, 1].tolist() == [384, 384] @@ -1625,7 +1627,9 @@ def group(specs): return KVCacheGroupSpec(list(specs), uniform) # Both rings overlay the second paged region, so the scratch group must - # address only that region while the paged group spans both. + # address only that region while the paged group spans both. The rings' + # distinct page geometry keeps them in their own transfer region rather + # than folding into the paged region they alias. tensor_regions = (("paged.0",), ("paged.1", "ring.0", "ring.1")) region_size, page_size = 512, 256 kv_cache_config = KVCacheConfig( @@ -1680,7 +1684,7 @@ def test_ring_scratch_without_mamba_registers_and_addresses_its_own_regions(): worker = _make_ring_worker() assert worker._ple_group_index is None - assert worker._scratch_region_indices == [1] + assert worker._scratch_region_indices == [2] desc_ids = worker._compute_desc_ids( block_ids=([1], [0]), @@ -1688,5 +1692,119 @@ def test_ring_scratch_without_mamba_registers_and_addresses_its_own_regions(): block_size_ratio=None, physical_blocks_per_logical=1, ) - # Paged block 1 in both regions; ring block 0 only in its scratch region. - assert sorted(desc_ids.tolist()) == sorted([0 * 2 + 1, 1 * 2 + 1, 1 * 2 + 0]) + # Paged block 1 in both paged regions; ring block 0 only in its own + # scratch region, which no longer folds into the overlaid paged region. + assert sorted(desc_ids.tolist()) == sorted([0 * 2 + 1, 1 * 2 + 1, 2 * 2 + 0]) + + +def _make_shared_base_ring_worker(): + """MLA swa group plus a compressor ring group aliasing the same + block-major allocation with a smaller page: the DeepSeek-V4.1-Flash + HMA shape where a compressor state cache shares its storage base with + the swa cache it serves.""" + from unittest.mock import MagicMock + + from vllm.config import set_current_vllm_config + from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw + from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import ( + NixlConnectorWorker, + ) + from vllm.v1.kv_cache_interface import ( + CircularBufferSpec, + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + MLAAttentionSpec, + UniformTypeKVCacheSpecs, + ) + + num_blocks = 2 + # Every 256-byte block holds one 256-byte MLA swa page (4 tokens x 64B) + # and a 128-byte compressor ring page (8 tokens x 16B) aliased at the + # same base address. + block_stride, swa_page, ring_page = 256, 256, 128 + raw = torch.zeros(num_blocks * block_stride, dtype=torch.uint8) + + swa_spec = MLAAttentionSpec( + block_size=4, num_kv_heads=1, head_size=64, dtype=torch.uint8 + ) + ring_spec = CircularBufferSpec( + block_size=8, + num_kv_heads=1, + head_size=16, + head_size_v=0, + dtype=torch.uint8, + ) + + def group(specs): + uniform = UniformTypeKVCacheSpecs.from_specs(specs) + assert uniform is not None + return KVCacheGroupSpec(list(specs), uniform) + + swa_name, ring_name = "layers.0.swa", "layers.2.compressor" + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=raw.nbytes, + layers=[name], + layer_stride=raw.nbytes, + block_stride=block_stride, + ) + for name in (swa_name, ring_name) + ], + kv_cache_groups=[ + group({swa_name: swa_spec}), + group({ring_name: ring_spec}), + ], + ) + + vllm_config = create_vllm_config(block_size=4) + vllm_config.cache_config.enable_prefix_caching = False + vllm_config.kv_transfer_config.kv_buffer_device = "cuda" + fake_backend = MagicMock() + fake_backend.get_supported_kernel_block_sizes.return_value = [4] + fake_backend.get_name.return_value = "TEST_ATTN" + fake_backend.full_cls_name.return_value = "test.AttentionBackend" + fake_platform = MagicMock() + fake_platform.device_type = "cuda" + fake_platform.get_nixl_memory_type.return_value = "VRAM" + + with ( + patch.object(bw, "NixlWrapper", _RecordingNixl), + patch.object(bw, "get_tensor_model_parallel_rank", return_value=0), + patch.object(bw, "get_tensor_model_parallel_world_size", return_value=1), + patch.object(bw, "get_current_attn_backends", return_value=[fake_backend]), + patch.object(bw, "current_platform", fake_platform), + set_current_vllm_config(vllm_config), + ): + worker = NixlConnectorWorker(vllm_config, "local-engine", kv_cache_config) + swa_cache = raw.view(num_blocks, swa_page) + ring_cache = raw.as_strided( + (num_blocks, 1, ring_page), (block_stride, block_stride, 1) + ) + worker.register_kv_caches({swa_name: swa_cache, ring_name: ring_cache}) + return worker + + +@pytest.mark.cpu_test +def test_shared_base_regions_with_distinct_block_lens_stay_separate(): + """A compressor ring aliasing an MLA region's base address must register + its own region instead of collapsing into the MLA region and inheriting + its block length, which truncated and misaligned its transferred state.""" + worker = _make_shared_base_ring_worker() + + base = worker.kv_caches_base_addr[worker.engine_id][0] + assert base[0] == base[1] + assert worker.num_regions == 2 + assert worker.block_len_per_layer == [256, 128] + assert worker.block_stride_per_layer == [256, 256] + assert worker._region_is_mla == [True, False] + assert worker._scratch_region_indices == [1] + assert worker.region_group_ids == [0, 1] + assert worker.num_descs == 4 + descs = worker.src_blocks_data + assert descs[:, 1].tolist() == [256, 256, 128, 128] + np.testing.assert_array_equal( + descs[:, 0], [base[0], base[0] + 256, base[0], base[0] + 256] + ) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 280088e886e5..ee261817013e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -1529,6 +1529,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): registration_ranges: dict[tuple[int, str], tuple[int, int, int]] = {} region_mem_types: list[str] = [] seen_base_addresses: list[int] = [] + seen_region_keys: list[tuple[int, int]] = [] self._ssm_region_indices = [] self._scratch_region_indices = [] self._ple_block_len = None @@ -1788,10 +1789,16 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ] for base_addr, block_len, block_stride in region_specs: - # A packed PP producer keeps one region per layer, so aliases at - # one address (a layer and its SWA view) stay distinct. - if base_addr in seen_base_addresses and not route_packed_layers: - region_index = seen_base_addresses.index(base_addr) + # Distinct transfer geometries must stay distinct regions even + # when they alias the same base address (e.g. an MLA swa_cache + # and a compressor state_cache sharing one HMA allocation): + # collapsing them would make the later layer silently transfer + # with the earlier layer's block length. A packed PP producer + # instead keeps one region per layer, so aliases at one address + # (a layer and its SWA view) stay distinct unconditionally. + region_key = (base_addr, block_len) + if region_key in seen_region_keys and not route_packed_layers: + region_index = seen_region_keys.index(region_key) assert region_mem_types[region_index] == mem_type self._region_is_mla[region_index] |= is_mla_region if not is_mla_region: @@ -1805,6 +1812,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): else: region_index = len(seen_base_addresses) seen_base_addresses.append(base_addr) + seen_region_keys.append(region_key) self.block_len_per_layer.append(block_len) self.block_stride_per_layer.append(block_stride) self.region_group_ids.append(group_id) @@ -1859,6 +1867,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): assert ( len(self.block_len_per_layer) == len(seen_base_addresses) + == len(seen_region_keys) == len(self._region_is_mla) == len(self.block_stride_per_layer) == len(self.region_group_ids) From 57978684fc7c5eda25698164e1f3b59442b117bf Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Wed, 7 Oct 2026 11:24:25 -0700 Subject: [PATCH 11/13] [KV Connector][NIXL] Opt-in whole-row regions for replicated BLHNC layers vllm#53781 made per-layer regions the default whenever a layer's page is contiguous, which a HiSparse decoder needs. For a device-resident BLHNC cache it costs one descriptor per layer per block (99 for GLM-5.3 instead of 1). Add kv_connector_extra_config.whole_row_regions (default false): every replicated (MLA) layer of a block-major row registers as one whole-row region. Head-sharded layers, such as a GQA draft under an MLA target, never take the whole-row branch: they keep per-layer SPLIT regions so each TP rank reads its own head slice of the prefill page. Host-staged buffers and host-resident groups are rejected. Builds on #58471 (draft head slicing) and #57662 (regions keyed by (base_addr, block_len)); a block-outermost layout starts every KV cache group at byte 0 of the shared row, so the first draft page aliases the first target page's base address. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Stu Cao --- docs/features/nixl_connector_usage.md | 1 + .../kv_connector/unit/test_nixl_connector.py | 187 ++++++++++++++++++ .../kv_connector/v1/nixl/base_worker.py | 98 +++++++-- 3 files changed, 267 insertions(+), 19 deletions(-) diff --git a/docs/features/nixl_connector_usage.md b/docs/features/nixl_connector_usage.md index dab0faaf9879..b663db74a70f 100644 --- a/docs/features/nixl_connector_usage.md +++ b/docs/features/nixl_connector_usage.md @@ -236,6 +236,7 @@ Additional configuration options in `kv_connector_extra_config`: | `bidirectional_kv_xfer` | `false` | Enable bidirectional D→P KV transfer. | | `kv_recompute_threshold` | `64` | Minimum number of remote tokens required to trigger a D→P pull. Below this threshold, P recomputes locally instead of pulling (to amortize transfer latency). | | `decoder_kv_blocks_ttl` | `480` | TTL (seconds) for KV blocks cached on D for bidirectional reuse. Blocks are released after this duration. Not renewed via heartbeats. | +| `whole_row_regions` | `false` | With a block-major (`BLHNC`) layout, register every replicated (MLA) layer of a block row as one NIXL region, so a block transfers as a single descriptor instead of one per layer. Head-sharded layers, such as a GQA draft under an MLA target, keep per-layer regions so each TP rank reads its head slice. Requires device-resident caches on both peers (incompatible with `kv_buffer_device=cpu` and HiSparse host-resident groups) and must be set on both P and D. | ### Multi-turn proxy setup diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 1cbacc2b3f12..0d298db05fb6 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -2760,6 +2760,193 @@ def test_register_packed_dsv4_mla_cache_as_single_region( ] +_BLHNC_BLOCK_SIZE = 16 +_BLHNC_NUM_BLOCKS = 3 +_BLHNC_MLA_LAYERS = 4 +_BLHNC_DRAFT_LAYERS = 2 +_BLHNC_DRAFT_KV_HEADS = 8 + + +def _register_blhnc_mla_target_with_gqa_draft( + engine_id: str, + tp_size: int, + tp_rank: int, + whole_row: bool = False, + kv_buffer_device: str | None = None, +): + """Register an MLA target (group 0) and a GQA draft (group 1) laid out as the + block-outermost allocator lays them out: one block row sized to the larger + group, with every group's pages starting at byte 0 of that row. The first + draft page therefore shares its base address with the first target page. + """ + from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend + + mla_spec = MLAAttentionSpec( + block_size=_BLHNC_BLOCK_SIZE, + num_kv_heads=1, + head_size=576, + dtype=torch.uint8, + ) + draft_spec = FullAttentionSpec( + block_size=_BLHNC_BLOCK_SIZE, + num_kv_heads=_BLHNC_DRAFT_KV_HEADS // tp_size, + head_size=32, + dtype=torch.bfloat16, + ) + target_names = [ + f"model.layers.{i}.self_attn.attn" for i in range(_BLHNC_MLA_LAYERS) + ] + draft_names = [ + f"draft_model.model.layers.{i}.self_attn.attn" + for i in range(_BLHNC_DRAFT_LAYERS) + ] + row = max( + _BLHNC_MLA_LAYERS * mla_spec.page_size_bytes, + _BLHNC_DRAFT_LAYERS * draft_spec.page_size_bytes, + ) + kv_cache_config = KVCacheConfig( + num_blocks=_BLHNC_NUM_BLOCKS, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(target_names, mla_spec), + KVCacheGroupSpec(draft_names, draft_spec), + ], + ) + vllm_config = create_vllm_config( + attention_backend="TRITON_ATTN", + kv_connector_extra_config={"whole_row_regions": True} if whole_row else None, + ) + vllm_config.cache_config.block_size = _BLHNC_BLOCK_SIZE + vllm_config.cache_config.kv_cache_layout = "BLHNC" + if kv_buffer_device is not None: + vllm_config.kv_transfer_config.kv_buffer_device = kv_buffer_device + + raw = torch.zeros( + row * _BLHNC_NUM_BLOCKS, dtype=torch.int8, device=current_platform.device_type + ) + + def page_views(names: list[str], page: int) -> dict[str, torch.Tensor]: + return { + name: raw.as_strided((_BLHNC_NUM_BLOCKS, page), (row, 1), idx * page) + for idx, name in enumerate(names) + } + + nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker" + with ( + patch(f"{nixl_worker}.NixlWrapper", FakeNixlWrapper), + patch(f"{nixl_worker}.threading.Event"), + patch(f"{nixl_worker}.threading.Thread"), + patch(f"{nixl_worker}.get_current_attn_backends") as mock_backends, + patch( + f"{nixl_worker}.get_tensor_model_parallel_world_size", + return_value=tp_size, + ), + patch(f"{nixl_worker}.get_tensor_model_parallel_rank", return_value=tp_rank), + ): + mock_backends.return_value = [TritonAttentionBackend] + worker = FakeNixlConnectorWorker( + vllm_config, + engine_id, + hand_shake_latency=0, + kv_cache_layout="BLHNC", + kv_cache_config=kv_cache_config, + ) + worker.use_mla = True + worker._head_sharded_draft_kv_heads = _BLHNC_DRAFT_KV_HEADS + worker.register_kv_caches( + page_views(target_names, mla_spec.page_size_bytes) + | page_views(draft_names, draft_spec.page_size_bytes) + ) + meta = NixlAgentMetadata( + engine_id=engine_id, + agent_metadata=FakeNixlWrapper.AGENT_METADATA, + kv_caches_base_addr=worker.kv_caches_base_addr[engine_id][worker.tp_rank], + device_id=worker.device_id, + num_blocks=worker.num_blocks, + block_lens=worker.block_len_per_layer, + block_strides=worker.block_stride_per_layer, + kv_cache_layout=worker.kv_cache_layout, + block_size=worker.block_size, + ssm_sizes=(0, 0), + attn_backend_name=worker.backend_name, + physical_blocks_per_logical_kv_block=1, + region_num_blocks=worker.region_num_blocks, + region_group_ids=worker.region_group_ids, + region_names=worker.region_names, + region_mem_types=worker.region_mem_types, + ) + return worker, meta, raw, row, mla_spec.page_size_bytes, draft_spec.page_size_bytes + + +@pytest.mark.parametrize("whole_row", [False, True]) +@pytest.mark.parametrize("tp_rank", [0, 3]) +def test_blhnc_gqa_draft_reads_own_head_slice_from_tp1_prefill( + default_vllm_config, dist_init, whole_row, tp_rank +): + """A block-outermost layout aliases the draft's first page onto the target's + first page. The draft page must still be its own head-sharded region, so a + TP4 decode rank reads exactly its head slice of the TP1 prefill's page.""" + tp_size = 4 + p_worker, p_meta, p_raw, p_row, mla_page, p_draft_page = ( + _register_blhnc_mla_target_with_gqa_draft("prefill", 1, 0, whole_row) + ) + d_worker, _, d_raw, d_row, _, d_draft_page = ( + _register_blhnc_mla_target_with_gqa_draft("decode", tp_size, tp_rank, whole_row) + ) + assert p_row == d_row + assert p_draft_page == d_draft_page * tp_size + + draft_regions = [ + i + for i, name in enumerate(d_worker.region_names) + if name.startswith("draft_model.") + ] + target_regions = [ + i for i in range(len(d_worker.region_names)) if i not in draft_regions + ] + assert len(draft_regions) == _BLHNC_DRAFT_LAYERS + assert len(target_regions) == (1 if whole_row else _BLHNC_MLA_LAYERS) + assert all(d_worker._is_head_sharded_draft_region(i) for i in draft_regions) + assert not any(d_worker._is_region_replicated(i) for i in draft_regions) + assert all(d_worker._is_region_replicated(i) for i in target_regions) + assert {d_worker.region_group_ids[i] for i in draft_regions} == {1} + assert {d_worker.region_group_ids[i] for i in target_regions} == {0} + d_bases = d_worker.kv_caches_base_addr["decode"][tp_rank] + assert [d_bases[i] - d_raw.data_ptr() for i in draft_regions] == [ + j * d_draft_page for j in range(_BLHNC_DRAFT_LAYERS) + ] + + nixl_worker = "vllm.distributed.kv_transfer.kv_connector.v1.nixl.base_worker" + with patch(f"{nixl_worker}.NixlWrapper", FakeNixlWrapper): + d_worker.nixl_wrapper = FakeNixlWrapper("decode") + d_worker.add_remote_agent(p_meta, remote_tp_size=1) + plan = d_worker.tp_mappings["prefill"] + rows = d_worker._build_fa_remote(plan, p_meta, block_size_ratio=1) + rows = rows.reshape(len(d_worker.region_names), _BLHNC_NUM_BLOCKS, 3) + p_base = p_raw.data_ptr() + for j, region in enumerate(draft_regions): + assert rows[region, :, :2].tolist() == [ + [ + p_base + j * p_draft_page + tp_rank * d_draft_page + b * p_row, + d_draft_page, + ] + for b in range(_BLHNC_NUM_BLOCKS) + ] + target_len = d_row if whole_row else mla_page + for region in target_regions: + offset = d_bases[region] - d_raw.data_ptr() + assert rows[region, :, :2].tolist() == [ + [p_base + offset + b * p_row, target_len] for b in range(_BLHNC_NUM_BLOCKS) + ] + + +def test_whole_row_regions_rejects_host_buffer(default_vllm_config, dist_init): + with pytest.raises(ValueError, match="whole_row_regions"): + _register_blhnc_mla_target_with_gqa_draft( + "decode", 1, 0, whole_row=True, kv_buffer_device="cpu" + ) + + class FakePlatform(Platform): device_type: str = "oot" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index ee261817013e..2ce5f2f22999 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -170,6 +170,11 @@ class NixlBaseConnectorWorker: # Overridden by NixlPushConnectorWorker. _TRANSFER_MODE: str = "pull" + # Set in __init__; these defaults keep registration usable on workers built + # without it (whole-row off, no head-sharded draft). + _whole_row_regions: bool = False + _head_sharded_draft_kv_heads: int | None = None + # Layer-name routing is supported only by NixlPushConnector for HMA or # packed KV caches under PP. _supports_pp_hma = False @@ -431,28 +436,43 @@ def _resolve_head_sharded_draft_kv_heads(self) -> int | None: return None return draft_config.get_total_num_kv_heads() + def _is_draft_layer_name(self, layer_name: str) -> bool: + """Whether ``layer_name`` belongs to the speculative draft. + + Draft layers may use a draft_model prefix or be numbered after the + target's layers (e.g. DFlash's model.layers.78 after 78 target layers). + """ + if layer_name.startswith("draft_model."): + return True + spec_config = self.vllm_config.speculative_config + draft_config = spec_config.draft_model_config if spec_config else None + if draft_config is None: + return False + match = re.search(r"(?:^|\.)layers\.(\d+)(?:\.|$)", layer_name) + if match is None: + return False + layer_idx = int(match.group(1)) + target_layers = self.vllm_config.model_config.get_total_num_hidden_layers() + draft_layers = draft_config.get_total_num_hidden_layers() + return target_layers <= layer_idx < target_layers + draft_layers + + def _is_head_sharded_draft_layer(self, layer_name: str, is_mla: bool) -> bool: + """Whether a layer is a head-sharded draft layer under an MLA target.""" + return ( + self._head_sharded_draft_kv_heads is not None + and not is_mla + and self._is_draft_layer_name(layer_name) + ) + def _is_head_sharded_draft_region(self, region_idx: int) -> bool: if self._head_sharded_draft_kv_heads is None or self._is_region_replicated( region_idx ): return False - # Draft layers may use a draft_model prefix or be numbered after the - # target's layers (e.g. DFlash's model.layers.78 after 78 target layers). # A target SPLIT region must keep the target mapping. region_name = self.region_names[region_idx] - if not region_name.startswith("draft_model."): - spec_config = self.vllm_config.speculative_config - draft_config = spec_config.draft_model_config if spec_config else None - if draft_config is None: - return False - match = re.search(r"(?:^|\.)layers\.(\d+)(?:\.|$)", region_name) - if match is None: - return False - layer_idx = int(match.group(1)) - target_layers = self.vllm_config.model_config.get_total_num_hidden_layers() - draft_layers = draft_config.get_total_num_hidden_layers() - if not target_layers <= layer_idx < target_layers + draft_layers: - return False + if not self._is_draft_layer_name(region_name): + return False region_kv_heads = self._region_num_kv_heads[region_idx] assert ( region_kv_heads is not None @@ -634,6 +654,19 @@ def __init__( "bidirectional_kv_xfer", False ) ) + # Block-major packed caches can move a block as one descriptor covering + # the complete storage row. Per-layer regions stay the default whenever + # a layer's page is contiguous, because a HiSparse decoder lands + # host-resident and device-only layers in different memory and must + # address them separately. Only replicated (MLA) layers join the row; a + # head-sharded draft keeps per-layer regions so each rank reads its head + # slice. Both P and D must set it or the handshake rejects the region + # count. + self._whole_row_regions = bool( + vllm_config.kv_transfer_config.get_from_extra_config( + "whole_row_regions", False + ) + ) self.kv_cache_config = kv_cache_config transfer_block_sizes = [ @@ -1504,6 +1537,18 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): "NIXL host staging does not preserve CSA-linear shared tensors." ) + if self._whole_row_regions and ( + self.use_host_buffer + or any( + group.host_resident for group in self.kv_cache_config.transfer_groups + ) + ): + raise ValueError( + "whole_row_regions requires device-resident KV caches; it cannot " + "be combined with kv_buffer_device=cpu or host-resident transfer " + "groups." + ) + if self.use_host_buffer: self.initialize_host_xfer_buffer(kv_caches=kv_caches) assert len(self.host_xfer_buffers) == len(kv_caches), ( @@ -1630,6 +1675,9 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): region_num_kv_heads = ( None if is_mla_region else getattr(layer_spec, "num_kv_heads", None) ) + head_sharded_draft = self._is_head_sharded_draft_layer( + layer_name, is_mla_region + ) logger.debug( "Registering layer %s with cache shape: %s", layer_name, cache.shape ) @@ -1748,12 +1796,24 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): offset = cache.data_ptr() - storage_addr assert offset >= 0 and offset + physical_page_size <= block_stride packed_member_layouts[layer_name] = (offset, physical_page_size) - elif storage_is_block_major and ( - not page_contiguous - and ((packed_storage and is_mla_region) or not self._is_csa_linear) + elif ( + storage_is_block_major + and not head_sharded_draft + and ( + ( + not page_contiguous + and ( + (packed_storage and is_mla_region) + or not self._is_csa_linear + ) + ) + or (self._whole_row_regions and is_mla_region) + ) ): # TODO(Lucas): handle TP slicing for packed_storage; for now - # restrict to MLA (DSv4) where kv is replicated. + # restrict to MLA (DSv4) where kv is replicated. A head-sharded + # draft page must stay its own region so each rank can read its + # head slice. # The complete storage row is one block stride. Dividing the # registered length instead would fold the padding tail in. storage_block_len = block_stride From edb4a3c77b87486b47585737c194d9c125cd81be Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Wed, 7 Oct 2026 23:45:32 -0700 Subject: [PATCH 12/13] [KV Connector][NIXL] Reject whole_row_regions when draft pages share the target row With the hybrid KV cache manager disabled, a block-outermost layout packs the head-sharded draft pages into the target's KV cache group and block row. Their size depends on the TP size, so a whole-row region would differ between a TP1 prefill and a TP4 decode and fail the handshake later ("KV cache sizes must match between P and D when replicated"). Fail at registration instead, naming the layer and group. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Stu Cao --- .../kv_connector/unit/test_nixl_connector.py | 58 +++++++++++++++---- .../kv_connector/v1/nixl/base_worker.py | 28 +++++++++ 2 files changed, 75 insertions(+), 11 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 0d298db05fb6..3ee6baabb3dd 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -73,6 +73,7 @@ KVCacheGroupSpec, KVCacheLayout, MLAAttentionSpec, + UniformTypeKVCacheSpecs, compute_layer_kv_cache_shape_bytes, ) from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput @@ -2773,11 +2774,16 @@ def _register_blhnc_mla_target_with_gqa_draft( tp_rank: int, whole_row: bool = False, kv_buffer_device: str | None = None, + packed_draft: bool = False, ): """Register an MLA target (group 0) and a GQA draft (group 1) laid out as the block-outermost allocator lays them out: one block row sized to the larger group, with every group's pages starting at byte 0 of that row. The first draft page therefore shares its base address with the first target page. + + With ``packed_draft``, target and draft share one group and the draft pages + follow the target pages in the row, as when the hybrid KV cache manager is + disabled. """ from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend @@ -2800,17 +2806,32 @@ def _register_blhnc_mla_target_with_gqa_draft( f"draft_model.model.layers.{i}.self_attn.attn" for i in range(_BLHNC_DRAFT_LAYERS) ] - row = max( - _BLHNC_MLA_LAYERS * mla_spec.page_size_bytes, - _BLHNC_DRAFT_LAYERS * draft_spec.page_size_bytes, - ) + target_bytes = _BLHNC_MLA_LAYERS * mla_spec.page_size_bytes + draft_bytes = _BLHNC_DRAFT_LAYERS * draft_spec.page_size_bytes + if packed_draft: + row = target_bytes + draft_bytes + draft_offset = target_bytes + kv_cache_groups = [ + KVCacheGroupSpec( + target_names + draft_names, + UniformTypeKVCacheSpecs( + block_size=_BLHNC_BLOCK_SIZE, + kv_cache_specs={name: mla_spec for name in target_names} + | {name: draft_spec for name in draft_names}, + ), + ) + ] + else: + row = max(target_bytes, draft_bytes) + draft_offset = 0 + kv_cache_groups = [ + KVCacheGroupSpec(target_names, mla_spec), + KVCacheGroupSpec(draft_names, draft_spec), + ] kv_cache_config = KVCacheConfig( num_blocks=_BLHNC_NUM_BLOCKS, kv_cache_tensors=[], - kv_cache_groups=[ - KVCacheGroupSpec(target_names, mla_spec), - KVCacheGroupSpec(draft_names, draft_spec), - ], + kv_cache_groups=kv_cache_groups, ) vllm_config = create_vllm_config( attention_backend="TRITON_ATTN", @@ -2825,9 +2846,13 @@ def _register_blhnc_mla_target_with_gqa_draft( row * _BLHNC_NUM_BLOCKS, dtype=torch.int8, device=current_platform.device_type ) - def page_views(names: list[str], page: int) -> dict[str, torch.Tensor]: + def page_views( + names: list[str], page: int, offset: int = 0 + ) -> dict[str, torch.Tensor]: return { - name: raw.as_strided((_BLHNC_NUM_BLOCKS, page), (row, 1), idx * page) + name: raw.as_strided( + (_BLHNC_NUM_BLOCKS, page), (row, 1), offset + idx * page + ) for idx, name in enumerate(names) } @@ -2855,7 +2880,7 @@ def page_views(names: list[str], page: int) -> dict[str, torch.Tensor]: worker._head_sharded_draft_kv_heads = _BLHNC_DRAFT_KV_HEADS worker.register_kv_caches( page_views(target_names, mla_spec.page_size_bytes) - | page_views(draft_names, draft_spec.page_size_bytes) + | page_views(draft_names, draft_spec.page_size_bytes, draft_offset) ) meta = NixlAgentMetadata( engine_id=engine_id, @@ -2947,6 +2972,17 @@ def test_whole_row_regions_rejects_host_buffer(default_vllm_config, dist_init): ) +def test_whole_row_regions_rejects_draft_packed_into_target_row( + default_vllm_config, dist_init +): + """With the hybrid KV cache manager off, BLHNC packs head-sharded draft + pages into the target's row, so a whole row differs between TP sizes.""" + with pytest.raises(ValueError, match="head-sharded draft pages"): + _register_blhnc_mla_target_with_gqa_draft( + "decode", 4, 0, whole_row=True, packed_draft=True + ) + + class FakePlatform(Platform): device_type: str = "oot" diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py index 2ce5f2f22999..8482576ed100 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py @@ -1615,6 +1615,20 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): ): compressed_region_owners.setdefault(cache.data_ptr(), cache) + # A whole row of a group that also holds head-sharded draft pages would + # carry TP-dependent bytes, so P and D rows could not match. + sharded_row_groups: set[int] = set() + if self._whole_row_regions: + group_by_layer = self.kv_cache_config.transfer_group_index_by_layer + sharded_row_groups = { + group_by_layer[name] + for name, spec in layer_specs.items() + if name in group_by_layer + and self._is_head_sharded_draft_layer( + name, isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)) + ) + } + track_region_layers = self._tracks_region_layers() region_layers: list[list[str]] = [] packed_member_layouts: dict[str, tuple[int, int]] = {} @@ -1776,6 +1790,20 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]): and storage_is_block_major and is_mla_region ) + if ( + storage_is_block_major + and self._whole_row_regions + and is_mla_region + and group_id in sharded_row_groups + ): + raise ValueError( + "whole_row_regions cannot register a block row that also " + f"holds head-sharded draft pages (layer {layer_name}, group " + f"{group_id}): the row size depends on the TP size, so P and " + "D rows would differ. Keep the draft in its own KV cache " + "group (do not combine with " + "--disable-hybrid-kv-cache-manager) or drop the flag." + ) if virtual_transfer_pages: # A compressed kernel row can contain multiple NIXL transfer pages. region_specs = [ From 9adc138cfcefdbb87349f1f685f4c51574658a67 Mon Sep 17 00:00:00 2001 From: Stu Cao Date: Thu, 8 Oct 2026 00:34:27 -0700 Subject: [PATCH 13/13] [KV Connector][NIXL] Let the fake NIXL worker take a mixed-spec group FakeNixlConnectorWorker derives its reference KV shape from the first KV cache group's spec, which has no head count when that group mixes specs (UniformTypeKVCacheSpecs). Use the group's first spec, as a packed block-outermost layout has, so the packed-draft rejection test can register. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Stu Cao --- tests/v1/kv_connector/unit/test_nixl_connector.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_nixl_connector.py b/tests/v1/kv_connector/unit/test_nixl_connector.py index 3ee6baabb3dd..f10233905e8a 100644 --- a/tests/v1/kv_connector/unit/test_nixl_connector.py +++ b/tests/v1/kv_connector/unit/test_nixl_connector.py @@ -553,6 +553,8 @@ def __init__( self.src_xfer_handles_by_block_size = {self.block_size: 1} self.src_blocks_data = np.empty((0, 3), dtype=np.uint64) rep_spec = self.kv_cache_config.kv_cache_groups[0].kv_cache_spec + if isinstance(rep_spec, UniformTypeKVCacheSpecs): + rep_spec = rep_spec.first_spec test_shape = compute_layer_kv_cache_shape_bytes(rep_spec, 1) self.transfer_topo = TransferTopology( tp_rank=self.tp_rank,