diff --git a/docs/source/user_guide/feature_guide/hisparse.md b/docs/source/user_guide/feature_guide/hisparse.md index f95853212955..8fb04256a91f 100644 --- a/docs/source/user_guide/feature_guide/hisparse.md +++ b/docs/source/user_guide/feature_guide/hisparse.md @@ -29,8 +29,9 @@ supported, with independent Main and Indexer KV dtype selection. | --- | --- | | Hardware and software | Ascend 950 with a SIMT-capable CANN 9.2 toolchain. | | Cache formats | C4, TurboQuant and INT8-format C8 are outside the current support scope. | -| Parallelism and transport | Context parallelism and external P/D transport are outside the current support scope. | +| Parallelism and transport | Context parallelism is unsupported. For P/D transport, use `HiSparsePDConnector` with Ascend Mooncake/HIXL and matching TP sizes. | | Host storage | Use private Host pools per tensor-parallel rank and Core eager Host mirroring. | +| P/D Host import | The D worker registers a separate pinned Host staging region for Mooncake, with capacity equal to the Host Source pool. Allow for this additional Host memory per TP rank. | | TopK | The existing SFA requirement of divisibility by 128 applies. | | Feature mutual exclusion | HiSparse and `additional_config.sparse_kv_offload_config.enabled` cannot be enabled together. | @@ -56,6 +57,31 @@ For native A5 C8, replace the two dtype options with: --attention-config '{"indexer_kv_dtype":"fp8","hisparse_config":{"eager_host_mirror":true}}' ``` +### P/D with a Memcache Prefill pool + +Run a Prefill instance without HiSparse and a Decode instance with HiSparse. +Configure both instances with the existing Mooncake transfer service. The +Prefill connector saves ordinary KV to Memcache; the Decode connector pulls +the Prefill KV directly. The Decode instance does not read Memcache. + +Prefill `--kv-transfer-config`: + +```json +{"kv_connector":"MultiConnector","kv_role":"kv_producer","kv_connector_extra_config":{"connectors":[{"kv_connector":"HiSparsePDConnector","kv_role":"kv_producer"},{"kv_connector":"AscendStoreConnector","kv_role":"kv_producer","kv_connector_extra_config":{"backend":"memcache","use_layerwise":false}}]}} +``` + +Decode `--kv-transfer-config`: + +```json +{"kv_connector":"MultiConnector","kv_role":"kv_consumer","kv_connector_extra_config":{"connectors":[{"kv_connector":"HiSparseConnector","kv_role":"kv_consumer","kv_connector_extra_config":{"host_pool_gib":4}},{"kv_connector":"HiSparsePDConnector","kv_role":"kv_consumer"}]}} +``` + +Set distinct Mooncake `kv_port` values for the P and D connectors and follow +the [Mooncake P/D setup](../../tutorials/features/pd_disaggregation_mooncake_single_node.md) +for its transfer service. Both instances must use the same model and KV format. +The D connector imports main KV into Host Source, writes Indexer KV to the NPU, +and notifies Core after transfer; Core restores the last page to Resident. + ### Verifying the Feature Send a request and check the server's metrics endpoint for the HiSparse miss diff --git a/release_notes/fragments/2.yaml b/release_notes/fragments/2.yaml new file mode 100644 index 000000000000..07c611cfe1e7 --- /dev/null +++ b/release_notes/fragments/2.yaml @@ -0,0 +1,4 @@ +pr: 2 +entries: + - category: feature + text: Added Ascend HiSparse P/D import using Mooncake transport, with Prefill KV imported into Decode Host Source and Indexer KV on the NPU. diff --git a/tests/ut/kv_offload/mooncake_v2/test_hisparse_pd.py b/tests/ut/kv_offload/mooncake_v2/test_hisparse_pd.py new file mode 100644 index 000000000000..fb02954e3152 --- /dev/null +++ b/tests/ut/kv_offload/mooncake_v2/test_hisparse_pd.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import patch + +import torch + +from vllm_ascend.distributed.kv_transfer.kv_p2p.hisparse_pd import _HiSparsePDReceivingThread, _HiSparsePDWorker +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker import MooncakePullConnectorWorker + + +def test_hisparse_pd_worker_selects_source_and_indexer() -> None: + worker = _HiSparsePDWorker.__new__(_HiSparsePDWorker) + worker.num_blocks = 8 + worker.kv_cache_config = SimpleNamespace( + hisparse_host_num_blocks=32, + kv_cache_groups=[ + SimpleNamespace(enable_kv_transfer=True, host_resident=True), + SimpleNamespace(enable_kv_transfer=True, host_resident=False), + SimpleNamespace(enable_kv_transfer=False, host_resident=False), + ], + ) + + assert [worker._is_transfer_group(i) for i in range(3)] == [True, True, False] + assert [worker._logical_num_blocks(i) for i in range(2)] == [32, 8] + + +def test_hisparse_pd_worker_slices_replicated_indexer_blocks() -> None: + worker = _HiSparsePDWorker.__new__(_HiSparsePDWorker) + worker.kv_transfer_config = SimpleNamespace(is_kv_consumer=True) + worker.kv_cache_config = SimpleNamespace( + hisparse_host_num_blocks=4, + hisparse_host_block_stride=4, + kv_cache_groups=[ + SimpleNamespace( + role="hisparse_source", layer_names=["source"], host_resident=True, enable_kv_transfer=True + ), + SimpleNamespace( + role="hisparse_indexer", layer_names=["indexer"], host_resident=False, enable_kv_transfer=True + ), + ], + ) + worker.vllm_config = SimpleNamespace( + compilation_config=SimpleNamespace(static_forward_context={"indexer": SimpleNamespace(hisparse_block_stride=2)}) + ) + indexer = torch.arange(16).reshape(4, 4) + source = torch.zeros((4, 4), dtype=torch.uint8) + with patch.object(MooncakePullConnectorWorker, "register_kv_caches") as register: + worker.register_kv_caches({"source": source, "indexer": indexer}) + + transferred = register.call_args.args[0]["indexer"][0] + assert torch.equal(transferred, indexer[::2]) + staged_source = register.call_args.args[0]["source"][0] + assert staged_source.data_ptr() != source.data_ptr() + staged_source.copy_(torch.arange(16, dtype=torch.uint8).reshape(4, 4)) + thread = _HiSparsePDReceivingThread.__new__(_HiSparsePDReceivingThread) + thread.host_stage_range = worker._host_stage_range + thread._after_transfer([staged_source.data_ptr()], [staged_source.nbytes]) + assert torch.equal(source, staged_source) + + +def test_hisparse_pd_imports_mla_components_by_token() -> None: + source = torch.zeros((2, 3), dtype=torch.uint8) + staging = torch.tensor([1, 2, 3, 4, 5, 6], dtype=torch.uint8) + thread = _HiSparsePDReceivingThread.__new__(_HiSparsePDReceivingThread) + thread.host_stage_range = ( + staging.data_ptr(), + staging.data_ptr() + staging.nbytes, + source.data_ptr(), + source.data_ptr() + source.nbytes, + ) + thread._packed_host_pages = [(staging.data_ptr(), 2, 2, 1)] + thread._after_transfer([staging.data_ptr(), staging.data_ptr() + 4], [4, 2]) + assert source.tolist() == [[1, 2, 5], [3, 4, 6]] diff --git a/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py index dbf35eff46bd..53eeb8083a4b 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py @@ -1061,6 +1061,37 @@ def test_attention_address_generation_handles_partial_head_overlap() -> None: assert lengths == [512] +def test_hisparse_packed_mla_destination_splits_prefill_components() -> None: + spec = MLAAttentionSpec(block_size=16, num_kv_heads=1, head_size=4, dtype=torch.float16) + thread = make_thread( + kv_cache_specs=[spec], + kv_caches_base_addr=[[1000]], + block_strides=[[128]], + block_lens=[[128]], + ) + thread._allow_packed_mla_destination = True + remote = make_pp_metadata( + block_lens=[[96, 32]], + block_strides=[[96, 32]], + tp_base_addrs={0: [[5000, 6000]]}, + ) + src: list[int] = [] + dst: list[int] = [] + lengths: list[int] = [] + + thread._append_spec_transfer_addresses( + 0, 0, 0, 1, 1, 1, {(0, 0): [("request", [1], [2])]}, remote, src, dst, lengths + ) + + assert src == [1128, 1224] + assert dst == [5192, 6064] + assert lengths == [96, 32] + + remote.block_lens[0] = [95, 32] + with pytest.raises(ValueError, match="packed Source length mismatch"): + thread._append_spec_transfer_addresses(0, 0, 0, 1, 1, 1, {(0, 0): [("request", [1], [2])]}, remote, [], [], []) + + @pytest.mark.parametrize(("local_tp", "local_pcp"), [(1, 4), (4, 1), (2, 2)]) @pytest.mark.parametrize(("remote_tp", "remote_pcp"), [(1, 4), (4, 1), (2, 2)]) @pytest.mark.parametrize("worker_rank", [1, 2, 3]) diff --git a/tests/ut/kv_offload/mooncake_v2/test_utils.py b/tests/ut/kv_offload/mooncake_v2/test_utils.py index 368e683c6c07..b0fa48b7ba3e 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_utils.py +++ b/tests/ut/kv_offload/mooncake_v2/test_utils.py @@ -99,6 +99,25 @@ def test_collect_register_regions_handles_independent_storages() -> None: } +def test_collect_register_regions_can_select_transfer_layers() -> None: + source = torch.empty(32, dtype=torch.uint8) + resident = torch.empty(48, dtype=torch.uint8) + config = SimpleNamespace( + kv_cache_tensors=[ + make_kv_cache_tensor(layers=["source"], layer_stride=32, block_stride=32, offset=0, size=32), + make_kv_cache_tensor(layers=["resident"], layer_stride=48, block_stride=48, offset=0, size=48), + ] + ) + + regions = collect_configured_register_regions( + config, {"source": source, "resident": resident}, included_layer_names={"source"} + ) + + assert regions.ptrs == [source.data_ptr()] + assert regions.lengths == [source.nbytes] + assert regions.logical_tensor_count == 1 + + def test_collect_register_regions_recovers_aligned_backing_before_view() -> None: alignment = 2 * 1024 * 1024 backing_size = 2 * alignment diff --git a/vllm_ascend/distributed/kv_transfer/__init__.py b/vllm_ascend/distributed/kv_transfer/__init__.py index 51dd5df0a514..fd66ff05f9a4 100644 --- a/vllm_ascend/distributed/kv_transfer/__init__.py +++ b/vllm_ascend/distributed/kv_transfer/__init__.py @@ -94,6 +94,12 @@ def register_connector(): "MooncakeConnector", ) + KVConnectorFactory.register_connector( + "HiSparsePDConnector", + "vllm_ascend.distributed.kv_transfer.kv_p2p.hisparse_pd", + "HiSparsePDConnector", + ) + KVConnectorFactory.register_connector( "MooncakePullConnector", "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector", diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/hisparse_pd.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/hisparse_pd.py new file mode 100644 index 000000000000..1ef3bfefced6 --- /dev/null +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/hisparse_pd.py @@ -0,0 +1,180 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM Ascend project +"""Pull P KV into HiSparse's Host Source and NPU Indexer groups.""" + +import ctypes +from typing import Any + +import torch +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole +from vllm.v1.kv_cache_interface import KVCacheConfig + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector import MooncakeBaseConnector +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler import MooncakePullConnectorScheduler +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker import ( + MooncakePullConnectorWorker, + MooncakePullRecvingThread, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils import as_kv_cache_tensors + +_HOST_ALIGNMENT = 2 * 1024 * 1024 + + +class _HiSparsePDReceivingThread(MooncakePullRecvingThread): + _allow_packed_mla_destination = True + + def _execute_tp_transfer_bucket(self, *args: Any, **kwargs: Any) -> None: + self._packed_host_pages: list[tuple[int, int, int, int]] = [] + try: + super()._execute_tp_transfer_bucket(*args, **kwargs) + finally: + self._packed_host_pages = [] + + def _append_spec_transfer_addresses(self, *args: Any, **kwargs: Any) -> None: + spec_index = kwargs["spec_index"] + remote_metadata = kwargs["remote_metadata"] + remote_tp_metadata = remote_metadata.metadata_by_pcp_rank[kwargs["remote_pcp_rank"]].metadata_by_tp_rank[ + kwargs["remote_tp_rank"] + ] + for (local_layer, remote_layer), entries in kwargs["transfer_entries_by_layer"].items(): + local_addrs = self.kv_caches_base_addr[local_layer] + remote_addrs = remote_tp_metadata.kv_caches_base_addr[remote_layer] + if len(local_addrs) != 1 or len(remote_addrs) != 2: + continue + spec = self.kv_cache_specs[spec_index] + lengths = remote_metadata.block_lens[remote_layer] + if len(lengths) != 2 or any(length % spec.block_size for length in lengths): + raise ValueError("HiSparse PD requires two whole-token MLA components") + row_lengths = (lengths[0] // spec.block_size, lengths[1] // spec.block_size) + if sum(lengths) != self.block_lens[local_layer][0]: + raise ValueError("HiSparse PD Source row size differs from producer MLA") + for _, local_blocks, _ in entries: + for block in local_blocks: + page = local_addrs[0] + block * self.block_strides[local_layer][0] + self._packed_host_pages.append((page, spec.block_size, *row_lengths)) + super()._append_spec_transfer_addresses(*args, **kwargs) + + def _after_transfer(self, local_addresses: list[int], lengths: list[int]) -> None: + staging_start, staging_end, source_start, source_end = self.host_stage_range + packed_pages = {page for page, _, _, _ in getattr(self, "_packed_host_pages", [])} + for address, length in zip(local_addresses, lengths, strict=True): + if not staging_start <= address < staging_end: + continue + destination = source_start + address - staging_start + if address + length > staging_end or destination + length > source_end: + raise ValueError("HiSparse PD Host transfer exceeds its registered Source range") + if any(page <= address < page + length for page in packed_pages): + continue + ctypes.memmove(destination, address, length) + for page, tokens, nope_bytes, rope_bytes in getattr(self, "_packed_host_pages", []): + destination = source_start + page - staging_start + page_bytes = tokens * (nope_bytes + rope_bytes) + if page < staging_start or page + page_bytes > staging_end or destination + page_bytes > source_end: + raise ValueError("HiSparse PD packed page exceeds its Source range") + for token in range(tokens): + row = destination + token * (nope_bytes + rope_bytes) + ctypes.memmove(row, page + token * nope_bytes, nope_bytes) + ctypes.memmove(row + nope_bytes, page + tokens * nope_bytes + token * rope_bytes, rope_bytes) + + +class _HiSparsePDWorker(MooncakePullConnectorWorker): + _receiving_thread_cls = _HiSparsePDReceivingThread + + def _is_transfer_group(self, group_index: int) -> bool: + return self.kv_cache_config.kv_cache_groups[group_index].enable_kv_transfer + + def _logical_num_blocks(self, group_index: int) -> int: + group = self.kv_cache_config.kv_cache_groups[group_index] + if group.host_resident: + assert self.kv_cache_config.hisparse_host_num_blocks is not None + return self.kv_cache_config.hisparse_host_num_blocks + return self.num_blocks + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]]) -> None: + if self.kv_transfer_config.is_kv_consumer: + context = self.vllm_config.compilation_config.static_forward_context + caches: dict[str, Any] = dict(kv_caches) + source_names = [ + name + for group in self.kv_cache_config.kv_cache_groups + if group.host_resident and group.enable_kv_transfer + for name in group.layer_names + ] + source_parts = [cache for name in source_names for cache in as_kv_cache_tensors(kv_caches[name])] + if not source_parts or any(cache.device.type != "cpu" for cache in source_parts): + raise ValueError("HiSparse PD Source must consist of Host tensors") + source_start = min(cache.data_ptr() for cache in source_parts) + source_end = max( + cache.data_ptr() + + cache.element_size() + + sum((size - 1) * stride * cache.element_size() for size, stride in zip(cache.shape, cache.stride())) + for cache in source_parts + ) + host_blocks = self.kv_cache_config.hisparse_host_num_blocks + host_stride = self.kv_cache_config.hisparse_host_block_stride + assert host_blocks is not None and host_stride is not None + staging_size = ( + (max(source_end - source_start, host_blocks * host_stride) + _HOST_ALIGNMENT - 1) // _HOST_ALIGNMENT + ) * _HOST_ALIGNMENT + owner = torch.empty(staging_size + _HOST_ALIGNMENT - 1, dtype=torch.uint8, pin_memory=True) + alignment_offset = (-owner.data_ptr()) % _HOST_ALIGNMENT + staging = owner.narrow(0, alignment_offset, staging_size) + self._source_staging_owner = owner + self._host_stage_range = (staging.data_ptr(), staging.data_ptr() + staging_size, source_start, source_end) + for name in source_names: + staged_parts = [] + for cache in as_kv_cache_tensors(kv_caches[name]): + offset = alignment_offset + cache.data_ptr() - source_start + if offset % cache.element_size(): + raise ValueError(f"Misaligned HiSparse Source layer {name!r}") + staged_parts.append( + staging.view(cache.dtype).as_strided( + cache.shape, cache.stride(), storage_offset=offset // cache.element_size() + ) + ) + caches[name] = tuple(staged_parts) + for group in self.kv_cache_config.kv_cache_groups: + if group.role != "hisparse_indexer": + continue + for name in group.layer_names: + stride = getattr(context[name], "hisparse_block_stride", 1) + if stride > 1: + caches[name] = tuple(cache[::stride] for cache in as_kv_cache_tensors(kv_caches[name])) + kv_caches = caches + super().register_kv_caches(kv_caches) + + def _configure_receiving_thread(self, thread: MooncakePullRecvingThread) -> None: + assert isinstance(thread, _HiSparsePDReceivingThread) + thread.host_stage_range = self._host_stage_range + + +class HiSparsePDConnector(MooncakeBaseConnector): + """Use Ascend Mooncake transport with Core HiSparse import ownership.""" + + def __init__( + self, + vllm_config: VllmConfig, + role: KVConnectorRole, + kv_cache_config: KVCacheConfig, + ) -> None: + super().__init__(vllm_config, role, kv_cache_config) + if self.kv_transfer_config.is_kv_consumer: + if kv_cache_config.hisparse_host_num_blocks is None: + raise ValueError("HiSparsePDConnector consumer requires HiSparse Host Source") + transferable = [group for group in kv_cache_config.kv_cache_groups if group.enable_kv_transfer] + if ( + len(transferable) != 2 + or not transferable[0].host_resident + or transferable[1].role != "hisparse_indexer" + ): + raise ValueError("HiSparsePDConnector requires Host Source and NPU Indexer transfer groups") + elif kv_cache_config.hisparse_host_num_blocks is not None: + raise ValueError("HiSparsePDConnector producer must use ordinary Prefill KV") + + if role == KVConnectorRole.SCHEDULER: + self.connector_scheduler = MooncakePullConnectorScheduler(vllm_config, str(self.engine_id), kv_cache_config) + elif role == KVConnectorRole.WORKER: + self.connector_worker = _HiSparsePDWorker(vllm_config, str(self.engine_id), kv_cache_config) + else: + raise ValueError(f"Unsupported KVConnectorRole: {role}") diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_worker.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_worker.py index 2cb09d319416..69a84487a2e8 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_worker.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_worker.py @@ -374,6 +374,12 @@ def __init__( logger.info("Initializing Mooncake worker %s", engine_id) + def _is_transfer_group(self, group_index: int) -> bool: + return True + + def _logical_num_blocks(self, group_index: int) -> int: + return self.num_blocks + def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None: """Return the registered cache descriptor for either transfer mode.""" return self.xfer_handshake_metadata @@ -414,6 +420,7 @@ def _build_kv_cache_spec_mappings(self) -> None: def _get_shared_page_metadata( self, caches: tuple[torch.Tensor, ...], + logical_num_blocks: int, ) -> tuple[int, int, tuple[int, ...], int] | None: """Describe views packed into one physical page allocation. @@ -433,7 +440,7 @@ def _get_shared_page_metadata( num_tensor_blocks = next(iter(tensor_num_blocks)) page_stride = next(iter(page_strides)) - if num_tensor_blocks <= 0 or page_stride <= 0 or num_tensor_blocks % self.num_blocks: + if num_tensor_blocks <= 0 or page_stride <= 0 or num_tensor_blocks % logical_num_blocks: return None page_base_addr = min(cache.data_ptr() for cache in caches) @@ -470,7 +477,7 @@ def _get_shared_page_metadata( page_base_addr, page_stride, tuple(selected_cache.shape[1:]), - num_tensor_blocks // self.num_blocks, + num_tensor_blocks // logical_num_blocks, ) @staticmethod @@ -524,6 +531,7 @@ def register_kv_caches( block_size_scales_per_layer: list[list[int]] = [] indexer_replication_sizes: list[int] = [] configured_layer_names: set[str] = set() + transfer_layer_names: set[str] = set() for tensor_config in self.kv_cache_config.kv_cache_tensors: for layer_name in get_kv_cache_tensor_layers(tensor_config): @@ -532,11 +540,15 @@ def register_kv_caches( if layer_name not in self.layer_name_to_group_index: raise ValueError(f"Configured KV cache layer {layer_name!r} does not belong to a KV cache group.") + configured_layer_names.add(layer_name) + group_index = self.layer_name_to_group_index[layer_name] + if not self._is_transfer_group(group_index): + continue cache_or_caches = kv_caches.get(layer_name) if cache_or_caches is None: raise ValueError(f"No KV cache was registered for configured layer {layer_name!r}.") - - configured_layer_names.add(layer_name) + transfer_layer_names.add(layer_name) + logical_num_blocks = self._logical_num_blocks(group_index) # Foreign target layers alias scratch. Publish persistent owners # only; MTP caches are absent from owners and remain replicated. owner = owners.get(layer_name) @@ -553,7 +565,7 @@ def register_kv_caches( spec = self.kv_cache_specs[spec_index] caches = as_kv_cache_tensors(cache_or_caches) shared_page_metadata = ( - self._get_shared_page_metadata(caches) + self._get_shared_page_metadata(caches, logical_num_blocks) if isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)) else None ) @@ -586,7 +598,12 @@ def register_kv_caches( block_strides.append(cache.stride(0) * element_size) block_lens.append(math.prod(block_shape) * element_size) block_shapes.append(block_shape) - block_size_scales.append(tensor_num_blocks // self.num_blocks) + if tensor_num_blocks % logical_num_blocks: + raise ValueError( + f"KV cache layer {layer_name!r} has {tensor_num_blocks} physical blocks " + f"for {logical_num_blocks} logical blocks." + ) + block_size_scales.append(tensor_num_blocks // logical_num_blocks) layer_names.append(layer_name) layer_block_size = spec.block_size @@ -611,7 +628,9 @@ def register_kv_caches( if unexpected_layers: raise ValueError(f"KV caches contain layers absent from kv_cache_tensors: {sorted(unexpected_layers)}.") - register_regions = collect_configured_register_regions(self.kv_cache_config, kv_caches) + register_regions = collect_configured_register_regions( + self.kv_cache_config, kv_caches, included_layer_names=transfer_layer_names + ) validate_register_region_count(register_regions) global_te.register_buffer(register_regions.ptrs, register_regions.lengths) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_worker.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_worker.py index 65eb14e063e0..7248d51d0424 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_worker.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_worker.py @@ -57,6 +57,8 @@ class MooncakePullRecvingThread(threading.Thread, MooncakeTransferAddressMixin): submission belong in _handle_requests. """ + _allow_packed_mla_destination = False + def __init__( self, engine: Any, @@ -143,6 +145,9 @@ def __init__( self.invalid_block_ids_lock = threading.Lock() assert self.local_metadata is not None + def _after_transfer(self, local_addresses: list[int], lengths: list[int]) -> None: + return + def add_requests( self, remote_engine_id: str, @@ -981,6 +986,25 @@ def _append_spec_transfer_addresses( for (local_layer_index, remote_layer_index), transfer_entries in transfer_entries_by_layer.items(): local_base_addrs = self.kv_caches_base_addr[local_layer_index] remote_base_addrs = remote_tp_metadata.kv_caches_base_addr[remote_layer_index] + local_block_strides = self.block_strides[local_layer_index] + local_block_lens = self.block_lens[local_layer_index] + if ( + self._allow_packed_mla_destination + and isinstance(spec, MLAAttentionSpec) + and len(local_base_addrs) == 1 + and len(remote_base_addrs) > 1 + ): + remote_lens = remote_metadata.block_lens[remote_layer_index] + if len(remote_lens) != len(remote_base_addrs) or sum(remote_lens) != local_block_lens[0]: + raise ValueError( + f"HiSparse packed Source length mismatch for {self.layer_names[local_layer_index]!r}" + ) + offsets = [0] + for length in remote_lens[:-1]: + offsets.append(offsets[-1] + length) + local_base_addrs = [local_base_addrs[0] + offset for offset in offsets] + local_block_strides = local_block_strides * len(remote_base_addrs) + local_block_lens = remote_lens if len(local_base_addrs) != len(remote_base_addrs): raise ValueError( f"Mooncake KV tensor count mismatch for layer {self.layer_names[local_layer_index]!r}: " @@ -998,9 +1022,9 @@ def _append_spec_transfer_addresses( f"{self.layer_names[local_layer_index]!r}, cache {cache_index}: " f"local={local_shape}, remote={remote_shape}" ) - local_block_stride = self.block_strides[local_layer_index][cache_index] + local_block_stride = local_block_strides[cache_index] remote_block_stride = remote_metadata.block_strides[remote_layer_index][cache_index] - local_block_len = self.block_lens[local_layer_index][cache_index] + local_block_len = local_block_lens[cache_index] remote_block_len = remote_metadata.block_lens[remote_layer_index][cache_index] if local_block_len != remote_block_len: raise ValueError( @@ -1203,6 +1227,7 @@ def _execute_tp_transfer_bucket( "Mooncake KV transfer failed for remote " f"PP/PCP/TP rank {remote_pp_rank}/{remote_pcp_rank}/{remote_tp_rank}, ret={ret}" ) + self._after_transfer(src_list, length_list) transfer_elapsed_ms = (time.perf_counter() - transfer_start_time) * 1000 logger.info( "KV cache transfer for requests %s took %.2f ms (%d ranges, %d bytes). " @@ -1324,6 +1349,8 @@ def _get_remote_metadata( class MooncakePullConnectorWorker(MooncakeBaseConnectorWorker): """Worker-side framework for Mooncake pull transfers.""" + _receiving_thread_cls = MooncakePullRecvingThread + def __init__( self, vllm_config: VllmConfig, @@ -1343,7 +1370,7 @@ def register_kv_caches( if not isinstance(self.xfer_handshake_metadata, MooncakeTransferMetadata): raise RuntimeError("Mooncake KV caches were registered without transfer metadata") ready_event = threading.Event() - self._recving_thread = MooncakePullRecvingThread( + self._recving_thread = self._receiving_thread_cls( engine=self.engine, vllm_config=self.vllm_config, kv_cache_config=self.kv_cache_config, @@ -1364,12 +1391,16 @@ def register_kv_caches( device=torch.npu.current_device(), ready_event=ready_event, ) + self._configure_receiving_thread(self._recving_thread) self._recving_thread.start() if not ready_event.wait(timeout=10): raise RuntimeError("Timed out starting Mooncake pull receiving thread") if not self._recving_thread.is_alive(): raise RuntimeError("Mooncake pull receiving thread failed to start") + def _configure_receiving_thread(self, thread: MooncakePullRecvingThread) -> None: + return + def start_load_kv(self, metadata: MooncakeConnectorMetadata) -> None: """Start the pull operations described by scheduler metadata.""" if not metadata.requests: diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py index 0d5e39610400..14d862a8a46b 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py @@ -177,6 +177,7 @@ def _get_tensor_span_nbytes(tensor: torch.Tensor) -> int: def collect_configured_register_regions( kv_cache_config: "KVCacheConfig", kv_caches: dict[str, Any], + included_layer_names: set[str] | None = None, ) -> RegisterRegions: """Collect one registration range per configured backing storage.""" ranges_by_storage: dict[int, tuple[int, int]] = {} @@ -198,6 +199,8 @@ def merge_storage_range( for tensor_config in kv_cache_config.kv_cache_tensors: layer_names = get_kv_cache_tensor_layers(tensor_config) + if included_layer_names is not None: + layer_names = [name for name in layer_names if name in included_layer_names] if not layer_names: continue