[Feature][HiSparse] Import P KV into Ascend Decode Host Source

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

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.

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]]

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])

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

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",

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}")