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 52ac746391ff..f10233905e8a 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 @@ -552,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, @@ -1383,6 +1386,231 @@ 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 + 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._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( + [(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, + ) + 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.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", + 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, + ) + @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_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, + ) + 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, @@ -2535,6 +2763,228 @@ 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, + 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 + + 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) + ] + 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=kv_cache_groups, + ) + 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, offset: int = 0 + ) -> dict[str, torch.Tensor]: + return { + name: raw.as_strided( + (_BLHNC_NUM_BLOCKS, page), (row, 1), offset + 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, draft_offset) + ) + 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" + ) + + +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/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py b/tests/v1/kv_connector/unit/test_nixl_desc_geometry.py index 9f6bc56deeb5..91f1a2f3b775 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 @@ -1560,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] @@ -1624,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( @@ -1679,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]), @@ -1687,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/tests/v1/kv_connector/unit/test_nixl_push_connector.py b/tests/v1/kv_connector/unit/test_nixl_push_connector.py index a36c7e5b2a71..0ea58892f1d2 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,8 @@ def fresh(cls) -> _StubWriterWorker: w._has_packed_cache = False 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/tests/v1/kv_connector/unit/test_tp_mapping.py b/tests/v1/kv_connector/unit/test_tp_mapping.py index 3ddc568a524d..34321b0731c5 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, + head_sharded_kv_heads: int | None = None, ) -> 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, + head_sharded_kv_heads=head_sharded_kv_heads, ) @@ -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..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 @@ -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 @@ -419,6 +424,76 @@ 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 draft under an MLA target, else None. + + 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 + 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_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 + # A target SPLIT region must keep the target mapping. + region_name = self.region_names[region_idx] + 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 + and region_kv_heads <= self._head_sharded_draft_kv_heads + ), ( + 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 + + 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)), ( @@ -579,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 = [ @@ -929,6 +1017,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 @@ -960,6 +1054,9 @@ def __init__( self._transfer_layer_group_ids = tuple[int, ...]() # Per-engine TP mappings. Generated during handshake. self.tp_mappings: dict[EngineId, TPMapping] = {} + # 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] = {} self.enforce_compat_hash = self.kv_transfer_config.get_from_extra_config( "enforce_handshake_compat", True @@ -1440,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), ( @@ -1465,6 +1574,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 @@ -1505,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]] = {} @@ -1562,6 +1686,12 @@ 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) + ) + 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 ) @@ -1660,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 = [ @@ -1680,12 +1824,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 @@ -1721,12 +1877,20 @@ 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: + 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 @@ -1736,12 +1900,14 @@ 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) 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: @@ -1789,10 +1955,12 @@ 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) == 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. @@ -2116,9 +2284,6 @@ 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]) region_num_blocks = nixl_agent_meta.region_num_blocks or [ nixl_agent_meta.num_blocks ] * len(nixl_agent_meta.kv_caches_base_addr) @@ -2139,11 +2304,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 `split_reads` remote ranks at a per-rank offset. - num_reads = 1 if replicated else split_reads - rank_offset = ( - 0 if replicated else 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) @@ -2360,6 +2531,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 +2728,24 @@ def _validate_remote_agent_handshake( assert not ( 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: + if self._TRANSFER_MODE == "push": + 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 " + "target when prefill TP exceeds decode TP or with DCP: " + 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}." + ) remote_physical_per_logical = ( nixl_agent_meta.physical_blocks_per_logical_kv_block @@ -2641,6 +2838,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 +2888,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 +4047,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..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 @@ -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,17 @@ 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. + + ``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 - 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 +142,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: