From 012682c18a3c99f0d70a3fede425c99b720092e7 Mon Sep 17 00:00:00 2001 From: LifeOfCoffee Date: Fri, 13 Nov 2026 11:19:26 +0800 Subject: [PATCH 1/2] feat(hisparse): import prefilling KV over Ascend PD transport Add a dedicated HiSparse PD connector that reuses Mooncake transfer while routing main KV through Host Source and Indexer KV to the NPU. Repack split MLA components by token before completing the Host import. Tested GLM 5.2 10-layer TP2 P and TP2 D on Ascend 950 with one-page and two-page prefix continuation; both matched local Decode output. The remote test used a mixed Python installation, so an exact branch build remains for CI. Signed-off-by: LifeOfCoffee --- .../user_guide/feature_guide/hisparse.md | 28 ++- .../mooncake_v2/test_hisparse_pd.py | 74 +++++++ .../mooncake_v2/test_pull_worker.py | 31 +++ tests/ut/kv_offload/mooncake_v2/test_utils.py | 19 ++ .../distributed/kv_transfer/__init__.py | 6 + .../kv_transfer/kv_p2p/hisparse_pd.py | 180 ++++++++++++++++++ .../kv_p2p/mooncake/base_worker.py | 33 +++- .../kv_p2p/mooncake/pull_worker.py | 37 +++- .../kv_transfer/kv_p2p/mooncake/utils.py | 3 + 9 files changed, 400 insertions(+), 11 deletions(-) create mode 100644 tests/ut/kv_offload/mooncake_v2/test_hisparse_pd.py create mode 100644 vllm_ascend/distributed/kv_transfer/kv_p2p/hisparse_pd.py 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/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 From d524e8ee6eec9cd60c1aff0d949dfbeb27d5582b Mon Sep 17 00:00:00 2001 From: LifeOfCoffee Date: Fri, 13 Nov 2026 11:20:50 +0800 Subject: [PATCH 2/2] docs(hisparse): add PD import release note Signed-off-by: LifeOfCoffee --- release_notes/fragments/2.yaml | 4 ++++ 1 file changed, 4 insertions(+) create mode 100644 release_notes/fragments/2.yaml 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.