From 9a4fb8f34f160cd367db33a8ddd4def07b5aa391 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Thu, 30 Jul 2026 10:38:56 +0800 Subject: [PATCH 01/48] add mooncake connector v2 base class Signed-off-by: nwpu-zxr --- .../kv_transfer/kv_p2p/mooncake/__init__.py | 0 .../kv_p2p/mooncake/base_scheduler.py | 273 ++++++++++++++++ .../kv_p2p/mooncake/base_worker.py | 66 ++++ .../kv_transfer/kv_p2p/mooncake/connector.py | 300 ++++++++++++++++++ .../kv_transfer/kv_p2p/mooncake/metadata.py | 16 + .../kv_p2p/mooncake/pull_scheduler.py | 191 +++++++++++ .../kv_p2p/mooncake/pull_worker.py | 14 + .../kv_transfer/kv_p2p/mooncake/stats.py | 177 +++++++++++ .../kv_transfer/kv_p2p/mooncake/utils.py | 0 9 files changed, 1037 insertions(+) create mode 100644 vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py create mode 100644 vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py create mode 100644 vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_worker.py create mode 100644 vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py create mode 100644 vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py create mode 100644 vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py create mode 100644 vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_worker.py create mode 100644 vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/stats.py create mode 100644 vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py new file mode 100644 index 000000000000..c402b5ffc3c2 --- /dev/null +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project +"""Common scheduler-side logic for Mooncake KV transfer connectors.""" + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorHandshakeMetadata, + KVConnectorMetadata, +) +from vllm.logger import init_logger +from vllm.utils.math_utils import cdiv +from vllm.utils.network_utils import get_ip +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.kv_cache_interface import ( + MambaSpec, + SlidingWindowSpec, + UniformTypeKVCacheSpecs, +) +from vllm.v1.outputs import KVConnectorOutput + +from vllm_ascend.ascend_config import ( + get_ascend_config, + init_ascend_config, +) + +if TYPE_CHECKING: + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.request import Request + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class GroupTransferInfo: + """Transfer-relevant properties of one KV cache group.""" + + tokens_per_block: int + blocks_per_window: int + is_state_group: bool + + +class MooncakeBaseConnectorScheduler: + """Scheduler logic shared by Mooncake transfer modes.""" + + def __init__( + self, + vllm_config: VllmConfig, + engine_id: str, + kv_cache_config: "KVCacheConfig", + ) -> None: + assert vllm_config.kv_transfer_config is not None + + self.vllm_config = vllm_config + self.kv_transfer_config = vllm_config.kv_transfer_config + self.kv_cache_config = kv_cache_config + self.engine_id = engine_id + self.block_size = vllm_config.cache_config.block_size + + init_ascend_config(vllm_config) + self.ascend_config = get_ascend_config() + + self.local_ip = get_ip() + self.side_channel_host = self.local_ip + self.pp_size = vllm_config.parallel_config.pipeline_parallel_size + self.tp_size = vllm_config.parallel_config.tensor_parallel_size + self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size + self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size + self.max_device_id = ( + self.tp_size + * vllm_config.parallel_config.data_parallel_size + * self.pcp_size + * vllm_config.parallel_config.pipeline_parallel_size + ) + + # Handshake base port for this DP group. + self.side_channel_port = ( + self.kv_transfer_config.kv_port + + vllm_config.parallel_config.data_parallel_rank + * self.tp_size + * vllm_config.parallel_config.pipeline_parallel_size + * self.pcp_size + ) + + # Worker metadata for a DP group that may span multiple nodes. + self.multi_nodes_meta_mapping: dict[str, dict[str, Any]] = {} + + self.kv_cache_groups = kv_cache_config.kv_cache_groups + self.use_compress = self._model_uses_compress() + self.group_transfer_info = [self._get_group_transfer_info(group) for group in self.kv_cache_groups] + self.need_truncate = self.use_compress or any(info.is_state_group for info in self.group_transfer_info) + + logger.info("Initializing Mooncake Scheduler %s", engine_id) + + def _model_uses_compress(self) -> bool: + hf_config = getattr(self.vllm_config.model_config, "hf_config", None) + compress_ratios = getattr(hf_config, "compress_ratios", None) + return isinstance(compress_ratios, (list, tuple, dict)) + + def _get_group_transfer_info(self, group: Any) -> GroupTransferInfo: + specs = self._get_group_unique_specs(group) + first_spec = specs[0] if specs else group.kv_cache_spec + block_size = getattr( + group.kv_cache_spec, + "block_size", + getattr(first_spec, "block_size", self.block_size), + ) + is_state_group = any(isinstance(spec, MambaSpec) for spec in specs) + sliding_window = 0 + compress_ratio = 1 + for spec in specs: + if isinstance(spec, SlidingWindowSpec): + sliding_window = spec.sliding_window + elif hasattr(spec, "compress_ratio"): + compress_ratio = spec.compress_ratio + + return GroupTransferInfo( + tokens_per_block=block_size * max(1, int(compress_ratio)), + blocks_per_window=(cdiv(sliding_window, block_size) + 1 if sliding_window else 0), + is_state_group=is_state_group, + ) + + @staticmethod + def _get_group_unique_specs(group: Any) -> list[Any]: + if not isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs): + return [group.kv_cache_spec] + + specs = [] + for layer_name in group.layer_names: + layer_spec = group.kv_cache_spec.kv_cache_specs[layer_name] + if layer_spec not in specs: + specs.append(layer_spec) + return specs + + def _get_transfer_block_ids(self, block_ids: BlockIds, prompt_len: int) -> BlockIds: + """Drop non-prompt attention blocks while retaining state groups.""" + if not block_ids: + return block_ids + + assert len(block_ids) == len(self.group_transfer_info), "Number of KV cache groups must match" + + transfer_block_ids = [] + cp_size = max(1, self.pcp_size * self.dcp_size) + for blocks, group_info in zip(block_ids, self.group_transfer_info): + if group_info.is_state_group: + transfer_block_ids.append(blocks) + else: + num_prompt_blocks = cdiv( + prompt_len, + group_info.tokens_per_block * cp_size, + ) + transfer_block_ids.append(blocks[:num_prompt_blocks]) + return tuple(transfer_block_ids) + + def _get_swa_transfer_block_ids(self, block_ids: BlockIds) -> BlockIds: + """Clip SWA groups to their window tail and drop block zero.""" + if not block_ids: + return block_ids + + assert len(block_ids) == len(self.group_transfer_info), "Number of KV cache groups must match" + + transfer_block_ids = [] + for blocks, group_info in zip(block_ids, self.group_transfer_info): + if group_info.is_state_group or group_info.blocks_per_window == 0: + transfer_block_ids.append(blocks) + else: + window_blocks = blocks[-group_info.blocks_per_window :] + transfer_block_ids.append([block_id for block_id in window_blocks if block_id != 0]) + return tuple(transfer_block_ids) + + def _state_prefill_token_count(self, num_prompt_tokens: int) -> int: + """Return the prompt token count transferred for stateful models.""" + if self.need_truncate and num_prompt_tokens > 1: + return num_prompt_tokens - 1 + return num_prompt_tokens + + def _truncate_request_for_prefill(self, request: "Request") -> None: + """Drop the last P-side prompt token for stateful model transfer.""" + params = request.kv_transfer_params + if params is None or params.get("_p_side_truncated") or request.num_prompt_tokens <= 1: + return + + if request.prompt_token_ids is not None: + request.prompt_token_ids.pop() + elif request.prompt_embeds is not None: + request.prompt_embeds = request.prompt_embeds[:-1] + else: + return + + request._all_token_ids.pop() + request.num_prompt_tokens -= 1 + request.max_tokens = 1 + params["_p_side_truncated"] = True + + def on_new_request(self, request: "Request") -> None: + """Mooncake currently requires no request-arrival bookkeeping.""" + + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: + """Mooncake currently requires no worker-output bookkeeping.""" + + def get_num_new_matched_tokens(self, request: "Request", num_computed_tokens: int) -> tuple[int | None, bool]: + raise NotImplementedError + + def update_state_after_alloc( + self, + request: "Request", + blocks: "KVCacheBlocks", + num_external_tokens: int, + ) -> None: + raise NotImplementedError + + def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata: + raise NotImplementedError + + def request_finished( + self, + request: "Request", + block_ids: BlockIds, + ) -> tuple[bool, dict[str, Any] | None]: + raise NotImplementedError + + def _port_offset_from_handshake_metadata( + self, + rank_metadata: KVConnectorHandshakeMetadata, + metadata_key: int | tuple[int, ...], + ) -> int: + handshake_port = getattr(rank_metadata, "handshake_port", 0) + if handshake_port > 0: + return handshake_port - self.kv_transfer_config.kv_port + if isinstance(metadata_key, int): + return metadata_key + raise ValueError(f"Mooncake handshake metadata is missing handshake_port for worker key {metadata_key}") + + def set_xfer_handshake_metadata_from_workers( + self, + metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], + ) -> None: + """Build the worker host mapping for a possibly multi-node DP group.""" + if not metadata: + return + + updated_mapping: dict[str, dict[str, Any]] = {} + kv_port = self.kv_transfer_config.kv_port + for metadata_key, rank_metadata in metadata.items(): + port_offset = self._port_offset_from_handshake_metadata(rank_metadata, metadata_key) + updated_mapping[str(port_offset)] = { + "host": rank_metadata.local_ip, + "engine_id": rank_metadata.engine_id, + "handshake_port": kv_port + port_offset, + } + + self.multi_nodes_meta_mapping.update(updated_mapping) + logger.info( + "MooncakeConnector set_xfer_handshake_metadata: worker_count=%d, updated=%s, multi_nodes_meta_mapping=%s", + len(metadata), + updated_mapping, + self.multi_nodes_meta_mapping, + ) + + def set_xfer_handshake_metadata( + self, + metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], + ) -> None: + """Handle the legacy port-offset-keyed handshake entry point.""" + self.set_xfer_handshake_metadata_from_workers(metadata) + + +__all__ = ["GroupTransferInfo", "MooncakeBaseConnectorScheduler"] 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 new file mode 100644 index 000000000000..a073daa57f03 --- /dev/null +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_worker.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project +"""Worker-side interface for Mooncake KV transfer connectors. + +This module intentionally contains no transfer implementation. Memory +registration, Mooncake engine lifecycle, rank mapping, and D2D execution belong +in concrete worker implementations. +""" + +from typing import TYPE_CHECKING + +import torch +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorHandshakeMetadata, +) +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + KVConnectorStats, +) + +from .metadata import MooncakeConnectorMetadata + +if TYPE_CHECKING: + from vllm.v1.kv_cache_interface import KVCacheConfig + + +class MooncakeBaseConnectorWorker: + """Common worker-side contract for Mooncake connectors.""" + + def __init__( + self, + vllm_config: VllmConfig, + engine_id: str, + kv_cache_config: "KVCacheConfig", + ) -> None: + self.vllm_config = vllm_config + self.engine_id = engine_id + self.kv_cache_config = kv_cache_config + self.xfer_handshake_metadata: ( + KVConnectorHandshakeMetadata | None + ) = None + + def register_kv_caches( + self, kv_caches: dict[str, torch.Tensor] + ) -> None: + """Register model KV cache tensors for D2D transfer.""" + raise NotImplementedError + + def get_finished(self) -> tuple[set[str], set[str]]: + """Return requests with completed receive and send operations.""" + raise NotImplementedError + + def get_block_ids_with_load_errors(self) -> set[int]: + """Return local block IDs whose KV load failed.""" + raise NotImplementedError + + def get_kv_connector_stats(self) -> KVConnectorStats | None: + """Return and reset transfer statistics for the current interval.""" + raise NotImplementedError + + def start_load_kv(self, metadata: MooncakeConnectorMetadata) -> None: + """Start D2D KV loading described by scheduler metadata.""" + raise NotImplementedError + + +__all__ = ["MooncakeBaseConnectorWorker"] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py new file mode 100644 index 000000000000..f03c14ba6394 --- /dev/null +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py @@ -0,0 +1,300 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project +"""Mooncake connector facades. + +The connector classes in this module bridge vLLM's KV connector interface to +the scheduler-side and worker-side Mooncake implementations. Transfer and +scheduling details belong in those implementations rather than in this facade. + +``MooncakeConnector`` remains an alias for the pull-based connector. The push +connector is declared as part of the new architecture but is not implemented. +""" + +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any + +import torch +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.utils import EngineId +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorBase_V1, + KVConnectorHandshakeMetadata, + KVConnectorMetadata, + KVConnectorRole, + SupportsHMA, +) +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + KVConnectorPromMetrics, + KVConnectorStats, + PromMetric, + PromMetricT, +) +from vllm.logger import init_logger +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.outputs import KVConnectorOutput + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler import ( + MooncakeBaseConnectorScheduler, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker import ( + MooncakeBaseConnectorWorker, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( + MooncakeConnectorMetadata, +) +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, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.stats import ( + MooncakeKVConnectorStats, + MooncakePromMetrics, +) + +if TYPE_CHECKING: + from vllm.forward_context import ForwardContext + from vllm.v1.attention.backend import AttentionMetadata + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.request import Request + +logger = init_logger(__name__) + + +class MooncakeBaseConnector(KVConnectorBase_V1, SupportsHMA): + """Common facade for the Mooncake scheduler and worker implementations.""" + + def __init__( + self, + vllm_config: VllmConfig, + role: KVConnectorRole, + kv_cache_config: "KVCacheConfig", + ) -> None: + super().__init__(vllm_config, role, kv_cache_config) + assert vllm_config.kv_transfer_config is not None + assert vllm_config.kv_transfer_config.engine_id is not None + + self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id + self.kv_cache_config = kv_cache_config + self.kv_transfer_config = vllm_config.kv_transfer_config + self.connector_scheduler: MooncakeBaseConnectorScheduler | None = None + self.connector_worker: MooncakeBaseConnectorWorker | None = None + + @classmethod + def get_required_kvcache_layout( + cls, vllm_config: VllmConfig + ) -> str | None: + """Use HND to avoid reformatting KV cache before D2D transfer.""" + if vllm_config.model_config is None: + logger.warning_once( + "Unable to detect the current vLLM model config; falling " + "back to the default KV cache layout." + ) + return None + + if vllm_config.model_config.use_mla: + # MLA has no separate head dimension, so HND versus NHD does not + # affect its transfer layout. + return None + + logger.info_once( + "MooncakeConnector is setting the KV cache layout to HND for " + "direct D2D transfer." + ) + return "HND" + + ############################################################ + # Scheduler-side methods + ############################################################ + + def get_num_new_matched_tokens( + self, request: "Request", num_computed_tokens: int + ) -> tuple[int | None, bool]: + assert self.connector_scheduler is not None + return self.connector_scheduler.get_num_new_matched_tokens( + request, num_computed_tokens + ) + + def update_state_after_alloc( + self, + request: "Request", + blocks: "KVCacheBlocks", + num_external_tokens: int, + ) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.update_state_after_alloc( + request, blocks, num_external_tokens + ) + + def build_connector_meta( + self, scheduler_output: SchedulerOutput + ) -> KVConnectorMetadata: + assert self.connector_scheduler is not None + return self.connector_scheduler.build_connector_meta(scheduler_output) + + def on_new_request(self, request: "Request") -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.on_new_request(request) + + def update_connector_output( + self, connector_output: KVConnectorOutput + ) -> None: + assert self.connector_scheduler is not None + self.connector_scheduler.update_connector_output(connector_output) + + def request_finished( + self, + request: "Request", + block_ids: list[int], + ) -> tuple[bool, dict[str, Any] | None]: + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished(request, (block_ids,)) + + def request_finished_all_groups( + self, + request: "Request", + block_ids: tuple[list[int], ...], + ) -> tuple[bool, dict[str, Any] | None]: + assert self.connector_scheduler is not None + return self.connector_scheduler.request_finished(request, block_ids) + + def set_xfer_handshake_metadata( + self, + metadata: Mapping[ + int | tuple[int, ...], KVConnectorHandshakeMetadata + ], + ) -> None: + """Set worker handshake metadata on the scheduler.""" + assert self.connector_scheduler is not None + self.connector_scheduler.set_xfer_handshake_metadata(metadata) + + def set_xfer_handshake_metadata_pp_aware( + self, + metadata: Mapping[ + int | tuple[int, ...], KVConnectorHandshakeMetadata + ], + ) -> None: + """Set PP-aware worker handshake metadata on the scheduler.""" + assert self.connector_scheduler is not None + self.connector_scheduler.set_xfer_handshake_metadata_from_workers(metadata) + + ############################################################ + # Worker-side methods + ############################################################ + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + assert self.connector_worker is not None + self.connector_worker.register_kv_caches(kv_caches) + + def get_finished( + self, finished_req_ids: set[str] + ) -> tuple[set[str], set[str]]: + """Return requests whose receive and send operations have finished.""" + assert self.connector_worker is not None + return self.connector_worker.get_finished() + + def get_block_ids_with_load_errors(self) -> set[int]: + """Return block IDs whose Mooncake KV load failed.""" + assert self.connector_worker is not None + return self.connector_worker.get_block_ids_with_load_errors() + + def get_kv_connector_stats(self) -> KVConnectorStats | None: + if self.connector_worker is None: + return None + return self.connector_worker.get_kv_connector_stats() + + @classmethod + def build_kv_connector_stats( + cls, data: dict[str, Any] | None = None + ) -> KVConnectorStats: + return MooncakeKVConnectorStats(data=data or {}) + + @classmethod + def build_prom_metrics( + cls, + vllm_config: VllmConfig, + metric_types: dict[type[PromMetric], type[PromMetricT]], + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ) -> KVConnectorPromMetrics: + return MooncakePromMetrics( + vllm_config, metric_types, labelnames, per_engine_labelvalues + ) + + def start_load_kv( + self, forward_context: "ForwardContext", **kwargs: Any + ) -> None: + assert self.connector_worker is not None + assert isinstance(self._connector_metadata, MooncakeConnectorMetadata) + self.connector_worker.start_load_kv(self._connector_metadata) + + def wait_for_layer_load(self, layer_name: str) -> None: + """MooncakeConnector does not load KV cache layer by layer.""" + + def save_kv_layer( + self, + layer_name: str, + kv_layer: torch.Tensor, + attn_metadata: "AttentionMetadata", + **kwargs: Any, + ) -> None: + """MooncakeConnector does not save KV cache explicitly.""" + + def wait_for_save(self) -> None: + """MooncakeConnector does not save KV cache explicitly.""" + + def get_handshake_metadata(self) -> KVConnectorHandshakeMetadata | None: + """Return this worker's out-of-band transfer handshake metadata.""" + assert self.connector_worker is not None + return self.connector_worker.xfer_handshake_metadata + + +class MooncakePullConnector(MooncakeBaseConnector): + """Pull-based Mooncake KV transfer connector.""" + + def __init__( + self, + vllm_config: VllmConfig, + role: KVConnectorRole, + kv_cache_config: "KVCacheConfig", + ) -> None: + super().__init__(vllm_config, role, kv_cache_config) + + if role == KVConnectorRole.SCHEDULER: + self.connector_scheduler = MooncakePullConnectorScheduler( + vllm_config, str(self.engine_id), kv_cache_config + ) + elif role == KVConnectorRole.WORKER: + self.connector_worker = MooncakePullConnectorWorker( + vllm_config, str(self.engine_id), kv_cache_config + ) + else: + raise ValueError(f"Unsupported KVConnectorRole: {role}") + + +class MooncakePushConnector(MooncakeBaseConnector): + """Placeholder for the not-yet-implemented push-based connector.""" + + def __init__( + self, + vllm_config: VllmConfig, + role: KVConnectorRole, + kv_cache_config: "KVCacheConfig", + ) -> None: + raise NotImplementedError( + "MooncakePushConnector is not implemented yet." + ) + + +# Backward compatibility: MooncakeConnector is the pull-based connector. +MooncakeConnector = MooncakePullConnector + + +__all__ = [ + "MooncakeBaseConnector", + "MooncakeConnector", + "MooncakePullConnector", + "MooncakePushConnector", +] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py new file mode 100644 index 000000000000..7e5bfe56647b --- /dev/null +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project +"""Metadata types for Mooncake KV transfer connectors.""" + +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorMetadata, +) + + +class MooncakeConnectorMetadata(KVConnectorMetadata): + """Scheduler-to-worker metadata for Mooncake KV transfers.""" + + pass + + +__all__ = ["MooncakeConnectorMetadata"] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py new file mode 100644 index 000000000000..6c1442f7b001 --- /dev/null +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project +"""Scheduler-side logic for Mooncake pull transfers.""" + +import math +import time +from typing import TYPE_CHECKING, Any + +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorMetadata, +) +from vllm.logger import init_logger +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.request import RequestStatus + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler import ( + MooncakeBaseConnectorScheduler, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( + MooncakeConnectorMetadata, +) + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.core.kv_cache_manager import KVCacheBlocks + from vllm.v1.kv_cache_interface import KVCacheConfig + from vllm.v1.request import Request + +logger = init_logger(__name__) + + +class MooncakePullConnectorScheduler(MooncakeBaseConnectorScheduler): + """Scheduler-side Mooncake pull connector implementation.""" + + def __init__( + self, + vllm_config: "VllmConfig", + engine_id: str, + kv_cache_config: "KVCacheConfig", + ) -> None: + super().__init__(vllm_config, engine_id, kv_cache_config) + + # Requests waiting for the worker to start a READ transfer. + self._reqs_need_recv: dict[str, tuple[Request, BlockIds, BlockIds, int]] = {} + # Producer requests whose blocks must remain allocated until read. + self._reqs_need_send: dict[str, float] = {} + self._reqs_in_batch: set[str] = set() + + def get_num_new_matched_tokens(self, request: "Request", num_computed_tokens: int) -> tuple[int, bool]: + """Return prompt tokens that will be loaded from a remote producer.""" + params = request.kv_transfer_params + logger.debug( + "MooncakeConnector get_num_new_matched_tokens: num_computed_tokens=%s, kv_transfer_params=%s", + num_computed_tokens, + params, + ) + + if params is not None and params.get("do_remote_prefill"): + token_ids = request.prompt_token_ids or [] + actual = self._state_prefill_token_count(len(token_ids)) + params["num_computed_tokens"] = num_computed_tokens + count = max(actual - num_computed_tokens, 0) + if count > 0: + return count, True + + if params is not None and params.get("do_remote_decode") and self.need_truncate: + self._truncate_request_for_prefill(request) + + return 0, False + + def update_state_after_alloc( + self, + request: "Request", + blocks: "KVCacheBlocks", + num_external_tokens: int, + ) -> None: + params = request.kv_transfer_params + logger.debug( + "MooncakeConnector update_state_after_alloc: num_external_tokens=%s, kv_transfer_params=%s", + num_external_tokens, + params, + ) + + if params is not None and (params.get("do_remote_prefill", False) or params.get("do_remote_decode", False)): + self._reqs_in_batch.add(request.request_id) + + if params is None or not params.get("do_remote_prefill"): + return + + if params.get("remote_block_ids"): + required_remote_fields = ( + "remote_engine_id", + "remote_host", + "remote_port", + "remote_request_id", + ) + if all(field in params for field in required_remote_fields): + local_block_ids = blocks.get_unhashed_block_ids_all_groups() if num_external_tokens > 0 else [] + local_full_block_ids = blocks.get_block_ids() if num_external_tokens > 0 else tuple() + self._reqs_need_recv[request.request_id] = ( + request, + local_block_ids, + local_full_block_ids, + num_external_tokens, + ) + else: + logger.warning("Got invalid KVTransferParams. params=%s.", params) + else: + assert num_external_tokens == 0 + + # Only trigger one transfer for a request. + params["do_remote_prefill"] = False + + def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata: + meta = MooncakeConnectorMetadata() + + for ( + req_id, + (req, block_ids, full_block_ids, num_external_tokens), + ) in self._reqs_need_recv.items(): + assert req.kv_transfer_params is not None + meta.add_new_req( + request_id=req_id, + local_block_ids=block_ids, + local_full_block_ids=full_block_ids, + num_external_tokens=num_external_tokens, + kv_transfer_params=req.kv_transfer_params, + ) + + self._reqs_need_recv.clear() + meta.requests_to_send = self._reqs_need_send + self._reqs_need_send = {} + meta.reqs_in_batch = self._reqs_in_batch + self._reqs_in_batch = set() + return meta + + def request_finished( + self, + request: "Request", + block_ids: BlockIds, + ) -> tuple[bool, dict[str, Any] | None]: + """Expose completed producer blocks for a remote READ transfer.""" + params = request.kv_transfer_params + logger.debug( + "MooncakeConnector request_finished: request_status=%s, kv_transfer_params=%s", + request.status, + params, + ) + + if ( + params is None + or not params.get("do_remote_decode") + or request.status != RequestStatus.FINISHED_LENGTH_CAPPED + ): + return False, None + + prompt_token_ids = request.prompt_token_ids or [] + prompt_len = len(prompt_token_ids) + num_prompt_blocks = math.ceil(prompt_len / self.block_size) + computed_block_ids = self._get_transfer_block_ids(block_ids, prompt_len) + computed_block_ids = self._get_swa_transfer_block_ids(computed_block_ids) + num_computed_blocks = sum(len(group_block_ids) for group_block_ids in computed_block_ids) + delay_free_blocks = num_computed_blocks > 0 + if delay_free_blocks: + logger.info( + "Delaying free of %d blocks for request %s", + num_computed_blocks, + request.request_id, + ) + self._reqs_need_send[request.request_id] = time.time() + + return delay_free_blocks, { + "do_remote_prefill": True, + "do_remote_decode": False, + "remote_block_ids": computed_block_ids, + "remote_engine_id": self.engine_id, + "remote_request_id": request.request_id, + "remote_host": self.side_channel_host, + "remote_port": self.side_channel_port, + "remote_pcp_size": self.pcp_size, + "remote_dcp_size": self.dcp_size, + "remote_ptp_size": self.tp_size, + "last_token_id": request.output_token_ids[-1], + "remote_multi_nodes_meta_mapping": (self.multi_nodes_meta_mapping), + "num_prompt_blocks": num_prompt_blocks, + "remote_block_size": self.block_size, + } + + +__all__ = ["MooncakePullConnectorScheduler"] 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 new file mode 100644 index 000000000000..3c7327e8096e --- /dev/null +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_worker.py @@ -0,0 +1,14 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project +"""Worker-side implementation entry point for Mooncake pull transfers.""" + +from .base_worker import MooncakeBaseConnectorWorker + + +class MooncakePullConnectorWorker(MooncakeBaseConnectorWorker): + """Worker-side Mooncake pull connector implementation.""" + + pass + + +__all__ = ["MooncakePullConnectorWorker"] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/stats.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/stats.py new file mode 100644 index 000000000000..495a6c2e896b --- /dev/null +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/stats.py @@ -0,0 +1,177 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project +"""Transfer statistics and Prometheus metrics for Mooncake.""" + +import copy +from dataclasses import dataclass +from typing import Any + +import numpy as np +from vllm.config import VllmConfig +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( + KVConnectorPromMetrics, + KVConnectorStats, + PromMetric, + PromMetricT, +) +from vllm.v1.metrics.utils import create_metric_per_engine + + +@dataclass +class MooncakeKVConnectorStats(KVConnectorStats): + """D2D transfer statistics collected by a Mooncake worker.""" + + def __post_init__(self) -> None: + if not self.data: + self.reset() + + def reset(self) -> None: + # Values must remain serializable because worker stats are aggregated + # outside of the worker process. + self.data: dict[str, list[float | int]] = { + "transfer_duration": [], + "bytes_transferred": [], + "num_failed_transfers": [], + } + + def record_transfer( + self, transfer_duration: float, bytes_transferred: int + ) -> None: + """Record one completed D2D transfer. + + ``transfer_duration`` is expressed in seconds. + """ + self.data["transfer_duration"].append(transfer_duration) + self.data["bytes_transferred"].append(bytes_transferred) + + def record_failed_transfer(self) -> None: + self.data["num_failed_transfers"].append(1) + + def clone_and_reset(self) -> "MooncakeKVConnectorStats": + previous = copy.copy(self) + self.reset() + return previous + + def is_empty(self) -> bool: + return ( + self.num_successful_transfers == 0 + and not self.data["num_failed_transfers"] + ) + + def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: + if not other.is_empty(): + for key, values in other.data.items(): + accumulator = self.data[key] + assert isinstance(accumulator, list) + accumulator.extend(values) + return self + + def reduce(self) -> dict[str, int | float]: + if self.num_successful_transfers == 0: + return { + "Num successful transfers": 0, + "Num failed transfers": len( + self.data["num_failed_transfers"] + ), + "Avg xfer time (ms)": 0, + "P90 xfer time (ms)": 0, + "Avg MB per transfer": 0, + "Throughput (MB/s)": 0, + } + + durations = np.asarray(self.data["transfer_duration"]) + megabytes = np.asarray(self.data["bytes_transferred"]) / 2**20 + total_duration = durations.sum() + throughput = ( + megabytes.sum() / total_duration if total_duration > 0 else 0 + ) + + return { + "Num successful transfers": self.num_successful_transfers, + "Num failed transfers": len(self.data["num_failed_transfers"]), + "Avg xfer time (ms)": round(durations.mean() * 1e3, 3), + "P90 xfer time (ms)": round( + np.percentile(durations, 90).item() * 1e3, 3 + ), + "Avg MB per transfer": round(megabytes.mean(), 3), + "Throughput (MB/s)": round(float(throughput), 3), + } + + @property + def num_successful_transfers(self) -> int: + return len(self.data["transfer_duration"]) + + +class MooncakePromMetrics(KVConnectorPromMetrics): + """Prometheus metrics for Mooncake D2D transfers.""" + + def __init__( + self, + vllm_config: VllmConfig, + metric_types: dict[type[PromMetric], type[PromMetricT]], + labelnames: list[str], + per_engine_labelvalues: dict[int, list[object]], + ) -> None: + super().__init__( + vllm_config, + metric_types, + labelnames, + per_engine_labelvalues, + ) + + transfer_duration = self._histogram_cls( + name="vllm:mooncake_xfer_time_seconds", + documentation=( + "Histogram of Mooncake KV cache D2D transfer duration." + ), + buckets=[ + 0.005, + 0.01, + 0.025, + 0.05, + 0.075, + 0.1, + 0.2, + 0.3, + 0.5, + 0.75, + 1.0, + 5.0, + ], + labelnames=labelnames, + ) + self.transfer_duration = create_metric_per_engine( + transfer_duration, self.per_engine_labelvalues + ) + + bytes_transferred = self._histogram_cls( + name="vllm:mooncake_bytes_transferred", + documentation=( + "Histogram of bytes transferred per Mooncake KV cache D2D " + "transfer." + ), + buckets=[2 ** (10 + index) for index in range(1, 25, 2)], + labelnames=labelnames, + ) + self.bytes_transferred = create_metric_per_engine( + bytes_transferred, self.per_engine_labelvalues + ) + + failed_transfers = self._counter_cls( + name="vllm:mooncake_num_failed_transfers", + documentation="Number of failed Mooncake KV cache D2D transfers.", + labelnames=labelnames, + ) + self.failed_transfers = create_metric_per_engine( + failed_transfers, self.per_engine_labelvalues + ) + + def observe( + self, transfer_stats_data: dict[str, Any], engine_idx: int = 0 + ) -> None: + for duration in transfer_stats_data["transfer_duration"]: + self.transfer_duration[engine_idx].observe(duration) + for num_bytes in transfer_stats_data["bytes_transferred"]: + self.bytes_transferred[engine_idx].observe(num_bytes) + for failure in transfer_stats_data["num_failed_transfers"]: + self.failed_transfers[engine_idx].inc(failure) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py new file mode 100644 index 000000000000..e69de29bb2d1 From 70f89c94de67b009ec776fba206293a03007cba5 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Thu, 30 Jul 2026 23:10:24 +0800 Subject: [PATCH 02/48] add base scheduler worker Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_scheduler.py | 110 ++++++--------- .../kv_p2p/mooncake/base_worker.py | 133 +++++++++++++++--- .../kv_transfer/kv_p2p/mooncake/metadata.py | 4 +- .../kv_p2p/mooncake/pull_scheduler.py | 1 - 4 files changed, 155 insertions(+), 93 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py index c402b5ffc3c2..958b3dff0c88 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py @@ -3,7 +3,6 @@ """Common scheduler-side logic for Mooncake KV transfer connectors.""" from collections.abc import Mapping -from dataclasses import dataclass from typing import TYPE_CHECKING, Any from vllm.config import VllmConfig @@ -18,7 +17,6 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( MambaSpec, - SlidingWindowSpec, UniformTypeKVCacheSpecs, ) from vllm.v1.outputs import KVConnectorOutput @@ -36,15 +34,6 @@ logger = init_logger(__name__) -@dataclass(frozen=True) -class GroupTransferInfo: - """Transfer-relevant properties of one KV cache group.""" - - tokens_per_block: int - blocks_per_window: int - is_state_group: bool - - class MooncakeBaseConnectorScheduler: """Scheduler logic shared by Mooncake transfer modes.""" @@ -61,6 +50,11 @@ def __init__( self.kv_cache_config = kv_cache_config self.engine_id = engine_id self.block_size = vllm_config.cache_config.block_size + self.num_speculative_tokens = ( + vllm_config.speculative_config.num_speculative_tokens + if vllm_config.speculative_config is not None + else 0 + ) init_ascend_config(vllm_config) self.ascend_config = get_ascend_config() @@ -81,49 +75,28 @@ def __init__( # Handshake base port for this DP group. self.side_channel_port = ( self.kv_transfer_config.kv_port - + vllm_config.parallel_config.data_parallel_rank - * self.tp_size - * vllm_config.parallel_config.pipeline_parallel_size - * self.pcp_size + + vllm_config.parallel_config.data_parallel_rank * self.tp_size * self.pp_size * self.pcp_size ) # Worker metadata for a DP group that may span multiple nodes. self.multi_nodes_meta_mapping: dict[str, dict[str, Any]] = {} self.kv_cache_groups = kv_cache_config.kv_cache_groups - self.use_compress = self._model_uses_compress() - self.group_transfer_info = [self._get_group_transfer_info(group) for group in self.kv_cache_groups] - self.need_truncate = self.use_compress or any(info.is_state_group for info in self.group_transfer_info) + self.group_block_size = [group.kv_cache_spec.block_size for group in self.kv_cache_groups] + self.group_unique_specs = [self._get_group_unique_specs(group) for group in self.kv_cache_groups] + self.need_truncate = self._needs_prefill_token_truncation() logger.info("Initializing Mooncake Scheduler %s", engine_id) - def _model_uses_compress(self) -> bool: + def _needs_prefill_token_truncation(self) -> bool: + """Return whether Prefill must leave the last prompt token to Decode.""" hf_config = getattr(self.vllm_config.model_config, "hf_config", None) compress_ratios = getattr(hf_config, "compress_ratios", None) - return isinstance(compress_ratios, (list, tuple, dict)) - - def _get_group_transfer_info(self, group: Any) -> GroupTransferInfo: - specs = self._get_group_unique_specs(group) - first_spec = specs[0] if specs else group.kv_cache_spec - block_size = getattr( - group.kv_cache_spec, - "block_size", - getattr(first_spec, "block_size", self.block_size), - ) - is_state_group = any(isinstance(spec, MambaSpec) for spec in specs) - sliding_window = 0 - compress_ratio = 1 - for spec in specs: - if isinstance(spec, SlidingWindowSpec): - sliding_window = spec.sliding_window - elif hasattr(spec, "compress_ratio"): - compress_ratio = spec.compress_ratio - - return GroupTransferInfo( - tokens_per_block=block_size * max(1, int(compress_ratio)), - blocks_per_window=(cdiv(sliding_window, block_size) + 1 if sliding_window else 0), - is_state_group=is_state_group, + uses_compressed_cache = isinstance(compress_ratios, (list, tuple, dict)) + has_state_group = any( + isinstance(spec, MambaSpec) for group_specs in self.group_unique_specs for spec in group_specs ) + return uses_compressed_cache or has_state_group @staticmethod def _get_group_unique_specs(group: Any) -> list[Any]: @@ -138,41 +111,33 @@ def _get_group_unique_specs(group: Any) -> list[Any]: return specs def _get_transfer_block_ids(self, block_ids: BlockIds, prompt_len: int) -> BlockIds: - """Drop non-prompt attention blocks while retaining state groups.""" + """Return prompt blocks while retaining evicted-block padding. + + Attention groups are clipped by their block size. Mamba groups are not + prompt-block aligned, so only their speculative tail blocks are removed. + """ if not block_ids: return block_ids - assert len(block_ids) == len(self.group_transfer_info), "Number of KV cache groups must match" + assert len(block_ids) == len(self.group_unique_specs), "Number of KV cache groups must match" - transfer_block_ids = [] + transfer_block_ids: list[list[int]] = [] cp_size = max(1, self.pcp_size * self.dcp_size) - for blocks, group_info in zip(block_ids, self.group_transfer_info): - if group_info.is_state_group: - transfer_block_ids.append(blocks) + for blocks, block_size, group_specs in zip( + block_ids, + self.group_block_size, + self.group_unique_specs, + ): + if any(isinstance(spec, MambaSpec) for spec in group_specs): + if self.num_speculative_tokens > 0: + transfer_block_ids.append(blocks[: -self.num_speculative_tokens]) + else: + transfer_block_ids.append(blocks) else: - num_prompt_blocks = cdiv( - prompt_len, - group_info.tokens_per_block * cp_size, - ) + num_prompt_blocks = cdiv(prompt_len, block_size * cp_size) transfer_block_ids.append(blocks[:num_prompt_blocks]) return tuple(transfer_block_ids) - def _get_swa_transfer_block_ids(self, block_ids: BlockIds) -> BlockIds: - """Clip SWA groups to their window tail and drop block zero.""" - if not block_ids: - return block_ids - - assert len(block_ids) == len(self.group_transfer_info), "Number of KV cache groups must match" - - transfer_block_ids = [] - for blocks, group_info in zip(block_ids, self.group_transfer_info): - if group_info.is_state_group or group_info.blocks_per_window == 0: - transfer_block_ids.append(blocks) - else: - window_blocks = blocks[-group_info.blocks_per_window :] - transfer_block_ids.append([block_id for block_id in window_blocks if block_id != 0]) - return tuple(transfer_block_ids) - def _state_prefill_token_count(self, num_prompt_tokens: int) -> int: """Return the prompt token count transferred for stateful models.""" if self.need_truncate and num_prompt_tokens > 1: @@ -182,7 +147,12 @@ def _state_prefill_token_count(self, num_prompt_tokens: int) -> int: def _truncate_request_for_prefill(self, request: "Request") -> None: """Drop the last P-side prompt token for stateful model transfer.""" params = request.kv_transfer_params - if params is None or params.get("_p_side_truncated") or request.num_prompt_tokens <= 1: + if ( + params is None + or not self.need_truncate + or params.get("_p_side_truncated") + or request.num_prompt_tokens <= 1 + ): return if request.prompt_token_ids is not None: @@ -270,4 +240,4 @@ def set_xfer_handshake_metadata( self.set_xfer_handshake_metadata_from_workers(metadata) -__all__ = ["GroupTransferInfo", "MooncakeBaseConnectorScheduler"] +__all__ = ["MooncakeBaseConnectorScheduler"] 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 a073daa57f03..3ebb2ceed474 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 @@ -1,31 +1,52 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project -"""Worker-side interface for Mooncake KV transfer connectors. +"""Common worker-side logic for Mooncake KV transfer connectors.""" -This module intentionally contains no transfer implementation. Memory -registration, Mooncake engine lifecycle, rank mapping, and D2D execution belong -in concrete worker implementations. -""" - -from typing import TYPE_CHECKING +import os +from typing import TYPE_CHECKING, Any import torch +import torch_npu # noqa: F401 from vllm.config import VllmConfig +from vllm.distributed import get_pcp_group from vllm.distributed.kv_transfer.kv_connector.v1.base import ( KVConnectorHandshakeMetadata, ) -from vllm.distributed.kv_transfer.kv_connector.v1.metrics import ( - KVConnectorStats, +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats +from vllm.distributed.parallel_state import ( + get_pp_group, + get_tensor_model_parallel_rank, + get_tp_group, ) +from vllm.logger import init_logger +from vllm.utils.network_utils import get_ip -from .metadata import MooncakeConnectorMetadata +from vllm_ascend.ascend_config import get_ascend_config, init_ascend_config +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( + MooncakeConnectorMetadata, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.stats import ( + MooncakeKVConnectorStats, +) +from vllm_ascend.distributed.kv_transfer.utils.mooncake_transfer_engine import ( + global_te, +) +from vllm_ascend.distributed.kv_transfer.utils.utils import ( + get_transfer_timeout_value, +) +from vllm_ascend.distributed.utils import ( + get_decode_context_model_parallel_rank, + get_decode_context_model_parallel_world_size, +) if TYPE_CHECKING: from vllm.v1.kv_cache_interface import KVCacheConfig +logger = init_logger(__name__) + class MooncakeBaseConnectorWorker: - """Common worker-side contract for Mooncake connectors.""" + """Worker implementation shared by Mooncake transfer modes.""" def __init__( self, @@ -34,16 +55,88 @@ def __init__( kv_cache_config: "KVCacheConfig", ) -> None: self.vllm_config = vllm_config + self.kv_transfer_config = vllm_config.kv_transfer_config self.engine_id = engine_id self.kv_cache_config = kv_cache_config - self.xfer_handshake_metadata: ( - KVConnectorHandshakeMetadata | None - ) = None + self.block_size = vllm_config.cache_config.block_size + self.num_blocks = kv_cache_config.num_blocks - def register_kv_caches( - self, kv_caches: dict[str, torch.Tensor] - ) -> None: - """Register model KV cache tensors for D2D transfer.""" + init_ascend_config(vllm_config) + self.ascend_config = get_ascend_config() + self._get_prefill_decode_size() + os.environ["ASCEND_TRANSFER_TIMEOUT"] = str(get_transfer_timeout_value()) + if self._prefill_tp_size < self._decode_tp_size: + raise ValueError( + f"prefill_tp_size: {self._prefill_tp_size} must be greater than " + f"or equal to decode_tp_size: {self._decode_tp_size}" + ) + + self.tp_rank = get_tensor_model_parallel_rank() + self.tp_size = vllm_config.parallel_config.tensor_parallel_size + self.tp_group = get_tp_group() + self.pp_rank = get_pp_group().rank_in_group + self.pp_size = vllm_config.parallel_config.pipeline_parallel_size + self.dp_rank = vllm_config.parallel_config.data_parallel_rank_local + self.dp_size = vllm_config.parallel_config.data_parallel_size_local + self.pcp_size = get_pcp_group().world_size + self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0 + self.dcp_size = get_decode_context_model_parallel_world_size() + self.dcp_rank = ( + get_decode_context_model_parallel_rank() if self.dcp_size > 1 else 0 + ) + if self.pp_size > 1 and self.pcp_size > 1: + raise ValueError("PP and PCP cannot be enabled at the same time") + + self.max_device_id = self.tp_size * self.dp_size * self.pcp_size * self.pp_size + self.kv_role = self.kv_transfer_config.kv_role + self.side_channel_host = get_ip() + self.side_channel_port = ( + self.kv_transfer_config.kv_port + + vllm_config.parallel_config.data_parallel_rank + * self.tp_size + * self.pp_size + * self.pcp_size + ) + device_index = ( + self.pp_rank * self.pcp_size + self.pcp_rank + ) * self.tp_size + self.tp_rank + self.handshake_port = self.side_channel_port + device_index + + device_name = str(torch.npu.current_device()) if self.pp_size > 1 else None + self.engine = global_te.get_transfer_engine( + self.side_channel_host, + device_name=device_name, + ) + self.te_rpc_port = self.engine.get_rpc_port() + self.xfer_handshake_metadata: KVConnectorHandshakeMetadata | None = None + self.xfer_stats = MooncakeKVConnectorStats() + + logger.info("Initializing Mooncake worker %s", engine_id) + + def _get_prefill_decode_size(self) -> None: + prefill_config: dict[str, Any] = ( + self.kv_transfer_config.get_from_extra_config("prefill", {}) + ) + decode_config: dict[str, Any] = ( + self.kv_transfer_config.get_from_extra_config("decode", {}) + ) + if "tp_size" not in prefill_config or "dp_size" not in prefill_config: + raise ValueError("Mooncake prefill config requires tp_size and dp_size") + if "tp_size" not in decode_config or "dp_size" not in decode_config: + raise ValueError("Mooncake decode config requires tp_size and dp_size") + + self._prefill_tp_size = prefill_config["tp_size"] + self._prefill_dp_size = prefill_config["dp_size"] + self._prefill_pp_size = prefill_config.get("pp_size", 1) + self._prefill_pp_layer_partition = prefill_config.get("pp_layer_partition") + self._decode_tp_size = decode_config["tp_size"] + self._decode_dp_size = decode_config["dp_size"] + self._decode_pp_size = decode_config.get("pp_size", 1) + if self._decode_pp_size != 1: + raise ValueError("Decode pipeline parallel size must be 1") + + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: + """Register model KV caches for D2D transfer.""" raise NotImplementedError def get_finished(self) -> tuple[set[str], set[str]]: @@ -56,7 +149,9 @@ def get_block_ids_with_load_errors(self) -> set[int]: def get_kv_connector_stats(self) -> KVConnectorStats | None: """Return and reset transfer statistics for the current interval.""" - raise NotImplementedError + if self.xfer_stats.is_empty(): + return None + return self.xfer_stats.clone_and_reset() def start_load_kv(self, metadata: MooncakeConnectorMetadata) -> None: """Start D2D KV loading described by scheduler metadata.""" diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 7e5bfe56647b..2bab6809601d 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -2,9 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project """Metadata types for Mooncake KV transfer connectors.""" -from vllm.distributed.kv_transfer.kv_connector.v1.base import ( - KVConnectorMetadata, -) +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata class MooncakeConnectorMetadata(KVConnectorMetadata): diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 6c1442f7b001..4db07c3b3eb6 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -159,7 +159,6 @@ def request_finished( prompt_len = len(prompt_token_ids) num_prompt_blocks = math.ceil(prompt_len / self.block_size) computed_block_ids = self._get_transfer_block_ids(block_ids, prompt_len) - computed_block_ids = self._get_swa_transfer_block_ids(computed_block_ids) num_computed_blocks = sum(len(group_block_ids) for group_block_ids in computed_block_ids) delay_free_blocks = num_computed_blocks > 0 if delay_free_blocks: From b6beface26f28861023a4748c1a70dc55c28f2ba Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Fri, 31 Jul 2026 22:56:05 +0800 Subject: [PATCH 03/48] add base worker Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_worker.py | 30 +---- .../kv_p2p/mooncake/pull_worker.py | 36 +++++- .../kv_transfer/kv_p2p/mooncake/utils.py | 111 ++++++++++++++++++ 3 files changed, 145 insertions(+), 32 deletions(-) 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 3ebb2ceed474..56891ba1436b 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 @@ -3,7 +3,7 @@ """Common worker-side logic for Mooncake KV transfer connectors.""" import os -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import torch import torch_npu # noqa: F401 @@ -63,13 +63,7 @@ def __init__( init_ascend_config(vllm_config) self.ascend_config = get_ascend_config() - self._get_prefill_decode_size() os.environ["ASCEND_TRANSFER_TIMEOUT"] = str(get_transfer_timeout_value()) - if self._prefill_tp_size < self._decode_tp_size: - raise ValueError( - f"prefill_tp_size: {self._prefill_tp_size} must be greater than " - f"or equal to decode_tp_size: {self._decode_tp_size}" - ) self.tp_rank = get_tensor_model_parallel_rank() self.tp_size = vllm_config.parallel_config.tensor_parallel_size @@ -113,28 +107,6 @@ def __init__( logger.info("Initializing Mooncake worker %s", engine_id) - def _get_prefill_decode_size(self) -> None: - prefill_config: dict[str, Any] = ( - self.kv_transfer_config.get_from_extra_config("prefill", {}) - ) - decode_config: dict[str, Any] = ( - self.kv_transfer_config.get_from_extra_config("decode", {}) - ) - if "tp_size" not in prefill_config or "dp_size" not in prefill_config: - raise ValueError("Mooncake prefill config requires tp_size and dp_size") - if "tp_size" not in decode_config or "dp_size" not in decode_config: - raise ValueError("Mooncake decode config requires tp_size and dp_size") - - self._prefill_tp_size = prefill_config["tp_size"] - self._prefill_dp_size = prefill_config["dp_size"] - self._prefill_pp_size = prefill_config.get("pp_size", 1) - self._prefill_pp_layer_partition = prefill_config.get("pp_layer_partition") - self._decode_tp_size = decode_config["tp_size"] - self._decode_dp_size = decode_config["dp_size"] - self._decode_pp_size = decode_config.get("pp_size", 1) - if self._decode_pp_size != 1: - raise ValueError("Decode pipeline parallel size must be 1") - def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: """Register model KV caches for D2D transfer.""" raise NotImplementedError 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 3c7327e8096e..dd3f7cf22866 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 @@ -2,13 +2,43 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project """Worker-side implementation entry point for Mooncake pull transfers.""" -from .base_worker import MooncakeBaseConnectorWorker +from typing import TYPE_CHECKING + +from vllm.config import VllmConfig + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker import ( + MooncakeBaseConnectorWorker, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( + MooncakeConnectorMetadata, +) + +if TYPE_CHECKING: + from vllm.v1.kv_cache_interface import KVCacheConfig class MooncakePullConnectorWorker(MooncakeBaseConnectorWorker): - """Worker-side Mooncake pull connector implementation.""" + """Worker-side framework for Mooncake pull transfers.""" + + def __init__( + self, + vllm_config: VllmConfig, + engine_id: str, + kv_cache_config: "KVCacheConfig", + ) -> None: + super().__init__(vllm_config, engine_id, kv_cache_config) + + def start_load_kv(self, metadata: MooncakeConnectorMetadata) -> None: + """Start the pull operations described by scheduler metadata.""" + raise NotImplementedError + + def get_finished(self) -> tuple[set[str], set[str]]: + """Return requests with completed receive and send operations.""" + raise NotImplementedError - pass + def get_block_ids_with_load_errors(self) -> set[int]: + """Return local block IDs whose pull operations failed.""" + raise NotImplementedError __all__ = ["MooncakePullConnectorWorker"] 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 e69de29bb2d1..bf9891d42d29 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py @@ -0,0 +1,111 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project +"""Shared utilities for Mooncake KV transfer connectors.""" + +import hashlib +import struct +import time + +import numpy as np +import numpy.typing as npt +import zmq +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def group_concurrent_contiguous( + src: list[int], + dst: list[int], + src_block_stride: int = 1, + dst_block_stride: int = 1, + block_len: int = 1, +) -> tuple[list[list[int]], list[list[int]]]: + """Group block ids that are contiguous in both id space and memory.""" + src_indices: npt.NDArray[np.int64] = np.array(src, dtype=np.int64) + dst_indices: npt.NDArray[np.int64] = np.array(dst, dtype=np.int64) + + if src_indices.size == 0: + return [], [] + + src_byte_contiguous = np.diff(src_indices) * src_block_stride == block_len + dst_byte_contiguous = np.diff(dst_indices) * dst_block_stride == block_len + brk = np.where(~(src_byte_contiguous & dst_byte_contiguous))[0] + 1 + src_groups = np.split(src_indices, brk) + dst_groups = np.split(dst_indices, brk) + + src_groups = [g.tolist() for g in src_groups] + dst_groups = [g.tolist() for g in dst_groups] + + return src_groups, dst_groups + + +def split_if_not_byte_contiguous( + src_groups: list[list[int]], + dst_groups: list[list[int]], + src_block_stride: int, + dst_block_stride: int, + block_len: int, +) -> tuple[list[list[int]], list[list[int]]]: + if src_block_stride == block_len and dst_block_stride == block_len: + return src_groups, dst_groups + + src = [bid for group in src_groups for bid in group] + dst = [bid for group in dst_groups for bid in group] + return group_concurrent_contiguous( + src, + dst, + src_block_stride=src_block_stride, + dst_block_stride=dst_block_stride, + block_len=block_len, + ) + + +def string_to_int64_hash(input_str): + """ + Hash the string using SHA-256 and convert it into an int64 integer. + """ + hashed_bytes = hashlib.sha256(input_str.encode("utf-8")).digest() + trunked_bytes = hashed_bytes[:8] + uint64_value = struct.unpack(" 0: + logger.warning("Send failed. error=%s, attempts_left=%d. ", e, retries_left) + time.sleep(0.1) + else: + logger.error("Send failed after all retries. error=%s. ", e) + raise RuntimeError(f"Failed to send data to {path} after {max_retries} retries: {e}") + + +def ensure_zmq_recv( + socket: zmq.Socket, # type: ignore + path: str, + max_retries: int = 3, +) -> bytes: + retries_left = max_retries + while True: + try: + return socket.recv() + except zmq.ZMQError as e: # type: ignore + retries_left -= 1 + if retries_left > 0: + logger.warning("Receive failed. error=%s, attempts_left=%d. ", e, retries_left) + time.sleep(0.1) + else: + logger.error("Receive failed after all retries. source=%s, error=%s. ", path, e) + raise RuntimeError(f"Failed to receive data after {max_retries} retries: {e}") From 18948f6353b80b320459385dfea49cfc24dcd2a1 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Sun, 2 Aug 2026 17:58:53 +0800 Subject: [PATCH 04/48] add basic worker Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_scheduler.py | 2 + .../kv_p2p/mooncake/base_worker.py | 157 ++++++++++++++++-- .../kv_transfer/kv_p2p/mooncake/metadata.py | 47 +++++- .../kv_p2p/mooncake/pull_scheduler.py | 23 +-- .../kv_transfer/kv_p2p/mooncake/utils.py | 34 ++++ 5 files changed, 238 insertions(+), 25 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py index 958b3dff0c88..6350cc323c86 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py @@ -169,9 +169,11 @@ def _truncate_request_for_prefill(self, request: "Request") -> None: def on_new_request(self, request: "Request") -> None: """Mooncake currently requires no request-arrival bookkeeping.""" + raise NotImplementedError def update_connector_output(self, connector_output: KVConnectorOutput) -> None: """Mooncake currently requires no worker-output bookkeeping.""" + raise NotImplementedError def get_num_new_matched_tokens(self, request: "Request", num_computed_tokens: int) -> tuple[int | None, bool]: raise NotImplementedError 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 56891ba1436b..8169a6d5accf 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 @@ -2,8 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project """Common worker-side logic for Mooncake KV transfer connectors.""" +import math import os -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Any import torch import torch_npu # noqa: F401 @@ -20,6 +21,10 @@ ) from vllm.logger import init_logger from vllm.utils.network_utils import get_ip +from vllm.v1.kv_cache_interface import ( + KVCacheSpec, + UniformTypeKVCacheSpecs, +) from vllm_ascend.ascend_config import get_ascend_config, init_ascend_config from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( @@ -75,9 +80,7 @@ def __init__( self.pcp_size = get_pcp_group().world_size self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0 self.dcp_size = get_decode_context_model_parallel_world_size() - self.dcp_rank = ( - get_decode_context_model_parallel_rank() if self.dcp_size > 1 else 0 - ) + self.dcp_rank = get_decode_context_model_parallel_rank() if self.dcp_size > 1 else 0 if self.pp_size > 1 and self.pcp_size > 1: raise ValueError("PP and PCP cannot be enabled at the same time") @@ -86,14 +89,9 @@ def __init__( self.side_channel_host = get_ip() self.side_channel_port = ( self.kv_transfer_config.kv_port - + vllm_config.parallel_config.data_parallel_rank - * self.tp_size - * self.pp_size - * self.pcp_size + + vllm_config.parallel_config.data_parallel_rank * self.tp_size * self.pp_size * self.pcp_size ) - device_index = ( - self.pp_rank * self.pcp_size + self.pcp_rank - ) * self.tp_size + self.tp_rank + device_index = (self.pp_rank * self.pcp_size + self.pcp_rank) * self.tp_size + self.tp_rank self.handshake_port = self.side_channel_port + device_index device_name = str(torch.npu.current_device()) if self.pp_size > 1 else None @@ -107,9 +105,142 @@ def __init__( logger.info("Initializing Mooncake worker %s", engine_id) + def _build_kv_cache_spec_mappings(self) -> None: + """Flatten group specs and map each layer to its group and spec. + + A regular KV cache group contributes one spec. A uniform-type group + contributes each distinct inner spec in layer order. Spec uniqueness is + scoped to a group because different groups have different block tables, + even when their specs compare equal. + """ + self.kv_cache_specs: list[KVCacheSpec] = [] + self.layer_name_to_group_index: dict[str, int] = {} + self.layer_name_to_spec_index: dict[str, int] = {} + + for group_index, group in enumerate(self.kv_cache_config.kv_cache_groups): + group_spec = group.kv_cache_spec + group_spec_indices: list[int] = [] + + for layer_name in group.layer_names: + if layer_name in self.layer_name_to_group_index: + raise ValueError(f"Layer {layer_name!r} belongs to more than one KV cache group.") + + if isinstance(group_spec, UniformTypeKVCacheSpecs): + try: + layer_spec = group_spec.kv_cache_specs[layer_name] + except KeyError as exc: + raise ValueError( + f"Uniform KV cache group {group_index} has no spec for layer {layer_name!r}." + ) from exc + else: + layer_spec = group_spec + + spec_index = next( + (index for index in group_spec_indices if self.kv_cache_specs[index] == layer_spec), + -1, + ) + if spec_index < 0: + spec_index = len(self.kv_cache_specs) + self.kv_cache_specs.append(layer_spec) + group_spec_indices.append(spec_index) + + self.layer_name_to_group_index[layer_name] = group_index + self.layer_name_to_spec_index[layer_name] = spec_index + + @staticmethod + def _as_kv_cache_tensors(cache_or_caches: Any) -> tuple[torch.Tensor, ...]: + if isinstance(cache_or_caches, torch.Tensor): + return (cache_or_caches,) + if isinstance(cache_or_caches, (list, tuple)) and all( + isinstance(cache, torch.Tensor) for cache in cache_or_caches + ): + return tuple(cache_or_caches) + raise TypeError( + f"A layer KV cache must be a tensor or a list/tuple of tensors, but got {type(cache_or_caches).__name__}." + ) + def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: - """Register model KV caches for D2D transfer.""" - raise NotImplementedError + """Collect per-layer KV cache metadata for D2D registration.""" + self.num_blocks = self.kv_cache_config.num_blocks + logger.info("num_blocks: %s", self.num_blocks) + self.kv_caches = kv_caches + self._build_kv_cache_spec_mappings() + + # All lists below use registered_layer_names order. Each inner list + # preserves the layer cache order, such as K/V or Mamba conv/SSM. + self.registered_layer_names: list[str] = [] + self.registered_group_indices: list[int] = [] + self.registered_spec_indices: list[int] = [] + self.kv_caches_base_addr: list[list[int]] = [] + self.block_stride_per_addr: list[list[int]] = [] + self.block_len_per_addr: list[list[int]] = [] + self.block_shape_per_addr: list[list[tuple[int, ...]]] = [] + self.block_size_scale: list[list[int]] = [] + + for layer_name in self.layer_name_to_group_index: + cache_or_caches = kv_caches.get(layer_name) + if cache_or_caches is None: + logger.debug( + "Skipping layer %s because it has no worker KV cache tensor.", + layer_name, + ) + continue + + base_addrs: list[int] = [] + block_strides: list[int] = [] + block_lens: list[int] = [] + block_shapes: list[tuple[int, ...]] = [] + block_size_scales: list[int] = [] + + for cache in self._as_kv_cache_tensors(cache_or_caches): + if cache.ndim == 0: + raise ValueError(f"KV cache tensor for layer {layer_name!r} must have a block dimension.") + + tensor_num_blocks = cache.shape[0] + if self.num_blocks <= 0 or tensor_num_blocks % self.num_blocks: + raise ValueError( + f"KV cache tensor for layer {layer_name!r} has " + f"{tensor_num_blocks} physical blocks, which is not a " + f"multiple of {self.num_blocks} logical blocks." + ) + + element_size = cache.element_size() + block_shape = tuple(cache.shape[1:]) + base_addrs.append(cache.data_ptr()) + 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) + + self.registered_layer_names.append(layer_name) + self.registered_group_indices.append(self.layer_name_to_group_index[layer_name]) + self.registered_spec_indices.append(self.layer_name_to_spec_index[layer_name]) + self.kv_caches_base_addr.append(base_addrs) + self.block_stride_per_addr.append(block_strides) + self.block_len_per_addr.append(block_lens) + self.block_shape_per_addr.append(block_shapes) + self.block_size_scale.append(block_size_scales) + + unknown_layers = kv_caches.keys() - self.layer_name_to_group_index.keys() + if unknown_layers: + logger.debug( + "Ignored KV cache tensors without a local KV cache group: %s", + sorted(unknown_layers), + ) + + logger.debug( + "Mooncake KV cache metadata: layers=%s, group_indices=%s, " + "spec_indices=%s, base_addrs=%s, block_strides=%s, block_lens=%s, " + "block_shapes=%s, block_size_scale=%s", + self.registered_layer_names, + self.registered_group_indices, + self.registered_spec_indices, + self.kv_caches_base_addr, + self.block_stride_per_addr, + self.block_len_per_addr, + self.block_shape_per_addr, + self.block_size_scale, + ) def get_finished(self) -> tuple[set[str], set[str]]: """Return requests with completed receive and send operations.""" diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 2bab6809601d..9463d8ff1d40 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -2,13 +2,56 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project """Metadata types for Mooncake KV transfer connectors.""" +from dataclasses import dataclass +from typing import Any + +from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata +@dataclass +class ReqMeta: + """Request metadata required by the initial Mooncake pull path.""" + + local_block_ids: BlockIds + num_external_tokens: int + num_computed_tokens: int + remote_block_ids: BlockIds + remote_host: str + remote_port: int + remote_engine_id: str + remote_request_id: str + num_prompt_blocks: int + remote_block_size: int + + class MooncakeConnectorMetadata(KVConnectorMetadata): """Scheduler-to-worker metadata for Mooncake KV transfers.""" - pass + def __init__(self) -> None: + self.requests: dict[str, ReqMeta] = {} + self.requests_to_send: dict[str, float] = {} + self.reqs_in_batch: set[str] = set() + + def add_new_req( + self, + request_id: str, + local_block_ids: BlockIds, + num_external_tokens: int, + kv_transfer_params: dict[str, Any], + ) -> None: + self.requests[request_id] = ReqMeta( + local_block_ids=local_block_ids, + num_external_tokens=num_external_tokens, + num_computed_tokens=kv_transfer_params.get("num_computed_tokens", 0), + remote_block_ids=kv_transfer_params["remote_block_ids"], + remote_host=kv_transfer_params["remote_host"], + remote_port=kv_transfer_params["remote_port"], + remote_engine_id=kv_transfer_params["remote_engine_id"], + remote_request_id=kv_transfer_params["remote_request_id"], + num_prompt_blocks=kv_transfer_params.get("num_prompt_blocks", 0), + remote_block_size=kv_transfer_params.get("remote_block_size", 0), + ) -__all__ = ["MooncakeConnectorMetadata"] +__all__ = ["MooncakeConnectorMetadata", "ReqMeta"] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 4db07c3b3eb6..9a410e5aa786 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -12,6 +12,7 @@ ) from vllm.logger import init_logger from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import RequestStatus from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler import ( @@ -42,7 +43,7 @@ def __init__( super().__init__(vllm_config, engine_id, kv_cache_config) # Requests waiting for the worker to start a READ transfer. - self._reqs_need_recv: dict[str, tuple[Request, BlockIds, BlockIds, int]] = {} + self._reqs_need_recv: dict[str, tuple[Request, BlockIds, int]] = {} # Producer requests whose blocks must remain allocated until read. self._reqs_need_send: dict[str, float] = {} self._reqs_in_batch: set[str] = set() @@ -96,12 +97,10 @@ def update_state_after_alloc( "remote_request_id", ) if all(field in params for field in required_remote_fields): - local_block_ids = blocks.get_unhashed_block_ids_all_groups() if num_external_tokens > 0 else [] - local_full_block_ids = blocks.get_block_ids() if num_external_tokens > 0 else tuple() + local_block_ids = blocks.get_unhashed_block_ids_all_groups() if num_external_tokens > 0 else tuple() self._reqs_need_recv[request.request_id] = ( request, local_block_ids, - local_full_block_ids, num_external_tokens, ) else: @@ -117,13 +116,12 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnector for ( req_id, - (req, block_ids, full_block_ids, num_external_tokens), + (req, block_ids, num_external_tokens), ) in self._reqs_need_recv.items(): assert req.kv_transfer_params is not None meta.add_new_req( request_id=req_id, local_block_ids=block_ids, - local_full_block_ids=full_block_ids, num_external_tokens=num_external_tokens, kv_transfer_params=req.kv_transfer_params, ) @@ -177,14 +175,19 @@ def request_finished( "remote_request_id": request.request_id, "remote_host": self.side_channel_host, "remote_port": self.side_channel_port, - "remote_pcp_size": self.pcp_size, - "remote_dcp_size": self.dcp_size, - "remote_ptp_size": self.tp_size, "last_token_id": request.output_token_ids[-1], - "remote_multi_nodes_meta_mapping": (self.multi_nodes_meta_mapping), "num_prompt_blocks": num_prompt_blocks, "remote_block_size": self.block_size, } + def on_new_request(self, request: "Request") -> None: + pass + + def update_connector_output( + self, + connector_output: KVConnectorOutput, + ) -> None: + pass + __all__ = ["MooncakePullConnectorScheduler"] 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 bf9891d42d29..033a8a0b7382 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py @@ -5,6 +5,8 @@ import hashlib import struct import time +from collections import OrderedDict +from dataclasses import dataclass import numpy as np import numpy.typing as npt @@ -14,6 +16,28 @@ logger = init_logger(__name__) +@dataclass +class SizedDict(OrderedDict): + """Insertion-ordered mapping with a bounded number of entries.""" + + def __init__(self, max_size=16000, *args, **kwargs): + self.max_size = max_size + super().__init__(*args, **kwargs) + + def __setitem__(self, key, value): + super().__setitem__(key, value) + if len(self) > self.max_size: + self.popitem(last=False) + + def __getitem__(self, key): + try: + return super().__getitem__(key) + except KeyError: + value: dict[int, list[int]] = {} + self[key] = value + return value + + def group_concurrent_contiguous( src: list[int], dst: list[int], @@ -109,3 +133,13 @@ def ensure_zmq_recv( else: logger.error("Receive failed after all retries. source=%s, error=%s. ", path, e) raise RuntimeError(f"Failed to receive data after {max_retries} retries: {e}") + + +__all__ = [ + "SizedDict", + "ensure_zmq_recv", + "ensure_zmq_send", + "group_concurrent_contiguous", + "split_if_not_byte_contiguous", + "string_to_int64_hash", +] From ccbcb67123f83833d5b24b3803ae1192606c5d4c Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Mon, 3 Aug 2026 10:12:18 +0800 Subject: [PATCH 05/48] add register_kv_caches Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_worker.py | 187 +++++++++--------- .../kv_transfer/kv_p2p/mooncake/metadata.py | 68 ++++++- .../kv_transfer/kv_p2p/mooncake/utils.py | 106 ++++++++++ 3 files changed, 262 insertions(+), 99 deletions(-) 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 8169a6d5accf..98ce307512bc 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 @@ -4,7 +4,7 @@ import math import os -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING import torch import torch_npu # noqa: F401 @@ -29,15 +29,21 @@ from vllm_ascend.ascend_config import get_ascend_config, init_ascend_config from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( MooncakeConnectorMetadata, + MooncakeTransferMetadata, ) from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.stats import ( MooncakeKVConnectorStats, ) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils import ( + as_kv_cache_tensors, + collect_configured_register_regions, +) from vllm_ascend.distributed.kv_transfer.utils.mooncake_transfer_engine import ( global_te, ) from vllm_ascend.distributed.kv_transfer.utils.utils import ( get_transfer_timeout_value, + validate_register_region_count, ) from vllm_ascend.distributed.utils import ( get_decode_context_model_parallel_rank, @@ -122,19 +128,10 @@ def _build_kv_cache_spec_mappings(self) -> None: group_spec_indices: list[int] = [] for layer_name in group.layer_names: - if layer_name in self.layer_name_to_group_index: - raise ValueError(f"Layer {layer_name!r} belongs to more than one KV cache group.") - if isinstance(group_spec, UniformTypeKVCacheSpecs): - try: - layer_spec = group_spec.kv_cache_specs[layer_name] - except KeyError as exc: - raise ValueError( - f"Uniform KV cache group {group_index} has no spec for layer {layer_name!r}." - ) from exc + layer_spec = group_spec.kv_cache_specs[layer_name] else: layer_spec = group_spec - spec_index = next( (index for index in group_spec_indices if self.kv_cache_specs[index] == layer_spec), -1, @@ -147,99 +144,95 @@ def _build_kv_cache_spec_mappings(self) -> None: self.layer_name_to_group_index[layer_name] = group_index self.layer_name_to_spec_index[layer_name] = spec_index - @staticmethod - def _as_kv_cache_tensors(cache_or_caches: Any) -> tuple[torch.Tensor, ...]: - if isinstance(cache_or_caches, torch.Tensor): - return (cache_or_caches,) - if isinstance(cache_or_caches, (list, tuple)) and all( - isinstance(cache, torch.Tensor) for cache in cache_or_caches - ): - return tuple(cache_or_caches) - raise TypeError( - f"A layer KV cache must be a tensor or a list/tuple of tensors, but got {type(cache_or_caches).__name__}." - ) - - def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: - """Collect per-layer KV cache metadata for D2D registration.""" + def register_kv_caches( + self, + kv_caches: dict[str, torch.Tensor | list[torch.Tensor]], + ) -> None: + """Register configured KV cache allocations and publish metadata.""" self.num_blocks = self.kv_cache_config.num_blocks logger.info("num_blocks: %s", self.num_blocks) self.kv_caches = kv_caches self._build_kv_cache_spec_mappings() - # All lists below use registered_layer_names order. Each inner list - # preserves the layer cache order, such as K/V or Mamba conv/SSM. - self.registered_layer_names: list[str] = [] - self.registered_group_indices: list[int] = [] - self.registered_spec_indices: list[int] = [] - self.kv_caches_base_addr: list[list[int]] = [] - self.block_stride_per_addr: list[list[int]] = [] - self.block_len_per_addr: list[list[int]] = [] - self.block_shape_per_addr: list[list[tuple[int, ...]]] = [] - self.block_size_scale: list[list[int]] = [] - - for layer_name in self.layer_name_to_group_index: - cache_or_caches = kv_caches.get(layer_name) - if cache_or_caches is None: - logger.debug( - "Skipping layer %s because it has no worker KV cache tensor.", - layer_name, - ) - continue - - base_addrs: list[int] = [] - block_strides: list[int] = [] - block_lens: list[int] = [] - block_shapes: list[tuple[int, ...]] = [] - block_size_scales: list[int] = [] - - for cache in self._as_kv_cache_tensors(cache_or_caches): - if cache.ndim == 0: - raise ValueError(f"KV cache tensor for layer {layer_name!r} must have a block dimension.") - - tensor_num_blocks = cache.shape[0] - if self.num_blocks <= 0 or tensor_num_blocks % self.num_blocks: - raise ValueError( - f"KV cache tensor for layer {layer_name!r} has " - f"{tensor_num_blocks} physical blocks, which is not a " - f"multiple of {self.num_blocks} logical blocks." - ) - - element_size = cache.element_size() - block_shape = tuple(cache.shape[1:]) - base_addrs.append(cache.data_ptr()) - 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) - - self.registered_layer_names.append(layer_name) - self.registered_group_indices.append(self.layer_name_to_group_index[layer_name]) - self.registered_spec_indices.append(self.layer_name_to_spec_index[layer_name]) - self.kv_caches_base_addr.append(base_addrs) - self.block_stride_per_addr.append(block_strides) - self.block_len_per_addr.append(block_lens) - self.block_shape_per_addr.append(block_shapes) - self.block_size_scale.append(block_size_scales) - - unknown_layers = kv_caches.keys() - self.layer_name_to_group_index.keys() - if unknown_layers: - logger.debug( - "Ignored KV cache tensors without a local KV cache group: %s", - sorted(unknown_layers), - ) + layer_names: list[str] = [] + group_indices: list[int] = [] + spec_indices: list[int] = [] + kv_caches_base_addr: list[list[int]] = [] + block_strides_per_layer: list[list[int]] = [] + block_lens_per_layer: list[list[int]] = [] + block_shapes_per_layer: list[list[tuple[int, ...]]] = [] + block_size_scales_per_layer: list[list[int]] = [] + configured_layer_names: set[str] = set() + + for tensor_config in self.kv_cache_config.kv_cache_tensors: + for layer_name in tensor_config.shared_by: + if layer_name in configured_layer_names: + raise ValueError(f"Layer {layer_name!r} is referenced by more than one configured KV cache tensor.") + 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.") + + 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}.") + + base_addrs: list[int] = [] + block_strides: list[int] = [] + block_lens: list[int] = [] + block_shapes: list[tuple[int, ...]] = [] + block_size_scales: list[int] = [] + + for cache in as_kv_cache_tensors(cache_or_caches): + tensor_num_blocks = cache.shape[0] + element_size = cache.element_size() + block_shape = tuple(cache.shape[1:]) + base_addrs.append(cache.data_ptr()) + 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) + + configured_layer_names.add(layer_name) + layer_names.append(layer_name) + group_indices.append(self.layer_name_to_group_index[layer_name]) + spec_indices.append(self.layer_name_to_spec_index[layer_name]) + kv_caches_base_addr.append(base_addrs) + block_strides_per_layer.append(block_strides) + block_lens_per_layer.append(block_lens) + block_shapes_per_layer.append(block_shapes) + block_size_scales_per_layer.append(block_size_scales) + + unexpected_layers = kv_caches.keys() - configured_layer_names + 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) + validate_register_region_count(register_regions) + global_te.register_buffer(register_regions.ptrs, register_regions.lengths) + + transfer_metadata = MooncakeTransferMetadata( + engine_id=self.engine_id, + te_rpc_port=self.te_rpc_port, + block_size=self.block_size, + num_blocks=self.num_blocks, + layer_names=layer_names, + group_indices=group_indices, + spec_indices=spec_indices, + kv_caches_base_addr=kv_caches_base_addr, + block_strides=block_strides_per_layer, + block_lens=block_lens_per_layer, + block_shapes=block_shapes_per_layer, + block_size_scales=block_size_scales_per_layer, + local_ip=self.side_channel_host, + handshake_port=self.handshake_port, + ) + self.transfer_metadata = transfer_metadata + self.xfer_handshake_metadata = transfer_metadata logger.debug( - "Mooncake KV cache metadata: layers=%s, group_indices=%s, " - "spec_indices=%s, base_addrs=%s, block_strides=%s, block_lens=%s, " - "block_shapes=%s, block_size_scale=%s", - self.registered_layer_names, - self.registered_group_indices, - self.registered_spec_indices, - self.kv_caches_base_addr, - self.block_stride_per_addr, - self.block_len_per_addr, - self.block_shape_per_addr, - self.block_size_scale, + "Mooncake KV cache transfer metadata: metadata=%s, register_ptrs=%s, register_lengths=%s", + transfer_metadata, + register_regions.ptrs, + register_regions.lengths, ) def get_finished(self) -> tuple[set[str], set[str]]: diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 9463d8ff1d40..9e90fbaaf720 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -6,7 +6,67 @@ from typing import Any from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds -from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata +from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorHandshakeMetadata, + KVConnectorMetadata, +) + + +@dataclass(frozen=True) +class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): + """Worker transfer information exchanged during the P/D handshake. + + All per-layer fields use ``layer_names`` order. Each nested list preserves + the cache tensor order for that layer, such as K/V or Mamba conv/SSM. + """ + + engine_id: str + te_rpc_port: int + block_size: int + num_blocks: int + layer_names: list[str] + group_indices: list[int] + spec_indices: list[int] + kv_caches_base_addr: list[list[int]] + block_strides: list[list[int]] + block_lens: list[list[int]] + block_shapes: list[list[tuple[int, ...]]] + block_size_scales: list[list[int]] + local_ip: str = "" + handshake_port: int = 0 + + def __post_init__(self) -> None: + num_layers = len(self.layer_names) + per_layer_fields = { + "group_indices": self.group_indices, + "spec_indices": self.spec_indices, + "kv_caches_base_addr": self.kv_caches_base_addr, + "block_strides": self.block_strides, + "block_lens": self.block_lens, + "block_shapes": self.block_shapes, + "block_size_scales": self.block_size_scales, + } + for field_name, values in per_layer_fields.items(): + if len(values) != num_layers: + raise ValueError( + f"Mooncake transfer metadata field {field_name!r} has " + f"{len(values)} layers, expected {num_layers}." + ) + + for layer_index, layer_name in enumerate(self.layer_names): + num_addrs = len(self.kv_caches_base_addr[layer_index]) + for field_name in ( + "block_strides", + "block_lens", + "block_shapes", + "block_size_scales", + ): + values = per_layer_fields[field_name][layer_index] + if len(values) != num_addrs: + raise ValueError( + f"Mooncake transfer metadata for layer {layer_name!r} " + f"has {len(values)} {field_name}, expected {num_addrs}." + ) @dataclass @@ -54,4 +114,8 @@ def add_new_req( ) -__all__ = ["MooncakeConnectorMetadata", "ReqMeta"] +__all__ = [ + "MooncakeConnectorMetadata", + "MooncakeTransferMetadata", + "ReqMeta", +] 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 033a8a0b7382..72d4b5b23abc 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py @@ -7,15 +7,119 @@ import time from collections import OrderedDict from dataclasses import dataclass +from typing import TYPE_CHECKING, Any import numpy as np import numpy.typing as npt +import torch import zmq from vllm.logger import init_logger +from vllm_ascend.distributed.kv_transfer.utils.utils import ( + RegisterRegions, + tensor_storage_key, +) + +if TYPE_CHECKING: + from vllm.v1.kv_cache_interface import KVCacheConfig + logger = init_logger(__name__) +def as_kv_cache_tensors(cache_or_caches: Any) -> tuple[torch.Tensor, ...]: + """Normalize one layer's KV cache into a tuple of tensors.""" + if isinstance(cache_or_caches, torch.Tensor): + return (cache_or_caches,) + if isinstance(cache_or_caches, (list, tuple)) and all(isinstance(cache, torch.Tensor) for cache in cache_or_caches): + return tuple(cache_or_caches) + raise TypeError( + f"A layer KV cache must be a tensor or a list/tuple of tensors, but got {type(cache_or_caches).__name__}." + ) + + +def _get_storage_nbytes(tensor: torch.Tensor) -> int: + """Return the byte size of the allocation backing tensor.""" + try: + return tensor.untyped_storage().nbytes() + except Exception: + try: + return tensor.storage().nbytes() + except Exception: + return tensor.nbytes + + +def collect_configured_register_regions( + kv_cache_config: "KVCacheConfig", + kv_caches: dict[str, Any], +) -> RegisterRegions: + """Collect one registration range per configured backing storage.""" + ranges_by_storage: dict[int, tuple[int, int]] = {} + configured_tensor_count = 0 + + def merge_storage_range( + storage_key: int, + register_start: int, + register_end: int, + ) -> None: + previous_range = ranges_by_storage.get(storage_key) + if previous_range is None: + ranges_by_storage[storage_key] = (register_start, register_end) + else: + ranges_by_storage[storage_key] = ( + min(previous_range[0], register_start), + max(previous_range[1], register_end), + ) + + for tensor_config in kv_cache_config.kv_cache_tensors: + if not tensor_config.shared_by: + continue + + cache_tensors: list[torch.Tensor] = [] + for layer_name in tensor_config.shared_by: + cache_tensors.extend(as_kv_cache_tensors(kv_caches.get(layer_name))) + + caches_by_storage: dict[int, list[torch.Tensor]] = {} + for cache in cache_tensors: + storage_key = tensor_storage_key(cache) + caches_by_storage.setdefault(storage_key, []).append(cache) + + if len(caches_by_storage) == 1: + # KVCacheTensor.size is authoritative when all layer views belong + # to its one configured allocation. Packed layouts expose each + # view at allocation_base + offset. + storage_key, storage_caches = next(iter(caches_by_storage.items())) + register_start = min(cache.data_ptr() for cache in storage_caches) - tensor_config.offset + merge_storage_range( + storage_key, + register_start, + register_start + tensor_config.size, + ) + else: + # Some cache types initialize component tensors independently. + # The config size cannot be applied to every storage, so register + # the actual range of each allocation. + for storage_key, storage_caches in caches_by_storage.items(): + register_start = min(cache.data_ptr() for cache in storage_caches) + register_end = max(storage_key + _get_storage_nbytes(cache) for cache in storage_caches) + if register_end <= register_start: + raise ValueError(f"Invalid KV cache storage range: start={register_start}, end={register_end}.") + merge_storage_range(storage_key, register_start, register_end) + + configured_tensor_count += 1 + + if not ranges_by_storage: + raise ValueError("KV cache config contains no registerable tensors.") + + ptrs = [region[0] for region in ranges_by_storage.values()] + lengths = [region[1] - region[0] for region in ranges_by_storage.values()] + return RegisterRegions( + ptrs=ptrs, + lengths=lengths, + logical_tensor_count=configured_tensor_count, + logical_total_bytes=sum(lengths), + ) + + @dataclass class SizedDict(OrderedDict): """Insertion-ordered mapping with a bounded number of entries.""" @@ -137,6 +241,8 @@ def ensure_zmq_recv( __all__ = [ "SizedDict", + "as_kv_cache_tensors", + "collect_configured_register_regions", "ensure_zmq_recv", "ensure_zmq_send", "group_concurrent_contiguous", From c1d03274e56d253eb6bcf5f585675f3c443c5259 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Mon, 3 Aug 2026 16:35:45 +0800 Subject: [PATCH 06/48] add scheduler send/recv thread Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_scheduler.py | 16 +- .../kv_transfer/kv_p2p/mooncake/connector.py | 68 ++--- .../kv_p2p/mooncake/pull_scheduler.py | 259 +++++++++++++++++- .../kv_transfer/kv_p2p/mooncake/utils.py | 24 ++ 4 files changed, 303 insertions(+), 64 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py index 6350cc323c86..781766fe7dcc 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py @@ -47,13 +47,12 @@ def __init__( self.vllm_config = vllm_config self.kv_transfer_config = vllm_config.kv_transfer_config + self.kv_role = self.kv_transfer_config.kv_role self.kv_cache_config = kv_cache_config self.engine_id = engine_id self.block_size = vllm_config.cache_config.block_size self.num_speculative_tokens = ( - vllm_config.speculative_config.num_speculative_tokens - if vllm_config.speculative_config is not None - else 0 + vllm_config.speculative_config.num_speculative_tokens if vllm_config.speculative_config is not None else 0 ) init_ascend_config(vllm_config) @@ -72,10 +71,14 @@ def __init__( * vllm_config.parallel_config.pipeline_parallel_size ) - # Handshake base port for this DP group. + # Worker handshake ports occupy + # [kv_port, kv_port + dp_size * pp_size * pcp_size * tp_size). + # Keep one scheduler control port per DP rank immediately after that + # range so it cannot collide with a worker handshake socket. self.side_channel_port = ( self.kv_transfer_config.kv_port - + vllm_config.parallel_config.data_parallel_rank * self.tp_size * self.pp_size * self.pcp_size + + self.max_device_id + + vllm_config.parallel_config.data_parallel_rank ) # Worker metadata for a DP group that may span multiple nodes. @@ -168,11 +171,9 @@ def _truncate_request_for_prefill(self, request: "Request") -> None: params["_p_side_truncated"] = True def on_new_request(self, request: "Request") -> None: - """Mooncake currently requires no request-arrival bookkeeping.""" raise NotImplementedError def update_connector_output(self, connector_output: KVConnectorOutput) -> None: - """Mooncake currently requires no worker-output bookkeeping.""" raise NotImplementedError def get_num_new_matched_tokens(self, request: "Request", num_computed_tokens: int) -> tuple[int | None, bool]: @@ -241,5 +242,4 @@ def set_xfer_handshake_metadata( """Handle the legacy port-offset-keyed handshake entry point.""" self.set_xfer_handshake_metadata_from_workers(metadata) - __all__ = ["MooncakeBaseConnectorScheduler"] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py index f03c14ba6394..166994a2c0cc 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py @@ -83,14 +83,11 @@ def __init__( self.connector_worker: MooncakeBaseConnectorWorker | None = None @classmethod - def get_required_kvcache_layout( - cls, vllm_config: VllmConfig - ) -> str | None: + def get_required_kvcache_layout(cls, vllm_config: VllmConfig) -> str | None: """Use HND to avoid reformatting KV cache before D2D transfer.""" if vllm_config.model_config is None: logger.warning_once( - "Unable to detect the current vLLM model config; falling " - "back to the default KV cache layout." + "Unable to detect the current vLLM model config; falling back to the default KV cache layout." ) return None @@ -99,23 +96,16 @@ def get_required_kvcache_layout( # affect its transfer layout. return None - logger.info_once( - "MooncakeConnector is setting the KV cache layout to HND for " - "direct D2D transfer." - ) + logger.info_once("MooncakeConnector is setting the KV cache layout to HND for direct D2D transfer.") return "HND" ############################################################ # Scheduler-side methods ############################################################ - def get_num_new_matched_tokens( - self, request: "Request", num_computed_tokens: int - ) -> tuple[int | None, bool]: + def get_num_new_matched_tokens(self, request: "Request", num_computed_tokens: int) -> tuple[int | None, bool]: assert self.connector_scheduler is not None - return self.connector_scheduler.get_num_new_matched_tokens( - request, num_computed_tokens - ) + return self.connector_scheduler.get_num_new_matched_tokens(request, num_computed_tokens) def update_state_after_alloc( self, @@ -124,13 +114,9 @@ def update_state_after_alloc( num_external_tokens: int, ) -> None: assert self.connector_scheduler is not None - self.connector_scheduler.update_state_after_alloc( - request, blocks, num_external_tokens - ) + self.connector_scheduler.update_state_after_alloc(request, blocks, num_external_tokens) - def build_connector_meta( - self, scheduler_output: SchedulerOutput - ) -> KVConnectorMetadata: + def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnectorMetadata: assert self.connector_scheduler is not None return self.connector_scheduler.build_connector_meta(scheduler_output) @@ -138,9 +124,7 @@ def on_new_request(self, request: "Request") -> None: assert self.connector_scheduler is not None self.connector_scheduler.on_new_request(request) - def update_connector_output( - self, connector_output: KVConnectorOutput - ) -> None: + def update_connector_output(self, connector_output: KVConnectorOutput) -> None: assert self.connector_scheduler is not None self.connector_scheduler.update_connector_output(connector_output) @@ -162,9 +146,7 @@ def request_finished_all_groups( def set_xfer_handshake_metadata( self, - metadata: Mapping[ - int | tuple[int, ...], KVConnectorHandshakeMetadata - ], + metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], ) -> None: """Set worker handshake metadata on the scheduler.""" assert self.connector_scheduler is not None @@ -172,9 +154,7 @@ def set_xfer_handshake_metadata( def set_xfer_handshake_metadata_pp_aware( self, - metadata: Mapping[ - int | tuple[int, ...], KVConnectorHandshakeMetadata - ], + metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], ) -> None: """Set PP-aware worker handshake metadata on the scheduler.""" assert self.connector_scheduler is not None @@ -188,9 +168,7 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]) -> None: assert self.connector_worker is not None self.connector_worker.register_kv_caches(kv_caches) - def get_finished( - self, finished_req_ids: set[str] - ) -> tuple[set[str], set[str]]: + def get_finished(self, finished_req_ids: set[str]) -> tuple[set[str], set[str]]: """Return requests whose receive and send operations have finished.""" assert self.connector_worker is not None return self.connector_worker.get_finished() @@ -206,9 +184,7 @@ def get_kv_connector_stats(self) -> KVConnectorStats | None: return self.connector_worker.get_kv_connector_stats() @classmethod - def build_kv_connector_stats( - cls, data: dict[str, Any] | None = None - ) -> KVConnectorStats: + def build_kv_connector_stats(cls, data: dict[str, Any] | None = None) -> KVConnectorStats: return MooncakeKVConnectorStats(data=data or {}) @classmethod @@ -219,13 +195,9 @@ def build_prom_metrics( labelnames: list[str], per_engine_labelvalues: dict[int, list[object]], ) -> KVConnectorPromMetrics: - return MooncakePromMetrics( - vllm_config, metric_types, labelnames, per_engine_labelvalues - ) + return MooncakePromMetrics(vllm_config, metric_types, labelnames, per_engine_labelvalues) - def start_load_kv( - self, forward_context: "ForwardContext", **kwargs: Any - ) -> None: + def start_load_kv(self, forward_context: "ForwardContext", **kwargs: Any) -> None: assert self.connector_worker is not None assert isinstance(self._connector_metadata, MooncakeConnectorMetadata) self.connector_worker.start_load_kv(self._connector_metadata) @@ -263,13 +235,9 @@ def __init__( super().__init__(vllm_config, role, kv_cache_config) if role == KVConnectorRole.SCHEDULER: - self.connector_scheduler = MooncakePullConnectorScheduler( - vllm_config, str(self.engine_id), kv_cache_config - ) + self.connector_scheduler = MooncakePullConnectorScheduler(vllm_config, str(self.engine_id), kv_cache_config) elif role == KVConnectorRole.WORKER: - self.connector_worker = MooncakePullConnectorWorker( - vllm_config, str(self.engine_id), kv_cache_config - ) + self.connector_worker = MooncakePullConnectorWorker(vllm_config, str(self.engine_id), kv_cache_config) else: raise ValueError(f"Unsupported KVConnectorRole: {role}") @@ -283,9 +251,7 @@ def __init__( role: KVConnectorRole, kv_cache_config: "KVCacheConfig", ) -> None: - raise NotImplementedError( - "MooncakePushConnector is not implemented yet." - ) + raise NotImplementedError("MooncakePushConnector is not implemented yet.") # Backward compatibility: MooncakeConnector is the pull-based connector. diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 9a410e5aa786..e09e3b09bc04 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -3,14 +3,23 @@ """Scheduler-side logic for Mooncake pull transfers.""" import math +import queue +import threading import time +from collections import OrderedDict +from collections.abc import Mapping from typing import TYPE_CHECKING, Any +import msgspec +import zmq +from vllm import envs from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds from vllm.distributed.kv_transfer.kv_connector.v1.base import ( + KVConnectorHandshakeMetadata, KVConnectorMetadata, ) from vllm.logger import init_logger +from vllm.utils.network_utils import make_zmq_path from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import RequestStatus @@ -21,6 +30,11 @@ from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( MooncakeConnectorMetadata, ) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils import ( + ensure_zmq_recv, + ensure_zmq_send, + zmq_ctx, +) if TYPE_CHECKING: from vllm.config import VllmConfig @@ -30,6 +44,168 @@ logger = init_logger(__name__) +GET_META_MSG = b"get_meta_msg" +DONE_RECVING_MSG = b"done_recving_msg" +ACK_MSG = b"ACK" + + +class MooncakeSchedulerSendingThread(threading.Thread): + """Serve worker metadata and collect D-side completion messages.""" + + def __init__( + self, + host: str, + port: int, + metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], + ready_event: threading.Event, + ) -> None: + super().__init__(daemon=True, name="MooncakeSchedulerSendingThread") + encoder = msgspec.msgpack.Encoder() + self.host = host + self.port = port + self.encoded_metadata = {key: encoder.encode(value) for key, value in metadata.items()} + self.ready_event = ready_event + self.delayed_free_requests: OrderedDict[str, float] = OrderedDict() + self.finished_requests: queue.SimpleQueue[str] = queue.SimpleQueue() + self.early_finished_requests: set[str] = set() + self.finished_request_ids: set[str] = set() + self.state_lock = threading.Lock() + + def add_delayed_request(self, request_id: str, delay_start_time: float) -> None: + with self.state_lock: + self.delayed_free_requests[request_id] = delay_start_time + if request_id in self.early_finished_requests: + self.early_finished_requests.remove(request_id) + self._mark_finished_locked(request_id) + + def get_and_clear_finished_requests(self) -> set[str]: + finished: set[str] = set() + with self.state_lock: + self._retrieve_expired_requests_locked() + while True: + try: + finished.add(self.finished_requests.get_nowait()) + except queue.Empty: + return finished + + def run(self) -> None: + path = make_zmq_path("tcp", self.host, self.port) + try: + logger.info("Mooncake scheduler sending thread listening on %s", path) + with zmq_ctx(zmq.ROUTER, path) as sock: + sock.setsockopt(zmq.RCVTIMEO, 1000) + self.ready_event.set() + self._run_busy_loop(sock) + except Exception: + self.ready_event.set() + logger.exception("Mooncake scheduler sending thread failed on %s", path) + + def _run_busy_loop(self, sock: zmq.Socket) -> None: + decoder = msgspec.msgpack.Decoder(type=tuple) + while True: + try: + frames = sock.recv_multipart() + except zmq.Again: + continue + + identity = frames[0] + payload = [frame for frame in frames[1:] if frame] + if len(payload) != 1: + logger.warning("Invalid Mooncake scheduler control frames: %s", frames) + continue + + try: + msg = decoder.decode(payload[0]) + if msg[0] == GET_META_MSG and len(msg) == 2: + metadata = self._get_metadata(msg[1]) + sock.send_multipart((identity, b"", metadata or b"")) + elif msg[0] == DONE_RECVING_MSG and len(msg) == 2: + self._handle_finished_request(str(msg[1])) + sock.send_multipart((identity, b"", ACK_MSG)) + else: + logger.warning("Unexpected Mooncake scheduler control message: %s", msg) + except Exception: + logger.exception("Failed to handle Mooncake scheduler control message") + + def _get_metadata(self, metadata_key: Any) -> bytes | None: + if isinstance(metadata_key, list): + metadata_key = tuple(metadata_key) + return self.encoded_metadata.get(metadata_key) + + def _handle_finished_request(self, request_id: str) -> None: + with self.state_lock: + if request_id in self.finished_request_ids: + return + if request_id in self.delayed_free_requests: + self._mark_finished_locked(request_id) + else: + self.early_finished_requests.add(request_id) + + def _mark_finished_locked(self, request_id: str) -> None: + if request_id in self.finished_request_ids: + return + self.finished_request_ids.add(request_id) + self.delayed_free_requests.pop(request_id, None) + self.finished_requests.put(request_id) + + def _retrieve_expired_requests_locked(self) -> None: + current_time = time.time() + while self.delayed_free_requests: + request_id = next(iter(self.delayed_free_requests)) + delay_start_time = self.delayed_free_requests[request_id] + if current_time - delay_start_time <= envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT: + break + + self._mark_finished_locked(request_id) + logger.error( + "Force freed expired Mooncake request %s after %s seconds", + request_id, + envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT, + ) + + +class MooncakeSchedulerRecvingThread(threading.Thread): + """Send scheduler-level completion messages from D to P.""" + + def __init__(self, ready_event: threading.Event) -> None: + super().__init__(daemon=True, name="MooncakeSchedulerRecvingThread") + self.ready_event = ready_event + self.request_queue: queue.Queue[tuple[str, int, str]] = queue.Queue() + self.encoder = msgspec.msgpack.Encoder() + + def add_request(self, remote_host: str, remote_port: int, request_id: str) -> None: + self.request_queue.put((remote_host, remote_port, request_id)) + + def run(self) -> None: + self.ready_event.set() + while True: + request = self.request_queue.get() + + try: + self._send_done_recving(*request) + except Exception: + logger.exception( + "Failed to send Mooncake scheduler completion for request %s", + request[2], + ) + self.request_queue.put(request) + finally: + self.request_queue.task_done() + + def _send_done_recving(self, remote_host: str, remote_port: int, request_id: str) -> None: + path = make_zmq_path("tcp", remote_host, remote_port) + with zmq_ctx(zmq.REQ, path) as sock: + sock.setsockopt(zmq.SNDTIMEO, 1000) + sock.setsockopt(zmq.RCVTIMEO, 1000) + ensure_zmq_send( + sock, + self.encoder.encode((DONE_RECVING_MSG, request_id)), + path, + ) + response = ensure_zmq_recv(sock, path) + if response != ACK_MSG: + raise RuntimeError(f"Unexpected Mooncake scheduler completion response: {response!r}") + class MooncakePullConnectorScheduler(MooncakeBaseConnectorScheduler): """Scheduler-side Mooncake pull connector implementation.""" @@ -47,6 +223,42 @@ def __init__( # Producer requests whose blocks must remain allocated until read. self._reqs_need_send: dict[str, float] = {} self._reqs_in_batch: set[str] = set() + # D request -> (P scheduler host, port, P request id). + self._reqs_recv_info: dict[str, tuple[str, int, str]] = {} + self._sending_thread: MooncakeSchedulerSendingThread | None = None + self._recving_thread: MooncakeSchedulerRecvingThread | None = None + + if self.kv_role == "kv_consumer": + recving_ready_event = threading.Event() + self._recving_thread = MooncakeSchedulerRecvingThread(recving_ready_event) + self._recving_thread.start() + recving_ready_event.wait() + elif self.kv_role != "kv_producer": + raise ValueError( + "Mooncake pull scheduler only supports kv_producer or " + f"kv_consumer, got {self.kv_role!r}" + ) + + def set_xfer_handshake_metadata_from_workers( + self, + metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], + ) -> None: + super().set_xfer_handshake_metadata_from_workers(metadata) + if self.kv_role != "kv_producer" or not metadata or self._sending_thread is not None: + return + + ready_event = threading.Event() + self._sending_thread = MooncakeSchedulerSendingThread( + self.side_channel_host, + self.side_channel_port, + metadata, + ready_event, + ) + self._sending_thread.start() + if not ready_event.wait(timeout=10): + raise RuntimeError("Timed out starting Mooncake scheduler sending thread") + if not self._sending_thread.is_alive(): + raise RuntimeError("Mooncake scheduler sending thread failed to start") def get_num_new_matched_tokens(self, request: "Request", num_computed_tokens: int) -> tuple[int, bool]: """Return prompt tokens that will be loaded from a remote producer.""" @@ -103,6 +315,11 @@ def update_state_after_alloc( local_block_ids, num_external_tokens, ) + self._reqs_recv_info[request.request_id] = ( + params["remote_host"], + params["remote_port"], + params["remote_request_id"], + ) else: logger.warning("Got invalid KVTransferParams. params=%s.", params) else: @@ -127,8 +344,6 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnector ) self._reqs_need_recv.clear() - meta.requests_to_send = self._reqs_need_send - self._reqs_need_send = {} meta.reqs_in_batch = self._reqs_in_batch self._reqs_in_batch = set() return meta @@ -165,7 +380,14 @@ def request_finished( num_computed_blocks, request.request_id, ) - self._reqs_need_send[request.request_id] = time.time() + delay_start_time = time.time() + self._reqs_need_send[request.request_id] = delay_start_time + if self._sending_thread is None: + raise RuntimeError("Mooncake scheduler metadata has not been initialized") + self._sending_thread.add_delayed_request( + request.request_id, + delay_start_time, + ) return delay_free_blocks, { "do_remote_prefill": True, @@ -187,7 +409,34 @@ def update_connector_output( self, connector_output: KVConnectorOutput, ) -> None: - pass + # D side: this output has already aggregated completion from all + # workers. Send one scheduler-to-scheduler ACK for the request. + for req_id in connector_output.finished_recving or (): + remote = self._reqs_recv_info.pop(req_id, None) + if remote is not None: + if self._recving_thread is None: + raise RuntimeError( + "Producer Mooncake scheduler received a receive-completion event" + ) + self._recving_thread.add_request(*remote) + + # P side: feed scheduler-received ACKs into vLLM's standard delayed + # free path. Scheduler._update_from_kv_xfer_finished reads this same + # KVConnectorOutput immediately after this hook returns. + finished_sending = ( + self._sending_thread.get_and_clear_finished_requests() if self._sending_thread is not None else set() + ) + if finished_sending: + for req_id in finished_sending: + self._reqs_need_send.pop(req_id, None) + if connector_output.finished_sending is None: + connector_output.finished_sending = finished_sending + else: + connector_output.finished_sending.update(finished_sending) -__all__ = ["MooncakePullConnectorScheduler"] +__all__ = [ + "MooncakePullConnectorScheduler", + "MooncakeSchedulerRecvingThread", + "MooncakeSchedulerSendingThread", +] 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 72d4b5b23abc..5f8dd53c2d89 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py @@ -2,10 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project """Shared utilities for Mooncake KV transfer connectors.""" +import contextlib import hashlib import struct import time from collections import OrderedDict +from collections.abc import Iterator from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -14,6 +16,7 @@ import torch import zmq from vllm.logger import init_logger +from vllm.utils.network_utils import make_zmq_socket from vllm_ascend.distributed.kv_transfer.utils.utils import ( RegisterRegions, @@ -26,6 +29,26 @@ logger = init_logger(__name__) +@contextlib.contextmanager +def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]: + """Create a Mooncake ROUTER or REQ socket and clean up its context.""" + if socket_type not in (zmq.ROUTER, zmq.REQ): + raise ValueError(f"Unexpected socket type: {socket_type}") + + context: zmq.Context | None = None + try: + context = zmq.Context() # type: ignore[attr-defined] + yield make_zmq_socket( + ctx=context, + path=addr, + socket_type=socket_type, + bind=socket_type == zmq.ROUTER, + ) + finally: + if context is not None: + context.destroy(linger=0) + + def as_kv_cache_tensors(cache_or_caches: Any) -> tuple[torch.Tensor, ...]: """Normalize one layer's KV cache into a tuple of tensors.""" if isinstance(cache_or_caches, torch.Tensor): @@ -248,4 +271,5 @@ def ensure_zmq_recv( "group_concurrent_contiguous", "split_if_not_byte_contiguous", "string_to_int64_hash", + "zmq_ctx", ] From 9b8b8475bd6a9c2bcfcdfabd580c85787a47c6e3 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Mon, 3 Aug 2026 17:05:44 +0800 Subject: [PATCH 07/48] add full block ids Signed-off-by: nwpu-zxr --- .../kv_transfer/kv_p2p/mooncake/metadata.py | 3 +++ .../kv_p2p/mooncake/pull_scheduler.py | 18 +++++++++++------- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 9e90fbaaf720..bbdb6ad96d1a 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -83,6 +83,7 @@ class ReqMeta: remote_request_id: str num_prompt_blocks: int remote_block_size: int + local_full_block_ids: BlockIds = tuple() class MooncakeConnectorMetadata(KVConnectorMetadata): @@ -99,6 +100,7 @@ def add_new_req( local_block_ids: BlockIds, num_external_tokens: int, kv_transfer_params: dict[str, Any], + local_full_block_ids: BlockIds | None = None, ) -> None: self.requests[request_id] = ReqMeta( local_block_ids=local_block_ids, @@ -111,6 +113,7 @@ def add_new_req( remote_request_id=kv_transfer_params["remote_request_id"], num_prompt_blocks=kv_transfer_params.get("num_prompt_blocks", 0), remote_block_size=kv_transfer_params.get("remote_block_size", 0), + local_full_block_ids=local_full_block_ids or tuple(), ) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index e09e3b09bc04..a9ca1bf6884c 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -2,7 +2,6 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project """Scheduler-side logic for Mooncake pull transfers.""" -import math import queue import threading import time @@ -219,7 +218,7 @@ def __init__( super().__init__(vllm_config, engine_id, kv_cache_config) # Requests waiting for the worker to start a READ transfer. - self._reqs_need_recv: dict[str, tuple[Request, BlockIds, int]] = {} + self._reqs_need_recv: dict[str, tuple[Request, BlockIds, BlockIds, int]] = {} # Producer requests whose blocks must remain allocated until read. self._reqs_need_send: dict[str, float] = {} self._reqs_in_batch: set[str] = set() @@ -309,10 +308,16 @@ def update_state_after_alloc( "remote_request_id", ) if all(field in params for field in required_remote_fields): - local_block_ids = blocks.get_unhashed_block_ids_all_groups() if num_external_tokens > 0 else tuple() + if num_external_tokens > 0: + local_block_ids = blocks.get_unhashed_block_ids_all_groups() + local_full_block_ids = blocks.get_block_ids() + else: + local_block_ids = tuple() + local_full_block_ids = tuple() self._reqs_need_recv[request.request_id] = ( request, local_block_ids, + local_full_block_ids, num_external_tokens, ) self._reqs_recv_info[request.request_id] = ( @@ -333,12 +338,13 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnector for ( req_id, - (req, block_ids, num_external_tokens), + (req, block_ids, full_block_ids, num_external_tokens), ) in self._reqs_need_recv.items(): assert req.kv_transfer_params is not None meta.add_new_req( request_id=req_id, local_block_ids=block_ids, + local_full_block_ids=full_block_ids, num_external_tokens=num_external_tokens, kv_transfer_params=req.kv_transfer_params, ) @@ -370,7 +376,6 @@ def request_finished( prompt_token_ids = request.prompt_token_ids or [] prompt_len = len(prompt_token_ids) - num_prompt_blocks = math.ceil(prompt_len / self.block_size) computed_block_ids = self._get_transfer_block_ids(block_ids, prompt_len) num_computed_blocks = sum(len(group_block_ids) for group_block_ids in computed_block_ids) delay_free_blocks = num_computed_blocks > 0 @@ -393,13 +398,12 @@ def request_finished( "do_remote_prefill": True, "do_remote_decode": False, "remote_block_ids": computed_block_ids, + "remote_num_prompt_tokens": prompt_len, "remote_engine_id": self.engine_id, "remote_request_id": request.request_id, "remote_host": self.side_channel_host, "remote_port": self.side_channel_port, "last_token_id": request.output_token_ids[-1], - "num_prompt_blocks": num_prompt_blocks, - "remote_block_size": self.block_size, } def on_new_request(self, request: "Request") -> None: From 918bc2ee28b6da8e5524911e337f8e8e21fa6a30 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Mon, 3 Aug 2026 17:41:42 +0800 Subject: [PATCH 08/48] add full-pp transfermetadata Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_scheduler.py | 38 +---------- .../kv_p2p/mooncake/base_worker.py | 3 + .../kv_transfer/kv_p2p/mooncake/metadata.py | 18 ++++-- .../kv_p2p/mooncake/pull_scheduler.py | 64 +++++++++++++++++-- 4 files changed, 77 insertions(+), 46 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py index 781766fe7dcc..46aa65cd02cb 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py @@ -81,9 +81,6 @@ def __init__( + vllm_config.parallel_config.data_parallel_rank ) - # Worker metadata for a DP group that may span multiple nodes. - self.multi_nodes_meta_mapping: dict[str, dict[str, Any]] = {} - self.kv_cache_groups = kv_cache_config.kv_cache_groups self.group_block_size = [group.kv_cache_spec.block_size for group in self.kv_cache_groups] self.group_unique_specs = [self._get_group_unique_specs(group) for group in self.kv_cache_groups] @@ -197,43 +194,12 @@ def request_finished( ) -> tuple[bool, dict[str, Any] | None]: raise NotImplementedError - def _port_offset_from_handshake_metadata( - self, - rank_metadata: KVConnectorHandshakeMetadata, - metadata_key: int | tuple[int, ...], - ) -> int: - handshake_port = getattr(rank_metadata, "handshake_port", 0) - if handshake_port > 0: - return handshake_port - self.kv_transfer_config.kv_port - if isinstance(metadata_key, int): - return metadata_key - raise ValueError(f"Mooncake handshake metadata is missing handshake_port for worker key {metadata_key}") - def set_xfer_handshake_metadata_from_workers( self, metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], ) -> None: - """Build the worker host mapping for a possibly multi-node DP group.""" - if not metadata: - return - - updated_mapping: dict[str, dict[str, Any]] = {} - kv_port = self.kv_transfer_config.kv_port - for metadata_key, rank_metadata in metadata.items(): - port_offset = self._port_offset_from_handshake_metadata(rank_metadata, metadata_key) - updated_mapping[str(port_offset)] = { - "host": rank_metadata.local_ip, - "engine_id": rank_metadata.engine_id, - "handshake_port": kv_port + port_offset, - } - - self.multi_nodes_meta_mapping.update(updated_mapping) - logger.info( - "MooncakeConnector set_xfer_handshake_metadata: worker_count=%d, updated=%s, multi_nodes_meta_mapping=%s", - len(metadata), - updated_mapping, - self.multi_nodes_meta_mapping, - ) + """Set worker transfer metadata on a concrete scheduler.""" + raise NotImplementedError def set_xfer_handshake_metadata( self, 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 98ce307512bc..05ee3d699d78 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 @@ -222,6 +222,9 @@ def register_kv_caches( block_lens=block_lens_per_layer, block_shapes=block_shapes_per_layer, block_size_scales=block_size_scales_per_layer, + pcp_size=self.pcp_size, + dcp_size=self.dcp_size, + tp_size=self.tp_size, local_ip=self.side_channel_host, handshake_port=self.handshake_port, ) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index bbdb6ad96d1a..3465e8c254fd 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -32,6 +32,9 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): block_lens: list[list[int]] block_shapes: list[list[tuple[int, ...]]] block_size_scales: list[list[int]] + pcp_size: int + dcp_size: int + tp_size: int local_ip: str = "" handshake_port: int = 0 @@ -69,6 +72,14 @@ def __post_init__(self) -> None: ) +@dataclass(frozen=True) +class MooncakeTransferMetadataGroups: + """All pipeline-stage metadata for one tensor-parallel rank.""" + + tp_rank: int + metadata_by_pp_rank: dict[int, MooncakeTransferMetadata] + + @dataclass class ReqMeta: """Request metadata required by the initial Mooncake pull path.""" @@ -81,8 +92,7 @@ class ReqMeta: remote_port: int remote_engine_id: str remote_request_id: str - num_prompt_blocks: int - remote_block_size: int + remote_num_prompt_tokens: int local_full_block_ids: BlockIds = tuple() @@ -111,8 +121,7 @@ def add_new_req( remote_port=kv_transfer_params["remote_port"], remote_engine_id=kv_transfer_params["remote_engine_id"], remote_request_id=kv_transfer_params["remote_request_id"], - num_prompt_blocks=kv_transfer_params.get("num_prompt_blocks", 0), - remote_block_size=kv_transfer_params.get("remote_block_size", 0), + remote_num_prompt_tokens=kv_transfer_params["remote_num_prompt_tokens"], local_full_block_ids=local_full_block_ids or tuple(), ) @@ -120,5 +129,6 @@ def add_new_req( __all__ = [ "MooncakeConnectorMetadata", "MooncakeTransferMetadata", + "MooncakeTransferMetadataGroups", "ReqMeta", ] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index a9ca1bf6884c..898ae7dc2a83 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -28,6 +28,8 @@ ) from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( MooncakeConnectorMetadata, + MooncakeTransferMetadata, + MooncakeTransferMetadataGroups, ) from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils import ( ensure_zmq_recv, @@ -56,13 +58,32 @@ def __init__( host: str, port: int, metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], + pp_size: int, ready_event: threading.Event, ) -> None: super().__init__(daemon=True, name="MooncakeSchedulerSendingThread") encoder = msgspec.msgpack.Encoder() self.host = host self.port = port - self.encoded_metadata = {key: encoder.encode(value) for key, value in metadata.items()} + metadata_by_tp_rank = self._group_metadata_by_tp_rank(metadata) + expected_pp_ranks = set(range(pp_size)) + for tp_rank, metadata_by_pp_rank in metadata_by_tp_rank.items(): + actual_pp_ranks = set(metadata_by_pp_rank) + if actual_pp_ranks != expected_pp_ranks: + raise ValueError( + f"Incomplete Mooncake metadata for TP rank {tp_rank}: " + f"got PP ranks {sorted(actual_pp_ranks)}, " + f"expected {sorted(expected_pp_ranks)}" + ) + self.encoded_metadata = { + tp_rank: encoder.encode( + MooncakeTransferMetadataGroups( + tp_rank=tp_rank, + metadata_by_pp_rank=dict(sorted(metadata_by_pp_rank.items())), + ) + ) + for tp_rank, metadata_by_pp_rank in metadata_by_tp_rank.items() + } self.ready_event = ready_event self.delayed_free_requests: OrderedDict[str, float] = OrderedDict() self.finished_requests: queue.SimpleQueue[str] = queue.SimpleQueue() @@ -70,6 +91,36 @@ def __init__( self.finished_request_ids: set[str] = set() self.state_lock = threading.Lock() + @staticmethod + def _group_metadata_by_tp_rank( + metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], + ) -> dict[int, dict[int, MooncakeTransferMetadata]]: + grouped: dict[int, dict[int, MooncakeTransferMetadata]] = {} + for metadata_key, rank_metadata in metadata.items(): + if isinstance(metadata_key, int): + pp_rank, tp_rank = 0, metadata_key + elif len(metadata_key) == 2: + pp_rank, tp_rank = metadata_key + else: + raise ValueError( + "Mooncake handshake metadata key must be tp_rank or " + f"(pp_rank, tp_rank), got {metadata_key!r}" + ) + + if not isinstance(rank_metadata, MooncakeTransferMetadata): + raise ValueError( + "Mooncake scheduler expects MooncakeTransferMetadata, " + f"got {type(rank_metadata).__name__} for key {metadata_key!r}" + ) + + metadata_by_pp_rank = grouped.setdefault(tp_rank, {}) + if pp_rank in metadata_by_pp_rank: + raise ValueError( + f"Duplicate Mooncake metadata for PP rank {pp_rank}, TP rank {tp_rank}" + ) + metadata_by_pp_rank[pp_rank] = rank_metadata + return grouped + def add_delayed_request(self, request_id: str, delay_start_time: float) -> None: with self.state_lock: self.delayed_free_requests[request_id] = delay_start_time @@ -126,10 +177,11 @@ def _run_busy_loop(self, sock: zmq.Socket) -> None: except Exception: logger.exception("Failed to handle Mooncake scheduler control message") - def _get_metadata(self, metadata_key: Any) -> bytes | None: - if isinstance(metadata_key, list): - metadata_key = tuple(metadata_key) - return self.encoded_metadata.get(metadata_key) + def _get_metadata(self, tp_rank: Any) -> bytes | None: + if not isinstance(tp_rank, int): + logger.warning("Invalid Mooncake metadata TP rank: %r", tp_rank) + return None + return self.encoded_metadata.get(tp_rank) def _handle_finished_request(self, request_id: str) -> None: with self.state_lock: @@ -242,7 +294,6 @@ def set_xfer_handshake_metadata_from_workers( self, metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], ) -> None: - super().set_xfer_handshake_metadata_from_workers(metadata) if self.kv_role != "kv_producer" or not metadata or self._sending_thread is not None: return @@ -251,6 +302,7 @@ def set_xfer_handshake_metadata_from_workers( self.side_channel_host, self.side_channel_port, metadata, + self.pp_size, ready_event, ) self._sending_thread.start() From 1b246d5ad0cc941b7f88d9f2b58411042f43c025 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Tue, 4 Aug 2026 16:10:27 +0800 Subject: [PATCH 09/48] add worker basic function Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_scheduler.py | 13 +- .../kv_p2p/mooncake/base_worker.py | 24 +- .../kv_p2p/mooncake/pull_scheduler.py | 13 +- .../kv_p2p/mooncake/pull_worker.py | 272 +++++++++++++++++- 4 files changed, 304 insertions(+), 18 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py index 46aa65cd02cb..387902721927 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py @@ -47,7 +47,14 @@ def __init__( self.vllm_config = vllm_config self.kv_transfer_config = vllm_config.kv_transfer_config - self.kv_role = self.kv_transfer_config.kv_role + if ( + self.kv_transfer_config.is_kv_consumer + == self.kv_transfer_config.is_kv_producer + ): + raise ValueError( + "Mooncake scheduler requires exactly one KV transfer role, " + f"got {self.kv_transfer_config.kv_role!r}" + ) self.kv_cache_config = kv_cache_config self.engine_id = engine_id self.block_size = vllm_config.cache_config.block_size @@ -63,6 +70,10 @@ def __init__( self.pp_size = vllm_config.parallel_config.pipeline_parallel_size self.tp_size = vllm_config.parallel_config.tensor_parallel_size self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size + assert self.pcp_size == 1, ( + "Mooncake temporarily requires prefill context parallel size 1, " + f"got {self.pcp_size}" + ) self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size self.max_device_id = ( self.tp_size 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 05ee3d699d78..15a9138c6b1c 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 @@ -65,8 +65,18 @@ def __init__( engine_id: str, kv_cache_config: "KVCacheConfig", ) -> None: + assert vllm_config.kv_transfer_config is not None + self.vllm_config = vllm_config self.kv_transfer_config = vllm_config.kv_transfer_config + if ( + self.kv_transfer_config.is_kv_consumer + == self.kv_transfer_config.is_kv_producer + ): + raise ValueError( + "Mooncake worker requires exactly one KV transfer role, " + f"got {self.kv_transfer_config.kv_role!r}" + ) self.engine_id = engine_id self.kv_cache_config = kv_cache_config self.block_size = vllm_config.cache_config.block_size @@ -82,16 +92,20 @@ def __init__( self.pp_rank = get_pp_group().rank_in_group self.pp_size = vllm_config.parallel_config.pipeline_parallel_size self.dp_rank = vllm_config.parallel_config.data_parallel_rank_local + if self.dp_rank is None: + raise ValueError("Mooncake worker requires a local DP rank") self.dp_size = vllm_config.parallel_config.data_parallel_size_local - self.pcp_size = get_pcp_group().world_size - self.pcp_rank = get_pcp_group().rank_in_group if self.pcp_size > 1 else 0 + pcp_group = get_pcp_group() + self.pcp_rank = pcp_group.rank_in_group + self.pcp_size = pcp_group.world_size + assert self.pcp_size == 1, ( + "Mooncake temporarily requires prefill context parallel size 1, " + f"got {self.pcp_size}" + ) self.dcp_size = get_decode_context_model_parallel_world_size() self.dcp_rank = get_decode_context_model_parallel_rank() if self.dcp_size > 1 else 0 - if self.pp_size > 1 and self.pcp_size > 1: - raise ValueError("PP and PCP cannot be enabled at the same time") self.max_device_id = self.tp_size * self.dp_size * self.pcp_size * self.pp_size - self.kv_role = self.kv_transfer_config.kv_role self.side_channel_host = get_ip() self.side_channel_port = ( self.kv_transfer_config.kv_port diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 898ae7dc2a83..27b80d4eb4a4 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -279,22 +279,21 @@ def __init__( self._sending_thread: MooncakeSchedulerSendingThread | None = None self._recving_thread: MooncakeSchedulerRecvingThread | None = None - if self.kv_role == "kv_consumer": + if self.kv_transfer_config.is_kv_consumer: recving_ready_event = threading.Event() self._recving_thread = MooncakeSchedulerRecvingThread(recving_ready_event) self._recving_thread.start() recving_ready_event.wait() - elif self.kv_role != "kv_producer": - raise ValueError( - "Mooncake pull scheduler only supports kv_producer or " - f"kv_consumer, got {self.kv_role!r}" - ) def set_xfer_handshake_metadata_from_workers( self, metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], ) -> None: - if self.kv_role != "kv_producer" or not metadata or self._sending_thread is not None: + if ( + not self.kv_transfer_config.is_kv_producer + or not metadata + or self._sending_thread is not None + ): return ready_event = threading.Event() 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 dd3f7cf22866..a7b259e86d24 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 @@ -2,20 +2,226 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project """Worker-side implementation entry point for Mooncake pull transfers.""" -from typing import TYPE_CHECKING +import queue +import threading +from typing import TYPE_CHECKING, Any +import msgspec +import torch +import zmq from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.utils.network_utils import make_zmq_path from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker import ( MooncakeBaseConnectorWorker, ) from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( MooncakeConnectorMetadata, + MooncakeTransferMetadata, + MooncakeTransferMetadataGroups, + ReqMeta, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils import ( + SizedDict, + ensure_zmq_recv, + ensure_zmq_send, + zmq_ctx, ) if TYPE_CHECKING: from vllm.v1.kv_cache_interface import KVCacheConfig +logger = init_logger(__name__) + + +class MooncakePullRecvingThread(threading.Thread): + """D-side execution thread for Mooncake pull requests. + + The queue/result lifecycle is intentionally implemented separately from + the transfer algorithm. Metadata lookup, block mapping, and Mooncake READ + submission belong in _handle_requests. + """ + + def __init__( + self, + engine: Any, + vllm_config: VllmConfig, + local_metadata: MooncakeTransferMetadata, + tp_rank: int, + tp_size: int, + dp_rank: int, + dp_size: int, + pcp_rank: int, + pcp_size: int, + dcp_rank: int, + dcp_size: int, + device: Any, + ready_event: threading.Event, + ) -> None: + super().__init__(daemon=True, name="MooncakePullRecvingThread") + self.engine = engine + self.vllm_config = vllm_config + self.local_metadata = local_metadata + self.tp_rank = tp_rank + self.tp_size = tp_size + self.dp_rank = dp_rank + self.dp_size = dp_size + self.pcp_rank = pcp_rank + self.pcp_size = pcp_size + self.dcp_rank = dcp_rank + self.dcp_size = dcp_size + assert self.pcp_size == 1, ( + "Mooncake pull worker temporarily requires pcp_size=1, " + f"got {self.pcp_size}" + ) + self.device = device + self.ready_event = ready_event + self.encoder = msgspec.msgpack.Encoder() + self.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) + self.remote_metadata: SizedDict = SizedDict() + self.request_queue: queue.Queue[tuple[str, int, dict[str, ReqMeta]]] = queue.Queue() + self.finished_requests: queue.SimpleQueue[str] = queue.SimpleQueue() + self.invalid_block_ids: set[int] = set() + self.invalid_block_ids_lock = threading.Lock() + assert self.local_metadata is not None + + def add_requests( + self, + remote_engine_id: str, + remote_port: int, + requests: dict[str, ReqMeta], + ) -> None: + """Queue requests for one remote engine and scheduler port.""" + if requests: + self.request_queue.put((remote_engine_id, remote_port, requests)) + + def get_and_clear_finished_requests(self) -> set[str]: + """Drain requests whose transfer attempt has completed.""" + finished: set[str] = set() + while True: + try: + finished.add(self.finished_requests.get_nowait()) + except queue.Empty: + return finished + + def get_and_clear_invalid_block_ids(self) -> set[int]: + """Drain local block IDs affected by failed pull attempts.""" + with self.invalid_block_ids_lock: + invalid_block_ids = self.invalid_block_ids + self.invalid_block_ids = set() + return invalid_block_ids + + def run(self) -> None: + """Consume queued requests and invoke the transfer implementation.""" + try: + torch.npu.set_device(self.device) + except Exception: + self.ready_event.set() + logger.exception("Failed to initialize the Mooncake pull thread device") + return + + self.ready_event.set() + while True: + remote_engine_id, remote_port, requests = self.request_queue.get() + try: + self._handle_requests(remote_engine_id, remote_port, requests) + except Exception: + for request_metadata in requests.values(): + self._mark_request_failed(request_metadata) + logger.exception( + "Mooncake pull failed for remote engine %s, port %d", + remote_engine_id, + remote_port, + ) + finally: + for request_id in requests: + self.finished_requests.put(request_id) + self.request_queue.task_done() + + def _mark_request_failed(self, request_metadata: ReqMeta) -> None: + with self.invalid_block_ids_lock: + for group_block_ids in request_metadata.local_block_ids: + self.invalid_block_ids.update(group_block_ids) + + def _handle_requests( + self, + remote_engine_id: str, + remote_port: int, + requests: dict[str, ReqMeta], + ) -> None: + """Fetch peer metadata and submit its aggregated requests.""" + remote_hosts = {request.remote_host for request in requests.values()} + if len(remote_hosts) != 1: + raise ValueError( + "Mooncake requests grouped for one remote engine and port must " + f"share one host, got {sorted(remote_hosts)}" + ) + remote_host = next(iter(remote_hosts)) + self._get_remote_metadata(remote_engine_id, remote_host, remote_port) + + # Transfer planning and Mooncake READ submission are implemented next. + # Until then, do not report an incomplete transfer as successful. + raise NotImplementedError("Mooncake pull transfer submission is not implemented") + + def _get_remote_metadata( + self, + remote_engine_id: str, + remote_host: str, + remote_port: int, + ) -> MooncakeTransferMetadataGroups: + cache_key = ( + remote_engine_id, + remote_host, + remote_port, + self.tp_rank, + ) + cached_metadata = self.remote_metadata.get(cache_key) + if cached_metadata is not None: + return cached_metadata + + path = make_zmq_path("tcp", remote_host, remote_port) + with zmq_ctx(zmq.REQ, path) as sock: + sock.setsockopt(zmq.SNDTIMEO, 1000) + sock.setsockopt(zmq.RCVTIMEO, 1000) + ensure_zmq_send( + sock, + self.encoder.encode((b"get_meta_msg", self.tp_rank)), + path, + ) + metadata_bytes = ensure_zmq_recv(sock, path) + + if not metadata_bytes: + raise RuntimeError( + "Mooncake producer scheduler returned no transfer metadata " + f"for TP rank {self.tp_rank} from {path}" + ) + + remote_metadata = self.decoder.decode(metadata_bytes) + if remote_metadata.tp_rank != self.tp_rank: + raise ValueError( + "Mooncake producer scheduler returned metadata for TP rank " + f"{remote_metadata.tp_rank}, expected {self.tp_rank}" + ) + if not remote_metadata.metadata_by_pp_rank: + raise ValueError( + "Mooncake producer scheduler returned no PP metadata for " + f"TP rank {self.tp_rank}" + ) + unexpected_engine_ids = { + metadata.engine_id + for metadata in remote_metadata.metadata_by_pp_rank.values() + if metadata.engine_id != remote_engine_id + } + if unexpected_engine_ids: + raise ValueError( + "Mooncake producer metadata engine ID mismatch: expected " + f"{remote_engine_id!r}, got {sorted(unexpected_engine_ids)!r}" + ) + + self.remote_metadata[cache_key] = remote_metadata + return remote_metadata + class MooncakePullConnectorWorker(MooncakeBaseConnectorWorker): """Worker-side framework for Mooncake pull transfers.""" @@ -27,18 +233,74 @@ def __init__( kv_cache_config: "KVCacheConfig", ) -> None: super().__init__(vllm_config, engine_id, kv_cache_config) + self._recving_thread: MooncakePullRecvingThread | None = None + + def register_kv_caches( + self, + kv_caches: dict[str, torch.Tensor | list[torch.Tensor]], + ) -> None: + """Register caches and start the D-side pull execution thread.""" + super().register_kv_caches(kv_caches) + if self.kv_transfer_config.is_kv_consumer: + ready_event = threading.Event() + self._recving_thread = MooncakePullRecvingThread( + engine=self.engine, + vllm_config=self.vllm_config, + local_metadata=self.xfer_handshake_metadata, + tp_rank=self.tp_rank, + tp_size=self.tp_size, + dp_rank=self.dp_rank, + dp_size=self.dp_size, + pcp_rank=self.pcp_rank, + pcp_size=self.pcp_size, + dcp_rank=self.dcp_rank, + dcp_size=self.dcp_size, + device=torch.npu.current_device(), + ready_event=ready_event, + ) + 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 start_load_kv(self, metadata: MooncakeConnectorMetadata) -> None: """Start the pull operations described by scheduler metadata.""" - raise NotImplementedError + if not metadata.requests: + return + if self.kv_transfer_config.is_kv_consumer: + request_groups: dict[str, dict[int, dict[str, ReqMeta]]] = {} + for request_id, request_metadata in metadata.requests.items(): + req_engine_id, req_remote_port = request_metadata.remote_engine_id, request_metadata.remote_port + requests_by_engine = request_groups.setdefault(req_engine_id, {}) + requests_by_port = requests_by_engine.setdefault(req_remote_port, {}) + requests_by_port[request_id] = request_metadata + + for remote_engine_id, requests_by_port in request_groups.items(): + for remote_port, requests_batch in requests_by_port.items(): + self._recving_thread.add_requests( + remote_engine_id, + remote_port, + requests_batch, + ) def get_finished(self) -> tuple[set[str], set[str]]: """Return requests with completed receive and send operations.""" - raise NotImplementedError + finished_recving = ( + self._recving_thread.get_and_clear_finished_requests() + if self._recving_thread is not None + else set() + ) + return set(), finished_recving def get_block_ids_with_load_errors(self) -> set[int]: """Return local block IDs whose pull operations failed.""" - raise NotImplementedError + if self._recving_thread is None: + return set() + return self._recving_thread.get_and_clear_invalid_block_ids() -__all__ = ["MooncakePullConnectorWorker"] +__all__ = [ + "MooncakePullConnectorWorker", + "MooncakePullRecvingThread", +] From 08b07b81bdaa293b1bdd27b8981bf33b176bbb6b Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Tue, 4 Aug 2026 23:46:35 +0800 Subject: [PATCH 10/48] add get_remote_metadata Signed-off-by: nwpu-zxr --- .../kv_transfer/kv_p2p/mooncake/metadata.py | 8 +- .../kv_p2p/mooncake/pull_scheduler.py | 52 ++++-- .../kv_p2p/mooncake/pull_worker.py | 151 ++++++++++-------- 3 files changed, 128 insertions(+), 83 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 3465e8c254fd..96571cbd44aa 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -74,10 +74,12 @@ def __post_init__(self) -> None: @dataclass(frozen=True) class MooncakeTransferMetadataGroups: - """All pipeline-stage metadata for one tensor-parallel rank.""" + """All worker transfer metadata exposed by one DP scheduler.""" - tp_rank: int - metadata_by_pp_rank: dict[int, MooncakeTransferMetadata] + engine_id: str + scheduler_host: str + scheduler_port: int + metadata_by_tp_rank: dict[int, dict[int, MooncakeTransferMetadata]] @dataclass diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 27b80d4eb4a4..b45c87c7eadd 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -57,7 +57,9 @@ def __init__( self, host: str, port: int, + engine_id: str, metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], + tp_size: int, pp_size: int, ready_event: threading.Event, ) -> None: @@ -65,21 +67,41 @@ def __init__( encoder = msgspec.msgpack.Encoder() self.host = host self.port = port + self.tp_size = tp_size + self.pp_size = pp_size + self.engine_id = engine_id metadata_by_tp_rank = self._group_metadata_by_tp_rank(metadata) - expected_pp_ranks = set(range(pp_size)) - for tp_rank, metadata_by_pp_rank in metadata_by_tp_rank.items(): - actual_pp_ranks = set(metadata_by_pp_rank) - if actual_pp_ranks != expected_pp_ranks: - raise ValueError( - f"Incomplete Mooncake metadata for TP rank {tp_rank}: " - f"got PP ranks {sorted(actual_pp_ranks)}, " - f"expected {sorted(expected_pp_ranks)}" - ) - self.encoded_metadata = { + + worker_engine_ids = { + worker_metadata.engine_id + for metadata_by_pp_rank in metadata_by_tp_rank.values() + for worker_metadata in metadata_by_pp_rank.values() + } + if worker_engine_ids != {engine_id}: + raise ValueError( + "Mooncake worker metadata engine IDs do not match scheduler " + f"engine ID {engine_id!r}: {sorted(worker_engine_ids)!r}" + ) + + metadata_by_tp_rank = { + tp_rank: dict(sorted(metadata_by_pp_rank.items())) + for tp_rank, metadata_by_pp_rank in sorted(metadata_by_tp_rank.items()) + } + self.encoded_metadata = encoder.encode( + MooncakeTransferMetadataGroups( + engine_id=engine_id, + scheduler_host=host, + scheduler_port=port, + metadata_by_tp_rank=metadata_by_tp_rank, + ) + ) + self.encoded_metadata_by_tp_rank = { tp_rank: encoder.encode( MooncakeTransferMetadataGroups( - tp_rank=tp_rank, - metadata_by_pp_rank=dict(sorted(metadata_by_pp_rank.items())), + engine_id=engine_id, + scheduler_host=host, + scheduler_port=port, + metadata_by_tp_rank={tp_rank: metadata_by_pp_rank}, ) ) for tp_rank, metadata_by_pp_rank in metadata_by_tp_rank.items() @@ -181,7 +203,9 @@ def _get_metadata(self, tp_rank: Any) -> bytes | None: if not isinstance(tp_rank, int): logger.warning("Invalid Mooncake metadata TP rank: %r", tp_rank) return None - return self.encoded_metadata.get(tp_rank) + if tp_rank == -1: + return self.encoded_metadata + return self.encoded_metadata_by_tp_rank.get(tp_rank) def _handle_finished_request(self, request_id: str) -> None: with self.state_lock: @@ -300,7 +324,9 @@ def set_xfer_handshake_metadata_from_workers( self._sending_thread = MooncakeSchedulerSendingThread( self.side_channel_host, self.side_channel_port, + self.engine_id, metadata, + self.tp_size, self.pp_size, ready_event, ) 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 a7b259e86d24..dd20ebd1e62f 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 @@ -80,7 +80,7 @@ def __init__( self.encoder = msgspec.msgpack.Encoder() self.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) self.remote_metadata: SizedDict = SizedDict() - self.request_queue: queue.Queue[tuple[str, int, dict[str, ReqMeta]]] = queue.Queue() + self.request_queue: queue.Queue[tuple[str, dict[str, ReqMeta]]] = queue.Queue() self.finished_requests: queue.SimpleQueue[str] = queue.SimpleQueue() self.invalid_block_ids: set[int] = set() self.invalid_block_ids_lock = threading.Lock() @@ -89,12 +89,11 @@ def __init__( def add_requests( self, remote_engine_id: str, - remote_port: int, requests: dict[str, ReqMeta], ) -> None: - """Queue requests for one remote engine and scheduler port.""" + """Queue requests for one remote engine.""" if requests: - self.request_queue.put((remote_engine_id, remote_port, requests)) + self.request_queue.put((remote_engine_id, requests)) def get_and_clear_finished_requests(self) -> set[str]: """Drain requests whose transfer attempt has completed.""" @@ -114,25 +113,32 @@ def get_and_clear_invalid_block_ids(self) -> set[int]: def run(self) -> None: """Consume queued requests and invoke the transfer implementation.""" - try: - torch.npu.set_device(self.device) - except Exception: - self.ready_event.set() - logger.exception("Failed to initialize the Mooncake pull thread device") - return - + torch.npu.set_device(self.device) self.ready_event.set() while True: - remote_engine_id, remote_port, requests = self.request_queue.get() + remote_engine_id, requests = self.request_queue.get() + request_endpoints = { + (request.remote_host, request.remote_port) + for request in requests.values() + } try: - self._handle_requests(remote_engine_id, remote_port, requests) - except Exception: + remote_host, remote_port = self._get_remote_endpoint( + request_endpoints + ) + self._handle_requests( + remote_engine_id, + remote_host, + remote_port, + requests, + ) + except Exception as exc: for request_metadata in requests.values(): self._mark_request_failed(request_metadata) logger.exception( - "Mooncake pull failed for remote engine %s, port %d", + "Mooncake pull failed for remote engine %s at %s: %s", remote_engine_id, - remote_port, + sorted(request_endpoints), + exc, ) finally: for request_id in requests: @@ -144,39 +150,74 @@ def _mark_request_failed(self, request_metadata: ReqMeta) -> None: for group_block_ids in request_metadata.local_block_ids: self.invalid_block_ids.update(group_block_ids) + @staticmethod + def _get_remote_endpoint( + endpoints: set[tuple[str, int]], + ) -> tuple[str, int]: + if len(endpoints) != 1: + raise ValueError( + "Requests for one remote engine must share one scheduler " + f"endpoint, got {sorted(endpoints)}" + ) + return next(iter(endpoints)) + def _handle_requests( self, remote_engine_id: str, + remote_host: str, remote_port: int, requests: dict[str, ReqMeta], ) -> None: """Fetch peer metadata and submit its aggregated requests.""" - remote_hosts = {request.remote_host for request in requests.values()} - if len(remote_hosts) != 1: - raise ValueError( - "Mooncake requests grouped for one remote engine and port must " - f"share one host, got {sorted(remote_hosts)}" - ) - remote_host = next(iter(remote_hosts)) self._get_remote_metadata(remote_engine_id, remote_host, remote_port) # Transfer planning and Mooncake READ submission are implemented next. # Until then, do not report an incomplete transfer as successful. raise NotImplementedError("Mooncake pull transfer submission is not implemented") + @staticmethod + def _validate_remote_metadata( + metadata: MooncakeTransferMetadataGroups, + expected_engine_id: str, + ) -> None: + if metadata.engine_id != expected_engine_id: + raise ValueError( + "Mooncake producer metadata engine ID mismatch: expected " + f"{expected_engine_id!r}, got {metadata.engine_id!r}" + ) + if not metadata.metadata_by_tp_rank: + raise ValueError("Mooncake producer scheduler returned no worker metadata") + + empty_tp_ranks = [ + tp_rank + for tp_rank, metadata_by_pp_rank in metadata.metadata_by_tp_rank.items() + if not metadata_by_pp_rank + ] + if empty_tp_ranks: + raise ValueError( + "Mooncake producer scheduler returned no PP metadata for TP ranks " + f"{sorted(empty_tp_ranks)}" + ) + + unexpected_engine_ids = { + worker_metadata.engine_id + for metadata_by_pp_rank in metadata.metadata_by_tp_rank.values() + for worker_metadata in metadata_by_pp_rank.values() + if worker_metadata.engine_id != expected_engine_id + } + if unexpected_engine_ids: + raise ValueError( + "Mooncake producer worker metadata engine ID mismatch: expected " + f"{expected_engine_id!r}, got {sorted(unexpected_engine_ids)!r}" + ) + def _get_remote_metadata( self, remote_engine_id: str, remote_host: str, remote_port: int, ) -> MooncakeTransferMetadataGroups: - cache_key = ( - remote_engine_id, - remote_host, - remote_port, - self.tp_rank, - ) - cached_metadata = self.remote_metadata.get(cache_key) + cached_metadata = self.remote_metadata.get(remote_engine_id) if cached_metadata is not None: return cached_metadata @@ -186,7 +227,7 @@ def _get_remote_metadata( sock.setsockopt(zmq.RCVTIMEO, 1000) ensure_zmq_send( sock, - self.encoder.encode((b"get_meta_msg", self.tp_rank)), + self.encoder.encode((b"get_meta_msg", -1)), path, ) metadata_bytes = ensure_zmq_recv(sock, path) @@ -194,32 +235,12 @@ def _get_remote_metadata( if not metadata_bytes: raise RuntimeError( "Mooncake producer scheduler returned no transfer metadata " - f"for TP rank {self.tp_rank} from {path}" + f"for engine {remote_engine_id!r} from {path}" ) remote_metadata = self.decoder.decode(metadata_bytes) - if remote_metadata.tp_rank != self.tp_rank: - raise ValueError( - "Mooncake producer scheduler returned metadata for TP rank " - f"{remote_metadata.tp_rank}, expected {self.tp_rank}" - ) - if not remote_metadata.metadata_by_pp_rank: - raise ValueError( - "Mooncake producer scheduler returned no PP metadata for " - f"TP rank {self.tp_rank}" - ) - unexpected_engine_ids = { - metadata.engine_id - for metadata in remote_metadata.metadata_by_pp_rank.values() - if metadata.engine_id != remote_engine_id - } - if unexpected_engine_ids: - raise ValueError( - "Mooncake producer metadata engine ID mismatch: expected " - f"{remote_engine_id!r}, got {sorted(unexpected_engine_ids)!r}" - ) - - self.remote_metadata[cache_key] = remote_metadata + self._validate_remote_metadata(remote_metadata, remote_engine_id) + self.remote_metadata[remote_engine_id] = remote_metadata return remote_metadata @@ -269,20 +290,16 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata) -> None: if not metadata.requests: return if self.kv_transfer_config.is_kv_consumer: - request_groups: dict[str, dict[int, dict[str, ReqMeta]]] = {} + assert self._recving_thread is not None + request_groups: dict[str, dict[str, ReqMeta]] = {} for request_id, request_metadata in metadata.requests.items(): - req_engine_id, req_remote_port = request_metadata.remote_engine_id, request_metadata.remote_port - requests_by_engine = request_groups.setdefault(req_engine_id, {}) - requests_by_port = requests_by_engine.setdefault(req_remote_port, {}) - requests_by_port[request_id] = request_metadata - - for remote_engine_id, requests_by_port in request_groups.items(): - for remote_port, requests_batch in requests_by_port.items(): - self._recving_thread.add_requests( - remote_engine_id, - remote_port, - requests_batch, - ) + requests = request_groups.setdefault( + request_metadata.remote_engine_id, {} + ) + requests[request_id] = request_metadata + + for remote_engine_id, requests in request_groups.items(): + self._recving_thread.add_requests(remote_engine_id, requests) def get_finished(self) -> tuple[set[str], set[str]]: """Return requests with completed receive and send operations.""" From 6186c97f9a641667c9d26a900c5f1fcd9a80b46d Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Wed, 5 Aug 2026 15:51:59 +0800 Subject: [PATCH 11/48] add transfer metadata Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/pull_worker.py | 50 +++++++++++++++++-- 1 file changed, 45 insertions(+), 5 deletions(-) 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 dd20ebd1e62f..be68e69f49a0 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 @@ -4,7 +4,7 @@ import queue import threading -from typing import TYPE_CHECKING, Any +from typing import Any import msgspec import torch @@ -12,7 +12,9 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.utils.network_utils import make_zmq_path +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheSpec +from vllm_ascend.ascend_config import AscendConfig from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker import ( MooncakeBaseConnectorWorker, ) @@ -28,9 +30,7 @@ ensure_zmq_send, zmq_ctx, ) - -if TYPE_CHECKING: - from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm_ascend.utils import enable_sfa_dcp_replicated_indexer logger = init_logger(__name__) @@ -47,6 +47,10 @@ def __init__( self, engine: Any, vllm_config: VllmConfig, + kv_cache_config: KVCacheConfig, + kv_cache_specs: list[KVCacheSpec], + layer_name_to_group_index: dict[str, int], + layer_name_to_spec_index: dict[str, int], local_metadata: MooncakeTransferMetadata, tp_rank: int, tp_size: int, @@ -62,7 +66,35 @@ def __init__( super().__init__(daemon=True, name="MooncakePullRecvingThread") self.engine = engine self.vllm_config = vllm_config + self.model_config = vllm_config.model_config + self.kv_cache_config = kv_cache_config + self.kv_cache_specs = kv_cache_specs + self.layer_name_to_group_index = layer_name_to_group_index + self.layer_name_to_spec_index = layer_name_to_spec_index self.local_metadata = local_metadata + + self.block_size = local_metadata.block_size + self.layer_names = local_metadata.layer_names + self.group_indices = local_metadata.group_indices + self.spec_indices = local_metadata.spec_indices + self.kv_caches_base_addr = local_metadata.kv_caches_base_addr + self.block_strides = local_metadata.block_strides + self.block_lens = local_metadata.block_lens + self.block_shapes = local_metadata.block_shapes + self.block_size_scales = local_metadata.block_size_scales + + self.use_mla = self.model_config.is_deepseek_mla + hf_text_config = self.model_config.hf_text_config + self.num_key_value_heads = getattr( + hf_text_config, "num_key_value_heads", 0 + ) + speculative_config = vllm_config.speculative_config + self.num_speculative_tokens = ( + speculative_config.num_speculative_tokens + if speculative_config is not None + else 0 + ) + self.tp_rank = tp_rank self.tp_size = tp_size self.dp_rank = dp_rank @@ -251,7 +283,7 @@ def __init__( self, vllm_config: VllmConfig, engine_id: str, - kv_cache_config: "KVCacheConfig", + kv_cache_config: KVCacheConfig, ) -> None: super().__init__(vllm_config, engine_id, kv_cache_config) self._recving_thread: MooncakePullRecvingThread | None = None @@ -267,9 +299,17 @@ def register_kv_caches( self._recving_thread = MooncakePullRecvingThread( engine=self.engine, vllm_config=self.vllm_config, + ascend_config=self.ascend_config, + kv_cache_config=self.kv_cache_config, + kv_cache_specs=self.kv_cache_specs, + layer_name_to_group_index=self.layer_name_to_group_index, + layer_name_to_spec_index=self.layer_name_to_spec_index, + kv_caches=self.kv_caches, local_metadata=self.xfer_handshake_metadata, tp_rank=self.tp_rank, tp_size=self.tp_size, + pp_rank=self.pp_rank, + pp_size=self.pp_size, dp_rank=self.dp_rank, dp_size=self.dp_size, pcp_rank=self.pcp_rank, From 36e87c7db4b74e5d9950ccb4e31e5623b112c70e Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Thu, 6 Aug 2026 11:54:13 +0800 Subject: [PATCH 12/48] add merge transfermetadata Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_worker.py | 29 +++++ .../kv_transfer/kv_p2p/mooncake/metadata.py | 35 ++++++ .../kv_p2p/mooncake/pull_worker.py | 100 ++++++++++++++++-- 3 files changed, 155 insertions(+), 9 deletions(-) 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 15a9138c6b1c..d248837e4f18 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 @@ -22,7 +22,10 @@ from vllm.logger import init_logger from vllm.utils.network_utils import get_ip from vllm.v1.kv_cache_interface import ( + AttentionSpec, + FullAttentionSpec, KVCacheSpec, + SlidingWindowSpec, UniformTypeKVCacheSpecs, ) @@ -158,6 +161,24 @@ def _build_kv_cache_spec_mappings(self) -> None: self.layer_name_to_group_index[layer_name] = group_index self.layer_name_to_spec_index[layer_name] = spec_index + @staticmethod + def _get_spec_transfer_properties( + spec: KVCacheSpec, + ) -> tuple[int, int | None]: + """Return the block size and head size required for KV transfer.""" + if not isinstance(spec, AttentionSpec): + return spec.block_size, None + + if isinstance(spec, (FullAttentionSpec, SlidingWindowSpec)): + head_size_v = getattr(spec, "head_size_v", spec.head_size) + if spec.head_size != head_size_v: + raise NotImplementedError( + "Mooncake does not support different K/V head sizes for " + f"{type(spec).__name__}: K={spec.head_size}, V={head_size_v}" + ) + + return spec.block_size, spec.head_size + def register_kv_caches( self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]], @@ -167,6 +188,12 @@ def register_kv_caches( logger.info("num_blocks: %s", self.num_blocks) self.kv_caches = kv_caches self._build_kv_cache_spec_mappings() + spec_properties = [ + self._get_spec_transfer_properties(spec) + for spec in self.kv_cache_specs + ] + spec_block_sizes = [properties[0] for properties in spec_properties] + spec_head_sizes = [properties[1] for properties in spec_properties] layer_names: list[str] = [] group_indices: list[int] = [] @@ -228,6 +255,8 @@ def register_kv_caches( te_rpc_port=self.te_rpc_port, block_size=self.block_size, num_blocks=self.num_blocks, + spec_block_sizes=spec_block_sizes, + spec_head_sizes=spec_head_sizes, layer_names=layer_names, group_indices=group_indices, spec_indices=spec_indices, diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 96571cbd44aa..22514e6d1165 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -24,6 +24,8 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): te_rpc_port: int block_size: int num_blocks: int + spec_block_sizes: list[int] + spec_head_sizes: list[int | None] layer_names: list[str] group_indices: list[int] spec_indices: list[int] @@ -72,6 +74,37 @@ def __post_init__(self) -> None: ) +@dataclass(frozen=True) +class MooncakeTPTransferMetadata: + """TP-private connection and KV-cache address information.""" + + te_rpc_port: int + kv_caches_base_addr: list[list[int]] + local_ip: str + handshake_port: int + + +@dataclass(frozen=True) +class MooncakePPTransferMetadata: + """Metadata shared by all TP workers belonging to one PP rank.""" + + block_size: int + num_blocks: int + spec_block_sizes: list[int] + spec_head_sizes: list[int | None] + layer_names: list[str] + group_indices: list[int] + spec_indices: list[int] + block_strides: list[list[int]] + block_lens: list[list[int]] + block_shapes: list[list[tuple[int, ...]]] + block_size_scales: list[list[int]] + pcp_size: int + dcp_size: int + tp_size: int + metadata_by_tp_rank: dict[int, MooncakeTPTransferMetadata] + + @dataclass(frozen=True) class MooncakeTransferMetadataGroups: """All worker transfer metadata exposed by one DP scheduler.""" @@ -130,6 +163,8 @@ def add_new_req( __all__ = [ "MooncakeConnectorMetadata", + "MooncakePPTransferMetadata", + "MooncakeTPTransferMetadata", "MooncakeTransferMetadata", "MooncakeTransferMetadataGroups", "ReqMeta", 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 be68e69f49a0..8483ed0d356a 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 @@ -12,14 +12,15 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.utils.network_utils import make_zmq_path -from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig, KVCacheSpec +from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec -from vllm_ascend.ascend_config import AscendConfig from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker import ( MooncakeBaseConnectorWorker, ) from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( MooncakeConnectorMetadata, + MooncakePPTransferMetadata, + MooncakeTPTransferMetadata, MooncakeTransferMetadata, MooncakeTransferMetadataGroups, ReqMeta, @@ -30,7 +31,6 @@ ensure_zmq_send, zmq_ctx, ) -from vllm_ascend.utils import enable_sfa_dcp_replicated_indexer logger = init_logger(__name__) @@ -74,6 +74,8 @@ def __init__( self.local_metadata = local_metadata self.block_size = local_metadata.block_size + self.spec_block_sizes = local_metadata.spec_block_sizes + self.spec_head_sizes = local_metadata.spec_head_sizes self.layer_names = local_metadata.layer_names self.group_indices = local_metadata.group_indices self.spec_indices = local_metadata.spec_indices @@ -111,7 +113,9 @@ def __init__( self.ready_event = ready_event self.encoder = msgspec.msgpack.Encoder() self.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) - self.remote_metadata: SizedDict = SizedDict() + self.remote_metadata: SizedDict[ + str, dict[int, MooncakePPTransferMetadata] + ] = SizedDict() self.request_queue: queue.Queue[tuple[str, dict[str, ReqMeta]]] = queue.Queue() self.finished_requests: queue.SimpleQueue[str] = queue.SimpleQueue() self.invalid_block_ids: set[int] = set() @@ -243,12 +247,91 @@ def _validate_remote_metadata( f"{expected_engine_id!r}, got {sorted(unexpected_engine_ids)!r}" ) + @staticmethod + def _aggregate_remote_metadata( + metadata: MooncakeTransferMetadataGroups, + ) -> dict[int, MooncakePPTransferMetadata]: + """Deduplicate TP-invariant metadata and group worker endpoints by PP.""" + workers_by_pp_rank: dict[int, dict[int, MooncakeTransferMetadata]] = {} + for tp_rank, workers_by_pp_rank_for_tp in metadata.metadata_by_tp_rank.items(): + for pp_rank, worker_metadata in workers_by_pp_rank_for_tp.items(): + workers_by_pp_rank.setdefault(pp_rank, {})[tp_rank] = worker_metadata + + common_fields = ( + "block_size", + "num_blocks", + "spec_block_sizes", + "spec_head_sizes", + "layer_names", + "group_indices", + "spec_indices", + "block_strides", + "block_lens", + "block_shapes", + "block_size_scales", + "pcp_size", + "dcp_size", + "tp_size", + ) + expected_tp_ranks = set(metadata.metadata_by_tp_rank) + aggregated: dict[int, MooncakePPTransferMetadata] = {} + for pp_rank, workers_by_tp_rank in workers_by_pp_rank.items(): + actual_tp_ranks = set(workers_by_tp_rank) + if actual_tp_ranks != expected_tp_ranks: + raise ValueError( + "Mooncake producer metadata has incomplete TP workers for " + f"PP rank {pp_rank}: expected {sorted(expected_tp_ranks)}, " + f"got {sorted(actual_tp_ranks)}" + ) + + reference_tp_rank = min(workers_by_tp_rank) + reference = workers_by_tp_rank[reference_tp_rank] + for tp_rank, worker_metadata in workers_by_tp_rank.items(): + mismatched_fields = [ + field_name + for field_name in common_fields + if getattr(worker_metadata, field_name) != getattr(reference, field_name) + ] + if mismatched_fields: + raise ValueError( + "Mooncake producer metadata differs across TP ranks " + f"for PP rank {pp_rank}: TP {reference_tp_rank} and " + f"TP {tp_rank} mismatch in {mismatched_fields}" + ) + + aggregated[pp_rank] = MooncakePPTransferMetadata( + block_size=reference.block_size, + num_blocks=reference.num_blocks, + spec_block_sizes=reference.spec_block_sizes, + spec_head_sizes=reference.spec_head_sizes, + layer_names=reference.layer_names, + group_indices=reference.group_indices, + spec_indices=reference.spec_indices, + block_strides=reference.block_strides, + block_lens=reference.block_lens, + block_shapes=reference.block_shapes, + block_size_scales=reference.block_size_scales, + pcp_size=reference.pcp_size, + dcp_size=reference.dcp_size, + tp_size=reference.tp_size, + metadata_by_tp_rank={ + tp_rank: MooncakeTPTransferMetadata( + te_rpc_port=worker_metadata.te_rpc_port, + kv_caches_base_addr=(worker_metadata.kv_caches_base_addr), + local_ip=worker_metadata.local_ip, + handshake_port=worker_metadata.handshake_port, + ) + for tp_rank, worker_metadata in sorted(workers_by_tp_rank.items()) + }, + ) + return aggregated + def _get_remote_metadata( self, remote_engine_id: str, remote_host: str, remote_port: int, - ) -> MooncakeTransferMetadataGroups: + ) -> dict[int, MooncakePPTransferMetadata]: cached_metadata = self.remote_metadata.get(remote_engine_id) if cached_metadata is not None: return cached_metadata @@ -270,8 +353,9 @@ def _get_remote_metadata( f"for engine {remote_engine_id!r} from {path}" ) - remote_metadata = self.decoder.decode(metadata_bytes) - self._validate_remote_metadata(remote_metadata, remote_engine_id) + transfer_metadata = self.decoder.decode(metadata_bytes) + self._validate_remote_metadata(transfer_metadata, remote_engine_id) + remote_metadata = self._aggregate_remote_metadata(transfer_metadata) self.remote_metadata[remote_engine_id] = remote_metadata return remote_metadata @@ -299,12 +383,10 @@ def register_kv_caches( self._recving_thread = MooncakePullRecvingThread( engine=self.engine, vllm_config=self.vllm_config, - ascend_config=self.ascend_config, kv_cache_config=self.kv_cache_config, kv_cache_specs=self.kv_cache_specs, layer_name_to_group_index=self.layer_name_to_group_index, layer_name_to_spec_index=self.layer_name_to_spec_index, - kv_caches=self.kv_caches, local_metadata=self.xfer_handshake_metadata, tp_rank=self.tp_rank, tp_size=self.tp_size, From bddcaec61a80de606c34bad013547e1293945c2e Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Thu, 6 Aug 2026 14:54:19 +0800 Subject: [PATCH 13/48] simplify metadata Signed-off-by: nwpu-zxr --- .../kv_transfer/kv_p2p/mooncake/metadata.py | 4 +- .../kv_p2p/mooncake/pull_scheduler.py | 162 ++++++++++++------ .../kv_p2p/mooncake/pull_worker.py | 127 +++----------- 3 files changed, 140 insertions(+), 153 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 22514e6d1165..2fb826aedfd9 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -107,12 +107,12 @@ class MooncakePPTransferMetadata: @dataclass(frozen=True) class MooncakeTransferMetadataGroups: - """All worker transfer metadata exposed by one DP scheduler.""" + """PP-grouped worker transfer metadata exposed by one DP scheduler.""" engine_id: str scheduler_host: str scheduler_port: int - metadata_by_tp_rank: dict[int, dict[int, MooncakeTransferMetadata]] + metadata_by_pp_rank: dict[int, MooncakePPTransferMetadata] @dataclass diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index b45c87c7eadd..f2970c6a4b41 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -28,6 +28,8 @@ ) from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( MooncakeConnectorMetadata, + MooncakePPTransferMetadata, + MooncakeTPTransferMetadata, MooncakeTransferMetadata, MooncakeTransferMetadataGroups, ) @@ -70,42 +72,15 @@ def __init__( self.tp_size = tp_size self.pp_size = pp_size self.engine_id = engine_id - metadata_by_tp_rank = self._group_metadata_by_tp_rank(metadata) - - worker_engine_ids = { - worker_metadata.engine_id - for metadata_by_pp_rank in metadata_by_tp_rank.values() - for worker_metadata in metadata_by_pp_rank.values() - } - if worker_engine_ids != {engine_id}: - raise ValueError( - "Mooncake worker metadata engine IDs do not match scheduler " - f"engine ID {engine_id!r}: {sorted(worker_engine_ids)!r}" - ) - - metadata_by_tp_rank = { - tp_rank: dict(sorted(metadata_by_pp_rank.items())) - for tp_rank, metadata_by_pp_rank in sorted(metadata_by_tp_rank.items()) - } + metadata_by_pp_rank = self._merge_metadata_by_pp_rank(metadata) self.encoded_metadata = encoder.encode( MooncakeTransferMetadataGroups( engine_id=engine_id, scheduler_host=host, scheduler_port=port, - metadata_by_tp_rank=metadata_by_tp_rank, + metadata_by_pp_rank=metadata_by_pp_rank, ) ) - self.encoded_metadata_by_tp_rank = { - tp_rank: encoder.encode( - MooncakeTransferMetadataGroups( - engine_id=engine_id, - scheduler_host=host, - scheduler_port=port, - metadata_by_tp_rank={tp_rank: metadata_by_pp_rank}, - ) - ) - for tp_rank, metadata_by_pp_rank in metadata_by_tp_rank.items() - } self.ready_event = ready_event self.delayed_free_requests: OrderedDict[str, float] = OrderedDict() self.finished_requests: queue.SimpleQueue[str] = queue.SimpleQueue() @@ -113,11 +88,12 @@ def __init__( self.finished_request_ids: set[str] = set() self.state_lock = threading.Lock() - @staticmethod - def _group_metadata_by_tp_rank( + def _merge_metadata_by_pp_rank( + self, metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], - ) -> dict[int, dict[int, MooncakeTransferMetadata]]: - grouped: dict[int, dict[int, MooncakeTransferMetadata]] = {} + ) -> dict[int, MooncakePPTransferMetadata]: + """Validate worker metadata and merge TP-invariant fields by PP.""" + workers_by_pp_rank: dict[int, dict[int, MooncakeTransferMetadata]] = {} for metadata_key, rank_metadata in metadata.items(): if isinstance(metadata_key, int): pp_rank, tp_rank = 0, metadata_key @@ -134,14 +110,109 @@ def _group_metadata_by_tp_rank( "Mooncake scheduler expects MooncakeTransferMetadata, " f"got {type(rank_metadata).__name__} for key {metadata_key!r}" ) + if rank_metadata.engine_id != self.engine_id: + raise ValueError( + "Mooncake worker metadata engine ID mismatch: expected " + f"{self.engine_id!r}, got {rank_metadata.engine_id!r}" + ) - metadata_by_pp_rank = grouped.setdefault(tp_rank, {}) - if pp_rank in metadata_by_pp_rank: + workers_by_tp_rank = workers_by_pp_rank.setdefault(pp_rank, {}) + if tp_rank in workers_by_tp_rank: raise ValueError( - f"Duplicate Mooncake metadata for PP rank {pp_rank}, TP rank {tp_rank}" + f"Duplicate Mooncake metadata for PP rank {pp_rank}, " + f"TP rank {tp_rank}" ) - metadata_by_pp_rank[pp_rank] = rank_metadata - return grouped + workers_by_tp_rank[tp_rank] = rank_metadata + + expected_pp_ranks = set(range(self.pp_size)) + if set(workers_by_pp_rank) != expected_pp_ranks: + raise ValueError( + "Mooncake worker metadata has incomplete PP ranks: expected " + f"{sorted(expected_pp_ranks)}, got " + f"{sorted(workers_by_pp_rank)}" + ) + + common_fields = ( + "block_size", + "num_blocks", + "spec_block_sizes", + "spec_head_sizes", + "layer_names", + "group_indices", + "spec_indices", + "block_strides", + "block_lens", + "block_shapes", + "block_size_scales", + "pcp_size", + "dcp_size", + "tp_size", + ) + expected_tp_ranks = set(range(self.tp_size)) + merged: dict[int, MooncakePPTransferMetadata] = {} + for pp_rank, workers_by_tp_rank in sorted( + workers_by_pp_rank.items() + ): + if set(workers_by_tp_rank) != expected_tp_ranks: + raise ValueError( + "Mooncake worker metadata has incomplete TP ranks for " + f"PP rank {pp_rank}: expected {sorted(expected_tp_ranks)}, " + f"got {sorted(workers_by_tp_rank)}" + ) + + reference_tp_rank = min(workers_by_tp_rank) + reference = workers_by_tp_rank[reference_tp_rank] + if reference.tp_size != self.tp_size: + raise ValueError( + "Mooncake worker metadata TP size mismatch for PP rank " + f"{pp_rank}: expected {self.tp_size}, " + f"got {reference.tp_size}" + ) + + for tp_rank, worker_metadata in workers_by_tp_rank.items(): + mismatched_fields = [ + field_name + for field_name in common_fields + if getattr(worker_metadata, field_name) + != getattr(reference, field_name) + ] + if mismatched_fields: + raise ValueError( + "Mooncake worker metadata differs across TP ranks for " + f"PP rank {pp_rank}: TP {reference_tp_rank} and " + f"TP {tp_rank} mismatch in {mismatched_fields}" + ) + + merged[pp_rank] = MooncakePPTransferMetadata( + block_size=reference.block_size, + num_blocks=reference.num_blocks, + spec_block_sizes=reference.spec_block_sizes, + spec_head_sizes=reference.spec_head_sizes, + layer_names=reference.layer_names, + group_indices=reference.group_indices, + spec_indices=reference.spec_indices, + block_strides=reference.block_strides, + block_lens=reference.block_lens, + block_shapes=reference.block_shapes, + block_size_scales=reference.block_size_scales, + pcp_size=reference.pcp_size, + dcp_size=reference.dcp_size, + tp_size=reference.tp_size, + metadata_by_tp_rank={ + tp_rank: MooncakeTPTransferMetadata( + te_rpc_port=worker_metadata.te_rpc_port, + kv_caches_base_addr=( + worker_metadata.kv_caches_base_addr + ), + local_ip=worker_metadata.local_ip, + handshake_port=worker_metadata.handshake_port, + ) + for tp_rank, worker_metadata in sorted( + workers_by_tp_rank.items() + ) + }, + ) + return merged def add_delayed_request(self, request_id: str, delay_start_time: float) -> None: with self.state_lock: @@ -188,9 +259,10 @@ def _run_busy_loop(self, sock: zmq.Socket) -> None: try: msg = decoder.decode(payload[0]) - if msg[0] == GET_META_MSG and len(msg) == 2: - metadata = self._get_metadata(msg[1]) - sock.send_multipart((identity, b"", metadata or b"")) + if msg and msg[0] == GET_META_MSG: + sock.send_multipart( + (identity, b"", self.encoded_metadata) + ) elif msg[0] == DONE_RECVING_MSG and len(msg) == 2: self._handle_finished_request(str(msg[1])) sock.send_multipart((identity, b"", ACK_MSG)) @@ -199,14 +271,6 @@ def _run_busy_loop(self, sock: zmq.Socket) -> None: except Exception: logger.exception("Failed to handle Mooncake scheduler control message") - def _get_metadata(self, tp_rank: Any) -> bytes | None: - if not isinstance(tp_rank, int): - logger.warning("Invalid Mooncake metadata TP rank: %r", tp_rank) - return None - if tp_rank == -1: - return self.encoded_metadata - return self.encoded_metadata_by_tp_rank.get(tp_rank) - def _handle_finished_request(self, request_id: str) -> None: with self.state_lock: if request_id in self.finished_request_ids: 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 8483ed0d356a..4b3e8f8c98fd 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 @@ -20,7 +20,6 @@ from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( MooncakeConnectorMetadata, MooncakePPTransferMetadata, - MooncakeTPTransferMetadata, MooncakeTransferMetadata, MooncakeTransferMetadataGroups, ReqMeta, @@ -54,6 +53,8 @@ def __init__( local_metadata: MooncakeTransferMetadata, tp_rank: int, tp_size: int, + pp_rank: int, + pp_size: int, dp_rank: int, dp_size: int, pcp_rank: int, @@ -99,6 +100,8 @@ def __init__( self.tp_rank = tp_rank self.tp_size = tp_size + self.pp_rank = pp_rank + self.pp_size = pp_size self.dp_rank = dp_rank self.dp_size = dp_size self.pcp_rank = pcp_rank @@ -207,8 +210,6 @@ def _handle_requests( """Fetch peer metadata and submit its aggregated requests.""" self._get_remote_metadata(remote_engine_id, remote_host, remote_port) - # Transfer planning and Mooncake READ submission are implemented next. - # Until then, do not report an incomplete transfer as successful. raise NotImplementedError("Mooncake pull transfer submission is not implemented") @staticmethod @@ -221,111 +222,33 @@ def _validate_remote_metadata( "Mooncake producer metadata engine ID mismatch: expected " f"{expected_engine_id!r}, got {metadata.engine_id!r}" ) - if not metadata.metadata_by_tp_rank: - raise ValueError("Mooncake producer scheduler returned no worker metadata") + if not metadata.metadata_by_pp_rank: + raise ValueError( + "Mooncake producer scheduler returned no PP metadata" + ) - empty_tp_ranks = [ - tp_rank - for tp_rank, metadata_by_pp_rank in metadata.metadata_by_tp_rank.items() - if not metadata_by_pp_rank + empty_pp_ranks = [ + pp_rank + for pp_rank, pp_metadata in metadata.metadata_by_pp_rank.items() + if not pp_metadata.metadata_by_tp_rank ] - if empty_tp_ranks: + if empty_pp_ranks: raise ValueError( - "Mooncake producer scheduler returned no PP metadata for TP ranks " - f"{sorted(empty_tp_ranks)}" + "Mooncake producer scheduler returned no TP metadata for " + f"PP ranks {sorted(empty_pp_ranks)}" ) - unexpected_engine_ids = { - worker_metadata.engine_id - for metadata_by_pp_rank in metadata.metadata_by_tp_rank.values() - for worker_metadata in metadata_by_pp_rank.values() - if worker_metadata.engine_id != expected_engine_id + invalid_tp_ranks = { + tp_rank + for pp_metadata in metadata.metadata_by_pp_rank.values() + for tp_rank in pp_metadata.metadata_by_tp_rank + if tp_rank < 0 or tp_rank >= pp_metadata.tp_size } - if unexpected_engine_ids: + if invalid_tp_ranks: raise ValueError( - "Mooncake producer worker metadata engine ID mismatch: expected " - f"{expected_engine_id!r}, got {sorted(unexpected_engine_ids)!r}" + "Mooncake producer scheduler returned invalid TP ranks " + f"{sorted(invalid_tp_ranks)}" ) - - @staticmethod - def _aggregate_remote_metadata( - metadata: MooncakeTransferMetadataGroups, - ) -> dict[int, MooncakePPTransferMetadata]: - """Deduplicate TP-invariant metadata and group worker endpoints by PP.""" - workers_by_pp_rank: dict[int, dict[int, MooncakeTransferMetadata]] = {} - for tp_rank, workers_by_pp_rank_for_tp in metadata.metadata_by_tp_rank.items(): - for pp_rank, worker_metadata in workers_by_pp_rank_for_tp.items(): - workers_by_pp_rank.setdefault(pp_rank, {})[tp_rank] = worker_metadata - - common_fields = ( - "block_size", - "num_blocks", - "spec_block_sizes", - "spec_head_sizes", - "layer_names", - "group_indices", - "spec_indices", - "block_strides", - "block_lens", - "block_shapes", - "block_size_scales", - "pcp_size", - "dcp_size", - "tp_size", - ) - expected_tp_ranks = set(metadata.metadata_by_tp_rank) - aggregated: dict[int, MooncakePPTransferMetadata] = {} - for pp_rank, workers_by_tp_rank in workers_by_pp_rank.items(): - actual_tp_ranks = set(workers_by_tp_rank) - if actual_tp_ranks != expected_tp_ranks: - raise ValueError( - "Mooncake producer metadata has incomplete TP workers for " - f"PP rank {pp_rank}: expected {sorted(expected_tp_ranks)}, " - f"got {sorted(actual_tp_ranks)}" - ) - - reference_tp_rank = min(workers_by_tp_rank) - reference = workers_by_tp_rank[reference_tp_rank] - for tp_rank, worker_metadata in workers_by_tp_rank.items(): - mismatched_fields = [ - field_name - for field_name in common_fields - if getattr(worker_metadata, field_name) != getattr(reference, field_name) - ] - if mismatched_fields: - raise ValueError( - "Mooncake producer metadata differs across TP ranks " - f"for PP rank {pp_rank}: TP {reference_tp_rank} and " - f"TP {tp_rank} mismatch in {mismatched_fields}" - ) - - aggregated[pp_rank] = MooncakePPTransferMetadata( - block_size=reference.block_size, - num_blocks=reference.num_blocks, - spec_block_sizes=reference.spec_block_sizes, - spec_head_sizes=reference.spec_head_sizes, - layer_names=reference.layer_names, - group_indices=reference.group_indices, - spec_indices=reference.spec_indices, - block_strides=reference.block_strides, - block_lens=reference.block_lens, - block_shapes=reference.block_shapes, - block_size_scales=reference.block_size_scales, - pcp_size=reference.pcp_size, - dcp_size=reference.dcp_size, - tp_size=reference.tp_size, - metadata_by_tp_rank={ - tp_rank: MooncakeTPTransferMetadata( - te_rpc_port=worker_metadata.te_rpc_port, - kv_caches_base_addr=(worker_metadata.kv_caches_base_addr), - local_ip=worker_metadata.local_ip, - handshake_port=worker_metadata.handshake_port, - ) - for tp_rank, worker_metadata in sorted(workers_by_tp_rank.items()) - }, - ) - return aggregated - def _get_remote_metadata( self, remote_engine_id: str, @@ -342,7 +265,7 @@ def _get_remote_metadata( sock.setsockopt(zmq.RCVTIMEO, 1000) ensure_zmq_send( sock, - self.encoder.encode((b"get_meta_msg", -1)), + self.encoder.encode((b"get_meta_msg",)), path, ) metadata_bytes = ensure_zmq_recv(sock, path) @@ -355,7 +278,7 @@ def _get_remote_metadata( transfer_metadata = self.decoder.decode(metadata_bytes) self._validate_remote_metadata(transfer_metadata, remote_engine_id) - remote_metadata = self._aggregate_remote_metadata(transfer_metadata) + remote_metadata = transfer_metadata.metadata_by_pp_rank self.remote_metadata[remote_engine_id] = remote_metadata return remote_metadata From 12cd4d6bc3c632d3f492cdcc79b8e39ba4e90e1c Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Thu, 6 Aug 2026 14:56:32 +0800 Subject: [PATCH 14/48] fix ruff format Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_scheduler.py | 18 ++---- .../kv_p2p/mooncake/base_worker.py | 18 ++---- .../kv_transfer/kv_p2p/mooncake/metadata.py | 3 +- .../kv_p2p/mooncake/pull_scheduler.py | 37 +++---------- .../kv_p2p/mooncake/pull_worker.py | 55 +++++-------------- .../kv_transfer/kv_p2p/mooncake/stats.py | 46 ++++------------ 6 files changed, 44 insertions(+), 133 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py index 387902721927..2a8c3c608ad0 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py @@ -47,13 +47,9 @@ def __init__( self.vllm_config = vllm_config self.kv_transfer_config = vllm_config.kv_transfer_config - if ( - self.kv_transfer_config.is_kv_consumer - == self.kv_transfer_config.is_kv_producer - ): + if self.kv_transfer_config.is_kv_consumer == self.kv_transfer_config.is_kv_producer: raise ValueError( - "Mooncake scheduler requires exactly one KV transfer role, " - f"got {self.kv_transfer_config.kv_role!r}" + f"Mooncake scheduler requires exactly one KV transfer role, got {self.kv_transfer_config.kv_role!r}" ) self.kv_cache_config = kv_cache_config self.engine_id = engine_id @@ -70,10 +66,7 @@ def __init__( self.pp_size = vllm_config.parallel_config.pipeline_parallel_size self.tp_size = vllm_config.parallel_config.tensor_parallel_size self.pcp_size = vllm_config.parallel_config.prefill_context_parallel_size - assert self.pcp_size == 1, ( - "Mooncake temporarily requires prefill context parallel size 1, " - f"got {self.pcp_size}" - ) + assert self.pcp_size == 1, f"Mooncake temporarily requires prefill context parallel size 1, got {self.pcp_size}" self.dcp_size = vllm_config.parallel_config.decode_context_parallel_size self.max_device_id = ( self.tp_size @@ -87,9 +80,7 @@ def __init__( # Keep one scheduler control port per DP rank immediately after that # range so it cannot collide with a worker handshake socket. self.side_channel_port = ( - self.kv_transfer_config.kv_port - + self.max_device_id - + vllm_config.parallel_config.data_parallel_rank + self.kv_transfer_config.kv_port + self.max_device_id + vllm_config.parallel_config.data_parallel_rank ) self.kv_cache_groups = kv_cache_config.kv_cache_groups @@ -219,4 +210,5 @@ def set_xfer_handshake_metadata( """Handle the legacy port-offset-keyed handshake entry point.""" self.set_xfer_handshake_metadata_from_workers(metadata) + __all__ = ["MooncakeBaseConnectorScheduler"] 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 d248837e4f18..98a298424fb6 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 @@ -72,13 +72,9 @@ def __init__( self.vllm_config = vllm_config self.kv_transfer_config = vllm_config.kv_transfer_config - if ( - self.kv_transfer_config.is_kv_consumer - == self.kv_transfer_config.is_kv_producer - ): + if self.kv_transfer_config.is_kv_consumer == self.kv_transfer_config.is_kv_producer: raise ValueError( - "Mooncake worker requires exactly one KV transfer role, " - f"got {self.kv_transfer_config.kv_role!r}" + f"Mooncake worker requires exactly one KV transfer role, got {self.kv_transfer_config.kv_role!r}" ) self.engine_id = engine_id self.kv_cache_config = kv_cache_config @@ -101,10 +97,7 @@ def __init__( pcp_group = get_pcp_group() self.pcp_rank = pcp_group.rank_in_group self.pcp_size = pcp_group.world_size - assert self.pcp_size == 1, ( - "Mooncake temporarily requires prefill context parallel size 1, " - f"got {self.pcp_size}" - ) + assert self.pcp_size == 1, f"Mooncake temporarily requires prefill context parallel size 1, got {self.pcp_size}" self.dcp_size = get_decode_context_model_parallel_world_size() self.dcp_rank = get_decode_context_model_parallel_rank() if self.dcp_size > 1 else 0 @@ -188,10 +181,7 @@ def register_kv_caches( logger.info("num_blocks: %s", self.num_blocks) self.kv_caches = kv_caches self._build_kv_cache_spec_mappings() - spec_properties = [ - self._get_spec_transfer_properties(spec) - for spec in self.kv_cache_specs - ] + spec_properties = [self._get_spec_transfer_properties(spec) for spec in self.kv_cache_specs] spec_block_sizes = [properties[0] for properties in spec_properties] spec_head_sizes = [properties[1] for properties in spec_properties] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 2fb826aedfd9..d3612a91fd13 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -54,8 +54,7 @@ def __post_init__(self) -> None: for field_name, values in per_layer_fields.items(): if len(values) != num_layers: raise ValueError( - f"Mooncake transfer metadata field {field_name!r} has " - f"{len(values)} layers, expected {num_layers}." + f"Mooncake transfer metadata field {field_name!r} has {len(values)} layers, expected {num_layers}." ) for layer_index, layer_name in enumerate(self.layer_names): diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index f2970c6a4b41..7b07db662995 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -101,8 +101,7 @@ def _merge_metadata_by_pp_rank( pp_rank, tp_rank = metadata_key else: raise ValueError( - "Mooncake handshake metadata key must be tp_rank or " - f"(pp_rank, tp_rank), got {metadata_key!r}" + f"Mooncake handshake metadata key must be tp_rank or (pp_rank, tp_rank), got {metadata_key!r}" ) if not isinstance(rank_metadata, MooncakeTransferMetadata): @@ -118,10 +117,7 @@ def _merge_metadata_by_pp_rank( workers_by_tp_rank = workers_by_pp_rank.setdefault(pp_rank, {}) if tp_rank in workers_by_tp_rank: - raise ValueError( - f"Duplicate Mooncake metadata for PP rank {pp_rank}, " - f"TP rank {tp_rank}" - ) + raise ValueError(f"Duplicate Mooncake metadata for PP rank {pp_rank}, TP rank {tp_rank}") workers_by_tp_rank[tp_rank] = rank_metadata expected_pp_ranks = set(range(self.pp_size)) @@ -150,9 +146,7 @@ def _merge_metadata_by_pp_rank( ) expected_tp_ranks = set(range(self.tp_size)) merged: dict[int, MooncakePPTransferMetadata] = {} - for pp_rank, workers_by_tp_rank in sorted( - workers_by_pp_rank.items() - ): + for pp_rank, workers_by_tp_rank in sorted(workers_by_pp_rank.items()): if set(workers_by_tp_rank) != expected_tp_ranks: raise ValueError( "Mooncake worker metadata has incomplete TP ranks for " @@ -173,8 +167,7 @@ def _merge_metadata_by_pp_rank( mismatched_fields = [ field_name for field_name in common_fields - if getattr(worker_metadata, field_name) - != getattr(reference, field_name) + if getattr(worker_metadata, field_name) != getattr(reference, field_name) ] if mismatched_fields: raise ValueError( @@ -201,15 +194,11 @@ def _merge_metadata_by_pp_rank( metadata_by_tp_rank={ tp_rank: MooncakeTPTransferMetadata( te_rpc_port=worker_metadata.te_rpc_port, - kv_caches_base_addr=( - worker_metadata.kv_caches_base_addr - ), + kv_caches_base_addr=(worker_metadata.kv_caches_base_addr), local_ip=worker_metadata.local_ip, handshake_port=worker_metadata.handshake_port, ) - for tp_rank, worker_metadata in sorted( - workers_by_tp_rank.items() - ) + for tp_rank, worker_metadata in sorted(workers_by_tp_rank.items()) }, ) return merged @@ -260,9 +249,7 @@ def _run_busy_loop(self, sock: zmq.Socket) -> None: try: msg = decoder.decode(payload[0]) if msg and msg[0] == GET_META_MSG: - sock.send_multipart( - (identity, b"", self.encoded_metadata) - ) + sock.send_multipart((identity, b"", self.encoded_metadata)) elif msg[0] == DONE_RECVING_MSG and len(msg) == 2: self._handle_finished_request(str(msg[1])) sock.send_multipart((identity, b"", ACK_MSG)) @@ -377,11 +364,7 @@ def set_xfer_handshake_metadata_from_workers( self, metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], ) -> None: - if ( - not self.kv_transfer_config.is_kv_producer - or not metadata - or self._sending_thread is not None - ): + if not self.kv_transfer_config.is_kv_producer or not metadata or self._sending_thread is not None: return ready_event = threading.Event() @@ -560,9 +543,7 @@ def update_connector_output( remote = self._reqs_recv_info.pop(req_id, None) if remote is not None: if self._recving_thread is None: - raise RuntimeError( - "Producer Mooncake scheduler received a receive-completion event" - ) + raise RuntimeError("Producer Mooncake scheduler received a receive-completion event") self._recving_thread.add_request(*remote) # P side: feed scheduler-received ACKs into vLLM's standard delayed 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 4b3e8f8c98fd..63c141dc4808 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 @@ -88,15 +88,9 @@ def __init__( self.use_mla = self.model_config.is_deepseek_mla hf_text_config = self.model_config.hf_text_config - self.num_key_value_heads = getattr( - hf_text_config, "num_key_value_heads", 0 - ) + self.num_key_value_heads = getattr(hf_text_config, "num_key_value_heads", 0) speculative_config = vllm_config.speculative_config - self.num_speculative_tokens = ( - speculative_config.num_speculative_tokens - if speculative_config is not None - else 0 - ) + self.num_speculative_tokens = speculative_config.num_speculative_tokens if speculative_config is not None else 0 self.tp_rank = tp_rank self.tp_size = tp_size @@ -108,17 +102,12 @@ def __init__( self.pcp_size = pcp_size self.dcp_rank = dcp_rank self.dcp_size = dcp_size - assert self.pcp_size == 1, ( - "Mooncake pull worker temporarily requires pcp_size=1, " - f"got {self.pcp_size}" - ) + assert self.pcp_size == 1, f"Mooncake pull worker temporarily requires pcp_size=1, got {self.pcp_size}" self.device = device self.ready_event = ready_event self.encoder = msgspec.msgpack.Encoder() self.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) - self.remote_metadata: SizedDict[ - str, dict[int, MooncakePPTransferMetadata] - ] = SizedDict() + self.remote_metadata: SizedDict[str, dict[int, MooncakePPTransferMetadata]] = SizedDict() self.request_queue: queue.Queue[tuple[str, dict[str, ReqMeta]]] = queue.Queue() self.finished_requests: queue.SimpleQueue[str] = queue.SimpleQueue() self.invalid_block_ids: set[int] = set() @@ -156,14 +145,9 @@ def run(self) -> None: self.ready_event.set() while True: remote_engine_id, requests = self.request_queue.get() - request_endpoints = { - (request.remote_host, request.remote_port) - for request in requests.values() - } + request_endpoints = {(request.remote_host, request.remote_port) for request in requests.values()} try: - remote_host, remote_port = self._get_remote_endpoint( - request_endpoints - ) + remote_host, remote_port = self._get_remote_endpoint(request_endpoints) self._handle_requests( remote_engine_id, remote_host, @@ -195,8 +179,7 @@ def _get_remote_endpoint( ) -> tuple[str, int]: if len(endpoints) != 1: raise ValueError( - "Requests for one remote engine must share one scheduler " - f"endpoint, got {sorted(endpoints)}" + f"Requests for one remote engine must share one scheduler endpoint, got {sorted(endpoints)}" ) return next(iter(endpoints)) @@ -223,9 +206,7 @@ def _validate_remote_metadata( f"{expected_engine_id!r}, got {metadata.engine_id!r}" ) if not metadata.metadata_by_pp_rank: - raise ValueError( - "Mooncake producer scheduler returned no PP metadata" - ) + raise ValueError("Mooncake producer scheduler returned no PP metadata") empty_pp_ranks = [ pp_rank @@ -234,8 +215,7 @@ def _validate_remote_metadata( ] if empty_pp_ranks: raise ValueError( - "Mooncake producer scheduler returned no TP metadata for " - f"PP ranks {sorted(empty_pp_ranks)}" + f"Mooncake producer scheduler returned no TP metadata for PP ranks {sorted(empty_pp_ranks)}" ) invalid_tp_ranks = { @@ -245,10 +225,8 @@ def _validate_remote_metadata( if tp_rank < 0 or tp_rank >= pp_metadata.tp_size } if invalid_tp_ranks: - raise ValueError( - "Mooncake producer scheduler returned invalid TP ranks " - f"{sorted(invalid_tp_ranks)}" - ) + raise ValueError(f"Mooncake producer scheduler returned invalid TP ranks {sorted(invalid_tp_ranks)}") + def _get_remote_metadata( self, remote_engine_id: str, @@ -272,8 +250,7 @@ def _get_remote_metadata( if not metadata_bytes: raise RuntimeError( - "Mooncake producer scheduler returned no transfer metadata " - f"for engine {remote_engine_id!r} from {path}" + f"Mooncake producer scheduler returned no transfer metadata for engine {remote_engine_id!r} from {path}" ) transfer_metadata = self.decoder.decode(metadata_bytes) @@ -338,9 +315,7 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata) -> None: assert self._recving_thread is not None request_groups: dict[str, dict[str, ReqMeta]] = {} for request_id, request_metadata in metadata.requests.items(): - requests = request_groups.setdefault( - request_metadata.remote_engine_id, {} - ) + requests = request_groups.setdefault(request_metadata.remote_engine_id, {}) requests[request_id] = request_metadata for remote_engine_id, requests in request_groups.items(): @@ -349,9 +324,7 @@ def start_load_kv(self, metadata: MooncakeConnectorMetadata) -> None: def get_finished(self) -> tuple[set[str], set[str]]: """Return requests with completed receive and send operations.""" finished_recving = ( - self._recving_thread.get_and_clear_finished_requests() - if self._recving_thread is not None - else set() + self._recving_thread.get_and_clear_finished_requests() if self._recving_thread is not None else set() ) return set(), finished_recving diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/stats.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/stats.py index 495a6c2e896b..07be8c62ee8d 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/stats.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/stats.py @@ -34,9 +34,7 @@ def reset(self) -> None: "num_failed_transfers": [], } - def record_transfer( - self, transfer_duration: float, bytes_transferred: int - ) -> None: + def record_transfer(self, transfer_duration: float, bytes_transferred: int) -> None: """Record one completed D2D transfer. ``transfer_duration`` is expressed in seconds. @@ -53,10 +51,7 @@ def clone_and_reset(self) -> "MooncakeKVConnectorStats": return previous def is_empty(self) -> bool: - return ( - self.num_successful_transfers == 0 - and not self.data["num_failed_transfers"] - ) + return self.num_successful_transfers == 0 and not self.data["num_failed_transfers"] def aggregate(self, other: KVConnectorStats) -> KVConnectorStats: if not other.is_empty(): @@ -70,9 +65,7 @@ def reduce(self) -> dict[str, int | float]: if self.num_successful_transfers == 0: return { "Num successful transfers": 0, - "Num failed transfers": len( - self.data["num_failed_transfers"] - ), + "Num failed transfers": len(self.data["num_failed_transfers"]), "Avg xfer time (ms)": 0, "P90 xfer time (ms)": 0, "Avg MB per transfer": 0, @@ -82,17 +75,13 @@ def reduce(self) -> dict[str, int | float]: durations = np.asarray(self.data["transfer_duration"]) megabytes = np.asarray(self.data["bytes_transferred"]) / 2**20 total_duration = durations.sum() - throughput = ( - megabytes.sum() / total_duration if total_duration > 0 else 0 - ) + throughput = megabytes.sum() / total_duration if total_duration > 0 else 0 return { "Num successful transfers": self.num_successful_transfers, "Num failed transfers": len(self.data["num_failed_transfers"]), "Avg xfer time (ms)": round(durations.mean() * 1e3, 3), - "P90 xfer time (ms)": round( - np.percentile(durations, 90).item() * 1e3, 3 - ), + "P90 xfer time (ms)": round(np.percentile(durations, 90).item() * 1e3, 3), "Avg MB per transfer": round(megabytes.mean(), 3), "Throughput (MB/s)": round(float(throughput), 3), } @@ -121,9 +110,7 @@ def __init__( transfer_duration = self._histogram_cls( name="vllm:mooncake_xfer_time_seconds", - documentation=( - "Histogram of Mooncake KV cache D2D transfer duration." - ), + documentation=("Histogram of Mooncake KV cache D2D transfer duration."), buckets=[ 0.005, 0.01, @@ -140,35 +127,24 @@ def __init__( ], labelnames=labelnames, ) - self.transfer_duration = create_metric_per_engine( - transfer_duration, self.per_engine_labelvalues - ) + self.transfer_duration = create_metric_per_engine(transfer_duration, self.per_engine_labelvalues) bytes_transferred = self._histogram_cls( name="vllm:mooncake_bytes_transferred", - documentation=( - "Histogram of bytes transferred per Mooncake KV cache D2D " - "transfer." - ), + documentation=("Histogram of bytes transferred per Mooncake KV cache D2D transfer."), buckets=[2 ** (10 + index) for index in range(1, 25, 2)], labelnames=labelnames, ) - self.bytes_transferred = create_metric_per_engine( - bytes_transferred, self.per_engine_labelvalues - ) + self.bytes_transferred = create_metric_per_engine(bytes_transferred, self.per_engine_labelvalues) failed_transfers = self._counter_cls( name="vllm:mooncake_num_failed_transfers", documentation="Number of failed Mooncake KV cache D2D transfers.", labelnames=labelnames, ) - self.failed_transfers = create_metric_per_engine( - failed_transfers, self.per_engine_labelvalues - ) + self.failed_transfers = create_metric_per_engine(failed_transfers, self.per_engine_labelvalues) - def observe( - self, transfer_stats_data: dict[str, Any], engine_idx: int = 0 - ) -> None: + def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0) -> None: for duration in transfer_stats_data["transfer_duration"]: self.transfer_duration[engine_idx].observe(duration) for num_bytes in transfer_stats_data["bytes_transferred"]: From b867c234518b79f27a127c174bc72db2a2d4206b Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Fri, 7 Aug 2026 15:19:19 +0800 Subject: [PATCH 15/48] add kernel block sizes Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_worker.py | 55 +++++++++++++++---- .../kv_transfer/kv_p2p/mooncake/metadata.py | 23 ++++++-- .../kv_p2p/mooncake/pull_scheduler.py | 6 +- 3 files changed, 67 insertions(+), 17 deletions(-) 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 98a298424fb6..20230a1d46ac 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 @@ -25,6 +25,9 @@ AttentionSpec, FullAttentionSpec, KVCacheSpec, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, SlidingWindowSpec, UniformTypeKVCacheSpecs, ) @@ -162,15 +165,21 @@ def _get_spec_transfer_properties( if not isinstance(spec, AttentionSpec): return spec.block_size, None - if isinstance(spec, (FullAttentionSpec, SlidingWindowSpec)): - head_size_v = getattr(spec, "head_size_v", spec.head_size) - if spec.head_size != head_size_v: - raise NotImplementedError( - "Mooncake does not support different K/V head sizes for " - f"{type(spec).__name__}: K={spec.head_size}, V={head_size_v}" - ) - - return spec.block_size, spec.head_size + if isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)): + return spec.block_size, 1 + elif isinstance(spec, (FullAttentionSpec, SlidingWindowSpec)): + return spec.block_size, spec.num_kv_heads + elif isinstance(spec, MambaSpec): + mamba_shapes = spec.shapes + if len(mamba_shapes) == 2: + head_size = mamba_shapes[0][1] + elif len(mamba_shapes) == 1: + head_size = mamba_shapes[0][1] if len(mamba_shapes[0]) == 2 else mamba_shapes[0][0] + else: + raise NotImplementedError("Mooncake Connector does not support current MambaSpec type now.") + return spec.block_size, head_size + else: + raise NotImplementedError("Mooncake Connector does not support {%s} now.", type(spec).__name__) def register_kv_caches( self, @@ -183,11 +192,12 @@ def register_kv_caches( self._build_kv_cache_spec_mappings() spec_properties = [self._get_spec_transfer_properties(spec) for spec in self.kv_cache_specs] spec_block_sizes = [properties[0] for properties in spec_properties] - spec_head_sizes = [properties[1] for properties in spec_properties] + spec_num_heads = [properties[1] for properties in spec_properties] layer_names: list[str] = [] group_indices: list[int] = [] spec_indices: list[int] = [] + kernel_block_sizes: list[int] = spec_block_sizes.copy() kv_caches_base_addr: list[list[int]] = [] block_strides_per_layer: list[list[int]] = [] block_lens_per_layer: list[list[int]] = [] @@ -231,6 +241,28 @@ def register_kv_caches( block_lens_per_layer.append(block_lens) block_shapes_per_layer.append(block_shapes) block_size_scales_per_layer.append(block_size_scales) + cur_layer_spec_idx = self.layer_name_to_spec_index[layer_name] + spec_block_size = spec_block_sizes[cur_layer_spec_idx] + layer_kernel_block_sizes = {spec_block_size // scale for scale in block_size_scales} + if len(layer_kernel_block_sizes) != 1: + raise ValueError( + f"Layer {layer_name!r} has inconsistent kernel block sizes {sorted(layer_kernel_block_sizes)}." + ) + kernel_block_size = layer_kernel_block_sizes.pop() + if ( + kernel_block_sizes[cur_layer_spec_idx] == spec_block_sizes[cur_layer_spec_idx] + and kernel_block_size < spec_block_sizes[cur_layer_spec_idx] + ): + kernel_block_sizes[cur_layer_spec_idx] = kernel_block_size + if ( + kernel_block_sizes[cur_layer_spec_idx] != spec_block_sizes[cur_layer_spec_idx] + and kernel_block_size != kernel_block_sizes[cur_layer_spec_idx] + ): + raise ValueError( + "Layer %s has kernel block size %s different from global kernel block size %s", + kernel_block_size, + kernel_block_sizes[cur_layer_spec_idx], + ) unexpected_layers = kv_caches.keys() - configured_layer_names if unexpected_layers: @@ -246,7 +278,8 @@ def register_kv_caches( block_size=self.block_size, num_blocks=self.num_blocks, spec_block_sizes=spec_block_sizes, - spec_head_sizes=spec_head_sizes, + kernel_block_sizes=kernel_block_sizes, + spec_num_heads=spec_num_heads, layer_names=layer_names, group_indices=group_indices, spec_indices=spec_indices, diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index d3612a91fd13..0bac9df99597 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -16,8 +16,10 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): """Worker transfer information exchanged during the P/D handshake. - All per-layer fields use ``layer_names`` order. Each nested list preserves - the cache tensor order for that layer, such as K/V or Mamba conv/SSM. + Per-spec fields use the flattened KV-cache spec order. In particular, + kernel_block_sizes records the block size seen by worker kernels after + expanding a logical/spec block. All per-layer fields use layer_names + order; each nested list preserves that layer's cache tensor order. """ engine_id: str @@ -25,7 +27,8 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): block_size: int num_blocks: int spec_block_sizes: list[int] - spec_head_sizes: list[int | None] + kernel_block_sizes: list[int] + spec_num_heads: list[int | None] layer_names: list[str] group_indices: list[int] spec_indices: list[int] @@ -41,6 +44,17 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): handshake_port: int = 0 def __post_init__(self) -> None: + num_specs = len(self.spec_block_sizes) + per_spec_fields = { + "kernel_block_sizes": self.kernel_block_sizes, + "spec_num_heads": self.spec_num_heads, + } + for field_name, values in per_spec_fields.items(): + if len(values) != num_specs: + raise ValueError( + f"Mooncake transfer metadata field {field_name!r} has {len(values)} specs, expected {num_specs}." + ) + num_layers = len(self.layer_names) per_layer_fields = { "group_indices": self.group_indices, @@ -90,7 +104,8 @@ class MooncakePPTransferMetadata: block_size: int num_blocks: int spec_block_sizes: list[int] - spec_head_sizes: list[int | None] + kernel_block_sizes: list[int] + spec_num_heads: list[int | None] layer_names: list[str] group_indices: list[int] spec_indices: list[int] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 7b07db662995..536a2cca42e9 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -132,7 +132,8 @@ def _merge_metadata_by_pp_rank( "block_size", "num_blocks", "spec_block_sizes", - "spec_head_sizes", + "kernel_block_sizes", + "spec_num_heads", "layer_names", "group_indices", "spec_indices", @@ -180,7 +181,8 @@ def _merge_metadata_by_pp_rank( block_size=reference.block_size, num_blocks=reference.num_blocks, spec_block_sizes=reference.spec_block_sizes, - spec_head_sizes=reference.spec_head_sizes, + kernel_block_sizes=reference.kernel_block_sizes, + spec_num_heads=reference.spec_num_heads, layer_names=reference.layer_names, group_indices=reference.group_indices, spec_indices=reference.spec_indices, From f8deee93124dd564c079f4a908962c0a489c4031 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Sat, 8 Aug 2026 01:57:36 +0800 Subject: [PATCH 16/48] add remote_tp_layout Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/pull_worker.py | 240 +++++++++++++++++- 1 file changed, 237 insertions(+), 3 deletions(-) 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 63c141dc4808..bd2e49679bcf 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 @@ -12,8 +12,17 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.utils.network_utils import make_zmq_path -from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheSpec, + MambaSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + SlidingWindowSpec, +) +from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker import ( MooncakeBaseConnectorWorker, ) @@ -76,7 +85,8 @@ def __init__( self.block_size = local_metadata.block_size self.spec_block_sizes = local_metadata.spec_block_sizes - self.spec_head_sizes = local_metadata.spec_head_sizes + self.kernel_block_sizes = local_metadata.kernel_block_sizes + self.spec_num_heads = local_metadata.spec_num_heads self.layer_names = local_metadata.layer_names self.group_indices = local_metadata.group_indices self.spec_indices = local_metadata.spec_indices @@ -108,6 +118,12 @@ def __init__( self.encoder = msgspec.msgpack.Encoder() self.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) self.remote_metadata: SizedDict[str, dict[int, MooncakePPTransferMetadata]] = SizedDict() + # Candidate producer TP ranks, cached independently for each engine: + # engine_id -> remote_pp_rank -> (local_spec_index, remote_spec_index) + # -> candidate_groups + self.remote_tp_rank_groups: SizedDict[str, dict[int, dict[tuple[int, int], list[list[int]]]]] = SizedDict() + # engine_id -> remote_pp_rank -> [[local_layer_index, remote_layer_index], ...] + self.remote_layer_index_pairs: SizedDict[str, dict[int, list[tuple[int, int]]]] = SizedDict() self.request_queue: queue.Queue[tuple[str, dict[str, ReqMeta]]] = queue.Queue() self.finished_requests: queue.SimpleQueue[str] = queue.SimpleQueue() self.invalid_block_ids: set[int] = set() @@ -183,6 +199,216 @@ def _get_remote_endpoint( ) return next(iter(endpoints)) + def _build_remote_transfer_layout( + self, + remote_metadata: dict[int, MooncakePPTransferMetadata], + ) -> tuple[ + dict[int, dict[tuple[int, int], list[list[int]]]], + dict[int, list[tuple[int, int]]], + ]: + """Build per-PP spec candidates and matching layer index pairs. + + TP candidates are constructed once for each local spec present under a + remote PP rank. Layer traversal is used only to record + [local_layer_index, remote_layer_index] pairs for later address + planning. + """ + groups_by_pp_rank: dict[int, dict[tuple[int, int], list[list[int]]]] = {} + layer_pairs_by_pp_rank: dict[int, list[tuple[int, int]]] = {} + local_layer_index_by_name = {layer_name: layer_index for layer_index, layer_name in enumerate(self.layer_names)} + matched_local_layer_indices: set[int] = set() + + for remote_pp_rank, pp_metadata in sorted(remote_metadata.items()): + if pp_metadata.pcp_size != 1: + raise ValueError( + "Mooncake pull temporarily requires remote pcp_size=1, " + f"got {pp_metadata.pcp_size} for PP rank {remote_pp_rank}" + ) + + layer_index_pairs: list[tuple[int, int]] = [] + spec_index_mapping: dict[int, int] = {} + for remote_layer_index, layer_name in enumerate(pp_metadata.layer_names): + local_layer_index = local_layer_index_by_name.get(layer_name) + if local_layer_index is None: + continue + + layer_index_pairs.append((local_layer_index, remote_layer_index)) + matched_local_layer_indices.add(local_layer_index) + local_spec_index = self.spec_indices[local_layer_index] + remote_spec_index = pp_metadata.spec_indices[remote_layer_index] + mapped_remote_spec_index = spec_index_mapping.get(local_spec_index) + if mapped_remote_spec_index is not None and mapped_remote_spec_index != remote_spec_index: + raise ValueError( + f"Mooncake local spec {local_spec_index} maps to multiple producer specs " + f"under PP rank {remote_pp_rank}: {mapped_remote_spec_index} and {remote_spec_index}" + ) + spec_index_mapping[local_spec_index] = remote_spec_index + + layer_pairs_by_pp_rank[remote_pp_rank] = layer_index_pairs + groups_by_pp_rank[remote_pp_rank] = { + (local_spec_index, remote_spec_index): self._get_spec_remote_tp_rank_groups( + local_spec_index, + remote_spec_index, + self.kv_cache_specs[local_spec_index], + pp_metadata, + ) + for local_spec_index, remote_spec_index in sorted(spec_index_mapping.items()) + } + + missing_local_layer_indices = set(range(len(self.layer_names))) - matched_local_layer_indices + if missing_local_layer_indices: + missing_local_layers = [ + self.layer_names[layer_index] for layer_index in sorted(missing_local_layer_indices) + ] + raise ValueError( + f"Mooncake producer metadata is missing layers required by this worker: {missing_local_layers}" + ) + return groups_by_pp_rank, layer_pairs_by_pp_rank + + def _get_spec_remote_tp_rank_groups( + self, + local_spec_index: int, + remote_spec_index: int, + spec: KVCacheSpec, + remote_metadata: MooncakePPTransferMetadata, + ) -> list[list[int]]: + """Infer the spec's TP strategy and build its remote rank groups.""" + if isinstance(spec, MambaSpec): + return self._get_mamba_remote_tp_rank_groups(remote_metadata.tp_size) + if remote_spec_index >= len(remote_metadata.spec_num_heads): + raise ValueError( + f"Mooncake local spec {local_spec_index} maps to missing producer spec {remote_spec_index}" + ) + + fixed_total_num_kv_heads = None + if isinstance(spec, (AscendSFAIndexerCacheSpec, SlidingWindowMLASpec)): + local_dcp_size = remote_dcp_size = 1 + fixed_total_num_kv_heads = 1 + elif isinstance(spec, MLAAttentionSpec): + local_dcp_size = self.dcp_size + remote_dcp_size = remote_metadata.dcp_size + fixed_total_num_kv_heads = 1 + elif isinstance(spec, SlidingWindowSpec): + local_dcp_size = remote_dcp_size = 1 + elif isinstance(spec, FullAttentionSpec): + local_dcp_size = self.dcp_size + remote_dcp_size = remote_metadata.dcp_size + else: + raise NotImplementedError(f"Mooncake pull has no TP grouping rule for KV cache spec {type(spec).__name__}") + + total_num_kv_heads = self._infer_total_num_kv_heads( + local_num_kv_heads=spec.num_kv_heads, + remote_num_kv_heads=remote_metadata.spec_num_heads[remote_spec_index], + remote_tp_size=remote_metadata.tp_size, + local_dcp_size=local_dcp_size, + remote_dcp_size=remote_dcp_size, + fixed_total_num_kv_heads=fixed_total_num_kv_heads, + ) + return self._get_attention_remote_tp_rank_groups( + remote_tp_size=remote_metadata.tp_size, + local_dcp_size=local_dcp_size, + remote_dcp_size=remote_dcp_size, + total_num_kv_heads=total_num_kv_heads, + ) + + def _get_mamba_remote_tp_rank_groups( + self, + remote_tp_size: int, + ) -> list[list[int]]: + """Map Mamba ranks like non-replicated FullAttention TP shards.""" + larger_tp_size = max(self.tp_size, remote_tp_size) + smaller_tp_size = min(self.tp_size, remote_tp_size) + if smaller_tp_size <= 0 or larger_tp_size % smaller_tp_size != 0: + raise ValueError( + f"Mooncake Mamba TP sizes must have an integer ratio, got local={self.tp_size}, remote={remote_tp_size}" + ) + tp_ratio = larger_tp_size // smaller_tp_size + if remote_tp_size >= self.tp_size: + start_rank = self.tp_rank * tp_ratio + return [[remote_tp_rank] for remote_tp_rank in range(start_rank, start_rank + tp_ratio)] + return [[self.tp_rank // tp_ratio]] + + def _infer_total_num_kv_heads( + self, + local_num_kv_heads: int, + remote_num_kv_heads: int | None, + remote_tp_size: int, + local_dcp_size: int, + remote_dcp_size: int, + fixed_total_num_kv_heads: int | None, + ) -> int: + """Infer and validate the total KV heads for one transfer spec.""" + if fixed_total_num_kv_heads is not None: + return fixed_total_num_kv_heads + assert max(self.tp_size, remote_tp_size) % min(self.tp_size, remote_tp_size) == 0 + assert max(local_num_kv_heads, remote_num_kv_heads) % min(local_num_kv_heads, remote_num_kv_heads) == 0 + local_head_tp_size = self.tp_size // local_dcp_size + remote_head_tp_size = remote_tp_size // remote_dcp_size + inferred_total_heads: set[int] = set() + if local_num_kv_heads > 1: + inferred_total_heads.add(local_num_kv_heads * self.tp_size) + if remote_num_kv_heads > 1: + inferred_total_heads.add(remote_num_kv_heads * remote_tp_size) + if local_dcp_size > 1: + inferred_total_heads.add(local_head_tp_size) + if remote_dcp_size > 1: + inferred_total_heads.add(remote_head_tp_size) + + if not inferred_total_heads: + inferred_total_heads.add(min(local_head_tp_size, remote_head_tp_size)) + if len(inferred_total_heads) != 1: + raise ValueError(f"Mooncake inferred inconsistent total KV head counts: {sorted(inferred_total_heads)}") + + return inferred_total_heads.pop() + + def _get_attention_remote_tp_rank_groups( + self, + remote_tp_size: int, + local_dcp_size: int, + remote_dcp_size: int, + total_num_kv_heads: int, + ) -> list[list[int]]: + """Build remote TP groups from an already inferred head topology.""" + local_head_tp_size = self.tp_size // local_dcp_size + remote_head_tp_size = remote_tp_size // remote_dcp_size + local_head_tp_rank = self.tp_rank // local_dcp_size + local_head_interval = self._get_head_interval(local_head_tp_rank, local_head_tp_size, total_num_kv_heads) + + ranks_by_head_piece: dict[tuple[int, int], list[int]] = {} + for remote_head_tp_rank in range(remote_head_tp_size): + remote_head_interval = self._get_head_interval(remote_head_tp_rank, remote_head_tp_size, total_num_kv_heads) + head_start = max(local_head_interval[0], remote_head_interval[0]) + head_end = min(local_head_interval[1], remote_head_interval[1]) + if head_start >= head_end: + continue + remote_rank_start = remote_head_tp_rank * remote_dcp_size + ranks_by_head_piece.setdefault((head_start, head_end), []).extend( + range(remote_rank_start, remote_rank_start + remote_dcp_size) + ) + + if not ranks_by_head_piece: + raise ValueError( + "MooncakeConnector found no remote TP group for local rank " + f"{self.tp_rank}, local_tp={self.tp_size}, remote_tp={remote_tp_size}, " + f"total_heads={total_num_kv_heads}, local_dcp={local_dcp_size}, remote_dcp={remote_dcp_size}" + ) + return [ranks for _, ranks in sorted(ranks_by_head_piece.items())] + + @staticmethod + def _get_head_interval( + tp_rank: int, + tp_size: int, + total_num_kv_heads: int, + ) -> tuple[int, int]: + if total_num_kv_heads >= tp_size: + heads_per_rank = total_num_kv_heads // tp_size + start = tp_rank * heads_per_rank + return start, start + heads_per_rank + + replication_size = tp_size // total_num_kv_heads + head_index = tp_rank // replication_size + return head_index, head_index + 1 + def _handle_requests( self, remote_engine_id: str, @@ -217,7 +443,6 @@ def _validate_remote_metadata( raise ValueError( f"Mooncake producer scheduler returned no TP metadata for PP ranks {sorted(empty_pp_ranks)}" ) - invalid_tp_ranks = { tp_rank for pp_metadata in metadata.metadata_by_pp_rank.values() @@ -256,7 +481,16 @@ def _get_remote_metadata( transfer_metadata = self.decoder.decode(metadata_bytes) self._validate_remote_metadata(transfer_metadata, remote_engine_id) remote_metadata = transfer_metadata.metadata_by_pp_rank + remote_tp_rank_groups, remote_layer_index_pairs = self._build_remote_transfer_layout(remote_metadata) self.remote_metadata[remote_engine_id] = remote_metadata + self.remote_tp_rank_groups[remote_engine_id] = remote_tp_rank_groups + self.remote_layer_index_pairs[remote_engine_id] = remote_layer_index_pairs + logger.debug( + "Mooncake remote transfer layout for engine %s: TP candidates=%s, layer index pairs=%s", + remote_engine_id, + remote_tp_rank_groups, + remote_layer_index_pairs, + ) return remote_metadata From 25e3744cdf747b7c6e31ec10d8be495a3f7f66ca Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Sat, 8 Aug 2026 18:48:24 +0800 Subject: [PATCH 17/48] change metadatagroups Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_worker.py | 3 -- .../kv_transfer/kv_p2p/mooncake/metadata.py | 10 ++--- .../kv_p2p/mooncake/pull_scheduler.py | 22 +++++----- .../kv_p2p/mooncake/pull_worker.py | 40 +++++++++---------- 4 files changed, 34 insertions(+), 41 deletions(-) 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 20230a1d46ac..20b84416ebb8 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 @@ -288,9 +288,6 @@ def register_kv_caches( block_lens=block_lens_per_layer, block_shapes=block_shapes_per_layer, block_size_scales=block_size_scales_per_layer, - pcp_size=self.pcp_size, - dcp_size=self.dcp_size, - tp_size=self.tp_size, local_ip=self.side_channel_host, handshake_port=self.handshake_port, ) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 0bac9df99597..fd5c9987d1f8 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -37,9 +37,6 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): block_lens: list[list[int]] block_shapes: list[list[tuple[int, ...]]] block_size_scales: list[list[int]] - pcp_size: int - dcp_size: int - tp_size: int local_ip: str = "" handshake_port: int = 0 @@ -113,9 +110,6 @@ class MooncakePPTransferMetadata: block_lens: list[list[int]] block_shapes: list[list[tuple[int, ...]]] block_size_scales: list[list[int]] - pcp_size: int - dcp_size: int - tp_size: int metadata_by_tp_rank: dict[int, MooncakeTPTransferMetadata] @@ -126,6 +120,10 @@ class MooncakeTransferMetadataGroups: engine_id: str scheduler_host: str scheduler_port: int + pp_size: int + pcp_size: int + dcp_size: int + tp_size: int metadata_by_pp_rank: dict[int, MooncakePPTransferMetadata] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 536a2cca42e9..76c669c2c9e3 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -63,6 +63,8 @@ def __init__( metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], tp_size: int, pp_size: int, + pcp_size: int, + dcp_size: int, ready_event: threading.Event, ) -> None: super().__init__(daemon=True, name="MooncakeSchedulerSendingThread") @@ -71,6 +73,8 @@ def __init__( self.port = port self.tp_size = tp_size self.pp_size = pp_size + self.pcp_size = pcp_size + self.dcp_size = dcp_size self.engine_id = engine_id metadata_by_pp_rank = self._merge_metadata_by_pp_rank(metadata) self.encoded_metadata = encoder.encode( @@ -78,6 +82,10 @@ def __init__( engine_id=engine_id, scheduler_host=host, scheduler_port=port, + pp_size=pp_size, + pcp_size=pcp_size, + dcp_size=dcp_size, + tp_size=tp_size, metadata_by_pp_rank=metadata_by_pp_rank, ) ) @@ -141,9 +149,6 @@ def _merge_metadata_by_pp_rank( "block_lens", "block_shapes", "block_size_scales", - "pcp_size", - "dcp_size", - "tp_size", ) expected_tp_ranks = set(range(self.tp_size)) merged: dict[int, MooncakePPTransferMetadata] = {} @@ -157,12 +162,6 @@ def _merge_metadata_by_pp_rank( reference_tp_rank = min(workers_by_tp_rank) reference = workers_by_tp_rank[reference_tp_rank] - if reference.tp_size != self.tp_size: - raise ValueError( - "Mooncake worker metadata TP size mismatch for PP rank " - f"{pp_rank}: expected {self.tp_size}, " - f"got {reference.tp_size}" - ) for tp_rank, worker_metadata in workers_by_tp_rank.items(): mismatched_fields = [ @@ -190,9 +189,6 @@ def _merge_metadata_by_pp_rank( block_lens=reference.block_lens, block_shapes=reference.block_shapes, block_size_scales=reference.block_size_scales, - pcp_size=reference.pcp_size, - dcp_size=reference.dcp_size, - tp_size=reference.tp_size, metadata_by_tp_rank={ tp_rank: MooncakeTPTransferMetadata( te_rpc_port=worker_metadata.te_rpc_port, @@ -377,6 +373,8 @@ def set_xfer_handshake_metadata_from_workers( metadata, self.tp_size, self.pp_size, + self.pcp_size, + self.dcp_size, ready_event, ) self._sending_thread.start() 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 bd2e49679bcf..b47e81e3d708 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 @@ -117,7 +117,7 @@ def __init__( self.ready_event = ready_event self.encoder = msgspec.msgpack.Encoder() self.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) - self.remote_metadata: SizedDict[str, dict[int, MooncakePPTransferMetadata]] = SizedDict() + self.remote_metadata: SizedDict[str, MooncakeTransferMetadataGroups] = SizedDict() # Candidate producer TP ranks, cached independently for each engine: # engine_id -> remote_pp_rank -> (local_spec_index, remote_spec_index) # -> candidate_groups @@ -201,7 +201,7 @@ def _get_remote_endpoint( def _build_remote_transfer_layout( self, - remote_metadata: dict[int, MooncakePPTransferMetadata], + remote_metadata: MooncakeTransferMetadataGroups, ) -> tuple[ dict[int, dict[tuple[int, int], list[list[int]]]], dict[int, list[tuple[int, int]]], @@ -218,13 +218,7 @@ def _build_remote_transfer_layout( local_layer_index_by_name = {layer_name: layer_index for layer_index, layer_name in enumerate(self.layer_names)} matched_local_layer_indices: set[int] = set() - for remote_pp_rank, pp_metadata in sorted(remote_metadata.items()): - if pp_metadata.pcp_size != 1: - raise ValueError( - "Mooncake pull temporarily requires remote pcp_size=1, " - f"got {pp_metadata.pcp_size} for PP rank {remote_pp_rank}" - ) - + for remote_pp_rank, pp_metadata in sorted(remote_metadata.metadata_by_pp_rank.items()): layer_index_pairs: list[tuple[int, int]] = [] spec_index_mapping: dict[int, int] = {} for remote_layer_index, layer_name in enumerate(pp_metadata.layer_names): @@ -251,6 +245,8 @@ def _build_remote_transfer_layout( remote_spec_index, self.kv_cache_specs[local_spec_index], pp_metadata, + remote_metadata.tp_size, + remote_metadata.dcp_size, ) for local_spec_index, remote_spec_index in sorted(spec_index_mapping.items()) } @@ -271,10 +267,12 @@ def _get_spec_remote_tp_rank_groups( remote_spec_index: int, spec: KVCacheSpec, remote_metadata: MooncakePPTransferMetadata, + remote_tp_size: int, + remote_dcp_size: int, ) -> list[list[int]]: """Infer the spec's TP strategy and build its remote rank groups.""" if isinstance(spec, MambaSpec): - return self._get_mamba_remote_tp_rank_groups(remote_metadata.tp_size) + return self._get_mamba_remote_tp_rank_groups(remote_tp_size) if remote_spec_index >= len(remote_metadata.spec_num_heads): raise ValueError( f"Mooncake local spec {local_spec_index} maps to missing producer spec {remote_spec_index}" @@ -286,26 +284,24 @@ def _get_spec_remote_tp_rank_groups( fixed_total_num_kv_heads = 1 elif isinstance(spec, MLAAttentionSpec): local_dcp_size = self.dcp_size - remote_dcp_size = remote_metadata.dcp_size fixed_total_num_kv_heads = 1 elif isinstance(spec, SlidingWindowSpec): local_dcp_size = remote_dcp_size = 1 elif isinstance(spec, FullAttentionSpec): local_dcp_size = self.dcp_size - remote_dcp_size = remote_metadata.dcp_size else: raise NotImplementedError(f"Mooncake pull has no TP grouping rule for KV cache spec {type(spec).__name__}") total_num_kv_heads = self._infer_total_num_kv_heads( local_num_kv_heads=spec.num_kv_heads, remote_num_kv_heads=remote_metadata.spec_num_heads[remote_spec_index], - remote_tp_size=remote_metadata.tp_size, + remote_tp_size=remote_tp_size, local_dcp_size=local_dcp_size, remote_dcp_size=remote_dcp_size, fixed_total_num_kv_heads=fixed_total_num_kv_heads, ) return self._get_attention_remote_tp_rank_groups( - remote_tp_size=remote_metadata.tp_size, + remote_tp_size=remote_tp_size, local_dcp_size=local_dcp_size, remote_dcp_size=remote_dcp_size, total_num_kv_heads=total_num_kv_heads, @@ -433,6 +429,11 @@ def _validate_remote_metadata( ) if not metadata.metadata_by_pp_rank: raise ValueError("Mooncake producer scheduler returned no PP metadata") + if metadata.pcp_size != 1: + raise ValueError( + "Mooncake pull temporarily requires remote pcp_size=1, " + f"got {metadata.pcp_size}" + ) empty_pp_ranks = [ pp_rank @@ -447,7 +448,7 @@ def _validate_remote_metadata( tp_rank for pp_metadata in metadata.metadata_by_pp_rank.values() for tp_rank in pp_metadata.metadata_by_tp_rank - if tp_rank < 0 or tp_rank >= pp_metadata.tp_size + if tp_rank < 0 or tp_rank >= metadata.tp_size } if invalid_tp_ranks: raise ValueError(f"Mooncake producer scheduler returned invalid TP ranks {sorted(invalid_tp_ranks)}") @@ -457,7 +458,7 @@ def _get_remote_metadata( remote_engine_id: str, remote_host: str, remote_port: int, - ) -> dict[int, MooncakePPTransferMetadata]: + ) -> MooncakeTransferMetadataGroups: cached_metadata = self.remote_metadata.get(remote_engine_id) if cached_metadata is not None: return cached_metadata @@ -480,9 +481,8 @@ def _get_remote_metadata( transfer_metadata = self.decoder.decode(metadata_bytes) self._validate_remote_metadata(transfer_metadata, remote_engine_id) - remote_metadata = transfer_metadata.metadata_by_pp_rank - remote_tp_rank_groups, remote_layer_index_pairs = self._build_remote_transfer_layout(remote_metadata) - self.remote_metadata[remote_engine_id] = remote_metadata + remote_tp_rank_groups, remote_layer_index_pairs = self._build_remote_transfer_layout(transfer_metadata) + self.remote_metadata[remote_engine_id] = transfer_metadata self.remote_tp_rank_groups[remote_engine_id] = remote_tp_rank_groups self.remote_layer_index_pairs[remote_engine_id] = remote_layer_index_pairs logger.debug( @@ -491,7 +491,7 @@ def _get_remote_metadata( remote_tp_rank_groups, remote_layer_index_pairs, ) - return remote_metadata + return transfer_metadata class MooncakePullConnectorWorker(MooncakeBaseConnectorWorker): From 414aea8230d0a744915c5486d22c3420d3758679 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Sat, 8 Aug 2026 23:07:08 +0800 Subject: [PATCH 18/48] add layerwise generate tp_layout Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/base_worker.py | 64 +--------------- .../kv_transfer/kv_p2p/mooncake/metadata.py | 22 +----- .../kv_p2p/mooncake/pull_scheduler.py | 4 - .../kv_p2p/mooncake/pull_worker.py | 73 ++++++++----------- 4 files changed, 35 insertions(+), 128 deletions(-) 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 20b84416ebb8..93ee6bcfbe13 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 @@ -21,16 +21,7 @@ ) from vllm.logger import init_logger from vllm.utils.network_utils import get_ip -from vllm.v1.kv_cache_interface import ( - AttentionSpec, - FullAttentionSpec, - KVCacheSpec, - MambaSpec, - MLAAttentionSpec, - SlidingWindowMLASpec, - SlidingWindowSpec, - UniformTypeKVCacheSpecs, -) +from vllm.v1.kv_cache_interface import KVCacheSpec, UniformTypeKVCacheSpecs from vllm_ascend.ascend_config import get_ascend_config, init_ascend_config from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( @@ -157,30 +148,6 @@ def _build_kv_cache_spec_mappings(self) -> None: self.layer_name_to_group_index[layer_name] = group_index self.layer_name_to_spec_index[layer_name] = spec_index - @staticmethod - def _get_spec_transfer_properties( - spec: KVCacheSpec, - ) -> tuple[int, int | None]: - """Return the block size and head size required for KV transfer.""" - if not isinstance(spec, AttentionSpec): - return spec.block_size, None - - if isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)): - return spec.block_size, 1 - elif isinstance(spec, (FullAttentionSpec, SlidingWindowSpec)): - return spec.block_size, spec.num_kv_heads - elif isinstance(spec, MambaSpec): - mamba_shapes = spec.shapes - if len(mamba_shapes) == 2: - head_size = mamba_shapes[0][1] - elif len(mamba_shapes) == 1: - head_size = mamba_shapes[0][1] if len(mamba_shapes[0]) == 2 else mamba_shapes[0][0] - else: - raise NotImplementedError("Mooncake Connector does not support current MambaSpec type now.") - return spec.block_size, head_size - else: - raise NotImplementedError("Mooncake Connector does not support {%s} now.", type(spec).__name__) - def register_kv_caches( self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]], @@ -190,14 +157,11 @@ def register_kv_caches( logger.info("num_blocks: %s", self.num_blocks) self.kv_caches = kv_caches self._build_kv_cache_spec_mappings() - spec_properties = [self._get_spec_transfer_properties(spec) for spec in self.kv_cache_specs] - spec_block_sizes = [properties[0] for properties in spec_properties] - spec_num_heads = [properties[1] for properties in spec_properties] + spec_block_sizes = [spec.block_size for spec in self.kv_cache_specs] layer_names: list[str] = [] group_indices: list[int] = [] spec_indices: list[int] = [] - kernel_block_sizes: list[int] = spec_block_sizes.copy() kv_caches_base_addr: list[list[int]] = [] block_strides_per_layer: list[list[int]] = [] block_lens_per_layer: list[list[int]] = [] @@ -241,28 +205,6 @@ def register_kv_caches( block_lens_per_layer.append(block_lens) block_shapes_per_layer.append(block_shapes) block_size_scales_per_layer.append(block_size_scales) - cur_layer_spec_idx = self.layer_name_to_spec_index[layer_name] - spec_block_size = spec_block_sizes[cur_layer_spec_idx] - layer_kernel_block_sizes = {spec_block_size // scale for scale in block_size_scales} - if len(layer_kernel_block_sizes) != 1: - raise ValueError( - f"Layer {layer_name!r} has inconsistent kernel block sizes {sorted(layer_kernel_block_sizes)}." - ) - kernel_block_size = layer_kernel_block_sizes.pop() - if ( - kernel_block_sizes[cur_layer_spec_idx] == spec_block_sizes[cur_layer_spec_idx] - and kernel_block_size < spec_block_sizes[cur_layer_spec_idx] - ): - kernel_block_sizes[cur_layer_spec_idx] = kernel_block_size - if ( - kernel_block_sizes[cur_layer_spec_idx] != spec_block_sizes[cur_layer_spec_idx] - and kernel_block_size != kernel_block_sizes[cur_layer_spec_idx] - ): - raise ValueError( - "Layer %s has kernel block size %s different from global kernel block size %s", - kernel_block_size, - kernel_block_sizes[cur_layer_spec_idx], - ) unexpected_layers = kv_caches.keys() - configured_layer_names if unexpected_layers: @@ -278,8 +220,6 @@ def register_kv_caches( block_size=self.block_size, num_blocks=self.num_blocks, spec_block_sizes=spec_block_sizes, - kernel_block_sizes=kernel_block_sizes, - spec_num_heads=spec_num_heads, layer_names=layer_names, group_indices=group_indices, spec_indices=spec_indices, diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index fd5c9987d1f8..2a4fdb2e9aba 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -16,10 +16,9 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): """Worker transfer information exchanged during the P/D handshake. - Per-spec fields use the flattened KV-cache spec order. In particular, - kernel_block_sizes records the block size seen by worker kernels after - expanding a logical/spec block. All per-layer fields use layer_names - order; each nested list preserves that layer's cache tensor order. + Per-spec fields use the flattened KV-cache spec order. All per-layer + fields use layer_names order; each nested list preserves that layer's + cache tensor order. """ engine_id: str @@ -27,8 +26,6 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): block_size: int num_blocks: int spec_block_sizes: list[int] - kernel_block_sizes: list[int] - spec_num_heads: list[int | None] layer_names: list[str] group_indices: list[int] spec_indices: list[int] @@ -41,17 +38,6 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): handshake_port: int = 0 def __post_init__(self) -> None: - num_specs = len(self.spec_block_sizes) - per_spec_fields = { - "kernel_block_sizes": self.kernel_block_sizes, - "spec_num_heads": self.spec_num_heads, - } - for field_name, values in per_spec_fields.items(): - if len(values) != num_specs: - raise ValueError( - f"Mooncake transfer metadata field {field_name!r} has {len(values)} specs, expected {num_specs}." - ) - num_layers = len(self.layer_names) per_layer_fields = { "group_indices": self.group_indices, @@ -101,8 +87,6 @@ class MooncakePPTransferMetadata: block_size: int num_blocks: int spec_block_sizes: list[int] - kernel_block_sizes: list[int] - spec_num_heads: list[int | None] layer_names: list[str] group_indices: list[int] spec_indices: list[int] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 76c669c2c9e3..6eb1ab21ae22 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -140,8 +140,6 @@ def _merge_metadata_by_pp_rank( "block_size", "num_blocks", "spec_block_sizes", - "kernel_block_sizes", - "spec_num_heads", "layer_names", "group_indices", "spec_indices", @@ -180,8 +178,6 @@ def _merge_metadata_by_pp_rank( block_size=reference.block_size, num_blocks=reference.num_blocks, spec_block_sizes=reference.spec_block_sizes, - kernel_block_sizes=reference.kernel_block_sizes, - spec_num_heads=reference.spec_num_heads, layer_names=reference.layer_names, group_indices=reference.group_indices, spec_indices=reference.spec_indices, 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 b47e81e3d708..d9c495e44d2a 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 @@ -85,8 +85,6 @@ def __init__( self.block_size = local_metadata.block_size self.spec_block_sizes = local_metadata.spec_block_sizes - self.kernel_block_sizes = local_metadata.kernel_block_sizes - self.spec_num_heads = local_metadata.spec_num_heads self.layer_names = local_metadata.layer_names self.group_indices = local_metadata.group_indices self.spec_indices = local_metadata.spec_indices @@ -119,8 +117,7 @@ def __init__( self.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) self.remote_metadata: SizedDict[str, MooncakeTransferMetadataGroups] = SizedDict() # Candidate producer TP ranks, cached independently for each engine: - # engine_id -> remote_pp_rank -> (local_spec_index, remote_spec_index) - # -> candidate_groups + # engine_id -> remote_pp_rank -> (local_layer_index, remote_layer_index) -> candidate_groups self.remote_tp_rank_groups: SizedDict[str, dict[int, dict[tuple[int, int], list[list[int]]]]] = SizedDict() # engine_id -> remote_pp_rank -> [[local_layer_index, remote_layer_index], ...] self.remote_layer_index_pairs: SizedDict[str, dict[int, list[tuple[int, int]]]] = SizedDict() @@ -206,50 +203,37 @@ def _build_remote_transfer_layout( dict[int, dict[tuple[int, int], list[list[int]]]], dict[int, list[tuple[int, int]]], ]: - """Build per-PP spec candidates and matching layer index pairs. - - TP candidates are constructed once for each local spec present under a - remote PP rank. Layer traversal is used only to record - [local_layer_index, remote_layer_index] pairs for later address - planning. - """ + """Build per-PP layer candidates and matching layer index pairs.""" groups_by_pp_rank: dict[int, dict[tuple[int, int], list[list[int]]]] = {} layer_pairs_by_pp_rank: dict[int, list[tuple[int, int]]] = {} - local_layer_index_by_name = {layer_name: layer_index for layer_index, layer_name in enumerate(self.layer_names)} matched_local_layer_indices: set[int] = set() for remote_pp_rank, pp_metadata in sorted(remote_metadata.metadata_by_pp_rank.items()): layer_index_pairs: list[tuple[int, int]] = [] - spec_index_mapping: dict[int, int] = {} - for remote_layer_index, layer_name in enumerate(pp_metadata.layer_names): - local_layer_index = local_layer_index_by_name.get(layer_name) - if local_layer_index is None: + groups_by_layer_pair: dict[tuple[int, int], list[list[int]]] = {} + remote_layer_index_by_name = { + layer_name: layer_index for layer_index, layer_name in enumerate(pp_metadata.layer_names) + } + for local_layer_index, layer_name in enumerate(self.layer_names): + remote_layer_index = remote_layer_index_by_name.get(layer_name) + if remote_layer_index is None: continue - layer_index_pairs.append((local_layer_index, remote_layer_index)) + layer_pair = (local_layer_index, remote_layer_index) + layer_index_pairs.append(layer_pair) matched_local_layer_indices.add(local_layer_index) local_spec_index = self.spec_indices[local_layer_index] - remote_spec_index = pp_metadata.spec_indices[remote_layer_index] - mapped_remote_spec_index = spec_index_mapping.get(local_spec_index) - if mapped_remote_spec_index is not None and mapped_remote_spec_index != remote_spec_index: - raise ValueError( - f"Mooncake local spec {local_spec_index} maps to multiple producer specs " - f"under PP rank {remote_pp_rank}: {mapped_remote_spec_index} and {remote_spec_index}" - ) - spec_index_mapping[local_spec_index] = remote_spec_index - - layer_pairs_by_pp_rank[remote_pp_rank] = layer_index_pairs - groups_by_pp_rank[remote_pp_rank] = { - (local_spec_index, remote_spec_index): self._get_spec_remote_tp_rank_groups( - local_spec_index, - remote_spec_index, + groups_by_layer_pair[layer_pair] = self._get_layer_remote_tp_rank_groups( + local_layer_index, + remote_layer_index, self.kv_cache_specs[local_spec_index], pp_metadata, remote_metadata.tp_size, remote_metadata.dcp_size, ) - for local_spec_index, remote_spec_index in sorted(spec_index_mapping.items()) - } + + layer_pairs_by_pp_rank[remote_pp_rank] = layer_index_pairs + groups_by_pp_rank[remote_pp_rank] = groups_by_layer_pair missing_local_layer_indices = set(range(len(self.layer_names))) - matched_local_layer_indices if missing_local_layer_indices: @@ -261,40 +245,43 @@ def _build_remote_transfer_layout( ) return groups_by_pp_rank, layer_pairs_by_pp_rank - def _get_spec_remote_tp_rank_groups( + def _get_layer_remote_tp_rank_groups( self, - local_spec_index: int, - remote_spec_index: int, + local_layer_index: int, + remote_layer_index: int, spec: KVCacheSpec, remote_metadata: MooncakePPTransferMetadata, remote_tp_size: int, remote_dcp_size: int, ) -> list[list[int]]: - """Infer the spec's TP strategy and build its remote rank groups.""" + """Infer one matched layer's TP strategy and remote rank groups.""" if isinstance(spec, MambaSpec): return self._get_mamba_remote_tp_rank_groups(remote_tp_size) - if remote_spec_index >= len(remote_metadata.spec_num_heads): - raise ValueError( - f"Mooncake local spec {local_spec_index} maps to missing producer spec {remote_spec_index}" - ) fixed_total_num_kv_heads = None if isinstance(spec, (AscendSFAIndexerCacheSpec, SlidingWindowMLASpec)): local_dcp_size = remote_dcp_size = 1 + local_num_kv_heads = remote_num_kv_heads = 1 fixed_total_num_kv_heads = 1 elif isinstance(spec, MLAAttentionSpec): local_dcp_size = self.dcp_size + local_num_kv_heads = remote_num_kv_heads = 1 fixed_total_num_kv_heads = 1 + # For FA or SWA with kv_heads > 1, must use HND kv_cache layout. elif isinstance(spec, SlidingWindowSpec): local_dcp_size = remote_dcp_size = 1 + local_num_kv_heads = self.block_shapes[local_layer_index][0][0] + remote_num_kv_heads = remote_metadata.block_shapes[remote_layer_index][0][0] elif isinstance(spec, FullAttentionSpec): local_dcp_size = self.dcp_size + local_num_kv_heads = self.block_shapes[local_layer_index][0][0] + remote_num_kv_heads = remote_metadata.block_shapes[remote_layer_index][0][0] else: raise NotImplementedError(f"Mooncake pull has no TP grouping rule for KV cache spec {type(spec).__name__}") total_num_kv_heads = self._infer_total_num_kv_heads( - local_num_kv_heads=spec.num_kv_heads, - remote_num_kv_heads=remote_metadata.spec_num_heads[remote_spec_index], + local_num_kv_heads=local_num_kv_heads, + remote_num_kv_heads=remote_num_kv_heads, remote_tp_size=remote_tp_size, local_dcp_size=local_dcp_size, remote_dcp_size=remote_dcp_size, From 75e338c5e2d29a9b4ab3fc80da1dd7c3614fed25 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Sun, 9 Aug 2026 01:51:23 +0800 Subject: [PATCH 19/48] add compute block ids Signed-off-by: nwpu-zxr --- .../kv_transfer/kv_p2p/mooncake/metadata.py | 5 +- .../kv_p2p/mooncake/pull_scheduler.py | 1 + .../kv_p2p/mooncake/pull_worker.py | 159 +++++++++++++++++- 3 files changed, 157 insertions(+), 8 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 2a4fdb2e9aba..1703e6062423 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -116,6 +116,7 @@ class ReqMeta: """Request metadata required by the initial Mooncake pull path.""" local_block_ids: BlockIds + local_num_prompt_tokens: int num_external_tokens: int num_computed_tokens: int remote_block_ids: BlockIds @@ -124,7 +125,7 @@ class ReqMeta: remote_engine_id: str remote_request_id: str remote_num_prompt_tokens: int - local_full_block_ids: BlockIds = tuple() + local_full_block_ids: BlockIds class MooncakeConnectorMetadata(KVConnectorMetadata): @@ -139,12 +140,14 @@ def add_new_req( self, request_id: str, local_block_ids: BlockIds, + local_num_prompt_tokens: int, num_external_tokens: int, kv_transfer_params: dict[str, Any], local_full_block_ids: BlockIds | None = None, ) -> None: self.requests[request_id] = ReqMeta( local_block_ids=local_block_ids, + local_num_prompt_tokens=local_num_prompt_tokens, num_external_tokens=num_external_tokens, num_computed_tokens=kv_transfer_params.get("num_computed_tokens", 0), remote_block_ids=kv_transfer_params["remote_block_ids"], diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 6eb1ab21ae22..d4bcdf78153c 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -465,6 +465,7 @@ def build_connector_meta(self, scheduler_output: SchedulerOutput) -> KVConnector request_id=req_id, local_block_ids=block_ids, local_full_block_ids=full_block_ids, + local_num_prompt_tokens=req.num_prompt_tokens, num_external_tokens=num_external_tokens, kv_transfer_params=req.kv_transfer_params, ) 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 d9c495e44d2a..39e7d90c1bb4 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 @@ -399,10 +399,158 @@ def _handle_requests( remote_port: int, requests: dict[str, ReqMeta], ) -> None: - """Fetch peer metadata and submit its aggregated requests.""" - self._get_remote_metadata(remote_engine_id, remote_host, remote_port) + """Build and submit the aggregated transfer plan for one producer.""" + remote_metadata = self._get_remote_metadata(remote_engine_id, remote_host, remote_port) + src_lists, dst_lists, length_lists = self._build_transfer_buckets(remote_engine_id, remote_metadata, requests) + self._submit_transfer_buckets(remote_metadata, src_lists, dst_lists, length_lists) - raise NotImplementedError("Mooncake pull transfer submission is not implemented") + def _build_transfer_buckets( + self, + remote_engine_id: str, + remote_metadata: MooncakeTransferMetadataGroups, + requests: dict[str, ReqMeta], + ) -> tuple[ + dict[int, dict[int, list[int]]], + dict[int, dict[int, list[int]]], + dict[int, dict[int, list[int]]], + ]: + """Build parallel source, destination, and length lists per remote worker.""" + src_lists: dict[int, dict[int, list[int]]] = {} + dst_lists: dict[int, dict[int, list[int]]] = {} + length_lists: dict[int, dict[int, list[int]]] = {} + tp_rank_groups_by_pp_rank = self.remote_tp_rank_groups[remote_engine_id] + layer_pairs_by_pp_rank = self.remote_layer_index_pairs[remote_engine_id] + + for request_id, request_metadata in requests.items(): + expanded_block_ids_by_spec: dict[int, tuple[list[int], list[int]]] = {} + for remote_pp_rank, layer_pairs in layer_pairs_by_pp_rank.items(): + pp_metadata = remote_metadata.metadata_by_pp_rank[remote_pp_rank] + src_lists_by_tp_rank = src_lists.setdefault(remote_pp_rank, {}) + dst_lists_by_tp_rank = dst_lists.setdefault(remote_pp_rank, {}) + length_lists_by_tp_rank = length_lists.setdefault(remote_pp_rank, {}) + tp_rank_groups_by_layer = tp_rank_groups_by_pp_rank[remote_pp_rank] + for local_layer_index, remote_layer_index in layer_pairs: + spec_index = self.spec_indices[local_layer_index] + spec = self.kv_cache_specs[spec_index] + expanded_block_ids = expanded_block_ids_by_spec.get(spec_index) + if expanded_block_ids is None: + group_index = self.group_indices[local_layer_index] + spec_index = self.spec_indices[local_layer_index] + remote_spec_index = pp_metadata.spec_indices[remote_layer_index] + local_kernel_block_size = ( + self.local_metadata.spec_block_sizes[spec_index] + // self.block_size_scales[local_layer_index][0] + ) + remote_kernel_block_size = ( + pp_metadata.spec_block_sizes[remote_spec_index] + // pp_metadata.block_size_scales[remote_layer_index][0] + ) + assert local_kernel_block_size == remote_kernel_block_size, ( + "MooncakeConnector does not support different local and remote kernel block size %s | %s.", + local_kernel_block_size, + remote_kernel_block_size, + ) + expanded_block_ids = self._expand_group_block_ids( + spec_index, + request_metadata.local_block_ids[group_index], + request_metadata.local_full_block_ids[group_index], + request_metadata.remote_block_ids[group_index], + request_metadata.local_num_prompt_tokens, + request_metadata.remote_num_prompt_tokens, + self.block_size_scales[local_layer_index][0], + pp_metadata.block_size_scales[remote_layer_index][0], + spec, + ) + expanded_block_ids_by_spec[spec_index] = expanded_block_ids + + layer_pair = (local_layer_index, remote_layer_index) + self._append_layer_transfer_info( + request_id=request_id, + local_layer_index=local_layer_index, + remote_layer_index=remote_layer_index, + spec=spec, + local_block_ids=expanded_block_ids[0], + remote_block_ids=expanded_block_ids[1], + remote_metadata=pp_metadata, + remote_tp_rank_groups=tp_rank_groups_by_layer[layer_pair], + src_lists=src_lists_by_tp_rank, + dst_lists=dst_lists_by_tp_rank, + length_lists=length_lists_by_tp_rank, + ) + + return src_lists, dst_lists, length_lists + + def _expand_group_block_ids( + self, + spec_index: int, + local_group_block_ids: list[int], + local_full_group_block_ids: list[int], + remote_group_block_ids: list[int], + local_num_prompt_tokens: int, + remote_num_prompt_tokens: int, + local_block_size_scale: int, + remote_block_size_scale: int, + spec: KVCacheSpec, + ) -> tuple[list[int], list[int]]: + """Generate the local and remote kernel block ID sequences for one spec.""" + if isinstance(spec, AscendSFAIndexerCacheSpec): + pass + if isinstance(spec, SlidingWindowMLASpec): + pass + if isinstance(spec, MLAAttentionSpec): + pass + if isinstance(spec, SlidingWindowSpec): + pass + if isinstance(spec, FullAttentionSpec): + pass + if isinstance(spec, MambaSpec): + pass + raise NotImplementedError(f"Mooncake block ID expansion does not support {type(spec).__name__}") + + def _append_layer_transfer_info( + self, + request_id: str, + local_layer_index: int, + remote_layer_index: int, + spec: KVCacheSpec, + local_block_ids: list[int], + remote_block_ids: list[int], + remote_metadata: MooncakePPTransferMetadata, + remote_tp_rank_groups: list[list[int]], + src_lists: dict[int, list[int]], + dst_lists: dict[int, list[int]], + length_lists: dict[int, list[int]], + ) -> None: + """Append one matched layer to the per-TP transfer lists.""" + raise NotImplementedError( + "Mooncake transfer address calculation is not implemented for " + f"request {request_id!r}, layer {self.layer_names[local_layer_index]!r}, " + f"remote layer index {remote_layer_index}, spec {type(spec).__name__}" + ) + + def _submit_transfer_buckets( + self, + remote_metadata: MooncakeTransferMetadataGroups, + src_lists: dict[int, dict[int, list[int]]], + dst_lists: dict[int, dict[int, list[int]]], + length_lists: dict[int, dict[int, list[int]]], + ) -> None: + """Submit each remote worker's aggregated address lists sequentially.""" + for remote_pp_rank, src_lists_by_tp_rank in src_lists.items(): + pp_metadata = remote_metadata.metadata_by_pp_rank[remote_pp_rank] + for remote_tp_rank, src_list in src_lists_by_tp_rank.items(): + if not src_list: + continue + tp_metadata = pp_metadata.metadata_by_tp_rank[remote_tp_rank] + session_id = f"{tp_metadata.local_ip}:{tp_metadata.te_rpc_port}" + dst_list = dst_lists[remote_pp_rank][remote_tp_rank] + length_list = length_lists[remote_pp_rank][remote_tp_rank] + ret = self.engine.batch_transfer_sync_read(session_id, src_list, dst_list, length_list) + if ret < 0: + raise RuntimeError( + "Mooncake KV transfer failed for remote " + f"PP rank {remote_pp_rank}, TP rank {remote_tp_rank}, ret={ret}" + ) @staticmethod def _validate_remote_metadata( @@ -417,10 +565,7 @@ def _validate_remote_metadata( if not metadata.metadata_by_pp_rank: raise ValueError("Mooncake producer scheduler returned no PP metadata") if metadata.pcp_size != 1: - raise ValueError( - "Mooncake pull temporarily requires remote pcp_size=1, " - f"got {metadata.pcp_size}" - ) + raise ValueError(f"Mooncake pull temporarily requires remote pcp_size=1, got {metadata.pcp_size}") empty_pp_ranks = [ pp_rank From 013a439fd92ee6cb98f689d74ecc9812010a1474 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Sun, 9 Aug 2026 17:51:49 +0800 Subject: [PATCH 20/48] add dcp=1 block ids compute Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/pull_worker.py | 55 +++++++++++++++---- 1 file changed, 44 insertions(+), 11 deletions(-) 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 39e7d90c1bb4..a72aca46c4d0 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 @@ -450,13 +450,16 @@ def _build_transfer_buckets( local_kernel_block_size, remote_kernel_block_size, ) - expanded_block_ids = self._expand_group_block_ids( + expanded_block_ids = self._compute_group_block_ids( spec_index, + self.local_metadata.spec_block_sizes[spec_index], + pp_metadata.spec_block_sizes[remote_spec_index], request_metadata.local_block_ids[group_index], request_metadata.local_full_block_ids[group_index], request_metadata.remote_block_ids[group_index], request_metadata.local_num_prompt_tokens, request_metadata.remote_num_prompt_tokens, + request_metadata.num_computed_tokens, self.block_size_scales[local_layer_index][0], pp_metadata.block_size_scales[remote_layer_index][0], spec, @@ -480,31 +483,61 @@ def _build_transfer_buckets( return src_lists, dst_lists, length_lists - def _expand_group_block_ids( + @staticmethod + def _expand_block_ids(block_ids: list[int], scale: int) -> list[int]: + """Expand logical block IDs into contiguous kernel block IDs.""" + if scale == 1: + return block_ids + return [block_id * scale + offset for block_id in block_ids for offset in range(scale)] + + def _compute_group_block_ids( self, spec_index: int, + local_block_size: int, + remote_block_size: int, local_group_block_ids: list[int], local_full_group_block_ids: list[int], remote_group_block_ids: list[int], local_num_prompt_tokens: int, remote_num_prompt_tokens: int, + num_computed_tokens: int, local_block_size_scale: int, remote_block_size_scale: int, spec: KVCacheSpec, ) -> tuple[list[int], list[int]]: """Generate the local and remote kernel block ID sequences for one spec.""" - if isinstance(spec, AscendSFAIndexerCacheSpec): - pass - if isinstance(spec, SlidingWindowMLASpec): - pass - if isinstance(spec, MLAAttentionSpec): - pass if isinstance(spec, SlidingWindowSpec): - pass + assert local_block_size == remote_block_size, "Mooncake SWA requires the same P/D logical block size." + local_unhashed_start_idx = len(local_full_group_block_ids) - len(local_group_block_ids) + local_kernel_block_ids = self._expand_block_ids(local_group_block_ids, local_block_size_scale) + remote_kernel_block_ids = self._expand_block_ids( + remote_group_block_ids[local_unhashed_start_idx:], remote_block_size_scale + ) + num_kernel_blocks = min(len(local_kernel_block_ids), len(remote_kernel_block_ids)) + return local_kernel_block_ids[:num_kernel_blocks], remote_kernel_block_ids[:num_kernel_blocks] if isinstance(spec, FullAttentionSpec): - pass + if self.dcp_size > 1 and not isinstance(spec, AscendSFAIndexerCacheSpec): + if local_block_size != remote_block_size: + raise NotImplementedError("Mooncake DCP with different P/D logical block sizes is not implemented") + remote_start_block_idx = num_computed_tokens // remote_block_size + local_kernel_block_ids = self._expand_block_ids(local_group_block_ids, local_block_size_scale) + remote_kernel_block_ids = self._expand_block_ids( + remote_group_block_ids[remote_start_block_idx:], remote_block_size_scale + ) + local_kernel_block_ids = local_kernel_block_ids[self.dcp_rank :: self.dcp_size] + remote_kernel_block_ids = remote_kernel_block_ids[self.dcp_rank :: self.dcp_size] + else: + local_kernel_block_ids = self._expand_block_ids(local_group_block_ids, local_block_size_scale) + remote_kernel_block_ids = self._expand_block_ids(remote_group_block_ids, remote_block_size_scale) + kernel_block_size = local_block_size // local_block_size_scale + remote_start_block_idx = num_computed_tokens // kernel_block_size + remote_kernel_block_ids = remote_kernel_block_ids[remote_start_block_idx:] + num_kernel_blocks = min(len(local_kernel_block_ids), len(remote_kernel_block_ids)) + return local_kernel_block_ids[:num_kernel_blocks], remote_kernel_block_ids[:num_kernel_blocks] if isinstance(spec, MambaSpec): - pass + local_state_block_id = local_group_block_ids[-self.num_speculative_tokens - 1] + remote_state_block_id = remote_group_block_ids[-1] + return [local_state_block_id], [remote_state_block_id] raise NotImplementedError(f"Mooncake block ID expansion does not support {type(spec).__name__}") def _append_layer_transfer_info( From 77d3cf89f57e33a5e0dbdadc0e4018fd3f9ad9e0 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Sun, 9 Aug 2026 22:05:11 +0800 Subject: [PATCH 21/48] add dcp>1 block ids compute Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/pull_scheduler.py | 26 +- .../kv_p2p/mooncake/pull_worker.py | 331 +++++++++++------- 2 files changed, 223 insertions(+), 134 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index d4bcdf78153c..d5f778611acf 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -428,23 +428,23 @@ def update_state_after_alloc( "remote_request_id", ) if all(field in params for field in required_remote_fields): - if num_external_tokens > 0: - local_block_ids = blocks.get_unhashed_block_ids_all_groups() - local_full_block_ids = blocks.get_block_ids() - else: - local_block_ids = tuple() - local_full_block_ids = tuple() - self._reqs_need_recv[request.request_id] = ( - request, - local_block_ids, - local_full_block_ids, - num_external_tokens, - ) - self._reqs_recv_info[request.request_id] = ( + remote = ( params["remote_host"], params["remote_port"], params["remote_request_id"], ) + if num_external_tokens > 0: + self._reqs_need_recv[request.request_id] = ( + request, + blocks.get_unhashed_block_ids_all_groups(), + blocks.get_block_ids(), + num_external_tokens, + ) + self._reqs_recv_info[request.request_id] = remote + else: + if self._recving_thread is None: + raise RuntimeError("Producer Mooncake scheduler cannot acknowledge a receive request") + self._recving_thread.add_request(*remote) else: logger.warning("Got invalid KVTransferParams. params=%s.", params) else: 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 a72aca46c4d0..34bef518557f 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 @@ -399,89 +399,89 @@ def _handle_requests( remote_port: int, requests: dict[str, ReqMeta], ) -> None: - """Build and submit the aggregated transfer plan for one producer.""" + """Build and execute transfer buckets for one producer engine.""" remote_metadata = self._get_remote_metadata(remote_engine_id, remote_host, remote_port) - src_lists, dst_lists, length_lists = self._build_transfer_buckets(remote_engine_id, remote_metadata, requests) - self._submit_transfer_buckets(remote_metadata, src_lists, dst_lists, length_lists) - - def _build_transfer_buckets( - self, - remote_engine_id: str, - remote_metadata: MooncakeTransferMetadataGroups, - requests: dict[str, ReqMeta], - ) -> tuple[ - dict[int, dict[int, list[int]]], - dict[int, dict[int, list[int]]], - dict[int, dict[int, list[int]]], - ]: - """Build parallel source, destination, and length lists per remote worker.""" - src_lists: dict[int, dict[int, list[int]]] = {} - dst_lists: dict[int, dict[int, list[int]]] = {} - length_lists: dict[int, dict[int, list[int]]] = {} tp_rank_groups_by_pp_rank = self.remote_tp_rank_groups[remote_engine_id] layer_pairs_by_pp_rank = self.remote_layer_index_pairs[remote_engine_id] + transfer_block_ids_by_spec: dict[str, dict[int, list[tuple[int, list[int], list[int]]]]] = {} + for remote_pp_rank, layer_pairs in layer_pairs_by_pp_rank.items(): + pp_metadata = remote_metadata.metadata_by_pp_rank[remote_pp_rank] + transfer_block_buckets = self._build_transfer_block_buckets( + pp_metadata, + layer_pairs, + tp_rank_groups_by_pp_rank[remote_pp_rank], + remote_metadata.dcp_size, + requests, + transfer_block_ids_by_spec, + ) + self._execute_transfer_block_buckets(remote_pp_rank, pp_metadata, transfer_block_buckets) - for request_id, request_metadata in requests.items(): - expanded_block_ids_by_spec: dict[int, tuple[list[int], list[int]]] = {} - for remote_pp_rank, layer_pairs in layer_pairs_by_pp_rank.items(): - pp_metadata = remote_metadata.metadata_by_pp_rank[remote_pp_rank] - src_lists_by_tp_rank = src_lists.setdefault(remote_pp_rank, {}) - dst_lists_by_tp_rank = dst_lists.setdefault(remote_pp_rank, {}) - length_lists_by_tp_rank = length_lists.setdefault(remote_pp_rank, {}) - tp_rank_groups_by_layer = tp_rank_groups_by_pp_rank[remote_pp_rank] - for local_layer_index, remote_layer_index in layer_pairs: - spec_index = self.spec_indices[local_layer_index] - spec = self.kv_cache_specs[spec_index] - expanded_block_ids = expanded_block_ids_by_spec.get(spec_index) - if expanded_block_ids is None: - group_index = self.group_indices[local_layer_index] - spec_index = self.spec_indices[local_layer_index] - remote_spec_index = pp_metadata.spec_indices[remote_layer_index] + def _build_transfer_block_buckets( + self, + remote_metadata: MooncakePPTransferMetadata, + layer_pairs: list[tuple[int, int]], + tp_rank_groups_by_layer: dict[tuple[int, int], list[list[int]]], + remote_dcp_size: int, + requests: dict[str, ReqMeta], + transfer_block_ids_by_spec: dict[str, dict[int, list[tuple[int, list[int], list[int]]]]], + ) -> dict[int, list[tuple[str, int, int, list[int], list[int]]]]: + """Bucket one PP rank's requests and layer block IDs by remote TP rank.""" + transfer_block_buckets: dict[int, list[tuple[str, int, int, list[int], list[int]]]] = {} + for selection_index, (request_id, request_metadata) in enumerate(requests.items()): + request_block_ids_by_spec = transfer_block_ids_by_spec.setdefault(request_id, {}) + for local_layer_index, remote_layer_index in layer_pairs: + spec_index = self.spec_indices[local_layer_index] + spec = self.kv_cache_specs[spec_index] + layer_pair = (local_layer_index, remote_layer_index) + transfer_block_ids = request_block_ids_by_spec.get(spec_index) + if transfer_block_ids is None: + group_index = self.group_indices[local_layer_index] + remote_spec_index = remote_metadata.spec_indices[remote_layer_index] + if not isinstance(spec, MambaSpec): local_kernel_block_size = ( self.local_metadata.spec_block_sizes[spec_index] // self.block_size_scales[local_layer_index][0] ) remote_kernel_block_size = ( - pp_metadata.spec_block_sizes[remote_spec_index] - // pp_metadata.block_size_scales[remote_layer_index][0] + remote_metadata.spec_block_sizes[remote_spec_index] + // remote_metadata.block_size_scales[remote_layer_index][0] ) assert local_kernel_block_size == remote_kernel_block_size, ( "MooncakeConnector does not support different local and remote kernel block size %s | %s.", local_kernel_block_size, remote_kernel_block_size, ) - expanded_block_ids = self._compute_group_block_ids( - spec_index, - self.local_metadata.spec_block_sizes[spec_index], - pp_metadata.spec_block_sizes[remote_spec_index], - request_metadata.local_block_ids[group_index], - request_metadata.local_full_block_ids[group_index], - request_metadata.remote_block_ids[group_index], - request_metadata.local_num_prompt_tokens, - request_metadata.remote_num_prompt_tokens, - request_metadata.num_computed_tokens, - self.block_size_scales[local_layer_index][0], - pp_metadata.block_size_scales[remote_layer_index][0], - spec, + transfer_block_ids = self._compute_group_block_ids( + request_id, + tp_rank_groups_by_layer[layer_pair], + remote_dcp_size, + spec_index, + self.local_metadata.spec_block_sizes[spec_index], + remote_metadata.spec_block_sizes[remote_spec_index], + request_metadata.local_block_ids[group_index], + request_metadata.local_full_block_ids[group_index], + request_metadata.remote_block_ids[group_index], + request_metadata.local_num_prompt_tokens, + request_metadata.remote_num_prompt_tokens, + request_metadata.num_computed_tokens, + self.block_size_scales[local_layer_index][0], + remote_metadata.block_size_scales[remote_layer_index][0], + spec, + selection_index, + ) + request_block_ids_by_spec[spec_index] = transfer_block_ids + + for remote_tp_rank, local_block_ids, remote_block_ids in transfer_block_ids: + transfer_block_buckets.setdefault(remote_tp_rank, []).append( + ( + request_id, + local_layer_index, + remote_layer_index, + local_block_ids, + remote_block_ids, ) - expanded_block_ids_by_spec[spec_index] = expanded_block_ids - - layer_pair = (local_layer_index, remote_layer_index) - self._append_layer_transfer_info( - request_id=request_id, - local_layer_index=local_layer_index, - remote_layer_index=remote_layer_index, - spec=spec, - local_block_ids=expanded_block_ids[0], - remote_block_ids=expanded_block_ids[1], - remote_metadata=pp_metadata, - remote_tp_rank_groups=tp_rank_groups_by_layer[layer_pair], - src_lists=src_lists_by_tp_rank, - dst_lists=dst_lists_by_tp_rank, - length_lists=length_lists_by_tp_rank, ) - - return src_lists, dst_lists, length_lists + return transfer_block_buckets @staticmethod def _expand_block_ids(block_ids: list[int], scale: int) -> list[int]: @@ -490,8 +490,16 @@ def _expand_block_ids(block_ids: list[int], scale: int) -> list[int]: return block_ids return [block_id * scale + offset for block_id in block_ids for offset in range(scale)] + @staticmethod + def _select_remote_tp_rank(candidate_tp_ranks: list[int], selection_index: int) -> int: + """Select a replicated producer TP rank using round-robin.""" + return candidate_tp_ranks[selection_index % len(candidate_tp_ranks)] + def _compute_group_block_ids( self, + request_id: str, + remote_tp_rank_groups: list[list[int]], + remote_dcp_size: int, spec_index: int, local_block_size: int, remote_block_size: int, @@ -504,8 +512,14 @@ def _compute_group_block_ids( local_block_size_scale: int, remote_block_size_scale: int, spec: KVCacheSpec, - ) -> tuple[list[int], list[int]]: - """Generate the local and remote kernel block ID sequences for one spec.""" + selection_index: int, + ) -> list[tuple[int, list[int], list[int]]]: + """Pair remote TP ranks with local and remote kernel block IDs.""" + is_dcp_transfer = ( + (self.dcp_size > 1 or remote_dcp_size > 1) + and isinstance(spec, FullAttentionSpec) + and not isinstance(spec, AscendSFAIndexerCacheSpec) + ) if isinstance(spec, SlidingWindowSpec): assert local_block_size == remote_block_size, "Mooncake SWA requires the same P/D logical block size." local_unhashed_start_idx = len(local_full_group_block_ids) - len(local_group_block_ids) @@ -513,77 +527,152 @@ def _compute_group_block_ids( remote_kernel_block_ids = self._expand_block_ids( remote_group_block_ids[local_unhashed_start_idx:], remote_block_size_scale ) - num_kernel_blocks = min(len(local_kernel_block_ids), len(remote_kernel_block_ids)) - return local_kernel_block_ids[:num_kernel_blocks], remote_kernel_block_ids[:num_kernel_blocks] - if isinstance(spec, FullAttentionSpec): - if self.dcp_size > 1 and not isinstance(spec, AscendSFAIndexerCacheSpec): + elif isinstance(spec, FullAttentionSpec): + if is_dcp_transfer: if local_block_size != remote_block_size: - raise NotImplementedError("Mooncake DCP with different P/D logical block sizes is not implemented") - remote_start_block_idx = num_computed_tokens // remote_block_size - local_kernel_block_ids = self._expand_block_ids(local_group_block_ids, local_block_size_scale) - remote_kernel_block_ids = self._expand_block_ids( - remote_group_block_ids[remote_start_block_idx:], remote_block_size_scale - ) - local_kernel_block_ids = local_kernel_block_ids[self.dcp_rank :: self.dcp_size] - remote_kernel_block_ids = remote_kernel_block_ids[self.dcp_rank :: self.dcp_size] - else: - local_kernel_block_ids = self._expand_block_ids(local_group_block_ids, local_block_size_scale) - remote_kernel_block_ids = self._expand_block_ids(remote_group_block_ids, remote_block_size_scale) - kernel_block_size = local_block_size // local_block_size_scale - remote_start_block_idx = num_computed_tokens // kernel_block_size - remote_kernel_block_ids = remote_kernel_block_ids[remote_start_block_idx:] - num_kernel_blocks = min(len(local_kernel_block_ids), len(remote_kernel_block_ids)) - return local_kernel_block_ids[:num_kernel_blocks], remote_kernel_block_ids[:num_kernel_blocks] - if isinstance(spec, MambaSpec): - local_state_block_id = local_group_block_ids[-self.num_speculative_tokens - 1] - remote_state_block_id = remote_group_block_ids[-1] - return [local_state_block_id], [remote_state_block_id] - raise NotImplementedError(f"Mooncake block ID expansion does not support {type(spec).__name__}") + raise NotImplementedError("Mooncake DCP with different P/D spec block sizes is not implemented") + + local_virtual_block_size = local_block_size * self.dcp_size + local_start_block_idx = num_computed_tokens // local_virtual_block_size + transfer_block_ids: list[tuple[int, list[int], list[int]]] = [] + for candidate_tp_ranks in remote_tp_rank_groups: + if remote_dcp_size == 1: + remote_tp_rank = ( + candidate_tp_ranks[0] + if len(candidate_tp_ranks) == 1 + else self._select_remote_tp_rank(candidate_tp_ranks, selection_index) + ) + remote_tp_ranks = [remote_tp_rank] + else: + remote_tp_ranks = candidate_tp_ranks + + block_ids_by_remote_tp_rank: dict[int, tuple[list[int], list[int]]] = {} + for local_block_offset, local_block_id in enumerate(local_group_block_ids): + local_block_idx = local_start_block_idx + local_block_offset + global_block_idx = local_block_idx * self.dcp_size + self.dcp_rank + remote_block_idx = global_block_idx // remote_dcp_size + if remote_block_idx >= len(remote_group_block_ids): + break + remote_dcp_rank = global_block_idx % remote_dcp_size + remote_tp_rank = remote_tp_ranks[remote_dcp_rank] + local_kernel_block_ids, remote_kernel_block_ids = block_ids_by_remote_tp_rank.setdefault( + remote_tp_rank, ([], []) + ) + local_kernel_block_ids.extend( + local_block_id * local_block_size_scale + offset for offset in range(local_block_size_scale) + ) + remote_block_id = remote_group_block_ids[remote_block_idx] + remote_kernel_block_ids.extend( + remote_block_id * remote_block_size_scale + offset + for offset in range(remote_block_size_scale) + ) + transfer_block_ids.extend( + (remote_tp_rank, local_kernel_block_ids, remote_kernel_block_ids) + for remote_tp_rank, ( + local_kernel_block_ids, + remote_kernel_block_ids, + ) in block_ids_by_remote_tp_rank.items() + ) + return transfer_block_ids - def _append_layer_transfer_info( + local_kernel_block_ids = self._expand_block_ids(local_group_block_ids, local_block_size_scale) + remote_kernel_block_ids = self._expand_block_ids(remote_group_block_ids, remote_block_size_scale) + kernel_block_size = local_block_size // local_block_size_scale + remote_start_block_idx = num_computed_tokens // kernel_block_size + remote_kernel_block_ids = remote_kernel_block_ids[remote_start_block_idx:] + elif isinstance(spec, MambaSpec): + local_kernel_block_ids = [local_group_block_ids[-self.num_speculative_tokens - 1]] + remote_kernel_block_ids = [remote_group_block_ids[-1]] + else: + raise NotImplementedError(f"Mooncake block ID expansion does not support {type(spec).__name__}") + + num_kernel_blocks = min(len(local_kernel_block_ids), len(remote_kernel_block_ids)) + local_kernel_block_ids = local_kernel_block_ids[:num_kernel_blocks] + remote_kernel_block_ids = remote_kernel_block_ids[:num_kernel_blocks] + return [ + ( + candidate_tp_ranks[0] + if len(candidate_tp_ranks) == 1 + else self._select_remote_tp_rank(candidate_tp_ranks, selection_index), + local_kernel_block_ids, + remote_kernel_block_ids, + ) + for candidate_tp_ranks in remote_tp_rank_groups + ] + + def _append_layer_transfer_addresses( self, request_id: str, local_layer_index: int, remote_layer_index: int, + remote_tp_rank: int, spec: KVCacheSpec, local_block_ids: list[int], remote_block_ids: list[int], remote_metadata: MooncakePPTransferMetadata, - remote_tp_rank_groups: list[list[int]], - src_lists: dict[int, list[int]], - dst_lists: dict[int, list[int]], - length_lists: dict[int, list[int]], + src_list: list[int], + dst_list: list[int], + length_list: list[int], ) -> None: - """Append one matched layer to the per-TP transfer lists.""" + """Append one layer's addresses to one remote TP transfer bucket.""" raise NotImplementedError( "Mooncake transfer address calculation is not implemented for " f"request {request_id!r}, layer {self.layer_names[local_layer_index]!r}, " - f"remote layer index {remote_layer_index}, spec {type(spec).__name__}" + f"remote PP layer index {remote_layer_index}, TP rank {remote_tp_rank}, " + f"spec {type(spec).__name__}" ) - def _submit_transfer_buckets( + def _execute_transfer_block_buckets( self, - remote_metadata: MooncakeTransferMetadataGroups, - src_lists: dict[int, dict[int, list[int]]], - dst_lists: dict[int, dict[int, list[int]]], - length_lists: dict[int, dict[int, list[int]]], + remote_pp_rank: int, + remote_metadata: MooncakePPTransferMetadata, + transfer_block_buckets: dict[int, list[tuple[str, int, int, list[int], list[int]]]], ) -> None: - """Submit each remote worker's aggregated address lists sequentially.""" - for remote_pp_rank, src_lists_by_tp_rank in src_lists.items(): - pp_metadata = remote_metadata.metadata_by_pp_rank[remote_pp_rank] - for remote_tp_rank, src_list in src_lists_by_tp_rank.items(): - if not src_list: - continue - tp_metadata = pp_metadata.metadata_by_tp_rank[remote_tp_rank] - session_id = f"{tp_metadata.local_ip}:{tp_metadata.te_rpc_port}" - dst_list = dst_lists[remote_pp_rank][remote_tp_rank] - length_list = length_lists[remote_pp_rank][remote_tp_rank] - ret = self.engine.batch_transfer_sync_read(session_id, src_list, dst_list, length_list) - if ret < 0: - raise RuntimeError( - "Mooncake KV transfer failed for remote " - f"PP rank {remote_pp_rank}, TP rank {remote_tp_rank}, ret={ret}" - ) + """Calculate addresses and execute one PP rank's TP buckets serially.""" + for remote_tp_rank, transfer_entries in transfer_block_buckets.items(): + self._execute_tp_transfer_bucket( + remote_pp_rank, + remote_tp_rank, + remote_metadata, + transfer_entries, + ) + + def _execute_tp_transfer_bucket( + self, + remote_pp_rank: int, + remote_tp_rank: int, + remote_metadata: MooncakePPTransferMetadata, + transfer_entries: list[tuple[str, int, int, list[int], list[int]]], + ) -> None: + """Calculate addresses and execute one remote PP/TP transfer bucket.""" + src_list: list[int] = [] + dst_list: list[int] = [] + length_list: list[int] = [] + for request_id, local_layer_index, remote_layer_index, local_block_ids, remote_block_ids in transfer_entries: + spec = self.kv_cache_specs[self.spec_indices[local_layer_index]] + self._append_layer_transfer_addresses( + request_id=request_id, + local_layer_index=local_layer_index, + remote_layer_index=remote_layer_index, + remote_tp_rank=remote_tp_rank, + spec=spec, + local_block_ids=local_block_ids, + remote_block_ids=remote_block_ids, + remote_metadata=remote_metadata, + src_list=src_list, + dst_list=dst_list, + length_list=length_list, + ) + + if not src_list: + return + tp_metadata = remote_metadata.metadata_by_tp_rank[remote_tp_rank] + session_id = f"{tp_metadata.local_ip}:{tp_metadata.te_rpc_port}" + ret = self.engine.batch_transfer_sync_read(session_id, src_list, dst_list, length_list) + if ret < 0: + raise RuntimeError( + f"Mooncake KV transfer failed for remote PP rank {remote_pp_rank}, TP rank {remote_tp_rank}, ret={ret}" + ) @staticmethod def _validate_remote_metadata( From 17d43c6fdac3d900038d166c3acd18b9e60a8740 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Mon, 10 Aug 2026 00:17:09 +0800 Subject: [PATCH 22/48] add compute address and transfer Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/pull_worker.py | 246 +++++++++++++++--- 1 file changed, 208 insertions(+), 38 deletions(-) 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 34bef518557f..99a925c9fe3b 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 @@ -414,7 +414,13 @@ def _handle_requests( requests, transfer_block_ids_by_spec, ) - self._execute_transfer_block_buckets(remote_pp_rank, pp_metadata, transfer_block_buckets) + self._execute_transfer_block_buckets( + remote_pp_rank, + pp_metadata, + remote_metadata.tp_size, + remote_metadata.dcp_size, + transfer_block_buckets, + ) def _build_transfer_block_buckets( self, @@ -424,9 +430,23 @@ def _build_transfer_block_buckets( remote_dcp_size: int, requests: dict[str, ReqMeta], transfer_block_ids_by_spec: dict[str, dict[int, list[tuple[int, list[int], list[int]]]]], - ) -> dict[int, list[tuple[str, int, int, list[int], list[int]]]]: - """Bucket one PP rank's requests and layer block IDs by remote TP rank.""" - transfer_block_buckets: dict[int, list[tuple[str, int, int, list[int], list[int]]]] = {} + ) -> dict[ + int, + dict[ + int, + dict[tuple[int, int], list[tuple[str, list[int], list[int]]]], + ], + ]: + """Bucket one PP rank's requests by remote TP, spec index, and layer pair.""" + # remote_tp_rank -> local_spec_index -> (local_layer_index, remote_layer_index) -> + # [(request_id, local_kernel_block_ids, remote_kernel_block_ids), ...] + transfer_block_buckets: dict[ + int, + dict[ + int, + dict[tuple[int, int], list[tuple[str, list[int], list[int]]]], + ], + ] = {} for selection_index, (request_id, request_metadata) in enumerate(requests.items()): request_block_ids_by_spec = transfer_block_ids_by_spec.setdefault(request_id, {}) for local_layer_index, remote_layer_index in layer_pairs: @@ -472,14 +492,10 @@ def _build_transfer_block_buckets( request_block_ids_by_spec[spec_index] = transfer_block_ids for remote_tp_rank, local_block_ids, remote_block_ids in transfer_block_ids: - transfer_block_buckets.setdefault(remote_tp_rank, []).append( - ( - request_id, - local_layer_index, - remote_layer_index, - local_block_ids, - remote_block_ids, - ) + transfer_entries_by_spec = transfer_block_buckets.setdefault(remote_tp_rank, {}) + transfer_entries_by_layer = transfer_entries_by_spec.setdefault(spec_index, {}) + transfer_entries_by_layer.setdefault(layer_pair, []).append( + (request_id, local_block_ids, remote_block_ids) ) return transfer_block_buckets @@ -600,64 +616,218 @@ def _compute_group_block_ids( for candidate_tp_ranks in remote_tp_rank_groups ] - def _append_layer_transfer_addresses( + def _append_spec_transfer_addresses( self, - request_id: str, - local_layer_index: int, - remote_layer_index: int, + spec_index: int, remote_tp_rank: int, - spec: KVCacheSpec, - local_block_ids: list[int], - remote_block_ids: list[int], + remote_tp_size: int, + remote_dcp_size: int, + transfer_entries_by_layer: dict[tuple[int, int], list[tuple[str, list[int], list[int]]]], remote_metadata: MooncakePPTransferMetadata, src_list: list[int], dst_list: list[int], length_list: list[int], ) -> None: - """Append one layer's addresses to one remote TP transfer bucket.""" - raise NotImplementedError( - "Mooncake transfer address calculation is not implemented for " - f"request {request_id!r}, layer {self.layer_names[local_layer_index]!r}, " - f"remote PP layer index {remote_layer_index}, TP rank {remote_tp_rank}, " - f"spec {type(spec).__name__}" + """Append all layers and requests belonging to one local spec.""" + for (local_layer_index, _), transfer_entries in transfer_entries_by_layer.items(): + for request_id, local_block_ids, remote_block_ids in transfer_entries: + if len(local_block_ids) != len(remote_block_ids): + raise ValueError( + f"Mooncake block ID count mismatch for request {request_id!r}, " + f"layer {self.layer_names[local_layer_index]!r}: " + f"local={len(local_block_ids)}, remote={len(remote_block_ids)}" + ) + if not any( + local_block_ids + for transfer_entries in transfer_entries_by_layer.values() + for _, local_block_ids, _ in transfer_entries + ): + return + + spec = self.kv_cache_specs[spec_index] + if isinstance(spec, MambaSpec): + raise NotImplementedError("Mooncake Mamba transfer address calculation is not implemented") + remote_tp_metadata = remote_metadata.metadata_by_tp_rank[remote_tp_rank] + transfer_whole_block = isinstance( + spec, + (MLAAttentionSpec, SlidingWindowMLASpec, AscendSFAIndexerCacheSpec), + ) + if transfer_whole_block: + 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] + 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}: " + f"local={len(local_base_addrs)}, remote={len(remote_base_addrs)}" + ) + for cache_index, (local_base_addr, remote_base_addr) in enumerate( + zip(local_base_addrs, remote_base_addrs) + ): + local_block_stride = self.block_strides[local_layer_index][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] + remote_block_len = remote_metadata.block_lens[remote_layer_index][cache_index] + if local_block_len != remote_block_len: + raise ValueError( + f"Mooncake whole-block length mismatch for layer " + f"{self.layer_names[local_layer_index]!r}, cache {cache_index}: " + f"local={local_block_len}, remote={remote_block_len}" + ) + for _, local_block_ids, remote_block_ids in transfer_entries: + for local_block_id, remote_block_id in zip(local_block_ids, remote_block_ids): + src_list.append(local_base_addr + local_block_id * local_block_stride) + dst_list.append(remote_base_addr + remote_block_id * remote_block_stride) + length_list.append(local_block_len) + return + + if not isinstance(spec, (FullAttentionSpec, SlidingWindowSpec)): + raise NotImplementedError(f"Mooncake transfer address calculation does not support {type(spec).__name__}") + + first_local_layer_index, first_remote_layer_index = next(iter(transfer_entries_by_layer)) + if isinstance(spec, SlidingWindowSpec): + local_attention_dcp_size = remote_attention_dcp_size = 1 + else: + local_attention_dcp_size = self.dcp_size + remote_attention_dcp_size = remote_dcp_size + local_num_heads = self.block_shapes[first_local_layer_index][0][0] + remote_num_heads = remote_metadata.block_shapes[first_remote_layer_index][0][0] + total_num_kv_heads = self._infer_total_num_kv_heads( + local_num_kv_heads=local_num_heads, + remote_num_kv_heads=remote_num_heads, + remote_tp_size=remote_tp_size, + local_dcp_size=local_attention_dcp_size, + remote_dcp_size=remote_attention_dcp_size, + fixed_total_num_kv_heads=None, + ) + local_head_tp_size = self.tp_size // local_attention_dcp_size + remote_head_tp_size = remote_tp_size // remote_attention_dcp_size + local_head_tp_rank = self.tp_rank // local_attention_dcp_size + remote_head_tp_rank = remote_tp_rank // remote_attention_dcp_size + local_head_start, local_head_end = self._get_head_interval( + local_head_tp_rank, local_head_tp_size, total_num_kv_heads + ) + remote_head_start, remote_head_end = self._get_head_interval( + remote_head_tp_rank, remote_head_tp_size, total_num_kv_heads ) + transfer_head_start = max(local_head_start, remote_head_start) + transfer_head_end = min(local_head_end, remote_head_end) + if transfer_head_start >= transfer_head_end: + raise ValueError( + f"Mooncake found no head overlap for local TP {self.tp_rank} and remote TP {remote_tp_rank}" + ) + if local_head_end - local_head_start != local_num_heads: + raise ValueError( + f"Mooncake local HND head shape does not match inferred topology for spec {spec_index}: " + f"shape={local_num_heads}, interval={local_head_end - local_head_start}" + ) + if remote_head_end - remote_head_start != remote_num_heads: + raise ValueError( + f"Mooncake remote HND head shape does not match inferred topology for spec {spec_index}: " + f"shape={remote_num_heads}, interval={remote_head_end - remote_head_start}" + ) + + local_head_offset = transfer_head_start - local_head_start + remote_head_offset = transfer_head_start - remote_head_start + num_transfer_heads = transfer_head_end - transfer_head_start + 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] + 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}: " + f"local={len(local_base_addrs)}, remote={len(remote_base_addrs)}" + ) + if self.block_shapes[local_layer_index][0][0] != local_num_heads: + raise ValueError(f"Mooncake local head count differs within spec {spec_index}") + if remote_metadata.block_shapes[remote_layer_index][0][0] != remote_num_heads: + raise ValueError(f"Mooncake remote head count differs within local spec {spec_index}") + + for cache_index, (local_base_addr, remote_base_addr) in enumerate(zip(local_base_addrs, remote_base_addrs)): + local_block_shape = self.block_shapes[local_layer_index][cache_index] + remote_block_shape = remote_metadata.block_shapes[remote_layer_index][cache_index] + if not local_block_shape or local_block_shape[0] != local_num_heads: + raise ValueError( + f"Mooncake local K/V head counts differ for layer {self.layer_names[local_layer_index]!r}" + ) + if not remote_block_shape or remote_block_shape[0] != remote_num_heads: + raise ValueError( + f"Mooncake remote K/V head counts differ for layer {self.layer_names[local_layer_index]!r}" + ) + + local_block_stride = self.block_strides[local_layer_index][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] + remote_block_len = remote_metadata.block_lens[remote_layer_index][cache_index] + if local_block_len % local_num_heads or remote_block_len % remote_num_heads: + raise ValueError( + f"Mooncake HND block length is not divisible by head count for layer " + f"{self.layer_names[local_layer_index]!r}, cache {cache_index}" + ) + local_head_len = local_block_len // local_num_heads + remote_head_len = remote_block_len // remote_num_heads + if local_head_len != remote_head_len: + raise ValueError( + f"Mooncake P/D head length mismatch for layer {self.layer_names[local_layer_index]!r}, " + f"cache {cache_index}: local={local_head_len}, remote={remote_head_len}" + ) + transfer_len = num_transfer_heads * local_head_len + local_inner_offset = local_head_offset * local_head_len + remote_inner_offset = remote_head_offset * remote_head_len + for _, local_block_ids, remote_block_ids in transfer_entries: + for local_block_id, remote_block_id in zip(local_block_ids, remote_block_ids): + src_list.append(local_base_addr + local_block_id * local_block_stride + local_inner_offset) + dst_list.append(remote_base_addr + remote_block_id * remote_block_stride + remote_inner_offset) + length_list.append(transfer_len) def _execute_transfer_block_buckets( self, remote_pp_rank: int, remote_metadata: MooncakePPTransferMetadata, - transfer_block_buckets: dict[int, list[tuple[str, int, int, list[int], list[int]]]], + remote_tp_size: int, + remote_dcp_size: int, + transfer_block_buckets: dict[ + int, + dict[ + int, + dict[tuple[int, int], list[tuple[str, list[int], list[int]]]], + ], + ], ) -> None: """Calculate addresses and execute one PP rank's TP buckets serially.""" - for remote_tp_rank, transfer_entries in transfer_block_buckets.items(): + for remote_tp_rank, transfer_entries_by_spec in transfer_block_buckets.items(): self._execute_tp_transfer_bucket( remote_pp_rank, remote_tp_rank, + remote_tp_size, + remote_dcp_size, remote_metadata, - transfer_entries, + transfer_entries_by_spec, ) def _execute_tp_transfer_bucket( self, remote_pp_rank: int, remote_tp_rank: int, + remote_tp_size: int, + remote_dcp_size: int, remote_metadata: MooncakePPTransferMetadata, - transfer_entries: list[tuple[str, int, int, list[int], list[int]]], + transfer_entries_by_spec: dict[ + int, + dict[tuple[int, int], list[tuple[str, list[int], list[int]]]], + ], ) -> None: """Calculate addresses and execute one remote PP/TP transfer bucket.""" src_list: list[int] = [] dst_list: list[int] = [] length_list: list[int] = [] - for request_id, local_layer_index, remote_layer_index, local_block_ids, remote_block_ids in transfer_entries: - spec = self.kv_cache_specs[self.spec_indices[local_layer_index]] - self._append_layer_transfer_addresses( - request_id=request_id, - local_layer_index=local_layer_index, - remote_layer_index=remote_layer_index, + for spec_index, transfer_entries_by_layer in transfer_entries_by_spec.items(): + self._append_spec_transfer_addresses( + spec_index=spec_index, remote_tp_rank=remote_tp_rank, - spec=spec, - local_block_ids=local_block_ids, - remote_block_ids=remote_block_ids, + remote_tp_size=remote_tp_size, + remote_dcp_size=remote_dcp_size, + transfer_entries_by_layer=transfer_entries_by_layer, remote_metadata=remote_metadata, src_list=src_list, dst_list=dst_list, From c2dcc200e75f0c15356df23114e9eef85800de3c Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Mon, 10 Aug 2026 00:46:19 +0800 Subject: [PATCH 23/48] add mamba address and transfer Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/pull_worker.py | 192 +++++++++++++++++- 1 file changed, 191 insertions(+), 1 deletion(-) 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 99a925c9fe3b..bfdcf6256e7c 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 @@ -12,6 +12,7 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.utils.network_utils import make_zmq_path +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, @@ -616,6 +617,185 @@ def _compute_group_block_ids( for candidate_tp_ranks in remote_tp_rank_groups ] + def _append_mamba_transfer_addresses( + self, + spec: MambaSpec, + remote_tp_rank: int, + remote_tp_size: int, + transfer_entries_by_layer: dict[tuple[int, int], list[tuple[str, list[int], list[int]]]], + remote_metadata: MooncakePPTransferMetadata, + src_list: list[int], + dst_list: list[int], + length_list: list[int], + ) -> None: + """Append Mamba conv/SSM slices for one producer TP rank.""" + larger_tp_size = max(self.tp_size, remote_tp_size) + smaller_tp_size = min(self.tp_size, remote_tp_size) + if larger_tp_size % smaller_tp_size != 0: + raise ValueError( + f"Mooncake Mamba TP sizes must have an integer ratio, got local={self.tp_size}, remote={remote_tp_size}" + ) + tp_ratio = larger_tp_size // smaller_tp_size + if remote_tp_size >= self.tp_size: + rank_offset = remote_tp_rank % tp_ratio + if remote_tp_rank // tp_ratio != self.tp_rank: + raise ValueError( + f"Mooncake Mamba remote TP {remote_tp_rank} does not belong to local TP {self.tp_rank}" + ) + else: + rank_offset = self.tp_rank % tp_ratio + if self.tp_rank // tp_ratio != remote_tp_rank: + raise ValueError( + f"Mooncake Mamba local TP {self.tp_rank} does not belong to remote TP {remote_tp_rank}" + ) + + remote_tp_metadata = remote_metadata.metadata_by_tp_rank[remote_tp_rank] + 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_shapes = self.block_shapes[local_layer_index] + remote_shapes = remote_metadata.block_shapes[remote_layer_index] + if not ( + len(local_base_addrs) + == len(remote_base_addrs) + == len(local_shapes) + == len(remote_shapes) + == len(spec.dtypes) + ): + raise ValueError( + f"Mooncake Mamba cache metadata count mismatch for layer {self.layer_names[local_layer_index]!r}" + ) + + for cache_index, (local_base_addr, remote_base_addr) in enumerate(zip(local_base_addrs, remote_base_addrs)): + local_block_stride = self.block_strides[local_layer_index][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] + remote_block_len = remote_metadata.block_lens[remote_layer_index][cache_index] + local_shape = local_shapes[cache_index] + remote_shape = remote_shapes[cache_index] + + if self.tp_size == remote_tp_size: + if remote_tp_rank != self.tp_rank or local_block_len != remote_block_len: + raise ValueError( + f"Mooncake Mamba equal-TP metadata mismatch for layer " + f"{self.layer_names[local_layer_index]!r}, cache {cache_index}" + ) + address_slices = [(0, 0, local_block_len)] + elif cache_index == 0 and len(local_shape) == len(remote_shape) == 2: + if local_shape[0] != remote_shape[0]: + raise ValueError( + f"Mooncake Mamba conv row mismatch for layer {self.layer_names[local_layer_index]!r}: " + f"local={local_shape}, remote={remote_shape}" + ) + dtype_size = torch.tensor([], dtype=spec.dtypes[cache_index]).element_size() + mamba_type = spec.mamba_type + if mamba_type == MambaAttentionBackendEnum.MAMBA1: + local_projection_widths = (local_shape[1],) + remote_projection_widths = (remote_shape[1],) + elif mamba_type in (MambaAttentionBackendEnum.MAMBA2, MambaAttentionBackendEnum.GDN_ATTN): + if len(local_shapes) < 2 or len(remote_shapes) < 2: + raise ValueError( + f"Mooncake {mamba_type.name} conv transfer requires a temporal state shape" + ) + local_state_shape = local_shapes[1] + remote_state_shape = remote_shapes[1] + local_state_width = local_state_shape[0] * local_state_shape[1] + remote_state_width = remote_state_shape[0] * remote_state_shape[1] + local_remainder = local_shape[1] - local_state_width + remote_remainder = remote_shape[1] - remote_state_width + if local_remainder <= 0 or remote_remainder <= 0 or local_remainder % 2 or remote_remainder % 2: + raise ValueError( + f"Mooncake {mamba_type.name} conv shape cannot be decomposed: " + f"local={local_shape}, remote={remote_shape}, " + f"local_state={local_state_shape}, remote_state={remote_state_shape}" + ) + if mamba_type == MambaAttentionBackendEnum.MAMBA2: + local_projection_widths = ( + local_state_width, + local_remainder // 2, + local_remainder // 2, + ) + remote_projection_widths = ( + remote_state_width, + remote_remainder // 2, + remote_remainder // 2, + ) + else: + local_projection_widths = ( + local_remainder // 2, + local_remainder // 2, + local_state_width, + ) + remote_projection_widths = ( + remote_remainder // 2, + remote_remainder // 2, + remote_state_width, + ) + else: + raise NotImplementedError( + f"Mooncake unequal-TP conv transfer does not support Mamba type {mamba_type!r}" + ) + if ( + sum(local_projection_widths) != local_shape[1] + or sum(remote_projection_widths) != remote_shape[1] + ): + raise ValueError( + f"Mooncake Mamba conv projections do not match cache shapes for layer " + f"{self.layer_names[local_layer_index]!r}: local={local_shape}, remote={remote_shape}" + ) + + address_slices: list[tuple[int, int, int]] = [] + local_projection_offset = 0 + remote_projection_offset = 0 + for local_projection_width, remote_projection_width in zip( + local_projection_widths, remote_projection_widths + ): + if remote_tp_size > self.tp_size: + if local_projection_width != remote_projection_width * tp_ratio: + raise ValueError("Mooncake Mamba conv projection TP ratio mismatch") + local_slice_offset = local_projection_offset + rank_offset * remote_projection_width + remote_slice_offset = remote_projection_offset + transfer_width = remote_projection_width + else: + if remote_projection_width != local_projection_width * tp_ratio: + raise ValueError("Mooncake Mamba conv projection TP ratio mismatch") + local_slice_offset = local_projection_offset + remote_slice_offset = remote_projection_offset + rank_offset * local_projection_width + transfer_width = local_projection_width + for row_index in range(local_shape[0]): + address_slices.append( + ( + (row_index * local_shape[1] + local_slice_offset) * dtype_size, + (row_index * remote_shape[1] + remote_slice_offset) * dtype_size, + transfer_width * dtype_size, + ) + ) + local_projection_offset += local_projection_width + remote_projection_offset += remote_projection_width + elif remote_tp_size > self.tp_size: + if local_block_len != remote_block_len * tp_ratio: + raise ValueError( + f"Mooncake Mamba state TP ratio mismatch for layer " + f"{self.layer_names[local_layer_index]!r}, cache {cache_index}" + ) + address_slices = [(rank_offset * remote_block_len, 0, remote_block_len)] + else: + if remote_block_len != local_block_len * tp_ratio: + raise ValueError( + f"Mooncake Mamba state TP ratio mismatch for layer " + f"{self.layer_names[local_layer_index]!r}, cache {cache_index}" + ) + address_slices = [(0, rank_offset * local_block_len, local_block_len)] + + for _, local_block_ids, remote_block_ids in transfer_entries: + for local_block_id, remote_block_id in zip(local_block_ids, remote_block_ids): + local_block_addr = local_base_addr + local_block_id * local_block_stride + remote_block_addr = remote_base_addr + remote_block_id * remote_block_stride + for local_inner_offset, remote_inner_offset, transfer_len in address_slices: + src_list.append(local_block_addr + local_inner_offset) + dst_list.append(remote_block_addr + remote_inner_offset) + length_list.append(transfer_len) + def _append_spec_transfer_addresses( self, spec_index: int, @@ -646,7 +826,17 @@ def _append_spec_transfer_addresses( spec = self.kv_cache_specs[spec_index] if isinstance(spec, MambaSpec): - raise NotImplementedError("Mooncake Mamba transfer address calculation is not implemented") + self._append_mamba_transfer_addresses( + spec, + remote_tp_rank, + remote_tp_size, + transfer_entries_by_layer, + remote_metadata, + src_list, + dst_list, + length_list, + ) + return remote_tp_metadata = remote_metadata.metadata_by_tp_rank[remote_tp_rank] transfer_whole_block = isinstance( spec, From ab6c0e4659fd6551204114dff1b7c3ab9d1b1ad6 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Mon, 10 Aug 2026 11:10:44 +0800 Subject: [PATCH 24/48] add dcp with different block size Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/pull_worker.py | 79 ++++++++++++++++++- 1 file changed, 78 insertions(+), 1 deletion(-) 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 bfdcf6256e7c..ac878345fdf4 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 @@ -547,7 +547,84 @@ def _compute_group_block_ids( elif isinstance(spec, FullAttentionSpec): if is_dcp_transfer: if local_block_size != remote_block_size: - raise NotImplementedError("Mooncake DCP with different P/D spec block sizes is not implemented") + local_kernel_block_size = local_block_size // local_block_size_scale + remote_kernel_block_size = remote_block_size // remote_block_size_scale + if local_kernel_block_size != remote_kernel_block_size: + raise ValueError( + "Mooncake DCP requires the same P/D kernel block size: " + f"local={local_kernel_block_size}, remote={remote_kernel_block_size}" + ) + + kernel_block_size = local_kernel_block_size + local_virtual_block_size = local_block_size * self.dcp_size + remote_virtual_block_size = remote_block_size * remote_dcp_size + transfer_start_token = num_computed_tokens + transfer_end_token = min(local_num_prompt_tokens - 1, remote_num_prompt_tokens) + if transfer_end_token <= transfer_start_token: + return [] + + local_start_block_idx = num_computed_tokens // local_virtual_block_size + local_virtual_kernel_blocks: list[tuple[int, int]] = [] + for local_block_offset, local_block_id in enumerate(local_group_block_ids): + local_block_idx = local_start_block_idx + local_block_offset + local_dcp_token_start = ( + local_block_idx * local_virtual_block_size + self.dcp_rank * local_block_size + ) + for kernel_offset in range(local_block_size_scale): + token_start = local_dcp_token_start + kernel_offset * kernel_block_size + if token_start + kernel_block_size <= transfer_start_token: + continue + if token_start >= transfer_end_token: + break + local_kernel_block_id = local_block_id * local_block_size_scale + kernel_offset + virtual_kernel_block_idx = token_start // kernel_block_size + local_virtual_kernel_blocks.append((local_kernel_block_id, virtual_kernel_block_idx)) + + transfer_block_ids: list[tuple[int, list[int], list[int]]] = [] + for candidate_tp_ranks in remote_tp_rank_groups: + if remote_dcp_size == 1: + remote_tp_rank = ( + candidate_tp_ranks[0] + if len(candidate_tp_ranks) == 1 + else self._select_remote_tp_rank(candidate_tp_ranks, selection_index) + ) + remote_tp_ranks = [remote_tp_rank] + else: + if len(candidate_tp_ranks) != remote_dcp_size: + raise ValueError( + "Mooncake DCP candidate rank count must equal remote DCP size: " + f"ranks={candidate_tp_ranks}, remote_dcp={remote_dcp_size}" + ) + remote_tp_ranks = candidate_tp_ranks + + block_ids_by_remote_tp_rank: dict[int, tuple[list[int], list[int]]] = {} + for local_kernel_block_id, virtual_kernel_block_idx in local_virtual_kernel_blocks: + token_start = virtual_kernel_block_idx * kernel_block_size + remote_block_idx = token_start // remote_virtual_block_size + if remote_block_idx >= len(remote_group_block_ids): + break + remote_virtual_offset = token_start % remote_virtual_block_size + remote_dcp_rank = remote_virtual_offset // remote_block_size + remote_kernel_offset = ( + remote_virtual_offset % remote_block_size + ) // remote_kernel_block_size + remote_tp_rank = remote_tp_ranks[remote_dcp_rank] + local_kernel_block_ids, remote_kernel_block_ids = block_ids_by_remote_tp_rank.setdefault( + remote_tp_rank, ([], []) + ) + local_kernel_block_ids.append(local_kernel_block_id) + remote_kernel_block_ids.append( + remote_group_block_ids[remote_block_idx] * remote_block_size_scale + + remote_kernel_offset + ) + transfer_block_ids.extend( + (remote_tp_rank, local_kernel_block_ids, remote_kernel_block_ids) + for remote_tp_rank, ( + local_kernel_block_ids, + remote_kernel_block_ids, + ) in block_ids_by_remote_tp_rank.items() + ) + return transfer_block_ids local_virtual_block_size = local_block_size * self.dcp_size local_start_block_idx = num_computed_tokens // local_virtual_block_size From 2cfcee3a813de1bb63e179943bd4f8be632a67de Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Mon, 10 Aug 2026 12:21:41 +0800 Subject: [PATCH 25/48] merge dcp with same/different block size code Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/pull_worker.py | 118 ++++++++---------- 1 file changed, 49 insertions(+), 69 deletions(-) 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 ac878345fdf4..2d33d6417f8b 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 @@ -546,24 +546,26 @@ def _compute_group_block_ids( ) elif isinstance(spec, FullAttentionSpec): if is_dcp_transfer: - if local_block_size != remote_block_size: - local_kernel_block_size = local_block_size // local_block_size_scale - remote_kernel_block_size = remote_block_size // remote_block_size_scale - if local_kernel_block_size != remote_kernel_block_size: - raise ValueError( - "Mooncake DCP requires the same P/D kernel block size: " - f"local={local_kernel_block_size}, remote={remote_kernel_block_size}" - ) + local_kernel_block_size = local_block_size // local_block_size_scale + remote_kernel_block_size = remote_block_size // remote_block_size_scale + if local_kernel_block_size != remote_kernel_block_size: + raise ValueError( + "Mooncake DCP requires the same P/D kernel block size: " + f"local={local_kernel_block_size}, remote={remote_kernel_block_size}" + ) + + local_virtual_block_size = local_block_size * self.dcp_size + remote_virtual_block_size = remote_block_size * remote_dcp_size + local_start_block_idx = num_computed_tokens // local_virtual_block_size + transfer_block_ids: list[tuple[int, list[int], list[int]]] = [] + if local_block_size != remote_block_size: kernel_block_size = local_kernel_block_size - local_virtual_block_size = local_block_size * self.dcp_size - remote_virtual_block_size = remote_block_size * remote_dcp_size transfer_start_token = num_computed_tokens transfer_end_token = min(local_num_prompt_tokens - 1, remote_num_prompt_tokens) if transfer_end_token <= transfer_start_token: return [] - local_start_block_idx = num_computed_tokens // local_virtual_block_size local_virtual_kernel_blocks: list[tuple[int, int]] = [] for local_block_offset, local_block_id in enumerate(local_group_block_ids): local_block_idx = local_start_block_idx + local_block_offset @@ -580,24 +582,24 @@ def _compute_group_block_ids( virtual_kernel_block_idx = token_start // kernel_block_size local_virtual_kernel_blocks.append((local_kernel_block_id, virtual_kernel_block_idx)) - transfer_block_ids: list[tuple[int, list[int], list[int]]] = [] - for candidate_tp_ranks in remote_tp_rank_groups: - if remote_dcp_size == 1: - remote_tp_rank = ( - candidate_tp_ranks[0] - if len(candidate_tp_ranks) == 1 - else self._select_remote_tp_rank(candidate_tp_ranks, selection_index) + for candidate_tp_ranks in remote_tp_rank_groups: + if remote_dcp_size == 1: + remote_tp_rank = ( + candidate_tp_ranks[0] + if len(candidate_tp_ranks) == 1 + else self._select_remote_tp_rank(candidate_tp_ranks, selection_index) + ) + remote_tp_ranks = [remote_tp_rank] + else: + if len(candidate_tp_ranks) != remote_dcp_size: + raise ValueError( + "Mooncake DCP candidate rank count must equal remote DCP size: " + f"ranks={candidate_tp_ranks}, remote_dcp={remote_dcp_size}" ) - remote_tp_ranks = [remote_tp_rank] - else: - if len(candidate_tp_ranks) != remote_dcp_size: - raise ValueError( - "Mooncake DCP candidate rank count must equal remote DCP size: " - f"ranks={candidate_tp_ranks}, remote_dcp={remote_dcp_size}" - ) - remote_tp_ranks = candidate_tp_ranks + remote_tp_ranks = candidate_tp_ranks - block_ids_by_remote_tp_rank: dict[int, tuple[list[int], list[int]]] = {} + block_ids_by_remote_tp_rank: dict[int, tuple[list[int], list[int]]] = {} + if local_block_size != remote_block_size: for local_kernel_block_id, virtual_kernel_block_idx in local_virtual_kernel_blocks: token_start = virtual_kernel_block_idx * kernel_block_size remote_block_idx = token_start // remote_virtual_block_size @@ -617,49 +619,28 @@ def _compute_group_block_ids( remote_group_block_ids[remote_block_idx] * remote_block_size_scale + remote_kernel_offset ) - transfer_block_ids.extend( - (remote_tp_rank, local_kernel_block_ids, remote_kernel_block_ids) - for remote_tp_rank, ( - local_kernel_block_ids, - remote_kernel_block_ids, - ) in block_ids_by_remote_tp_rank.items() - ) - return transfer_block_ids - - local_virtual_block_size = local_block_size * self.dcp_size - local_start_block_idx = num_computed_tokens // local_virtual_block_size - transfer_block_ids: list[tuple[int, list[int], list[int]]] = [] - for candidate_tp_ranks in remote_tp_rank_groups: - if remote_dcp_size == 1: - remote_tp_rank = ( - candidate_tp_ranks[0] - if len(candidate_tp_ranks) == 1 - else self._select_remote_tp_rank(candidate_tp_ranks, selection_index) - ) - remote_tp_ranks = [remote_tp_rank] else: - remote_tp_ranks = candidate_tp_ranks + for local_block_offset, local_block_id in enumerate(local_group_block_ids): + local_block_idx = local_start_block_idx + local_block_offset + global_block_idx = local_block_idx * self.dcp_size + self.dcp_rank + remote_block_idx = global_block_idx // remote_dcp_size + if remote_block_idx >= len(remote_group_block_ids): + break + remote_dcp_rank = global_block_idx % remote_dcp_size + remote_tp_rank = remote_tp_ranks[remote_dcp_rank] + local_kernel_block_ids, remote_kernel_block_ids = block_ids_by_remote_tp_rank.setdefault( + remote_tp_rank, ([], []) + ) + local_kernel_block_ids.extend( + local_block_id * local_block_size_scale + offset + for offset in range(local_block_size_scale) + ) + remote_block_id = remote_group_block_ids[remote_block_idx] + remote_kernel_block_ids.extend( + remote_block_id * remote_block_size_scale + offset + for offset in range(remote_block_size_scale) + ) - block_ids_by_remote_tp_rank: dict[int, tuple[list[int], list[int]]] = {} - for local_block_offset, local_block_id in enumerate(local_group_block_ids): - local_block_idx = local_start_block_idx + local_block_offset - global_block_idx = local_block_idx * self.dcp_size + self.dcp_rank - remote_block_idx = global_block_idx // remote_dcp_size - if remote_block_idx >= len(remote_group_block_ids): - break - remote_dcp_rank = global_block_idx % remote_dcp_size - remote_tp_rank = remote_tp_ranks[remote_dcp_rank] - local_kernel_block_ids, remote_kernel_block_ids = block_ids_by_remote_tp_rank.setdefault( - remote_tp_rank, ([], []) - ) - local_kernel_block_ids.extend( - local_block_id * local_block_size_scale + offset for offset in range(local_block_size_scale) - ) - remote_block_id = remote_group_block_ids[remote_block_idx] - remote_kernel_block_ids.extend( - remote_block_id * remote_block_size_scale + offset - for offset in range(remote_block_size_scale) - ) transfer_block_ids.extend( (remote_tp_rank, local_kernel_block_ids, remote_kernel_block_ids) for remote_tp_rank, ( @@ -668,7 +649,6 @@ def _compute_group_block_ids( ) in block_ids_by_remote_tp_rank.items() ) return transfer_block_ids - local_kernel_block_ids = self._expand_block_ids(local_group_block_ids, local_block_size_scale) remote_kernel_block_ids = self._expand_block_ids(remote_group_block_ids, remote_block_size_scale) kernel_block_size = local_block_size // local_block_size_scale From 24572e93327d8dbcda0e6e44bcbc0d6d58931c55 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Mon, 10 Aug 2026 16:23:46 +0800 Subject: [PATCH 26/48] add threadpool Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/pull_worker.py | 130 +++++++++++------- 1 file changed, 77 insertions(+), 53 deletions(-) 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 2d33d6417f8b..064d3e7460c5 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 @@ -4,6 +4,7 @@ import queue import threading +from concurrent.futures import Future, ThreadPoolExecutor, as_completed from typing import Any import msgspec @@ -114,6 +115,12 @@ def __init__( assert self.pcp_size == 1, f"Mooncake pull worker temporarily requires pcp_size=1, got {self.pcp_size}" self.device = device self.ready_event = ready_event + self.executor = ThreadPoolExecutor( + max_workers=32, + initializer=torch.npu.set_device, + initargs=(self.device,), + ) + self.can_report_invalid_block_ids = len(kv_cache_config.kv_cache_groups) == 1 self.encoder = msgspec.msgpack.Encoder() self.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) self.remote_metadata: SizedDict[str, MooncakeTransferMetadataGroups] = SizedDict() @@ -162,15 +169,18 @@ def run(self) -> None: request_endpoints = {(request.remote_host, request.remote_port) for request in requests.values()} try: remote_host, remote_port = self._get_remote_endpoint(request_endpoints) - self._handle_requests( + failed_request_ids = self._handle_requests( remote_engine_id, remote_host, remote_port, requests, ) + for request_id in failed_request_ids: + self._mark_request_failed(requests[request_id]) except Exception as exc: - for request_metadata in requests.values(): - self._mark_request_failed(request_metadata) + if self.can_report_invalid_block_ids: + for request_metadata in requests.values(): + self._mark_request_failed(request_metadata) logger.exception( "Mooncake pull failed for remote engine %s at %s: %s", remote_engine_id, @@ -183,6 +193,8 @@ def run(self) -> None: self.request_queue.task_done() def _mark_request_failed(self, request_metadata: ReqMeta) -> None: + if not self.can_report_invalid_block_ids: + return with self.invalid_block_ids_lock: for group_block_ids in request_metadata.local_block_ids: self.invalid_block_ids.update(group_block_ids) @@ -399,29 +411,67 @@ def _handle_requests( remote_host: str, remote_port: int, requests: dict[str, ReqMeta], - ) -> None: + ) -> set[str]: """Build and execute transfer buckets for one producer engine.""" remote_metadata = self._get_remote_metadata(remote_engine_id, remote_host, remote_port) tp_rank_groups_by_pp_rank = self.remote_tp_rank_groups[remote_engine_id] layer_pairs_by_pp_rank = self.remote_layer_index_pairs[remote_engine_id] transfer_block_ids_by_spec: dict[str, dict[int, list[tuple[int, list[int], list[int]]]]] = {} - for remote_pp_rank, layer_pairs in layer_pairs_by_pp_rank.items(): - pp_metadata = remote_metadata.metadata_by_pp_rank[remote_pp_rank] - transfer_block_buckets = self._build_transfer_block_buckets( - pp_metadata, - layer_pairs, - tp_rank_groups_by_pp_rank[remote_pp_rank], - remote_metadata.dcp_size, - requests, - transfer_block_ids_by_spec, - ) - self._execute_transfer_block_buckets( - remote_pp_rank, - pp_metadata, - remote_metadata.tp_size, - remote_metadata.dcp_size, - transfer_block_buckets, - ) + future_to_task: dict[Future[None], tuple[int, int, set[str]]] = {} + submission_error: Exception | None = None + try: + for remote_pp_rank, layer_pairs in layer_pairs_by_pp_rank.items(): + pp_metadata = remote_metadata.metadata_by_pp_rank[remote_pp_rank] + transfer_block_buckets, request_ids_by_remote_tp_rank = self._build_transfer_block_buckets( + pp_metadata, + layer_pairs, + tp_rank_groups_by_pp_rank[remote_pp_rank], + remote_metadata.dcp_size, + requests, + transfer_block_ids_by_spec, + ) + for remote_tp_rank, transfer_entries_by_spec in transfer_block_buckets.items(): + future = self.executor.submit( + self._execute_tp_transfer_bucket, + remote_pp_rank, + remote_tp_rank, + remote_metadata.tp_size, + remote_metadata.dcp_size, + pp_metadata, + transfer_entries_by_spec, + ) + future_to_task[future] = ( + remote_pp_rank, + remote_tp_rank, + request_ids_by_remote_tp_rank[remote_tp_rank], + ) + except Exception as exc: + submission_error = exc + + failed_request_ids: set[str] = set() + for future in as_completed(future_to_task): + remote_pp_rank, remote_tp_rank, request_ids = future_to_task[future] + try: + future.result() + except Exception: + logger.exception( + "Mooncake transfer task failed for remote PP rank %s, TP rank %s, requests=%s", + remote_pp_rank, + remote_tp_rank, + sorted(request_ids), + ) + if self.can_report_invalid_block_ids: + failed_request_ids.update(request_ids) + else: + logger.warning( + "Ignoring Mooncake transfer failure for hybrid KV cache requests %s because " + "vLLM invalid block reporting currently supports only one KV cache group", + sorted(request_ids), + ) + + if submission_error is not None: + raise submission_error + return failed_request_ids def _build_transfer_block_buckets( self, @@ -431,12 +481,9 @@ def _build_transfer_block_buckets( remote_dcp_size: int, requests: dict[str, ReqMeta], transfer_block_ids_by_spec: dict[str, dict[int, list[tuple[int, list[int], list[int]]]]], - ) -> dict[ - int, - dict[ - int, - dict[tuple[int, int], list[tuple[str, list[int], list[int]]]], - ], + ) -> tuple[ + dict[int, dict[int, dict[tuple[int, int], list[tuple[str, list[int], list[int]]]]]], + dict[int, set[str]], ]: """Bucket one PP rank's requests by remote TP, spec index, and layer pair.""" # remote_tp_rank -> local_spec_index -> (local_layer_index, remote_layer_index) -> @@ -448,6 +495,7 @@ def _build_transfer_block_buckets( dict[tuple[int, int], list[tuple[str, list[int], list[int]]]], ], ] = {} + request_ids_by_remote_tp_rank: dict[int, set[str]] = {} for selection_index, (request_id, request_metadata) in enumerate(requests.items()): request_block_ids_by_spec = transfer_block_ids_by_spec.setdefault(request_id, {}) for local_layer_index, remote_layer_index in layer_pairs: @@ -493,12 +541,13 @@ def _build_transfer_block_buckets( request_block_ids_by_spec[spec_index] = transfer_block_ids for remote_tp_rank, local_block_ids, remote_block_ids in transfer_block_ids: + request_ids_by_remote_tp_rank.setdefault(remote_tp_rank, set()).add(request_id) transfer_entries_by_spec = transfer_block_buckets.setdefault(remote_tp_rank, {}) transfer_entries_by_layer = transfer_entries_by_spec.setdefault(spec_index, {}) transfer_entries_by_layer.setdefault(layer_pair, []).append( (request_id, local_block_ids, remote_block_ids) ) - return transfer_block_buckets + return transfer_block_buckets, request_ids_by_remote_tp_rank @staticmethod def _expand_block_ids(block_ids: list[int], scale: int) -> list[int]: @@ -1027,31 +1076,6 @@ def _append_spec_transfer_addresses( dst_list.append(remote_base_addr + remote_block_id * remote_block_stride + remote_inner_offset) length_list.append(transfer_len) - def _execute_transfer_block_buckets( - self, - remote_pp_rank: int, - remote_metadata: MooncakePPTransferMetadata, - remote_tp_size: int, - remote_dcp_size: int, - transfer_block_buckets: dict[ - int, - dict[ - int, - dict[tuple[int, int], list[tuple[str, list[int], list[int]]]], - ], - ], - ) -> None: - """Calculate addresses and execute one PP rank's TP buckets serially.""" - for remote_tp_rank, transfer_entries_by_spec in transfer_block_buckets.items(): - self._execute_tp_transfer_bucket( - remote_pp_rank, - remote_tp_rank, - remote_tp_size, - remote_dcp_size, - remote_metadata, - transfer_entries_by_spec, - ) - def _execute_tp_transfer_bucket( self, remote_pp_rank: int, From 6414193883a2179f6a80f944dd65a473b5321614 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Wed, 12 Aug 2026 17:12:19 +0800 Subject: [PATCH 27/48] some bugfix Signed-off-by: nwpu-zxr --- vllm_ascend/distributed/kv_transfer/__init__.py | 12 ++++++++++++ .../kv_transfer/kv_p2p/mooncake/base_scheduler.py | 3 +-- .../kv_transfer/kv_p2p/mooncake/base_worker.py | 3 +-- .../kv_transfer/kv_p2p/mooncake/connector.py | 3 +-- .../kv_transfer/kv_p2p/mooncake/pull_scheduler.py | 3 +-- .../kv_transfer/kv_p2p/mooncake/pull_worker.py | 4 +--- .../distributed/kv_transfer/kv_p2p/mooncake/utils.py | 3 +-- vllm_ascend/utils.py | 2 +- 8 files changed, 19 insertions(+), 14 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/__init__.py b/vllm_ascend/distributed/kv_transfer/__init__.py index 91c013d5ded4..6f140dc3dc76 100644 --- a/vllm_ascend/distributed/kv_transfer/__init__.py +++ b/vllm_ascend/distributed/kv_transfer/__init__.py @@ -46,6 +46,18 @@ def register_connector(): "MooncakeConnectorV1", "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_connector", "MooncakeConnector" ) + KVConnectorFactory.register_connector( + "MooncakeConnectorV2", + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake", + "MooncakeConnector", + ) + + KVConnectorFactory.register_connector( + "MooncakePullConnector", + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake", + "MooncakePullConnector", + ) + KVConnectorFactory.register_connector( "MooncakeHybridConnector", "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake_hybrid_connector", diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py index 2a8c3c608ad0..58474688fc83 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py @@ -11,7 +11,7 @@ KVConnectorHandshakeMetadata, KVConnectorMetadata, ) -from vllm.logger import init_logger +from vllm.logger import logger from vllm.utils.math_utils import cdiv from vllm.utils.network_utils import get_ip from vllm.v1.core.sched.output import SchedulerOutput @@ -31,7 +31,6 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request -logger = init_logger(__name__) class MooncakeBaseConnectorScheduler: 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 93ee6bcfbe13..f42944e1add0 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 @@ -19,7 +19,7 @@ get_tensor_model_parallel_rank, get_tp_group, ) -from vllm.logger import init_logger +from vllm.logger import logger from vllm.utils.network_utils import get_ip from vllm.v1.kv_cache_interface import KVCacheSpec, UniformTypeKVCacheSpecs @@ -50,7 +50,6 @@ if TYPE_CHECKING: from vllm.v1.kv_cache_interface import KVCacheConfig -logger = init_logger(__name__) class MooncakeBaseConnectorWorker: diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py index 166994a2c0cc..c528b84b33cf 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py @@ -29,7 +29,7 @@ PromMetric, PromMetricT, ) -from vllm.logger import init_logger +from vllm.logger import logger from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import KVConnectorOutput @@ -60,7 +60,6 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request -logger = init_logger(__name__) class MooncakeBaseConnector(KVConnectorBase_V1, SupportsHMA): diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index d5f778611acf..53d1a43493c7 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -17,7 +17,7 @@ KVConnectorHandshakeMetadata, KVConnectorMetadata, ) -from vllm.logger import init_logger +from vllm.logger import logger from vllm.utils.network_utils import make_zmq_path from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import KVConnectorOutput @@ -45,7 +45,6 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.request import Request -logger = init_logger(__name__) GET_META_MSG = b"get_meta_msg" DONE_RECVING_MSG = b"done_recving_msg" 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 064d3e7460c5..d4b9d34c4c13 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 @@ -11,7 +11,7 @@ import torch import zmq from vllm.config import VllmConfig -from vllm.logger import init_logger +from vllm.logger import logger from vllm.utils.network_utils import make_zmq_path from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_interface import ( @@ -42,8 +42,6 @@ zmq_ctx, ) -logger = init_logger(__name__) - class MooncakePullRecvingThread(threading.Thread): """D-side execution thread for Mooncake pull 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 5f8dd53c2d89..a3ce0b581ad8 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py @@ -15,7 +15,7 @@ import numpy.typing as npt import torch import zmq -from vllm.logger import init_logger +from vllm.logger import logger from vllm.utils.network_utils import make_zmq_socket from vllm_ascend.distributed.kv_transfer.utils.utils import ( @@ -26,7 +26,6 @@ if TYPE_CHECKING: from vllm.v1.kv_cache_interface import KVCacheConfig -logger = init_logger(__name__) @contextlib.contextmanager diff --git a/vllm_ascend/utils.py b/vllm_ascend/utils.py index b8497f0d6abc..155b304593a1 100644 --- a/vllm_ascend/utils.py +++ b/vllm_ascend/utils.py @@ -1331,7 +1331,7 @@ def is_gqa_backend(vllm_config: VllmConfig) -> bool: def uses_mooncake_connector(kv_transfer_config: Any) -> bool: - mooncake_connector_names = {"MooncakeConnector", "MooncakeConnectorV1"} + mooncake_connector_names = {"MooncakeConnector", "MooncakeConnectorV1", "MooncakeConnectorV2"} return bool(_collect_kv_connector_names(kv_transfer_config) & mooncake_connector_names) From d94f206c0442f37d9baba3b6431bfc43d0b105d3 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Wed, 12 Aug 2026 20:49:16 +0800 Subject: [PATCH 28/48] change connector path Signed-off-by: nwpu-zxr --- vllm_ascend/distributed/kv_transfer/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/__init__.py b/vllm_ascend/distributed/kv_transfer/__init__.py index 6f140dc3dc76..7eda355384f8 100644 --- a/vllm_ascend/distributed/kv_transfer/__init__.py +++ b/vllm_ascend/distributed/kv_transfer/__init__.py @@ -48,13 +48,13 @@ def register_connector(): KVConnectorFactory.register_connector( "MooncakeConnectorV2", - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake", + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.MooncakeConnector", "MooncakeConnector", ) KVConnectorFactory.register_connector( "MooncakePullConnector", - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake", + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.MooncakePullConnector", "MooncakePullConnector", ) From 0d87a1bad61e1a2212bba65d260ecb19254e7083 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Wed, 12 Aug 2026 20:59:29 +0800 Subject: [PATCH 29/48] add __init__ Signed-off-by: nwpu-zxr --- .../kv_transfer/kv_p2p/mooncake/__init__.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py index e69de29bb2d1..a1e33b2f2db7 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py @@ -0,0 +1,41 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vllm_ascend project +"""Mooncake KV-cache transfer connector (disaggregated prefill / decode).""" + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler import ( + MooncakeBaseConnectorScheduler, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker import ( + MooncakeBaseConnectorWorker, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector import ( + MooncakeBaseConnector, + MooncakeConnector, + MooncakePullConnector, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( + MooncakeTransferMetadata, + MooncakeConnectorMetadata, +) +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, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.stats import ( + MooncakeKVConnectorStats, +) + +__all__ = [ + "MooncakeTransferMetadata", + "MooncakeBaseConnector", + "MooncakeBaseConnectorScheduler", + "MooncakeBaseConnectorWorker", + "MooncakeConnector", + "MooncakeConnectorMetadata", + "MooncakeKVConnectorStats", + "MooncakePullConnector", + "MooncakePullConnectorScheduler", + "MooncakePullConnectorWorker", +] From f5c182d67cbb7c00798dc76ecb80f701bbaa617c Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Wed, 12 Aug 2026 21:10:38 +0800 Subject: [PATCH 30/48] fix import Signed-off-by: nwpu-zxr --- vllm_ascend/distributed/kv_transfer/__init__.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/__init__.py b/vllm_ascend/distributed/kv_transfer/__init__.py index 7eda355384f8..aad8e1a3a27d 100644 --- a/vllm_ascend/distributed/kv_transfer/__init__.py +++ b/vllm_ascend/distributed/kv_transfer/__init__.py @@ -48,13 +48,13 @@ def register_connector(): KVConnectorFactory.register_connector( "MooncakeConnectorV2", - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.MooncakeConnector", + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector.MooncakeConnector", "MooncakeConnector", ) KVConnectorFactory.register_connector( "MooncakePullConnector", - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.MooncakePullConnector", + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector.MooncakePullConnector", "MooncakePullConnector", ) From 351d6c06c38e7671d48acfeea268ca483f8ba161 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Thu, 13 Aug 2026 12:53:12 +0800 Subject: [PATCH 31/48] fix ruff Signed-off-by: nwpu-zxr --- vllm_ascend/distributed/kv_transfer/__init__.py | 8 ++++---- .../kv_transfer/kv_p2p/mooncake/base_scheduler.py | 1 - .../kv_transfer/kv_p2p/mooncake/base_worker.py | 1 - .../distributed/kv_transfer/kv_p2p/mooncake/connector.py | 1 - .../distributed/kv_transfer/kv_p2p/mooncake/utils.py | 1 - 5 files changed, 4 insertions(+), 8 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/__init__.py b/vllm_ascend/distributed/kv_transfer/__init__.py index aad8e1a3a27d..00af1dde8d93 100644 --- a/vllm_ascend/distributed/kv_transfer/__init__.py +++ b/vllm_ascend/distributed/kv_transfer/__init__.py @@ -48,14 +48,14 @@ def register_connector(): KVConnectorFactory.register_connector( "MooncakeConnectorV2", - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector.MooncakeConnector", + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector", "MooncakeConnector", ) KVConnectorFactory.register_connector( - "MooncakePullConnector", - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector.MooncakePullConnector", - "MooncakePullConnector", + "MooncakePullConnector", + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector", + "MooncakePullConnector", ) KVConnectorFactory.register_connector( diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py index 58474688fc83..4a258036f76c 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/base_scheduler.py @@ -32,7 +32,6 @@ from vllm.v1.request import Request - class MooncakeBaseConnectorScheduler: """Scheduler logic shared by Mooncake transfer modes.""" 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 f42944e1add0..1a58b1dd902a 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 @@ -51,7 +51,6 @@ from vllm.v1.kv_cache_interface import KVCacheConfig - class MooncakeBaseConnectorWorker: """Worker implementation shared by Mooncake transfer modes.""" diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py index c528b84b33cf..5e761dbb9d52 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/connector.py @@ -61,7 +61,6 @@ from vllm.v1.request import Request - class MooncakeBaseConnector(KVConnectorBase_V1, SupportsHMA): """Common facade for the Mooncake scheduler and worker implementations.""" 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 a3ce0b581ad8..b52de6c19207 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py @@ -27,7 +27,6 @@ from vllm.v1.kv_cache_interface import KVCacheConfig - @contextlib.contextmanager def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]: """Create a Mooncake ROUTER or REQ socket and clean up its context.""" From 489bfcb1d47b922224d6cadd1dc2b1fe7c64d9cc Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Sun, 16 Aug 2026 18:27:51 +0800 Subject: [PATCH 32/48] fix ruff Signed-off-by: nwpu-zxr --- vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py index a1e33b2f2db7..06c3295b5918 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/__init__.py @@ -14,8 +14,8 @@ MooncakePullConnector, ) from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( - MooncakeTransferMetadata, MooncakeConnectorMetadata, + MooncakeTransferMetadata, ) from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler import ( MooncakePullConnectorScheduler, From cab1d3b5fb0ef8b59d4221798a72af2e03cdbe36 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Fri, 21 Aug 2026 12:23:27 +0800 Subject: [PATCH 33/48] fix mypy Signed-off-by: nwpu-zxr --- .../kv_transfer/kv_p2p/mooncake/metadata.py | 36 ++++++++++--------- .../kv_p2p/mooncake/pull_scheduler.py | 14 ++++---- .../kv_p2p/mooncake/pull_worker.py | 15 +++++--- .../kv_transfer/kv_p2p/mooncake/utils.py | 24 +++++++------ 4 files changed, 51 insertions(+), 38 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 1703e6062423..0d974c0c6b4e 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Ascend project """Metadata types for Mooncake KV transfer connectors.""" +from collections.abc import Sequence from dataclasses import dataclass from typing import Any @@ -39,30 +40,31 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): def __post_init__(self) -> None: num_layers = len(self.layer_names) - per_layer_fields = { - "group_indices": self.group_indices, - "spec_indices": self.spec_indices, - "kv_caches_base_addr": self.kv_caches_base_addr, - "block_strides": self.block_strides, - "block_lens": self.block_lens, - "block_shapes": self.block_shapes, - "block_size_scales": self.block_size_scales, - } - for field_name, values in per_layer_fields.items(): + per_layer_fields: tuple[tuple[str, Sequence[object]], ...] = ( + ("group_indices", self.group_indices), + ("spec_indices", self.spec_indices), + ("kv_caches_base_addr", self.kv_caches_base_addr), + ("block_strides", self.block_strides), + ("block_lens", self.block_lens), + ("block_shapes", self.block_shapes), + ("block_size_scales", self.block_size_scales), + ) + for field_name, values in per_layer_fields: if len(values) != num_layers: raise ValueError( f"Mooncake transfer metadata field {field_name!r} has {len(values)} layers, expected {num_layers}." ) + nested_per_layer_fields: tuple[tuple[str, Sequence[Sequence[object]]], ...] = ( + ("block_strides", self.block_strides), + ("block_lens", self.block_lens), + ("block_shapes", self.block_shapes), + ("block_size_scales", self.block_size_scales), + ) for layer_index, layer_name in enumerate(self.layer_names): num_addrs = len(self.kv_caches_base_addr[layer_index]) - for field_name in ( - "block_strides", - "block_lens", - "block_shapes", - "block_size_scales", - ): - values = per_layer_fields[field_name][layer_index] + for field_name, per_layer_values in nested_per_layer_fields: + values = per_layer_values[layer_index] if len(values) != num_addrs: raise ValueError( f"Mooncake transfer metadata for layer {layer_name!r} " diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 53d1a43493c7..0d468a60c74e 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -217,20 +217,20 @@ def run(self) -> None: path = make_zmq_path("tcp", self.host, self.port) try: logger.info("Mooncake scheduler sending thread listening on %s", path) - with zmq_ctx(zmq.ROUTER, path) as sock: - sock.setsockopt(zmq.RCVTIMEO, 1000) + with zmq_ctx(zmq.ROUTER, path) as sock: # type: ignore[attr-defined] + sock.setsockopt(zmq.RCVTIMEO, 1000) # type: ignore[attr-defined] self.ready_event.set() self._run_busy_loop(sock) except Exception: self.ready_event.set() logger.exception("Mooncake scheduler sending thread failed on %s", path) - def _run_busy_loop(self, sock: zmq.Socket) -> None: + def _run_busy_loop(self, sock: Any) -> None: decoder = msgspec.msgpack.Decoder(type=tuple) while True: try: frames = sock.recv_multipart() - except zmq.Again: + except zmq.Again: # type: ignore[attr-defined] continue identity = frames[0] @@ -313,9 +313,9 @@ def run(self) -> None: def _send_done_recving(self, remote_host: str, remote_port: int, request_id: str) -> None: path = make_zmq_path("tcp", remote_host, remote_port) - with zmq_ctx(zmq.REQ, path) as sock: - sock.setsockopt(zmq.SNDTIMEO, 1000) - sock.setsockopt(zmq.RCVTIMEO, 1000) + with zmq_ctx(zmq.REQ, path) as sock: # type: ignore[attr-defined] + sock.setsockopt(zmq.SNDTIMEO, 1000) # type: ignore[attr-defined] + sock.setsockopt(zmq.RCVTIMEO, 1000) # type: ignore[attr-defined] ensure_zmq_send( sock, self.encoder.encode((DONE_RECVING_MSG, request_id)), 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 d4b9d34c4c13..d6da079c5277 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 @@ -334,6 +334,8 @@ def _infer_total_num_kv_heads( """Infer and validate the total KV heads for one transfer spec.""" if fixed_total_num_kv_heads is not None: return fixed_total_num_kv_heads + if remote_num_kv_heads is None: + raise ValueError("Remote KV head count is required when it cannot be inferred from the cache spec") assert max(self.tp_size, remote_tp_size) % min(self.tp_size, remote_tp_size) == 0 assert max(local_num_kv_heads, remote_num_kv_heads) % min(local_num_kv_heads, remote_num_kv_heads) == 0 local_head_tp_size = self.tp_size // local_dcp_size @@ -778,6 +780,7 @@ def _append_mamba_transfer_addresses( local_shape = local_shapes[cache_index] remote_shape = remote_shapes[cache_index] + address_slices: list[tuple[int, int, int]] if self.tp_size == remote_tp_size: if remote_tp_rank != self.tp_rank or local_block_len != remote_block_len: raise ValueError( @@ -793,6 +796,8 @@ def _append_mamba_transfer_addresses( ) dtype_size = torch.tensor([], dtype=spec.dtypes[cache_index]).element_size() mamba_type = spec.mamba_type + local_projection_widths: tuple[int, ...] + remote_projection_widths: tuple[int, ...] if mamba_type == MambaAttentionBackendEnum.MAMBA1: local_projection_widths = (local_shape[1],) remote_projection_widths = (remote_shape[1],) @@ -848,7 +853,7 @@ def _append_mamba_transfer_addresses( f"{self.layer_names[local_layer_index]!r}: local={local_shape}, remote={remote_shape}" ) - address_slices: list[tuple[int, int, int]] = [] + address_slices = [] local_projection_offset = 0 remote_projection_offset = 0 for local_projection_width, remote_projection_width in zip( @@ -1157,9 +1162,9 @@ def _get_remote_metadata( return cached_metadata path = make_zmq_path("tcp", remote_host, remote_port) - with zmq_ctx(zmq.REQ, path) as sock: - sock.setsockopt(zmq.SNDTIMEO, 1000) - sock.setsockopt(zmq.RCVTIMEO, 1000) + with zmq_ctx(zmq.REQ, path) as sock: # type: ignore[attr-defined] + sock.setsockopt(zmq.SNDTIMEO, 1000) # type: ignore[attr-defined] + sock.setsockopt(zmq.RCVTIMEO, 1000) # type: ignore[attr-defined] ensure_zmq_send( sock, self.encoder.encode((b"get_meta_msg",)), @@ -1206,6 +1211,8 @@ def register_kv_caches( """Register caches and start the D-side pull execution thread.""" super().register_kv_caches(kv_caches) if self.kv_transfer_config.is_kv_consumer: + 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( engine=self.engine, 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 b52de6c19207..108952d59421 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py @@ -9,7 +9,7 @@ from collections import OrderedDict from collections.abc import Iterator from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, TypeVar, cast import numpy as np import numpy.typing as npt @@ -28,19 +28,19 @@ @contextlib.contextmanager -def zmq_ctx(socket_type: Any, addr: str) -> Iterator[zmq.Socket]: +def zmq_ctx(socket_type: Any, addr: str) -> Iterator[Any]: """Create a Mooncake ROUTER or REQ socket and clean up its context.""" - if socket_type not in (zmq.ROUTER, zmq.REQ): + if socket_type not in (zmq.ROUTER, zmq.REQ): # type: ignore[attr-defined] raise ValueError(f"Unexpected socket type: {socket_type}") - context: zmq.Context | None = None + context: Any | None = None try: context = zmq.Context() # type: ignore[attr-defined] yield make_zmq_socket( ctx=context, path=addr, socket_type=socket_type, - bind=socket_type == zmq.ROUTER, + bind=socket_type == zmq.ROUTER, # type: ignore[attr-defined] ) finally: if context is not None: @@ -141,24 +141,28 @@ def merge_storage_range( ) +_KT = TypeVar("_KT") +_VT = TypeVar("_VT") + + @dataclass -class SizedDict(OrderedDict): +class SizedDict(OrderedDict[_KT, _VT]): """Insertion-ordered mapping with a bounded number of entries.""" - def __init__(self, max_size=16000, *args, **kwargs): + def __init__(self, max_size: int = 16000, *args: Any, **kwargs: Any) -> None: self.max_size = max_size super().__init__(*args, **kwargs) - def __setitem__(self, key, value): + def __setitem__(self, key: _KT, value: _VT) -> None: super().__setitem__(key, value) if len(self) > self.max_size: self.popitem(last=False) - def __getitem__(self, key): + def __getitem__(self, key: _KT) -> _VT: try: return super().__getitem__(key) except KeyError: - value: dict[int, list[int]] = {} + value = cast(_VT, {}) self[key] = value return value From 5d54a5d30d905553fb0ba9dae0cda09a55df7a34 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Fri, 21 Aug 2026 14:33:22 +0800 Subject: [PATCH 34/48] add reuseable zmq_context for pull scheduler Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/pull_scheduler.py | 40 ++++++++++++++++--- 1 file changed, 35 insertions(+), 5 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 0d468a60c74e..4caa8d6897da 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -5,7 +5,7 @@ import queue import threading import time -from collections import OrderedDict +from collections import OrderedDict, defaultdict, deque from collections.abc import Mapping from typing import TYPE_CHECKING, Any @@ -18,7 +18,7 @@ KVConnectorMetadata, ) from vllm.logger import logger -from vllm.utils.network_utils import make_zmq_path +from vllm.utils.network_utils import make_zmq_path, make_zmq_socket from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.outputs import KVConnectorOutput from vllm.v1.request import RequestStatus @@ -291,6 +291,9 @@ def __init__(self, ready_event: threading.Event) -> None: self.ready_event = ready_event self.request_queue: queue.Queue[tuple[str, int, str]] = queue.Queue() self.encoder = msgspec.msgpack.Encoder() + self.remote_sockets: dict[str, deque[Any]] = defaultdict(deque) + self.remote_sockets_lock = threading.Lock() + self.zmq_context: Any | None = None def add_request(self, remote_host: str, remote_port: int, request_id: str) -> None: self.request_queue.put((remote_host, remote_port, request_id)) @@ -313,9 +316,8 @@ def run(self) -> None: def _send_done_recving(self, remote_host: str, remote_port: int, request_id: str) -> None: path = make_zmq_path("tcp", remote_host, remote_port) - with zmq_ctx(zmq.REQ, path) as sock: # type: ignore[attr-defined] - sock.setsockopt(zmq.SNDTIMEO, 1000) # type: ignore[attr-defined] - sock.setsockopt(zmq.RCVTIMEO, 1000) # type: ignore[attr-defined] + sock = self._get_remote_socket(path) + try: ensure_zmq_send( sock, self.encoder.encode((DONE_RECVING_MSG, request_id)), @@ -324,6 +326,34 @@ def _send_done_recving(self, remote_host: str, remote_port: int, request_id: str response = ensure_zmq_recv(sock, path) if response != ACK_MSG: raise RuntimeError(f"Unexpected Mooncake scheduler completion response: {response!r}") + except Exception: + sock.close(linger=0) + raise + self._return_remote_socket(path, sock) + + def _get_remote_socket(self, path: str) -> Any: + """Borrow a persistent REQ socket for one producer scheduler.""" + with self.remote_sockets_lock: + sockets = self.remote_sockets[path] + if sockets: + return sockets.popleft() + + if self.zmq_context is None: + self.zmq_context = zmq.Context() # type: ignore[attr-defined] + sock = make_zmq_socket( + ctx=self.zmq_context, + path=path, + socket_type=zmq.REQ, # type: ignore[attr-defined] + bind=False, + ) + sock.setsockopt(zmq.SNDTIMEO, 1000) # type: ignore[attr-defined] + sock.setsockopt(zmq.RCVTIMEO, 1000) # type: ignore[attr-defined] + return sock + + def _return_remote_socket(self, path: str, sock: Any) -> None: + """Return a healthy REQ socket to its endpoint pool.""" + with self.remote_sockets_lock: + self.remote_sockets[path].append(sock) class MooncakePullConnectorScheduler(MooncakeBaseConnectorScheduler): From 1bd655397db4f958c3f9816a13d962d3f94daa1f Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Fri, 21 Aug 2026 14:46:11 +0800 Subject: [PATCH 35/48] add mooncake connector v2 UT Signed-off-by: nwpu-zxr --- tests/ut/kv_offload/mooncake_v2/__init__.py | 1 + tests/ut/kv_offload/mooncake_v2/helpers.py | 162 ++++++ .../mooncake_v2/test_base_worker.py | 110 ++++ .../kv_offload/mooncake_v2/test_connector.py | 105 ++++ .../kv_offload/mooncake_v2/test_metadata.py | 78 +++ .../mooncake_v2/test_pull_scheduler.py | 301 +++++++++++ .../mooncake_v2/test_pull_worker.py | 488 ++++++++++++++++++ tests/ut/kv_offload/mooncake_v2/test_stats.py | 37 ++ tests/ut/kv_offload/mooncake_v2/test_utils.py | 142 +++++ 9 files changed, 1424 insertions(+) create mode 100644 tests/ut/kv_offload/mooncake_v2/__init__.py create mode 100644 tests/ut/kv_offload/mooncake_v2/helpers.py create mode 100644 tests/ut/kv_offload/mooncake_v2/test_base_worker.py create mode 100644 tests/ut/kv_offload/mooncake_v2/test_connector.py create mode 100644 tests/ut/kv_offload/mooncake_v2/test_metadata.py create mode 100644 tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py create mode 100644 tests/ut/kv_offload/mooncake_v2/test_pull_worker.py create mode 100644 tests/ut/kv_offload/mooncake_v2/test_stats.py create mode 100644 tests/ut/kv_offload/mooncake_v2/test_utils.py diff --git a/tests/ut/kv_offload/mooncake_v2/__init__.py b/tests/ut/kv_offload/mooncake_v2/__init__.py new file mode 100644 index 000000000000..b294a3e672ff --- /dev/null +++ b/tests/ut/kv_offload/mooncake_v2/__init__.py @@ -0,0 +1 @@ +"""Unit tests for the refactored Mooncake connector.""" diff --git a/tests/ut/kv_offload/mooncake_v2/helpers.py b/tests/ut/kv_offload/mooncake_v2/helpers.py new file mode 100644 index 000000000000..0e4afe6c6225 --- /dev/null +++ b/tests/ut/kv_offload/mooncake_v2/helpers.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Shared builders for refactored Mooncake connector tests.""" + +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import torch +from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, SlidingWindowSpec + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( + MooncakePPTransferMetadata, + MooncakeTPTransferMetadata, + MooncakeTransferMetadata, + MooncakeTransferMetadataGroups, +) + + +def make_full_spec(block_size: int = 16, num_kv_heads: int = 1) -> FullAttentionSpec: + return FullAttentionSpec( + block_size=block_size, + num_kv_heads=num_kv_heads, + head_size=8, + head_size_v=8, + dtype=torch.float16, + ) + + +def make_sliding_spec(block_size: int = 16) -> SlidingWindowSpec: + return SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=8, + dtype=torch.float16, + sliding_window=64, + ) + + +def make_mamba_spec(block_size: int = 16) -> MambaSpec: + return MambaSpec( + block_size=block_size, + shapes=((3, 16), (2, 4, 4)), + dtypes=(torch.float16, torch.float16), + ) + + +def make_transfer_metadata( + *, + engine_id: str = "engine-p", + te_rpc_port: int = 9000, + local_ip: str = "10.0.0.1", + handshake_port: int = 5000, + layer_names: list[str] | None = None, + group_indices: list[int] | None = None, + spec_indices: list[int] | None = None, + spec_block_sizes: list[int] | None = None, + base_addrs: list[list[int]] | None = None, + block_strides: list[list[int]] | None = None, + block_lens: list[list[int]] | None = None, + block_shapes: list[list[tuple[int, ...]]] | None = None, + block_size_scales: list[list[int]] | None = None, +) -> MooncakeTransferMetadata: + layer_names = layer_names or ["model.layers.0.self_attn"] + num_layers = len(layer_names) + return MooncakeTransferMetadata( + engine_id=engine_id, + te_rpc_port=te_rpc_port, + block_size=16, + num_blocks=32, + spec_block_sizes=spec_block_sizes or [16], + layer_names=layer_names, + group_indices=group_indices or [0] * num_layers, + spec_indices=spec_indices or [0] * num_layers, + kv_caches_base_addr=base_addrs or [[1000 + index * 1000] for index in range(num_layers)], + block_strides=block_strides or [[128] for _ in range(num_layers)], + block_lens=block_lens or [[128] for _ in range(num_layers)], + block_shapes=block_shapes or [[(1, 16, 4)] for _ in range(num_layers)], + block_size_scales=block_size_scales or [[1] for _ in range(num_layers)], + local_ip=local_ip, + handshake_port=handshake_port, + ) + + +def make_pp_metadata( + *, + layer_names: list[str] | None = None, + spec_indices: list[int] | None = None, + spec_block_sizes: list[int] | None = None, + block_shapes: list[list[tuple[int, ...]]] | None = None, + block_strides: list[list[int]] | None = None, + block_lens: list[list[int]] | None = None, + block_size_scales: list[list[int]] | None = None, + tp_base_addrs: dict[int, list[list[int]]] | None = None, +) -> MooncakePPTransferMetadata: + layer_names = layer_names or ["model.layers.0.self_attn"] + num_layers = len(layer_names) + tp_base_addrs = tp_base_addrs or {0: [[5000 + index * 1000] for index in range(num_layers)]} + return MooncakePPTransferMetadata( + block_size=16, + num_blocks=32, + spec_block_sizes=spec_block_sizes or [16], + layer_names=layer_names, + group_indices=[0] * num_layers, + spec_indices=spec_indices or [0] * num_layers, + block_strides=block_strides or [[128] for _ in range(num_layers)], + block_lens=block_lens or [[128] for _ in range(num_layers)], + block_shapes=block_shapes or [[(1, 16, 4)] for _ in range(num_layers)], + block_size_scales=block_size_scales or [[1] for _ in range(num_layers)], + metadata_by_tp_rank={ + tp_rank: MooncakeTPTransferMetadata( + te_rpc_port=9000 + tp_rank, + kv_caches_base_addr=base_addrs, + local_ip=f"10.0.0.{tp_rank + 1}", + handshake_port=5000 + tp_rank, + ) + for tp_rank, base_addrs in tp_base_addrs.items() + }, + ) + + +def make_metadata_groups( + *, + engine_id: str = "engine-p", + tp_size: int = 1, + pp_metadata: MooncakePPTransferMetadata | None = None, +) -> MooncakeTransferMetadataGroups: + return MooncakeTransferMetadataGroups( + engine_id=engine_id, + scheduler_host="10.0.0.10", + scheduler_port=6000, + pp_size=1, + pcp_size=1, + dcp_size=1, + tp_size=tp_size, + metadata_by_pp_rank={0: pp_metadata or make_pp_metadata()}, + ) + + +def make_request(**overrides: Any) -> SimpleNamespace: + values: dict[str, Any] = { + "request_id": "request-0", + "prompt_token_ids": list(range(32)), + "prompt_embeds": None, + "num_prompt_tokens": 32, + "kv_transfer_params": {}, + "status": "running", + "output_token_ids": [100], + "_all_token_ids": list(range(32)), + "max_tokens": 16, + } + values.update(overrides) + return SimpleNamespace(**values) + + +def make_blocks( + unhashed: tuple[list[int], ...] = ([10, 11],), + full: tuple[list[int], ...] = ([1, 2, 10, 11],), +) -> MagicMock: + blocks = MagicMock() + blocks.get_unhashed_block_ids_all_groups.return_value = unhashed + blocks.get_block_ids.return_value = full + return blocks diff --git a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py new file mode 100644 index 000000000000..01718f78a0c1 --- /dev/null +++ b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 + +from unittest.mock import MagicMock + +import torch +from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + UniformTypeKVCacheSpecs, +) + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker import ( + MooncakeBaseConnectorWorker, +) + +from .helpers import make_full_spec, make_sliding_spec + + +def test_build_spec_mappings_expands_uniform_group_by_layer_spec() -> None: + full = make_full_spec() + sliding = make_sliding_spec() + uniform = UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={"layer.0": full, "layer.1": sliding}, + ) + worker = MooncakeBaseConnectorWorker.__new__(MooncakeBaseConnectorWorker) + worker.kv_cache_config = MagicMock( + kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0", "layer.1"], kv_cache_spec=uniform)] + ) + + worker._build_kv_cache_spec_mappings() + + assert worker.kv_cache_specs == [full, sliding] + assert worker.layer_name_to_group_index == {"layer.0": 0, "layer.1": 0} + assert worker.layer_name_to_spec_index == {"layer.0": 0, "layer.1": 1} + + +def test_register_kv_caches_uses_config_order_and_publishes_tensor_metadata(monkeypatch) -> None: + spec = make_full_spec() + k_cache = torch.empty((4, 1, 16, 8), dtype=torch.float16) + v_cache = torch.empty((4, 1, 16, 8), dtype=torch.float16) + config = KVCacheConfig( + num_blocks=2, + kv_cache_tensors=[ + KVCacheTensor( + size=k_cache.nbytes + v_cache.nbytes, + shared_by=["layer.0"], + ) + ], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=spec)], + ) + worker = MooncakeBaseConnectorWorker.__new__(MooncakeBaseConnectorWorker) + worker.kv_cache_config = config + worker.engine_id = "engine-d" + worker.te_rpc_port = 9000 + worker.block_size = 16 + worker.side_channel_host = "10.0.0.1" + worker.handshake_port = 5000 + transfer_engine = MagicMock() + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.global_te", + transfer_engine, + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.validate_register_region_count", + MagicMock(), + ) + + worker.register_kv_caches({"layer.0": [k_cache, v_cache]}) + + metadata = worker.xfer_handshake_metadata + assert metadata is not None + assert metadata.layer_names == ["layer.0"] + assert metadata.kv_caches_base_addr == [[k_cache.data_ptr(), v_cache.data_ptr()]] + assert metadata.block_shapes == [[(1, 16, 8), (1, 16, 8)]] + assert metadata.block_size_scales == [[2, 2]] + assert metadata.block_strides == [[k_cache.stride(0) * 2, v_cache.stride(0) * 2]] + transfer_engine.register_buffer.assert_called_once() + + +def test_register_kv_caches_rejects_missing_and_unconfigured_layers() -> None: + spec = make_full_spec() + config = KVCacheConfig( + num_blocks=2, + kv_cache_tensors=[KVCacheTensor(size=64, shared_by=["layer.0"])], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=spec)], + ) + worker = MooncakeBaseConnectorWorker.__new__(MooncakeBaseConnectorWorker) + worker.kv_cache_config = config + worker.engine_id = "engine" + worker.te_rpc_port = 9000 + worker.block_size = 16 + worker.side_channel_host = "host" + worker.handshake_port = 5000 + + try: + worker.register_kv_caches({}) + except ValueError as error: + assert "No KV cache" in str(error) + else: + raise AssertionError("missing configured layer must fail") + + cache = torch.empty((2, 1, 16, 8), dtype=torch.float16) + try: + worker.register_kv_caches({"layer.0": cache, "unexpected": cache}) + except ValueError as error: + assert "absent from kv_cache_tensors" in str(error) + else: + raise AssertionError("unexpected layer must fail") diff --git a/tests/ut/kv_offload/mooncake_v2/test_connector.py b/tests/ut/kv_offload/mooncake_v2/test_connector.py new file mode 100644 index 000000000000..7262041da091 --- /dev/null +++ b/tests/ut/kv_offload/mooncake_v2/test_connector.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector import ( + MooncakeBaseConnector, + MooncakeConnector, + MooncakePullConnector, + MooncakePushConnector, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import MooncakeConnectorMetadata + + +def make_facade() -> MooncakeBaseConnector: + connector = MooncakeBaseConnector.__new__(MooncakeBaseConnector) + connector.connector_scheduler = MagicMock() + connector.connector_worker = MagicMock() + connector._connector_metadata = MooncakeConnectorMetadata() + return connector + + +def test_mooncake_connector_aliases_pull_connector() -> None: + assert MooncakeConnector is MooncakePullConnector + + +def test_required_layout_is_hnd_for_regular_attention_and_none_for_mla() -> None: + config = SimpleNamespace(model_config=SimpleNamespace(use_mla=False)) + assert MooncakeBaseConnector.get_required_kvcache_layout(config) == "HND" # type: ignore[arg-type] + + config.model_config.use_mla = True + assert MooncakeBaseConnector.get_required_kvcache_layout(config) is None # type: ignore[arg-type] + + config.model_config = None + assert MooncakeBaseConnector.get_required_kvcache_layout(config) is None # type: ignore[arg-type] + + +def test_facade_delegates_scheduler_methods() -> None: + connector = make_facade() + request = MagicMock() + blocks = MagicMock() + output = MagicMock() + connector.connector_scheduler.get_num_new_matched_tokens.return_value = (8, True) + connector.connector_scheduler.request_finished.return_value = (True, {"remote": True}) + + assert connector.get_num_new_matched_tokens(request, 4) == (8, True) + connector.update_state_after_alloc(request, blocks, 8) + connector.update_connector_output(output) + assert connector.request_finished(request, [1, 2]) == (True, {"remote": True}) + + connector.connector_scheduler.update_state_after_alloc.assert_called_once_with(request, blocks, 8) + connector.connector_scheduler.request_finished.assert_called_once_with(request, ([1, 2],)) + + +def test_facade_delegates_worker_methods_and_metadata() -> None: + connector = make_facade() + connector.connector_worker.get_finished.return_value = ({"sent"}, {"received"}) + connector.connector_worker.get_block_ids_with_load_errors.return_value = {10} + connector.connector_worker.xfer_handshake_metadata = MagicMock() + connector._connector_metadata.requests["request"] = MagicMock() + + connector.register_kv_caches({}) + connector.start_load_kv(MagicMock()) + + assert connector.get_finished(set()) == ({"sent"}, {"received"}) + assert connector.get_block_ids_with_load_errors() == {10} + assert connector.get_handshake_metadata() is connector.connector_worker.xfer_handshake_metadata + connector.connector_worker.start_load_kv.assert_called_once_with(connector._connector_metadata) + + +def test_pull_connector_selects_implementation_by_role() -> None: + config = MagicMock() + config.kv_transfer_config.engine_id = "engine" + config.kv_transfer_config.kv_role = "kv_producer" + kv_cache_config = MagicMock() + + def initialize_base(self, *_args, **_kwargs) -> None: + self.engine_id = "engine" + self.connector_scheduler = None + self.connector_worker = None + + with ( + patch.object(MooncakeBaseConnector, "__init__", initialize_base), + patch( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector.MooncakePullConnectorScheduler", + return_value=MagicMock(), + ) as scheduler_cls, + patch( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.connector.MooncakePullConnectorWorker", + return_value=MagicMock(), + ) as worker_cls, + ): + scheduler_connector = MooncakePullConnector(config, KVConnectorRole.SCHEDULER, kv_cache_config) + worker_connector = MooncakePullConnector(config, KVConnectorRole.WORKER, kv_cache_config) + + assert scheduler_connector.connector_scheduler is scheduler_cls.return_value + assert worker_connector.connector_worker is worker_cls.return_value + + +def test_push_connector_is_explicitly_unimplemented() -> None: + with pytest.raises(NotImplementedError, match="not implemented"): + MooncakePushConnector(MagicMock(), KVConnectorRole.SCHEDULER, MagicMock()) diff --git a/tests/ut/kv_offload/mooncake_v2/test_metadata.py b/tests/ut/kv_offload/mooncake_v2/test_metadata.py new file mode 100644 index 000000000000..69550eeb1605 --- /dev/null +++ b/tests/ut/kv_offload/mooncake_v2/test_metadata.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 + +import msgspec +import pytest + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( + MooncakeConnectorMetadata, + MooncakeTransferMetadataGroups, +) + +from .helpers import make_metadata_groups, make_transfer_metadata + + +def test_transfer_metadata_accepts_aligned_layer_fields() -> None: + metadata = make_transfer_metadata( + layer_names=["layer.0", "layer.1"], + base_addrs=[[1000, 2000], [3000]], + block_strides=[[128, 128], [256]], + block_lens=[[64, 64], [256]], + block_shapes=[[(1, 16, 2), (1, 16, 2)], [(2, 16, 4)]], + block_size_scales=[[1, 1], [2]], + ) + + assert metadata.layer_names == ["layer.0", "layer.1"] + assert metadata.kv_caches_base_addr[0] == [1000, 2000] + + +def test_transfer_metadata_rejects_misaligned_layer_count() -> None: + with pytest.raises(ValueError, match="group_indices.*1 layers, expected 2"): + make_transfer_metadata( + layer_names=["layer.0", "layer.1"], + group_indices=[0], + ) + + +def test_transfer_metadata_rejects_misaligned_tensor_count() -> None: + with pytest.raises(ValueError, match="layer 'layer.0'.*block_strides"): + make_transfer_metadata( + base_addrs=[[1000, 2000]], + block_strides=[[128]], + ) + + +def test_transfer_metadata_groups_msgpack_round_trip() -> None: + groups = make_metadata_groups() + + encoded = msgspec.msgpack.encode(groups) + decoded = msgspec.msgpack.decode(encoded, type=MooncakeTransferMetadataGroups) + + assert decoded == groups + assert decoded.metadata_by_pp_rank[0].metadata_by_tp_rank[0].te_rpc_port == 9000 + + +def test_connector_metadata_adds_complete_request() -> None: + metadata = MooncakeConnectorMetadata() + metadata.add_new_req( + request_id="request-d", + local_block_ids=([10, 11],), + local_full_block_ids=([1, 2, 10, 11],), + local_num_prompt_tokens=32, + num_external_tokens=16, + kv_transfer_params={ + "num_computed_tokens": 16, + "remote_block_ids": ([20, 21],), + "remote_host": "10.0.0.1", + "remote_port": 6000, + "remote_engine_id": "engine-p", + "remote_request_id": "request-p", + "remote_num_prompt_tokens": 31, + }, + ) + + request = metadata.requests["request-d"] + assert request.local_block_ids == ([10, 11],) + assert request.local_full_block_ids == ([1, 2, 10, 11],) + assert request.local_num_prompt_tokens == 32 + assert request.num_computed_tokens == 16 + assert request.remote_request_id == "request-p" diff --git a/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py b/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py new file mode 100644 index 000000000000..c8104043d2ab --- /dev/null +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py @@ -0,0 +1,301 @@ +# SPDX-License-Identifier: Apache-2.0 + +import threading +import time +from dataclasses import replace +from types import SimpleNamespace +from unittest.mock import MagicMock + +import msgspec +import pytest +import zmq +from vllm.v1.request import RequestStatus + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler import ( + MooncakeBaseConnectorScheduler, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( + MooncakeTransferMetadataGroups, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler import ( + ACK_MSG, + MooncakePullConnectorScheduler, + MooncakeSchedulerRecvingThread, + MooncakeSchedulerSendingThread, +) + +from .helpers import make_blocks, make_full_spec, make_mamba_spec, make_request, make_transfer_metadata + + +def make_sending_thread( + metadata: dict[int | tuple[int, int], object] | None = None, + *, + tp_size: int = 1, + pp_size: int = 1, +) -> MooncakeSchedulerSendingThread: + return MooncakeSchedulerSendingThread( + host="127.0.0.1", + port=6000, + engine_id="engine-p", + metadata=metadata or {0: make_transfer_metadata()}, # type: ignore[arg-type] + tp_size=tp_size, + pp_size=pp_size, + pcp_size=1, + dcp_size=1, + ready_event=threading.Event(), + ) + + +def make_pull_scheduler() -> MooncakePullConnectorScheduler: + scheduler = MooncakePullConnectorScheduler.__new__(MooncakePullConnectorScheduler) + scheduler.need_truncate = False + scheduler.num_speculative_tokens = 0 + scheduler.pcp_size = 1 + scheduler.dcp_size = 1 + scheduler.group_block_size = [16] + scheduler.group_unique_specs = [[make_full_spec()]] + scheduler.engine_id = "engine-p" + scheduler.side_channel_host = "10.0.0.10" + scheduler.side_channel_port = 6000 + scheduler._reqs_need_recv = {} + scheduler._reqs_need_send = {} + scheduler._reqs_in_batch = set() + scheduler._reqs_recv_info = {} + scheduler._sending_thread = None + scheduler._recving_thread = None + return scheduler + + +def test_sending_thread_merges_tp_private_and_pp_common_metadata() -> None: + first = make_transfer_metadata(te_rpc_port=9000, local_ip="10.0.0.1", base_addrs=[[1000]]) + second = make_transfer_metadata(te_rpc_port=9001, local_ip="10.0.0.2", base_addrs=[[2000]]) + + thread = make_sending_thread({0: first, 1: second}, tp_size=2) + decoded = msgspec.msgpack.decode(thread.encoded_metadata, type=MooncakeTransferMetadataGroups) + + pp_metadata = decoded.metadata_by_pp_rank[0] + assert pp_metadata.layer_names == first.layer_names + assert pp_metadata.block_shapes == first.block_shapes + assert pp_metadata.metadata_by_tp_rank[0].kv_caches_base_addr == [[1000]] + assert pp_metadata.metadata_by_tp_rank[1].kv_caches_base_addr == [[2000]] + assert pp_metadata.metadata_by_tp_rank[1].te_rpc_port == 9001 + + +def test_sending_thread_accepts_pp_aware_keys() -> None: + pp0 = make_transfer_metadata(layer_names=["pp0.layer"]) + pp1 = make_transfer_metadata(layer_names=["pp1.layer"], base_addrs=[[3000]]) + + thread = make_sending_thread({(0, 0): pp0, (1, 0): pp1}, pp_size=2) + decoded = msgspec.msgpack.decode(thread.encoded_metadata, type=MooncakeTransferMetadataGroups) + + assert decoded.metadata_by_pp_rank[0].layer_names == ["pp0.layer"] + assert decoded.metadata_by_pp_rank[1].layer_names == ["pp1.layer"] + + +def test_sending_thread_rejects_incomplete_or_inconsistent_workers() -> None: + metadata = make_transfer_metadata() + with pytest.raises(ValueError, match="incomplete TP ranks"): + make_sending_thread({0: metadata}, tp_size=2) + + mismatched = replace(metadata, block_lens=[[256]], te_rpc_port=9001) + with pytest.raises(ValueError, match="mismatch in.*block_lens"): + make_sending_thread({0: metadata, 1: mismatched}, tp_size=2) + + +def test_sending_thread_handles_early_and_normal_completion_once() -> None: + thread = make_sending_thread() + + thread._handle_finished_request("early") + assert thread.get_and_clear_finished_requests() == set() + thread.add_delayed_request("early", time.time()) + assert thread.get_and_clear_finished_requests() == {"early"} + + thread.add_delayed_request("normal", time.time()) + thread._handle_finished_request("normal") + thread._handle_finished_request("normal") + assert thread.get_and_clear_finished_requests() == {"normal"} + + +def test_sending_thread_force_frees_expired_request(monkeypatch: pytest.MonkeyPatch) -> None: + thread = make_sending_thread() + thread.add_delayed_request("expired", 10.0) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", + 5, + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.time.time", + MagicMock(return_value=20.0), + ) + + assert thread.get_and_clear_finished_requests() == {"expired"} + assert not thread.delayed_free_requests + + +def test_recving_thread_reuses_socket_after_ack_and_discards_it_on_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + thread = MooncakeSchedulerRecvingThread(threading.Event()) + socket = MagicMock() + thread._get_remote_socket = MagicMock(return_value=socket) # type: ignore[method-assign] + thread._return_remote_socket = MagicMock() # type: ignore[method-assign] + send = MagicMock() + recv = MagicMock(return_value=ACK_MSG) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.ensure_zmq_send", + send, + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.ensure_zmq_recv", + recv, + ) + + thread._send_done_recving("10.0.0.1", 6000, "request-p") + + path = "tcp://10.0.0.1:6000" + thread._get_remote_socket.assert_called_once_with(path) + thread._return_remote_socket.assert_called_once_with(path, socket) + socket.close.assert_not_called() + + recv.return_value = b"not-ack" + with pytest.raises(RuntimeError, match="Unexpected.*completion response"): + thread._send_done_recving("10.0.0.1", 6000, "request-p") + + socket.close.assert_called_once_with(linger=0) + assert thread._return_remote_socket.call_count == 1 + + +def test_recving_thread_socket_pool_creates_once_and_reuses_by_endpoint( + monkeypatch: pytest.MonkeyPatch, +) -> None: + thread = MooncakeSchedulerRecvingThread(threading.Event()) + context = MagicMock() + socket = MagicMock() + context_cls = MagicMock(return_value=context) + make_socket = MagicMock(return_value=socket) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.zmq.Context", + context_cls, + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.make_zmq_socket", + make_socket, + ) + path = "tcp://10.0.0.1:6000" + + first = thread._get_remote_socket(path) + thread._return_remote_socket(path, first) + second = thread._get_remote_socket(path) + + assert first is second is socket + context_cls.assert_called_once_with() + make_socket.assert_called_once_with(ctx=context, path=path, socket_type=zmq.REQ, bind=False) + assert socket.setsockopt.call_count == 2 + + +def test_base_scheduler_clips_attention_and_keeps_mamba_state_blocks() -> None: + scheduler = MooncakeBaseConnectorScheduler.__new__(MooncakeBaseConnectorScheduler) + scheduler.pcp_size = 1 + scheduler.dcp_size = 1 + scheduler.num_speculative_tokens = 2 + scheduler.group_block_size = [16, 16] + scheduler.group_unique_specs = [[make_full_spec()], [make_mamba_spec()]] + + result = scheduler._get_transfer_block_ids( + ([10, 11, 12, 13], [20, 21, 22, 23]), + prompt_len=17, + ) + + assert result == ([10, 11], [20, 21]) + + +def test_get_num_new_matched_tokens_records_local_prefix() -> None: + scheduler = make_pull_scheduler() + request = make_request(kv_transfer_params={"do_remote_prefill": True}) + + count, is_async = scheduler.get_num_new_matched_tokens(request, num_computed_tokens=16) + + assert (count, is_async) == (16, True) + assert request.kv_transfer_params["num_computed_tokens"] == 16 + + +def test_update_after_alloc_and_build_metadata() -> None: + scheduler = make_pull_scheduler() + request = make_request( + request_id="request-d", + kv_transfer_params={ + "do_remote_prefill": True, + "remote_block_ids": ([20, 21],), + "remote_engine_id": "engine-p", + "remote_host": "10.0.0.1", + "remote_port": 6000, + "remote_request_id": "request-p", + "remote_num_prompt_tokens": 31, + "num_computed_tokens": 16, + }, + ) + + scheduler.update_state_after_alloc(request, make_blocks(), num_external_tokens=16) + metadata = scheduler.build_connector_meta(MagicMock()) + + assert request.kv_transfer_params["do_remote_prefill"] is False + assert metadata.reqs_in_batch == {"request-d"} + assert metadata.requests["request-d"].local_block_ids == ([10, 11],) + assert metadata.requests["request-d"].local_full_block_ids == ([1, 2, 10, 11],) + assert scheduler._reqs_need_recv == {} + + +def test_zero_external_tokens_acknowledges_without_worker_transfer() -> None: + scheduler = make_pull_scheduler() + scheduler._recving_thread = MagicMock() + request = make_request( + kv_transfer_params={ + "do_remote_prefill": True, + "remote_block_ids": ([20],), + "remote_engine_id": "engine-p", + "remote_host": "10.0.0.1", + "remote_port": 6000, + "remote_request_id": "request-p", + } + ) + + scheduler.update_state_after_alloc(request, make_blocks(), num_external_tokens=0) + + assert scheduler._reqs_need_recv == {} + scheduler._recving_thread.add_request.assert_called_once_with("10.0.0.1", 6000, "request-p") + + +def test_request_finished_delays_blocks_and_builds_remote_params() -> None: + scheduler = make_pull_scheduler() + scheduler._sending_thread = MagicMock() + request = make_request( + request_id="request-p", + status=RequestStatus.FINISHED_LENGTH_CAPPED, + kv_transfer_params={"do_remote_decode": True}, + output_token_ids=[123], + ) + + delay_free, params = scheduler.request_finished(request, ([10, 11, 12],)) + + assert delay_free is True + assert params is not None + assert params["remote_block_ids"] == ([10, 11],) + assert params["remote_request_id"] == "request-p" + assert params["last_token_id"] == 123 + scheduler._sending_thread.add_delayed_request.assert_called_once() + + +def test_update_connector_output_routes_worker_completion_and_scheduler_ack() -> None: + scheduler = make_pull_scheduler() + scheduler._recving_thread = MagicMock() + scheduler._sending_thread = MagicMock() + scheduler._sending_thread.get_and_clear_finished_requests.return_value = {"request-p"} + scheduler._reqs_recv_info["request-d"] = ("10.0.0.1", 6000, "request-p") + scheduler._reqs_need_send["request-p"] = time.time() + output = SimpleNamespace(finished_recving={"request-d"}, finished_sending=None) + + scheduler.update_connector_output(output) # type: ignore[arg-type] + + scheduler._recving_thread.add_request.assert_called_once_with("10.0.0.1", 6000, "request-p") + assert output.finished_sending == {"request-p"} + assert scheduler._reqs_need_send == {} diff --git a/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py new file mode 100644 index 000000000000..6c73a661a585 --- /dev/null +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py @@ -0,0 +1,488 @@ +# SPDX-License-Identifier: Apache-2.0 + +import queue +import threading +from concurrent.futures import Future +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( + MooncakeConnectorMetadata, + ReqMeta, +) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker import ( + MooncakePullConnectorWorker, + MooncakePullRecvingThread, +) + +from .helpers import ( + make_full_spec, + make_mamba_spec, + make_metadata_groups, + make_pp_metadata, + make_sliding_spec, + make_transfer_metadata, +) + + +def make_thread(**overrides: object) -> MooncakePullRecvingThread: + thread = MooncakePullRecvingThread.__new__(MooncakePullRecvingThread) + thread.tp_rank = 0 + thread.tp_size = 1 + thread.dcp_rank = 0 + thread.dcp_size = 1 + thread.num_speculative_tokens = 0 + thread.layer_names = ["model.layers.0.self_attn"] + thread.group_indices = [0] + thread.spec_indices = [0] + thread.kv_cache_specs = [make_full_spec()] + thread.local_metadata = make_transfer_metadata() + thread.kv_caches_base_addr = [[1000]] + thread.block_strides = [[128]] + thread.block_lens = [[128]] + thread.block_shapes = [[(1, 16, 4)]] + thread.can_report_invalid_block_ids = True + thread.invalid_block_ids = set() + thread.invalid_block_ids_lock = threading.Lock() + thread.finished_requests = queue.SimpleQueue() + thread.request_queue = queue.Queue() + for name, value in overrides.items(): + setattr(thread, name, value) + return thread + + +def make_req_meta( + *, + local: tuple[list[int], ...] = ([10, 11],), + remote: tuple[list[int], ...] = ([20, 21],), + engine_id: str = "engine-p", +) -> ReqMeta: + return ReqMeta( + local_block_ids=local, + local_num_prompt_tokens=32, + num_external_tokens=16, + num_computed_tokens=0, + remote_block_ids=remote, + remote_host="10.0.0.1", + remote_port=6000, + remote_engine_id=engine_id, + remote_request_id="request-p", + remote_num_prompt_tokens=32, + local_full_block_ids=local, + ) + + +def test_remote_endpoint_requires_one_endpoint() -> None: + assert MooncakePullRecvingThread._get_remote_endpoint({("host", 1)}) == ("host", 1) + with pytest.raises(ValueError, match="must share one scheduler endpoint"): + MooncakePullRecvingThread._get_remote_endpoint({("a", 1), ("b", 2)}) + + +def test_result_queues_are_drained_atomically() -> None: + thread = make_thread() + thread.finished_requests.put("request-a") + thread.finished_requests.put("request-b") + thread.invalid_block_ids = {10, 11} + + assert thread.get_and_clear_finished_requests() == {"request-a", "request-b"} + assert thread.get_and_clear_finished_requests() == set() + assert thread.get_and_clear_invalid_block_ids() == {10, 11} + assert thread.get_and_clear_invalid_block_ids() == set() + + +def test_failed_request_marks_all_local_groups_for_non_hybrid_cache() -> None: + thread = make_thread() + request = make_req_meta(local=([10, 11], [20])) + + thread._mark_request_failed(request) + + assert thread.invalid_block_ids == {10, 11, 20} + + +def test_hybrid_cache_failure_is_not_reported_as_invalid_blocks() -> None: + thread = make_thread(can_report_invalid_block_ids=False) + + thread._mark_request_failed(make_req_meta(local=([10], [20]))) + + assert thread.invalid_block_ids == set() + + +def test_attention_tp_groups_cover_replication_split_and_dcp() -> None: + replicated = make_thread(tp_size=4, tp_rank=1) + assert replicated._get_attention_remote_tp_rank_groups(8, 1, 1, 4) == [[2, 3]] + + split = make_thread(tp_size=2, tp_rank=0) + assert split._get_attention_remote_tp_rank_groups(4, 1, 1, 8) == [[0], [1]] + + mla_dcp = make_thread(tp_size=4, tp_rank=2, dcp_size=4, dcp_rank=2) + assert mla_dcp._get_attention_remote_tp_rank_groups(8, 4, 8, 1) == [list(range(8))] + + unequal_dcp = make_thread(tp_size=8, tp_rank=3, dcp_size=2, dcp_rank=1) + assert unequal_dcp._get_attention_remote_tp_rank_groups(8, 2, 4, 4) == [[0, 1, 2, 3]] + + +def test_mamba_tp_groups_follow_non_replicated_tp_ratio() -> None: + local_smaller = make_thread(tp_size=2, tp_rank=1) + assert local_smaller._get_mamba_remote_tp_rank_groups(4) == [[2], [3]] + + local_larger = make_thread(tp_size=4, tp_rank=3) + assert local_larger._get_mamba_remote_tp_rank_groups(2) == [[1]] + + with pytest.raises(ValueError, match="integer ratio"): + local_larger._get_mamba_remote_tp_rank_groups(3) + + +def test_build_remote_layout_matches_layers_across_pp_ranks() -> None: + thread = make_thread( + layer_names=["layer.0", "layer.1"], + spec_indices=[0, 0], + kv_cache_specs=[make_full_spec()], + ) + thread._get_layer_remote_tp_rank_groups = MagicMock(return_value=[[0]]) # type: ignore[method-assign] + pp0 = make_pp_metadata(layer_names=["layer.0"]) + pp1 = make_pp_metadata(layer_names=["layer.1"], tp_base_addrs={0: [[6000]]}) + groups = make_metadata_groups() + groups = SimpleNamespace( + tp_size=1, + dcp_size=1, + metadata_by_pp_rank={0: pp0, 1: pp1}, + ) + + rank_groups, pairs = thread._build_remote_transfer_layout(groups) # type: ignore[arg-type] + + assert pairs == {0: [(0, 0)], 1: [(1, 0)]} + assert rank_groups == {0: {(0, 0): [[0]]}, 1: {(1, 0): [[0]]}} + + +def test_build_remote_layout_rejects_missing_local_layer() -> None: + thread = make_thread(layer_names=["layer.0", "missing"], spec_indices=[0, 0]) + thread._get_layer_remote_tp_rank_groups = MagicMock(return_value=[[0]]) # type: ignore[method-assign] + groups = make_metadata_groups(pp_metadata=make_pp_metadata(layer_names=["layer.0"])) + + with pytest.raises(ValueError, match="missing layers.*missing"): + thread._build_remote_transfer_layout(groups) + + +def test_expand_block_ids_and_round_robin_candidate_selection() -> None: + assert MooncakePullRecvingThread._expand_block_ids([3, 5], 2) == [6, 7, 10, 11] + assert MooncakePullRecvingThread._select_remote_tp_rank([2, 3, 4], 4) == 3 + + +def test_compute_full_attention_blocks_skips_remote_prefix_and_balances_replica() -> None: + thread = make_thread() + + result = thread._compute_group_block_ids( + request_id="request", + remote_tp_rank_groups=[[2, 3]], + remote_dcp_size=1, + spec_index=0, + local_block_size=16, + remote_block_size=16, + local_group_block_ids=[10, 11], + local_full_group_block_ids=[10, 11], + remote_group_block_ids=[20, 21, 22], + local_num_prompt_tokens=32, + remote_num_prompt_tokens=32, + num_computed_tokens=16, + local_block_size_scale=1, + remote_block_size_scale=1, + spec=make_full_spec(), + selection_index=1, + ) + + assert result == [(3, [10, 11], [21, 22])] + + +def test_compute_sliding_window_blocks_uses_unhashed_suffix() -> None: + thread = make_thread() + + result = thread._compute_group_block_ids( + "request", + [[0]], + 1, + 0, + 16, + 16, + [10, 11], + [0, 0, 10, 11], + [0, 0, 20, 21], + 64, + 64, + 32, + 1, + 1, + make_sliding_spec(), + 0, + ) + + assert result == [(0, [10, 11], [20, 21])] + + +def test_compute_mamba_state_selects_pre_speculative_local_block() -> None: + thread = make_thread(num_speculative_tokens=2) + + result = thread._compute_group_block_ids( + "request", + [[0]], + 1, + 0, + 16, + 16, + [1, 2, 3, 4], + [1, 2, 3, 4], + [8, 9, 10], + 32, + 31, + 0, + 1, + 1, + make_mamba_spec(), + 0, + ) + + assert result == [(0, [2], [10])] + + +def test_compute_dcp_blocks_maps_local_shards_to_remote_dcp_ranks() -> None: + thread = make_thread(tp_size=4, tp_rank=1, dcp_size=2, dcp_rank=1) + + result = thread._compute_group_block_ids( + "request", + [[0, 1, 2, 3]], + 4, + 0, + 16, + 16, + [10, 11], + [10, 11], + [20, 21], + 64, + 64, + 0, + 1, + 1, + make_full_spec(), + 0, + ) + + assert result == [(1, [10], [20]), (3, [11], [20])] + + +def test_compute_dcp_blocks_supports_different_logical_block_sizes() -> None: + thread = make_thread(dcp_size=2, dcp_rank=0) + + result = thread._compute_group_block_ids( + "request", + [[5]], + 1, + 0, + 16, + 32, + [10, 11], + [10, 11], + [20, 21], + 65, + 64, + 0, + 1, + 2, + make_full_spec(), + 0, + ) + + assert result == [(5, [10, 11], [40, 42])] + + +def test_transfer_bucket_reuses_block_mapping_for_layers_of_same_spec() -> None: + thread = make_thread( + layer_names=["layer.0", "layer.1"], + group_indices=[0, 0], + spec_indices=[0, 0], + block_shapes=[[(1, 16, 4)], [(1, 16, 4)]], + block_strides=[[128], [128]], + block_lens=[[128], [128]], + kv_caches_base_addr=[[1000], [2000]], + ) + thread.local_metadata = make_transfer_metadata( + layer_names=["layer.0", "layer.1"], + spec_indices=[0, 0], + spec_block_sizes=[16], + ) + thread._compute_group_block_ids = MagicMock(return_value=[(0, [10], [20])]) # type: ignore[method-assign] + remote = make_pp_metadata( + layer_names=["layer.0", "layer.1"], + spec_indices=[0, 0], + spec_block_sizes=[16], + tp_base_addrs={0: [[5000], [6000]]}, + ) + + buckets, request_ids = thread._build_transfer_block_buckets( + remote, + [(0, 0), (1, 1)], + {(0, 0): [[0]], (1, 1): [[0]]}, + 1, + {"request": make_req_meta()}, + {}, + ) + + thread._compute_group_block_ids.assert_called_once() + assert set(buckets[0][0]) == {(0, 0), (1, 1)} + assert request_ids == {0: {"request"}} + + +def test_attention_address_generation_handles_partial_head_overlap() -> None: + thread = make_thread( + tp_size=1, + tp_rank=0, + kv_cache_specs=[make_full_spec(num_kv_heads=4)], + block_shapes=[[(4, 16, 8)]], + block_lens=[[128]], + block_strides=[[128]], + kv_caches_base_addr=[[1000]], + ) + remote = make_pp_metadata( + block_shapes=[[(2, 16, 8)]], + block_lens=[[64]], + block_strides=[[64]], + tp_base_addrs={0: [[5000]], 1: [[6000]]}, + ) + src: list[int] = [] + dst: list[int] = [] + lengths: list[int] = [] + + thread._append_spec_transfer_addresses( + 0, + remote_tp_rank=1, + remote_tp_size=2, + remote_dcp_size=1, + transfer_entries_by_layer={(0, 0): [("request", [1], [2])]}, + remote_metadata=remote, + src_list=src, + dst_list=dst, + length_list=lengths, + ) + + assert src == [1000 + 128 + 64] + assert dst == [6000 + 2 * 64] + assert lengths == [64] + + +def test_mamba_equal_tp_address_generation_transfers_each_cache() -> None: + thread = make_thread( + kv_caches_base_addr=[[1000, 2000]], + block_shapes=[[(3, 16), (2, 4, 4)]], + block_strides=[[128, 256]], + block_lens=[[96, 128]], + ) + remote = make_pp_metadata( + block_shapes=[[(3, 16), (2, 4, 4)]], + block_strides=[[128, 256]], + block_lens=[[96, 128]], + tp_base_addrs={0: [[5000, 6000]]}, + ) + src: list[int] = [] + dst: list[int] = [] + lengths: list[int] = [] + + thread._append_mamba_transfer_addresses( + make_mamba_spec(), + 0, + 1, + {(0, 0): [("request", [1], [2])]}, + remote, + src, + dst, + lengths, + ) + + assert src == [1128, 2256] + assert dst == [5256, 6512] + assert lengths == [96, 128] + + +def test_execute_bucket_calls_mooncake_and_raises_on_negative_return() -> None: + thread = make_thread(engine=MagicMock()) + thread._append_spec_transfer_addresses = MagicMock( # type: ignore[method-assign] + side_effect=lambda **kwargs: ( + kwargs["src_list"].append(1000), + kwargs["dst_list"].append(2000), + kwargs["length_list"].append(128), + ) + ) + remote = make_pp_metadata() + entries = {0: {(0, 0): [("request", [1], [2])]}} + + thread.engine.batch_transfer_sync_read.return_value = 0 + thread._execute_tp_transfer_bucket(0, 0, 1, 1, remote, entries) + thread.engine.batch_transfer_sync_read.assert_called_once_with("10.0.0.1:9000", [1000], [2000], [128]) + + thread.engine.batch_transfer_sync_read.return_value = -1 + with pytest.raises(RuntimeError, match="transfer failed"): + thread._execute_tp_transfer_bucket(0, 0, 1, 1, remote, entries) + + +@pytest.mark.parametrize(("can_report", "expected_failed"), [(True, {"request-b"}), (False, set())]) +def test_handle_requests_attributes_failed_tp_to_affected_requests( + can_report: bool, + expected_failed: set[str], +) -> None: + thread = make_thread(can_report_invalid_block_ids=can_report) + remote_pp = make_pp_metadata(tp_base_addrs={0: [[5000]], 1: [[6000]]}) + remote_groups = make_metadata_groups(tp_size=2, pp_metadata=remote_pp) + thread._get_remote_metadata = MagicMock(return_value=remote_groups) # type: ignore[method-assign] + thread.remote_tp_rank_groups = {"engine-p": {0: {(0, 0): [[0], [1]]}}} + thread.remote_layer_index_pairs = {"engine-p": {0: [(0, 0)]}} + thread._build_transfer_block_buckets = MagicMock( # type: ignore[method-assign] + return_value=( + { + 0: {0: {(0, 0): [("request-a", [1], [2])]}}, + 1: {0: {(0, 0): [("request-b", [3], [4])]}}, + }, + {0: {"request-a"}, 1: {"request-b"}}, + ) + ) + + def submit(_func: object, _pp: int, tp_rank: int, *_args: object) -> Future[None]: + future: Future[None] = Future() + if tp_rank == 1: + future.set_exception(RuntimeError("remote TP failed")) + else: + future.set_result(None) + return future + + thread.executor = MagicMock() + thread.executor.submit.side_effect = submit + requests = {"request-a": make_req_meta(), "request-b": make_req_meta()} + + assert thread._handle_requests("engine-p", "10.0.0.1", 6000, requests) == expected_failed + + +def test_connector_worker_groups_start_load_by_remote_engine() -> None: + worker = MooncakePullConnectorWorker.__new__(MooncakePullConnectorWorker) + worker.kv_transfer_config = SimpleNamespace(is_kv_consumer=True) + worker._recving_thread = MagicMock() + metadata = MooncakeConnectorMetadata() + metadata.requests = { + "a": make_req_meta(engine_id="engine-1"), + "b": make_req_meta(engine_id="engine-2"), + "c": make_req_meta(engine_id="engine-1"), + } + + worker.start_load_kv(metadata) + + calls = worker._recving_thread.add_requests.call_args_list + assert len(calls) == 2 + grouped = {call.args[0]: set(call.args[1]) for call in calls} + assert grouped == {"engine-1": {"a", "c"}, "engine-2": {"b"}} + + +def test_connector_worker_exposes_finished_and_invalid_blocks() -> None: + worker = MooncakePullConnectorWorker.__new__(MooncakePullConnectorWorker) + worker._recving_thread = MagicMock() + worker._recving_thread.get_and_clear_finished_requests.return_value = {"request"} + worker._recving_thread.get_and_clear_invalid_block_ids.return_value = {10, 11} + + assert worker.get_finished() == (set(), {"request"}) + assert worker.get_block_ids_with_load_errors() == {10, 11} diff --git a/tests/ut/kv_offload/mooncake_v2/test_stats.py b/tests/ut/kv_offload/mooncake_v2/test_stats.py new file mode 100644 index 000000000000..26876e1a15e2 --- /dev/null +++ b/tests/ut/kv_offload/mooncake_v2/test_stats.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.stats import MooncakeKVConnectorStats + + +def test_stats_record_reduce_and_reset() -> None: + stats = MooncakeKVConnectorStats(data={}) + stats.record_transfer(0.01, 2**20) + stats.record_transfer(0.03, 3 * 2**20) + stats.record_failed_transfer() + + reduced = stats.reduce() + + assert reduced["Num successful transfers"] == 2 + assert reduced["Num failed transfers"] == 1 + assert reduced["Avg xfer time (ms)"] == 20.0 + assert reduced["Avg MB per transfer"] == 2.0 + assert reduced["Throughput (MB/s)"] == 100.0 + + previous = stats.clone_and_reset() + assert previous.num_successful_transfers == 2 + assert stats.is_empty() + + +def test_stats_aggregate_ignores_empty_and_merges_nonempty() -> None: + stats = MooncakeKVConnectorStats(data={}) + empty = MooncakeKVConnectorStats(data={}) + other = MooncakeKVConnectorStats(data={}) + other.record_transfer(0.1, 1024) + other.record_failed_transfer() + + assert stats.aggregate(empty) is stats + stats.aggregate(other) + + assert stats.data["transfer_duration"] == [0.1] + assert stats.data["bytes_transferred"] == [1024] + assert stats.data["num_failed_transfers"] == [1] diff --git a/tests/ut/kv_offload/mooncake_v2/test_utils.py b/tests/ut/kv_offload/mooncake_v2/test_utils.py new file mode 100644 index 000000000000..0df3fc977abf --- /dev/null +++ b/tests/ut/kv_offload/mooncake_v2/test_utils.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +import zmq + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils import ( + SizedDict, + as_kv_cache_tensors, + collect_configured_register_regions, + ensure_zmq_recv, + ensure_zmq_send, + group_concurrent_contiguous, + split_if_not_byte_contiguous, + string_to_int64_hash, + zmq_ctx, +) + + +def test_as_kv_cache_tensors_normalizes_supported_inputs() -> None: + first = torch.empty(1) + second = torch.empty(1) + + assert as_kv_cache_tensors(first) == (first,) + assert as_kv_cache_tensors([first, second]) == (first, second) + with pytest.raises(TypeError, match="must be a tensor"): + as_kv_cache_tensors([first, object()]) + + +def test_sized_dict_evicts_oldest_entry_and_initializes_missing_value() -> None: + values: SizedDict[str, dict[int, list[int]]] = SizedDict(max_size=2) + values["a"] = {0: [1]} + values["b"] = {0: [2]} + values["c"] = {0: [3]} + + assert list(values) == ["b", "c"] + assert values["missing"] == {} + assert list(values) == ["c", "missing"] + + +def test_collect_register_regions_deduplicates_shared_storage() -> None: + backing = torch.empty(128, dtype=torch.uint8) + config = SimpleNamespace( + kv_cache_tensors=[SimpleNamespace(shared_by=["layer.0", "layer.1"], offset=0, size=backing.nbytes)] + ) + + regions = collect_configured_register_regions( + config, + {"layer.0": backing[:64], "layer.1": backing[64:]}, + ) + + assert regions.ptrs == [backing.data_ptr()] + assert regions.lengths == [backing.nbytes] + assert regions.logical_tensor_count == 1 + + +def test_collect_register_regions_handles_independent_storages() -> None: + first = torch.empty(32, dtype=torch.uint8) + second = torch.empty(48, dtype=torch.uint8) + config = SimpleNamespace(kv_cache_tensors=[SimpleNamespace(shared_by=["layer.0", "layer.1"], offset=0, size=80)]) + + regions = collect_configured_register_regions(config, {"layer.0": first, "layer.1": second}) + + assert set(zip(regions.ptrs, regions.lengths)) == { + (first.data_ptr(), first.nbytes), + (second.data_ptr(), second.nbytes), + } + + +def test_group_concurrent_contiguous_honors_byte_strides() -> None: + src_groups, dst_groups = group_concurrent_contiguous( + [1, 2, 4], + [10, 11, 13], + src_block_stride=64, + dst_block_stride=64, + block_len=64, + ) + + assert src_groups == [[1, 2], [4]] + assert dst_groups == [[10, 11], [13]] + + +def test_split_if_not_byte_contiguous_preserves_fast_path_identity() -> None: + src = [[1, 2]] + dst = [[3, 4]] + + result_src, result_dst = split_if_not_byte_contiguous(src, dst, 64, 64, 64) + + assert result_src is src + assert result_dst is dst + + +def test_string_hash_is_stable_and_distinguishes_inputs() -> None: + assert string_to_int64_hash("request") == string_to_int64_hash("request") + assert string_to_int64_hash("request") != string_to_int64_hash("other") + + +def test_zmq_ctx_rejects_unknown_socket_type() -> None: + with pytest.raises(ValueError, match="Unexpected socket type"), zmq_ctx(-1, "tcp://127.0.0.1:1"): + pass + + +def test_zmq_ctx_destroys_context(monkeypatch: pytest.MonkeyPatch) -> None: + context = MagicMock() + socket = MagicMock() + monkeypatch.setattr(zmq, "Context", MagicMock(return_value=context)) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils.make_zmq_socket", + MagicMock(return_value=socket), + ) + + with zmq_ctx(zmq.REQ, "tcp://127.0.0.1:1") as result: + assert result is socket + + context.destroy.assert_called_once_with(linger=0) + + +def test_zmq_retry_helpers_succeed_after_transient_error(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils.time.sleep", MagicMock()) + socket = MagicMock() + socket.send.side_effect = [zmq.ZMQError("transient"), None] + socket.recv.side_effect = [zmq.ZMQError("transient"), b"response"] + + ensure_zmq_send(socket, b"request", "tcp://peer", max_retries=2) + assert ensure_zmq_recv(socket, "tcp://peer", max_retries=2) == b"response" + assert socket.send.call_count == 2 + assert socket.recv.call_count == 2 + + +def test_zmq_retry_helpers_raise_after_limit(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils.time.sleep", MagicMock()) + socket = MagicMock() + socket.send.side_effect = zmq.ZMQError("failed") + socket.recv.side_effect = zmq.ZMQError("failed") + + with pytest.raises(RuntimeError, match="Failed to send"): + ensure_zmq_send(socket, b"request", "tcp://peer", max_retries=2) + with pytest.raises(RuntimeError, match="Failed to receive"): + ensure_zmq_recv(socket, "tcp://peer", max_retries=2) From e414ca6ec129b130f2d9ab59f45b771e50fdea34 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Fri, 21 Aug 2026 15:11:03 +0800 Subject: [PATCH 36/48] fix UT Signed-off-by: nwpu-zxr --- tests/ut/kv_offload/mooncake_v2/test_metadata.py | 1 + tests/ut/kv_offload/mooncake_v2/test_pull_worker.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/tests/ut/kv_offload/mooncake_v2/test_metadata.py b/tests/ut/kv_offload/mooncake_v2/test_metadata.py index 69550eeb1605..892118eaa5ee 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_metadata.py +++ b/tests/ut/kv_offload/mooncake_v2/test_metadata.py @@ -36,6 +36,7 @@ def test_transfer_metadata_rejects_misaligned_layer_count() -> None: def test_transfer_metadata_rejects_misaligned_tensor_count() -> None: with pytest.raises(ValueError, match="layer 'layer.0'.*block_strides"): make_transfer_metadata( + layer_names=["layer.0"], base_addrs=[[1000, 2000]], block_strides=[[128]], ) 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 6c73a661a585..044012685ead 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py @@ -43,6 +43,7 @@ def make_thread(**overrides: object) -> MooncakePullRecvingThread: thread.block_strides = [[128]] thread.block_lens = [[128]] thread.block_shapes = [[(1, 16, 4)]] + thread.block_size_scales = [[1]] thread.can_report_invalid_block_ids = True thread.invalid_block_ids = set() thread.invalid_block_ids_lock = threading.Lock() @@ -303,6 +304,7 @@ def test_transfer_bucket_reuses_block_mapping_for_layers_of_same_spec() -> None: block_shapes=[[(1, 16, 4)], [(1, 16, 4)]], block_strides=[[128], [128]], block_lens=[[128], [128]], + block_size_scales=[[1], [1]], kv_caches_base_addr=[[1000], [2000]], ) thread.local_metadata = make_transfer_metadata( From f000e2937fccb9daa067128b52dc5d69a4254de9 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Fri, 21 Aug 2026 15:32:44 +0800 Subject: [PATCH 37/48] add more UT Signed-off-by: nwpu-zxr --- .../mooncake_v2/test_pull_scheduler.py | 149 +++- .../mooncake_v2/test_pull_worker.py | 691 ++++++++++++++++++ tests/ut/kv_offload/mooncake_v2/test_utils.py | 10 +- 3 files changed, 843 insertions(+), 7 deletions(-) diff --git a/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py b/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py index c8104043d2ab..74cb3f4b5c42 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py @@ -4,10 +4,11 @@ import time from dataclasses import replace from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, call import msgspec import pytest +import torch import zmq from vllm.v1.request import RequestStatus @@ -189,7 +190,12 @@ def test_recving_thread_socket_pool_creates_once_and_reuses_by_endpoint( assert first is second is socket context_cls.assert_called_once_with() - make_socket.assert_called_once_with(ctx=context, path=path, socket_type=zmq.REQ, bind=False) + make_socket.assert_called_once_with( + ctx=context, + path=path, + socket_type=zmq.REQ, # type: ignore[attr-defined] + bind=False, + ) assert socket.setsockopt.call_count == 2 @@ -299,3 +305,142 @@ def test_update_connector_output_routes_worker_completion_and_scheduler_ack() -> scheduler._recving_thread.add_request.assert_called_once_with("10.0.0.1", 6000, "request-p") assert output.finished_sending == {"request-p"} assert scheduler._reqs_need_send == {} + + +class _StopLoop(BaseException): + """Stop an otherwise unbounded thread loop after one test iteration.""" + + +def test_base_scheduler_detects_state_and_compressed_prefill_truncation() -> None: + scheduler = MooncakeBaseConnectorScheduler.__new__(MooncakeBaseConnectorScheduler) + scheduler.vllm_config = SimpleNamespace(model_config=SimpleNamespace(hf_config=SimpleNamespace())) + scheduler.group_unique_specs = [[make_full_spec()]] + assert scheduler._needs_prefill_token_truncation() is False + + scheduler.vllm_config.model_config.hf_config.compress_ratios = [2] + assert scheduler._needs_prefill_token_truncation() is True + + scheduler.vllm_config.model_config.hf_config.compress_ratios = None + scheduler.group_unique_specs = [[make_full_spec()], [make_mamba_spec()]] + assert scheduler._needs_prefill_token_truncation() is True + + +def test_truncate_request_for_prefill_is_idempotent_for_token_ids() -> None: + scheduler = MooncakeBaseConnectorScheduler.__new__(MooncakeBaseConnectorScheduler) + scheduler.need_truncate = True + request = make_request( + prompt_token_ids=[1, 2, 3], + num_prompt_tokens=3, + _all_token_ids=[1, 2, 3], + kv_transfer_params={}, + max_tokens=8, + ) + + scheduler._truncate_request_for_prefill(request) + scheduler._truncate_request_for_prefill(request) + + assert request.prompt_token_ids == [1, 2] + assert request._all_token_ids == [1, 2] + assert request.num_prompt_tokens == 2 + assert request.max_tokens == 1 + assert request.kv_transfer_params["_p_side_truncated"] is True + + +def test_truncate_request_for_prefill_supports_prompt_embeddings() -> None: + scheduler = MooncakeBaseConnectorScheduler.__new__(MooncakeBaseConnectorScheduler) + scheduler.need_truncate = True + prompt_embeds = torch.arange(12).reshape(3, 4) + request = make_request( + prompt_token_ids=None, + prompt_embeds=prompt_embeds, + num_prompt_tokens=3, + _all_token_ids=[1, 2, 3], + kv_transfer_params={}, + ) + + scheduler._truncate_request_for_prefill(request) + + assert torch.equal(request.prompt_embeds, prompt_embeds[:2]) + assert request.num_prompt_tokens == 2 + + +def test_sending_busy_loop_serves_metadata_and_completion_ack() -> None: + thread = make_sending_thread() + thread._handle_finished_request = MagicMock() # type: ignore[method-assign] + encoder = msgspec.msgpack.Encoder() + socket = MagicMock() + socket.recv_multipart.side_effect = [ + [b"worker", b"", encoder.encode((b"get_meta_msg",))], + [b"worker", b"", encoder.encode((b"done_recving_msg", "request-p"))], + _StopLoop(), + ] + + with pytest.raises(_StopLoop): + thread._run_busy_loop(socket) + + assert socket.send_multipart.call_args_list == [ + call((b"worker", b"", thread.encoded_metadata)), + call((b"worker", b"", ACK_MSG)), + ] + thread._handle_finished_request.assert_called_once_with("request-p") + + +def test_recving_run_requeues_failed_completion() -> None: + ready_event = threading.Event() + thread = MooncakeSchedulerRecvingThread(ready_event) + request = ("10.0.0.1", 6000, "request-p") + request_queue = MagicMock() + request_queue.get.side_effect = [request, _StopLoop()] + thread.request_queue = request_queue + thread._send_done_recving = MagicMock(side_effect=RuntimeError("network error")) # type: ignore[method-assign] + + with pytest.raises(_StopLoop): + thread.run() + + assert ready_event.is_set() + request_queue.put.assert_called_once_with(request) + request_queue.task_done.assert_called_once_with() + + +def test_set_worker_metadata_starts_only_one_producer_sending_thread( + monkeypatch: pytest.MonkeyPatch, +) -> None: + scheduler = make_pull_scheduler() + scheduler.kv_transfer_config = SimpleNamespace(is_kv_producer=True) + scheduler.side_channel_host = "127.0.0.1" + scheduler.side_channel_port = 6000 + scheduler.tp_size = 1 + scheduler.pp_size = 1 + scheduler.pcp_size = 1 + scheduler.dcp_size = 1 + created: list[MagicMock] = [] + + def make_fake_thread(*args, **_kwargs): + fake = MagicMock() + ready_event = args[-1] + fake.start.side_effect = ready_event.set + fake.is_alive.return_value = True + created.append(fake) + return fake + + thread_cls = MagicMock(side_effect=make_fake_thread) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.MooncakeSchedulerSendingThread", + thread_cls, + ) + metadata = {0: make_transfer_metadata()} + + scheduler.set_xfer_handshake_metadata_from_workers(metadata) + scheduler.set_xfer_handshake_metadata_from_workers(metadata) + + thread_cls.assert_called_once() + created[0].start.assert_called_once_with() + + +def test_set_worker_metadata_is_ignored_on_consumer() -> None: + scheduler = make_pull_scheduler() + scheduler.kv_transfer_config = SimpleNamespace(is_kv_producer=False) + + scheduler.set_xfer_handshake_metadata_from_workers({0: make_transfer_metadata()}) + + assert scheduler._sending_thread is None 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 044012685ead..69519820ecbd 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py @@ -3,19 +3,26 @@ import queue import threading from concurrent.futures import Future +from dataclasses import replace from types import SimpleNamespace from unittest.mock import MagicMock +import msgspec import pytest +import torch +from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +from vllm.v1.kv_cache_interface import MLAAttentionSpec from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( MooncakeConnectorMetadata, + MooncakeTransferMetadataGroups, ReqMeta, ) 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 SizedDict from .helpers import ( make_full_spec, @@ -488,3 +495,687 @@ def test_connector_worker_exposes_finished_and_invalid_blocks() -> None: assert worker.get_finished() == (set(), {"request"}) assert worker.get_block_ids_with_load_errors() == {10, 11} + + +class _StopWorkerLoop(BaseException): + """Stop the worker's unbounded queue loop after one test iteration.""" + + +def test_get_remote_metadata_fetches_once_and_populates_layout_cache( + monkeypatch: pytest.MonkeyPatch, +) -> None: + thread = make_thread() + thread.encoder = msgspec.msgpack.Encoder() + thread.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) + thread.remote_metadata = SizedDict() + thread.remote_tp_rank_groups = SizedDict() + thread.remote_layer_index_pairs = SizedDict() + expected = make_metadata_groups() + socket = MagicMock() + context = MagicMock() + context.__enter__.return_value = socket + send = MagicMock() + recv = MagicMock(return_value=msgspec.msgpack.encode(expected)) + thread._build_remote_transfer_layout = MagicMock( # type: ignore[method-assign] + return_value=({0: {(0, 0): [[0]]}}, {0: [(0, 0)]}) + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.zmq_ctx", + MagicMock(return_value=context), + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_send", + send, + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_recv", + recv, + ) + + first = thread._get_remote_metadata("engine-p", "10.0.0.1", 6000) + second = thread._get_remote_metadata("engine-p", "10.0.0.1", 6000) + + assert first is second + send.assert_called_once() + recv.assert_called_once() + thread._build_remote_transfer_layout.assert_called_once_with(expected) + assert thread.remote_tp_rank_groups["engine-p"] == {0: {(0, 0): [[0]]}} + assert thread.remote_layer_index_pairs["engine-p"] == {0: [(0, 0)]} + + +def test_get_remote_metadata_rejects_empty_response(monkeypatch: pytest.MonkeyPatch) -> None: + thread = make_thread() + thread.encoder = msgspec.msgpack.Encoder() + thread.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) + thread.remote_metadata = SizedDict() + socket = MagicMock() + context = MagicMock() + context.__enter__.return_value = socket + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.zmq_ctx", + MagicMock(return_value=context), + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_send", + MagicMock(), + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_recv", + MagicMock(return_value=b""), + ) + + with pytest.raises(RuntimeError, match="returned no transfer metadata"): + thread._get_remote_metadata("engine-p", "10.0.0.1", 6000) + + +@pytest.mark.parametrize( + ("metadata", "expected_error"), + [ + (replace(make_metadata_groups(), engine_id="wrong"), "engine ID mismatch"), + (replace(make_metadata_groups(), pcp_size=2), "requires remote pcp_size=1"), + (replace(make_metadata_groups(), metadata_by_pp_rank={}), "no PP metadata"), + ( + replace( + make_metadata_groups(), + metadata_by_pp_rank={0: replace(make_pp_metadata(), metadata_by_tp_rank={})}, + ), + "no TP metadata", + ), + ( + make_metadata_groups( + pp_metadata=make_pp_metadata(tp_base_addrs={2: [[5000]]}), + ), + "invalid TP ranks", + ), + ], +) +def test_validate_remote_metadata_rejects_invalid_topology( + metadata: MooncakeTransferMetadataGroups, + expected_error: str, +) -> None: + with pytest.raises(ValueError, match=expected_error): + MooncakePullRecvingThread._validate_remote_metadata(metadata, "engine-p") + + +def test_worker_run_marks_only_failed_requests_and_finishes_the_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + thread = make_thread(device=0, ready_event=threading.Event()) + requests = { + "request-a": make_req_meta(local=([10],)), + "request-b": make_req_meta(local=([20],)), + } + request_queue = MagicMock() + request_queue.get.side_effect = [("engine-p", requests), _StopWorkerLoop()] + thread.request_queue = request_queue + thread._handle_requests = MagicMock(return_value={"request-b"}) # type: ignore[method-assign] + set_device = MagicMock() + monkeypatch.setattr(torch.npu, "set_device", set_device) + + with pytest.raises(_StopWorkerLoop): + thread.run() + + set_device.assert_called_once_with(0) + assert thread.ready_event.is_set() + assert thread.invalid_block_ids == {20} + assert thread.get_and_clear_finished_requests() == {"request-a", "request-b"} + request_queue.task_done.assert_called_once_with() + + +def test_whole_block_mla_address_generation_uses_independent_stride() -> None: + spec = MLAAttentionSpec(block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float16) + thread = make_thread( + kv_cache_specs=[spec], + kv_caches_base_addr=[[1000]], + block_strides=[[256]], + block_lens=[[128]], + ) + remote = make_pp_metadata( + block_strides=[[512]], + block_lens=[[128]], + tp_base_addrs={0: [[5000]]}, + ) + src: list[int] = [] + dst: list[int] = [] + lengths: list[int] = [] + + thread._append_spec_transfer_addresses( + 0, + 0, + 1, + 1, + {(0, 0): [("request", [1, 2], [3, 4])]}, + remote, + src, + dst, + lengths, + ) + + assert src == [1256, 1512] + assert dst == [6536, 7048] + assert lengths == [128, 128] + + +def test_mamba_unequal_tp_slices_conv_projections_and_state() -> None: + spec = SimpleNamespace( + dtypes=(torch.float16, torch.float16), + mamba_type=MambaAttentionBackendEnum.MAMBA2, + ) + thread = make_thread( + tp_size=1, + tp_rank=0, + kv_caches_base_addr=[[1000, 2000]], + block_shapes=[[(3, 16), (2, 4, 4)]], + block_strides=[[96, 64]], + block_lens=[[96, 64]], + ) + remote = make_pp_metadata( + block_shapes=[[(3, 8), (1, 4, 4)]], + block_strides=[[48, 32]], + block_lens=[[48, 32]], + tp_base_addrs={0: [[5000, 6000]], 1: [[7000, 8000]]}, + ) + src: list[int] = [] + dst: list[int] = [] + lengths: list[int] = [] + + thread._append_mamba_transfer_addresses( + spec, # type: ignore[arg-type] + remote_tp_rank=1, + remote_tp_size=2, + transfer_entries_by_layer={(0, 0): [("request", [1], [2])]}, + remote_metadata=remote, + src_list=src, + dst_list=dst, + length_list=lengths, + ) + + assert src == [1104, 1136, 1168, 1116, 1148, 1180, 1124, 1156, 1188, 2096] + assert dst == [7096, 7112, 7128, 7104, 7120, 7136, 7108, 7124, 7140, 8064] + assert lengths == [8, 8, 8, 4, 4, 4, 4, 4, 4, 32] + + +@pytest.mark.parametrize( + ( + "local_dcp_size", + "local_dcp_rank", + "remote_dcp_size", + "candidate_ranks", + "remote_blocks", + "expected", + ), + [ + (1, 0, 4, [0, 1, 2, 3], [20], [(0, [10], [20]), (1, [11], [20])]), + (4, 2, 1, [0], list(range(20, 28)), [(0, [10, 11], [22, 26])]), + (4, 2, 4, [0, 1, 2, 3], [20, 21], [(2, [10, 11], [20, 21])]), + ], +) +def test_compute_same_block_size_dcp_matrix( + local_dcp_size: int, + local_dcp_rank: int, + remote_dcp_size: int, + candidate_ranks: list[int], + remote_blocks: list[int], + expected: list[tuple[int, list[int], list[int]]], +) -> None: + thread = make_thread( + tp_size=max(local_dcp_size, 1), + tp_rank=local_dcp_rank, + dcp_size=local_dcp_size, + dcp_rank=local_dcp_rank, + ) + + result = thread._compute_group_block_ids( + "request", + [candidate_ranks], + remote_dcp_size, + 0, + 16, + 16, + [10, 11], + [10, 11], + remote_blocks, + 128, + 128, + 0, + 1, + 1, + make_full_spec(), + 0, + ) + + assert result == expected + + +def test_compute_dcp_stops_when_remote_prompt_has_fewer_blocks() -> None: + thread = make_thread(tp_size=4, tp_rank=2, dcp_size=4, dcp_rank=2) + + result = thread._compute_group_block_ids( + "request", + [[0]], + 1, + 0, + 16, + 16, + [10, 11], + [10, 11], + [20, 21, 22], + 128, + 47, + 0, + 1, + 1, + make_full_spec(), + 0, + ) + + assert result == [(0, [10], [22])] + + +def test_get_remote_metadata_propagates_network_error_without_caching( + monkeypatch: pytest.MonkeyPatch, +) -> None: + thread = make_thread() + thread.encoder = msgspec.msgpack.Encoder() + thread.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) + thread.remote_metadata = SizedDict() + context = MagicMock() + context.__enter__.return_value = MagicMock() + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.zmq_ctx", + MagicMock(return_value=context), + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_send", + MagicMock(), + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_recv", + MagicMock(side_effect=RuntimeError("timed out")), + ) + + with pytest.raises(RuntimeError, match="timed out"): + thread._get_remote_metadata("engine-p", "10.0.0.1", 6000) + + assert thread.remote_metadata.get("engine-p") is None + + +@pytest.mark.parametrize(("is_consumer", "expected_thread_count"), [(True, 1), (False, 0)]) +def test_register_kv_caches_starts_receiving_thread_only_for_consumer( + monkeypatch: pytest.MonkeyPatch, + is_consumer: bool, + expected_thread_count: int, +) -> None: + worker = MooncakePullConnectorWorker.__new__(MooncakePullConnectorWorker) + worker.kv_transfer_config = SimpleNamespace(is_kv_consumer=is_consumer) + worker.engine = MagicMock() + worker.vllm_config = MagicMock() + worker.kv_cache_config = MagicMock() + worker.kv_cache_specs = [make_full_spec()] + worker.layer_name_to_group_index = {"model.layers.0.self_attn": 0} + worker.layer_name_to_spec_index = {"model.layers.0.self_attn": 0} + worker.tp_rank = worker.pp_rank = worker.dp_rank = worker.pcp_rank = worker.dcp_rank = 0 + worker.tp_size = worker.pp_size = worker.dp_size = worker.pcp_size = worker.dcp_size = 1 + worker._recving_thread = None + + def fake_register(instance: MooncakePullConnectorWorker, _kv_caches: object) -> None: + instance.xfer_handshake_metadata = make_transfer_metadata() + + created_threads: list[MagicMock] = [] + + def fake_thread(**kwargs: object) -> MagicMock: + ready_event = kwargs["ready_event"] + assert isinstance(ready_event, threading.Event) + thread = MagicMock() + thread.start.side_effect = ready_event.set + thread.is_alive.return_value = True + created_threads.append(thread) + return thread + + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.MooncakeBaseConnectorWorker.register_kv_caches", + fake_register, + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.MooncakePullRecvingThread", + fake_thread, + ) + monkeypatch.setattr(torch.npu, "current_device", MagicMock(return_value=0)) + + worker.register_kv_caches({}) + + assert len(created_threads) == expected_thread_count + if is_consumer: + created_threads[0].start.assert_called_once_with() + assert worker._recving_thread is created_threads[0] + else: + assert worker._recving_thread is None + + +def test_mamba_unequal_tp_slices_into_wider_remote_cache() -> None: + spec = SimpleNamespace( + dtypes=(torch.float16, torch.float16), + mamba_type=MambaAttentionBackendEnum.MAMBA2, + ) + thread = make_thread( + tp_size=2, + tp_rank=1, + kv_caches_base_addr=[[1000, 2000]], + block_shapes=[[(3, 8), (1, 4, 4)]], + block_strides=[[48, 32]], + block_lens=[[48, 32]], + ) + remote = make_pp_metadata( + block_shapes=[[(3, 16), (2, 4, 4)]], + block_strides=[[96, 64]], + block_lens=[[96, 64]], + tp_base_addrs={0: [[5000, 6000]]}, + ) + src: list[int] = [] + dst: list[int] = [] + lengths: list[int] = [] + + thread._append_mamba_transfer_addresses( + spec, # type: ignore[arg-type] + remote_tp_rank=0, + remote_tp_size=1, + transfer_entries_by_layer={(0, 0): [("request", [1], [2])]}, + remote_metadata=remote, + src_list=src, + dst_list=dst, + length_list=lengths, + ) + + assert src == [1048, 1064, 1080, 1056, 1072, 1088, 1060, 1076, 1092, 2032] + assert dst == [5200, 5232, 5264, 5212, 5244, 5276, 5220, 5252, 5284, 6160] + assert lengths == [8, 8, 8, 4, 4, 4, 4, 4, 4, 32] + +@pytest.mark.parametrize( + ( + "local_tp_size", + "remote_tp_size", + "local_heads", + "remote_heads", + "local_dcp_size", + "remote_dcp_size", + "fixed_heads", + "expected_heads", + ), + [ + (8, 4, 1, 1, 1, 1, None, 4), + (4, 8, 2, 1, 1, 1, None, 8), + (16, 8, 1, 1, 4, 2, None, 4), + (8, 16, 1, 1, 2, 4, None, 4), + (16, 8, 1, 1, 4, 2, 1, 1), + ], +) +def test_infer_total_kv_heads_across_tp_and_dcp_strategies( + local_tp_size: int, + remote_tp_size: int, + local_heads: int, + remote_heads: int, + local_dcp_size: int, + remote_dcp_size: int, + fixed_heads: int | None, + expected_heads: int, +) -> None: + thread = make_thread(tp_size=local_tp_size) + + assert ( + thread._infer_total_num_kv_heads( + local_num_kv_heads=local_heads, + remote_num_kv_heads=remote_heads, + remote_tp_size=remote_tp_size, + local_dcp_size=local_dcp_size, + remote_dcp_size=remote_dcp_size, + fixed_total_num_kv_heads=fixed_heads, + ) + == expected_heads + ) + + +def test_infer_total_kv_heads_rejects_inconsistent_topology() -> None: + thread = make_thread(tp_size=4) + + with pytest.raises(ValueError, match="inconsistent total KV head counts"): + thread._infer_total_num_kv_heads( + local_num_kv_heads=2, + remote_num_kv_heads=2, + remote_tp_size=8, + local_dcp_size=1, + remote_dcp_size=1, + fixed_total_num_kv_heads=None, + ) + + +@pytest.mark.parametrize( + ( + "local_tp_size", + "local_tp_rank", + "local_dcp_size", + "remote_tp_size", + "remote_dcp_size", + "total_heads", + "expected_groups", + ), + [ + (4, 2, 1, 8, 1, 4, [[4, 5]]), + (2, 0, 1, 4, 1, 8, [[0], [1]]), + (16, 10, 4, 8, 2, 4, [[4, 5]]), + (8, 5, 2, 16, 4, 4, [[8, 9, 10, 11]]), + (8, 6, 4, 16, 4, 4, [[8, 9, 10, 11], [12, 13, 14, 15]]), + (16, 13, 4, 8, 4, 4, [[4, 5, 6, 7]]), + (4, 3, 4, 8, 8, 1, [list(range(8))]), + (8, 6, 8, 4, 4, 1, [list(range(4))]), + ], +) +def test_attention_remote_tp_rank_groups_cover_tp_and_dcp_matrix( + local_tp_size: int, + local_tp_rank: int, + local_dcp_size: int, + remote_tp_size: int, + remote_dcp_size: int, + total_heads: int, + expected_groups: list[list[int]], +) -> None: + thread = make_thread(tp_size=local_tp_size, tp_rank=local_tp_rank, dcp_size=local_dcp_size) + + assert ( + thread._get_attention_remote_tp_rank_groups( + remote_tp_size=remote_tp_size, + local_dcp_size=local_dcp_size, + remote_dcp_size=remote_dcp_size, + total_num_kv_heads=total_heads, + ) + == expected_groups + ) + + +def test_layer_remote_tp_rank_groups_apply_spec_specific_dcp_rules() -> None: + full_thread = make_thread( + tp_size=16, + tp_rank=10, + dcp_size=4, + block_shapes=[[(1, 16, 4)]], + ) + full_remote = make_pp_metadata(block_shapes=[[(1, 16, 4)]]) + assert full_thread._get_layer_remote_tp_rank_groups( + 0, 0, make_full_spec(), full_remote, remote_tp_size=8, remote_dcp_size=2 + ) == [[4, 5]] + + mla_thread = make_thread(tp_size=4, tp_rank=2, dcp_size=4) + mla_spec = MLAAttentionSpec(block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float16) + assert mla_thread._get_layer_remote_tp_rank_groups( + 0, 0, mla_spec, make_pp_metadata(), remote_tp_size=8, remote_dcp_size=8 + ) == [list(range(8))] + + swa_thread = make_thread( + tp_size=4, + tp_rank=2, + dcp_size=4, + block_shapes=[[(1, 16, 4)]], + ) + swa_remote = make_pp_metadata(block_shapes=[[(1, 16, 4)]]) + assert swa_thread._get_layer_remote_tp_rank_groups( + 0, 0, make_sliding_spec(), swa_remote, remote_tp_size=4, remote_dcp_size=4 + ) == [[2]] + + +@pytest.mark.parametrize( + ( + "local_dcp_size", + "local_dcp_rank", + "remote_dcp_size", + "candidate_ranks", + "local_blocks", + "remote_blocks", + "expected", + ), + [ + ( + 1, + 0, + 4, + [0, 1, 2, 3], + [10, 11, 12, 13], + [20], + [(0, [10], [20]), (1, [11], [20]), (2, [12], [20]), (3, [13], [20])], + ), + (4, 3, 2, [4, 5], [10, 11], [20, 21, 22, 23], [(5, [10, 11], [21, 23])]), + ( + 2, + 1, + 4, + [4, 5, 6, 7], + [10, 11, 12, 13], + [20, 21], + [(5, [10, 12], [20, 21]), (7, [11, 13], [20, 21])], + ), + (4, 2, 4, [4, 5, 6, 7], [10, 11], [20, 21], [(6, [10, 11], [20, 21])]), + ], +) +def test_compute_block_ids_across_equal_block_size_dcp_matrix( + local_dcp_size: int, + local_dcp_rank: int, + remote_dcp_size: int, + candidate_ranks: list[int], + local_blocks: list[int], + remote_blocks: list[int], + expected: list[tuple[int, list[int], list[int]]], +) -> None: + thread = make_thread( + tp_size=max(local_dcp_size, 1), + tp_rank=local_dcp_rank, + dcp_size=local_dcp_size, + dcp_rank=local_dcp_rank, + ) + + result = thread._compute_group_block_ids( + request_id="request", + remote_tp_rank_groups=[candidate_ranks], + remote_dcp_size=remote_dcp_size, + spec_index=0, + local_block_size=16, + remote_block_size=16, + local_group_block_ids=local_blocks, + local_full_group_block_ids=local_blocks, + remote_group_block_ids=remote_blocks, + local_num_prompt_tokens=256, + remote_num_prompt_tokens=256, + num_computed_tokens=0, + local_block_size_scale=1, + remote_block_size_scale=1, + spec=make_full_spec(), + selection_index=0, + ) + + assert result == expected + + +@pytest.mark.parametrize( + ( + "local_dcp_size", + "local_dcp_rank", + "remote_dcp_size", + "candidate_ranks", + "local_block_size", + "remote_block_size", + "local_scale", + "remote_scale", + "local_blocks", + "remote_blocks", + "prompt_tokens", + "expected", + ), + [ + ( + 2, + 0, + 4, + [0, 1, 2, 3], + 16, + 32, + 1, + 2, + [10, 11, 12, 13], + [20], + 128, + [(0, [10], [40]), (1, [11], [40]), (2, [12], [40]), (3, [13], [40])], + ), + ( + 4, + 1, + 2, + [4, 5], + 32, + 16, + 2, + 1, + [10, 11], + [30, 31, 32, 33, 34, 35], + 256, + [(4, [20, 22], [31, 35]), (5, [21, 23], [31, 35])], + ), + ], +) +def test_compute_block_ids_across_unequal_block_size_dcp_matrix( + local_dcp_size: int, + local_dcp_rank: int, + remote_dcp_size: int, + candidate_ranks: list[int], + local_block_size: int, + remote_block_size: int, + local_scale: int, + remote_scale: int, + local_blocks: list[int], + remote_blocks: list[int], + prompt_tokens: int, + expected: list[tuple[int, list[int], list[int]]], +) -> None: + thread = make_thread( + tp_size=max(local_dcp_size, 1), + tp_rank=local_dcp_rank, + dcp_size=local_dcp_size, + dcp_rank=local_dcp_rank, + ) + + result = thread._compute_group_block_ids( + request_id="request", + remote_tp_rank_groups=[candidate_ranks], + remote_dcp_size=remote_dcp_size, + spec_index=0, + local_block_size=local_block_size, + remote_block_size=remote_block_size, + local_group_block_ids=local_blocks, + local_full_group_block_ids=local_blocks, + remote_group_block_ids=remote_blocks, + local_num_prompt_tokens=prompt_tokens, + remote_num_prompt_tokens=prompt_tokens, + num_computed_tokens=0, + local_block_size_scale=local_scale, + remote_block_size_scale=remote_scale, + spec=make_full_spec(local_block_size), + selection_index=0, + ) + + assert result == expected diff --git a/tests/ut/kv_offload/mooncake_v2/test_utils.py b/tests/ut/kv_offload/mooncake_v2/test_utils.py index 0df3fc977abf..2bdb3940c07a 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_utils.py +++ b/tests/ut/kv_offload/mooncake_v2/test_utils.py @@ -112,7 +112,7 @@ def test_zmq_ctx_destroys_context(monkeypatch: pytest.MonkeyPatch) -> None: MagicMock(return_value=socket), ) - with zmq_ctx(zmq.REQ, "tcp://127.0.0.1:1") as result: + with zmq_ctx(zmq.REQ, "tcp://127.0.0.1:1") as result: # type: ignore[attr-defined] assert result is socket context.destroy.assert_called_once_with(linger=0) @@ -121,8 +121,8 @@ def test_zmq_ctx_destroys_context(monkeypatch: pytest.MonkeyPatch) -> None: def test_zmq_retry_helpers_succeed_after_transient_error(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils.time.sleep", MagicMock()) socket = MagicMock() - socket.send.side_effect = [zmq.ZMQError("transient"), None] - socket.recv.side_effect = [zmq.ZMQError("transient"), b"response"] + socket.send.side_effect = [zmq.ZMQError("transient"), None] # type: ignore[attr-defined] + socket.recv.side_effect = [zmq.ZMQError("transient"), b"response"] # type: ignore[attr-defined] ensure_zmq_send(socket, b"request", "tcp://peer", max_retries=2) assert ensure_zmq_recv(socket, "tcp://peer", max_retries=2) == b"response" @@ -133,8 +133,8 @@ def test_zmq_retry_helpers_succeed_after_transient_error(monkeypatch: pytest.Mon def test_zmq_retry_helpers_raise_after_limit(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr("vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils.time.sleep", MagicMock()) socket = MagicMock() - socket.send.side_effect = zmq.ZMQError("failed") - socket.recv.side_effect = zmq.ZMQError("failed") + socket.send.side_effect = zmq.ZMQError("failed") # type: ignore[attr-defined] + socket.recv.side_effect = zmq.ZMQError("failed") # type: ignore[attr-defined] with pytest.raises(RuntimeError, match="Failed to send"): ensure_zmq_send(socket, b"request", "tcp://peer", max_retries=2) From 1f971c05735d0dda0f50e655acc432244f35885e Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Fri, 21 Aug 2026 15:56:37 +0800 Subject: [PATCH 38/48] fix ut ruff Signed-off-by: nwpu-zxr --- tests/ut/kv_offload/mooncake_v2/test_pull_worker.py | 1 + 1 file changed, 1 insertion(+) 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 69519820ecbd..b6a9b3150520 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py @@ -890,6 +890,7 @@ def test_mamba_unequal_tp_slices_into_wider_remote_cache() -> None: assert dst == [5200, 5232, 5264, 5212, 5244, 5276, 5220, 5252, 5284, 6160] assert lengths == [8, 8, 8, 4, 4, 4, 4, 4, 4, 32] + @pytest.mark.parametrize( ( "local_tp_size", From 930c862ba8158aeac31d665b53b12025a6c8488b Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Fri, 21 Aug 2026 16:22:02 +0800 Subject: [PATCH 39/48] add base class UT Signed-off-by: nwpu-zxr --- .../mooncake_v2/test_base_scheduler.py | 155 ++++++++++++++++++ .../mooncake_v2/test_base_worker.py | 145 ++++++++++++++++ 2 files changed, 300 insertions(+) create mode 100644 tests/ut/kv_offload/mooncake_v2/test_base_scheduler.py diff --git a/tests/ut/kv_offload/mooncake_v2/test_base_scheduler.py b/tests/ut/kv_offload/mooncake_v2/test_base_scheduler.py new file mode 100644 index 000000000000..fa53aa842610 --- /dev/null +++ b/tests/ut/kv_offload/mooncake_v2/test_base_scheduler.py @@ -0,0 +1,155 @@ +# SPDX-License-Identifier: Apache-2.0 + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs + +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler import ( + MooncakeBaseConnectorScheduler, +) + +from .helpers import make_full_spec, make_mamba_spec, make_sliding_spec + + +def make_scheduler_config( + *, + is_consumer: bool = True, + is_producer: bool = False, + pcp_size: int = 1, +) -> SimpleNamespace: + return SimpleNamespace( + kv_transfer_config=SimpleNamespace( + is_kv_consumer=is_consumer, + is_kv_producer=is_producer, + kv_role="kv_consumer" if is_consumer else "kv_producer", + kv_port=6000, + ), + cache_config=SimpleNamespace(block_size=16), + speculative_config=SimpleNamespace(num_speculative_tokens=3), + parallel_config=SimpleNamespace( + pipeline_parallel_size=2, + tensor_parallel_size=4, + prefill_context_parallel_size=pcp_size, + decode_context_parallel_size=2, + data_parallel_size=3, + data_parallel_rank=1, + ), + model_config=SimpleNamespace(hf_config=SimpleNamespace(compress_ratios=None)), + ) + + +def test_base_scheduler_initializes_parallel_layout_and_control_port( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = make_scheduler_config() + group = KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=make_full_spec()) + kv_cache_config = SimpleNamespace(kv_cache_groups=[group]) + ascend_config = object() + init_ascend_config = MagicMock() + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.init_ascend_config", + init_ascend_config, + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.get_ascend_config", + MagicMock(return_value=ascend_config), + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.get_ip", + MagicMock(return_value="10.0.0.1"), + ) + + scheduler = MooncakeBaseConnectorScheduler(config, "engine-d", kv_cache_config) # type: ignore[arg-type] + + assert scheduler.engine_id == "engine-d" + assert scheduler.block_size == 16 + assert scheduler.num_speculative_tokens == 3 + assert scheduler.ascend_config is ascend_config + assert scheduler.side_channel_host == "10.0.0.1" + assert scheduler.max_device_id == 24 + assert scheduler.side_channel_port == 6025 + assert scheduler.group_block_size == [16] + assert scheduler.group_unique_specs == [[group.kv_cache_spec]] + assert scheduler.need_truncate is False + init_ascend_config.assert_called_once_with(config) + + +@pytest.mark.parametrize(("is_consumer", "is_producer"), [(False, False), (True, True)]) +def test_base_scheduler_rejects_invalid_transfer_roles( + is_consumer: bool, + is_producer: bool, +) -> None: + config = make_scheduler_config(is_consumer=is_consumer, is_producer=is_producer) + + with pytest.raises(ValueError, match="exactly one KV transfer role"): + MooncakeBaseConnectorScheduler(config, "engine", SimpleNamespace()) # type: ignore[arg-type] + + +def test_base_scheduler_rejects_unsupported_pcp(monkeypatch: pytest.MonkeyPatch) -> None: + config = make_scheduler_config(pcp_size=2) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.init_ascend_config", + MagicMock(), + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.get_ascend_config", + MagicMock(), + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.get_ip", + MagicMock(return_value="10.0.0.1"), + ) + + with pytest.raises(AssertionError, match="prefill context parallel size 1"): + MooncakeBaseConnectorScheduler(config, "engine", SimpleNamespace()) # type: ignore[arg-type] + + +def test_base_scheduler_expands_unique_uniform_specs_in_layer_order() -> None: + full = make_full_spec() + sliding = make_sliding_spec() + uniform = UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={"layer.0": full, "layer.1": full, "layer.2": sliding}, + ) + group = KVCacheGroupSpec( + layer_names=["layer.1", "layer.0", "layer.2"], + kv_cache_spec=uniform, + ) + + assert MooncakeBaseConnectorScheduler._get_group_unique_specs(group) == [full, sliding] + + +def test_base_scheduler_transfer_blocks_handles_empty_and_state_without_speculation() -> None: + scheduler = MooncakeBaseConnectorScheduler.__new__(MooncakeBaseConnectorScheduler) + scheduler.pcp_size = 1 + scheduler.dcp_size = 1 + scheduler.num_speculative_tokens = 0 + scheduler.group_block_size = [16] + scheduler.group_unique_specs = [[make_mamba_spec()]] + + assert scheduler._get_transfer_block_ids((), prompt_len=32) == () + assert scheduler._get_transfer_block_ids(([10, 11],), prompt_len=32) == ([10, 11],) + + +def test_base_scheduler_abstract_contract_and_legacy_metadata_delegation() -> None: + scheduler = MooncakeBaseConnectorScheduler.__new__(MooncakeBaseConnectorScheduler) + + with pytest.raises(NotImplementedError): + scheduler.on_new_request(MagicMock()) + with pytest.raises(NotImplementedError): + scheduler.update_connector_output(MagicMock()) + with pytest.raises(NotImplementedError): + scheduler.get_num_new_matched_tokens(MagicMock(), 0) + with pytest.raises(NotImplementedError): + scheduler.update_state_after_alloc(MagicMock(), MagicMock(), 0) + with pytest.raises(NotImplementedError): + scheduler.build_connector_meta(MagicMock()) + with pytest.raises(NotImplementedError): + scheduler.request_finished(MagicMock(), ()) + + scheduler.set_xfer_handshake_metadata_from_workers = MagicMock() # type: ignore[method-assign] + metadata = {0: MagicMock()} + scheduler.set_xfer_handshake_metadata(metadata) + scheduler.set_xfer_handshake_metadata_from_workers.assert_called_once_with(metadata) diff --git a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py index 01718f78a0c1..3a80789666be 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py @@ -1,7 +1,9 @@ # SPDX-License-Identifier: Apache-2.0 +from types import SimpleNamespace from unittest.mock import MagicMock +import pytest import torch from vllm.v1.kv_cache_interface import ( KVCacheConfig, @@ -13,6 +15,7 @@ from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker import ( MooncakeBaseConnectorWorker, ) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.stats import MooncakeKVConnectorStats from .helpers import make_full_spec, make_sliding_spec @@ -108,3 +111,145 @@ def test_register_kv_caches_rejects_missing_and_unconfigured_layers() -> None: assert "absent from kv_cache_tensors" in str(error) else: raise AssertionError("unexpected layer must fail") + + +def make_worker_config( + *, + is_consumer: bool = True, + is_producer: bool = False, + local_dp_rank: int | None = 1, +) -> SimpleNamespace: + return SimpleNamespace( + kv_transfer_config=SimpleNamespace( + is_kv_consumer=is_consumer, + is_kv_producer=is_producer, + kv_role="kv_consumer" if is_consumer else "kv_producer", + kv_port=6000, + ), + cache_config=SimpleNamespace(block_size=16), + parallel_config=SimpleNamespace( + tensor_parallel_size=4, + pipeline_parallel_size=2, + data_parallel_rank_local=local_dp_rank, + data_parallel_size_local=3, + data_parallel_rank=2, + ), + ) + + +def patch_worker_runtime( + monkeypatch: pytest.MonkeyPatch, + *, + pcp_size: int = 1, +) -> tuple[MagicMock, object]: + monkeypatch.setenv("ASCEND_TRANSFER_TIMEOUT", "original") + tp_group = object() + transfer_engine = MagicMock() + transfer_engine.get_rpc_port.return_value = 9000 + global_te = MagicMock() + global_te.get_transfer_engine.return_value = transfer_engine + patches = { + "init_ascend_config": MagicMock(), + "get_ascend_config": MagicMock(return_value=object()), + "get_transfer_timeout_value": MagicMock(return_value=30), + "get_tensor_model_parallel_rank": MagicMock(return_value=2), + "get_tp_group": MagicMock(return_value=tp_group), + "get_pp_group": MagicMock(return_value=SimpleNamespace(rank_in_group=1)), + "get_pcp_group": MagicMock(return_value=SimpleNamespace(rank_in_group=0, world_size=pcp_size)), + "get_decode_context_model_parallel_world_size": MagicMock(return_value=2), + "get_decode_context_model_parallel_rank": MagicMock(return_value=1), + "get_ip": MagicMock(return_value="10.0.0.2"), + "global_te": global_te, + } + for name, value in patches.items(): + monkeypatch.setattr( + f"vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.{name}", + value, + ) + monkeypatch.setattr(torch.npu, "current_device", MagicMock(return_value=7)) + return global_te, tp_group + + +def test_base_worker_initializes_parallel_ranks_and_handshake_port( + monkeypatch: pytest.MonkeyPatch, +) -> None: + config = make_worker_config() + kv_cache_config = SimpleNamespace(num_blocks=32) + global_te, tp_group = patch_worker_runtime(monkeypatch) + + worker = MooncakeBaseConnectorWorker(config, "engine-d", kv_cache_config) # type: ignore[arg-type] + + assert worker.engine_id == "engine-d" + assert worker.block_size == 16 + assert worker.num_blocks == 32 + assert worker.tp_rank == 2 + assert worker.tp_size == 4 + assert worker.tp_group is tp_group + assert worker.pp_rank == 1 + assert worker.pp_size == 2 + assert worker.dp_rank == 1 + assert worker.dp_size == 3 + assert worker.pcp_rank == 0 + assert worker.pcp_size == 1 + assert worker.dcp_rank == 1 + assert worker.dcp_size == 2 + assert worker.max_device_id == 24 + assert worker.side_channel_host == "10.0.0.2" + assert worker.side_channel_port == 6016 + assert worker.handshake_port == 6022 + assert worker.te_rpc_port == 9000 + assert worker.xfer_handshake_metadata is None + assert worker.xfer_stats.is_empty() + global_te.get_transfer_engine.assert_called_once_with("10.0.0.2", device_name="7") + + +@pytest.mark.parametrize(("is_consumer", "is_producer"), [(False, False), (True, True)]) +def test_base_worker_rejects_invalid_transfer_roles( + is_consumer: bool, + is_producer: bool, +) -> None: + config = make_worker_config(is_consumer=is_consumer, is_producer=is_producer) + + with pytest.raises(ValueError, match="exactly one KV transfer role"): + MooncakeBaseConnectorWorker(config, "engine", SimpleNamespace(num_blocks=1)) # type: ignore[arg-type] + + +def test_base_worker_requires_local_dp_rank(monkeypatch: pytest.MonkeyPatch) -> None: + config = make_worker_config(local_dp_rank=None) + patch_worker_runtime(monkeypatch) + + with pytest.raises(ValueError, match="requires a local DP rank"): + MooncakeBaseConnectorWorker(config, "engine", SimpleNamespace(num_blocks=1)) # type: ignore[arg-type] + + +def test_base_worker_rejects_unsupported_pcp(monkeypatch: pytest.MonkeyPatch) -> None: + config = make_worker_config() + patch_worker_runtime(monkeypatch, pcp_size=2) + + with pytest.raises(AssertionError, match="prefill context parallel size 1"): + MooncakeBaseConnectorWorker(config, "engine", SimpleNamespace(num_blocks=1)) # type: ignore[arg-type] + + +def test_base_worker_stats_are_returned_and_reset() -> None: + worker = MooncakeBaseConnectorWorker.__new__(MooncakeBaseConnectorWorker) + worker.xfer_stats = MooncakeKVConnectorStats(data={}) + + assert worker.get_kv_connector_stats() is None + + worker.xfer_stats.record_transfer(0.1, 1024) + result = worker.get_kv_connector_stats() + + assert result is not None + assert result.data["transfer_duration"] == [0.1] + assert worker.xfer_stats.is_empty() + + +def test_base_worker_abstract_contract() -> None: + worker = MooncakeBaseConnectorWorker.__new__(MooncakeBaseConnectorWorker) + + with pytest.raises(NotImplementedError): + worker.get_finished() + with pytest.raises(NotImplementedError): + worker.get_block_ids_with_load_errors() + with pytest.raises(NotImplementedError): + worker.start_load_kv(MagicMock()) From cbbb4e391a92e507648f26049056f6c426ffae01 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Sat, 22 Aug 2026 17:59:57 +0800 Subject: [PATCH 40/48] support kvpp Signed-off-by: nwpu-zxr --- tests/ut/kv_offload/mooncake_v2/helpers.py | 17 +-- .../mooncake_v2/test_base_worker.py | 1 + .../kv_offload/mooncake_v2/test_metadata.py | 8 ++ .../mooncake_v2/test_pull_scheduler.py | 87 +++++++++++- .../mooncake_v2/test_pull_worker.py | 90 +++++++++++- .../kv_p2p/mooncake/base_worker.py | 10 +- .../kv_transfer/kv_p2p/mooncake/metadata.py | 18 +-- .../kv_p2p/mooncake/pull_scheduler.py | 134 +++++++++++++----- .../kv_p2p/mooncake/pull_worker.py | 82 ++++++++--- 9 files changed, 365 insertions(+), 82 deletions(-) diff --git a/tests/ut/kv_offload/mooncake_v2/helpers.py b/tests/ut/kv_offload/mooncake_v2/helpers.py index 0e4afe6c6225..61721bd1adc4 100644 --- a/tests/ut/kv_offload/mooncake_v2/helpers.py +++ b/tests/ut/kv_offload/mooncake_v2/helpers.py @@ -52,8 +52,7 @@ def make_transfer_metadata( handshake_port: int = 5000, layer_names: list[str] | None = None, group_indices: list[int] | None = None, - spec_indices: list[int] | None = None, - spec_block_sizes: list[int] | None = None, + layer_block_sizes: list[int] | None = None, base_addrs: list[list[int]] | None = None, block_strides: list[list[int]] | None = None, block_lens: list[list[int]] | None = None, @@ -67,10 +66,9 @@ def make_transfer_metadata( te_rpc_port=te_rpc_port, block_size=16, num_blocks=32, - spec_block_sizes=spec_block_sizes or [16], layer_names=layer_names, + layer_block_sizes=layer_block_sizes or [16] * num_layers, group_indices=group_indices or [0] * num_layers, - spec_indices=spec_indices or [0] * num_layers, kv_caches_base_addr=base_addrs or [[1000 + index * 1000] for index in range(num_layers)], block_strides=block_strides or [[128] for _ in range(num_layers)], block_lens=block_lens or [[128] for _ in range(num_layers)], @@ -84,24 +82,24 @@ def make_transfer_metadata( def make_pp_metadata( *, layer_names: list[str] | None = None, - spec_indices: list[int] | None = None, - spec_block_sizes: list[int] | None = None, + layer_block_sizes: list[int] | None = None, block_shapes: list[list[tuple[int, ...]]] | None = None, block_strides: list[list[int]] | None = None, block_lens: list[list[int]] | None = None, block_size_scales: list[list[int]] | None = None, tp_base_addrs: dict[int, list[list[int]]] | None = None, + tp_layer_indices: dict[int, list[int]] | None = None, ) -> MooncakePPTransferMetadata: layer_names = layer_names or ["model.layers.0.self_attn"] num_layers = len(layer_names) tp_base_addrs = tp_base_addrs or {0: [[5000 + index * 1000] for index in range(num_layers)]} + tp_layer_indices = tp_layer_indices or {} return MooncakePPTransferMetadata( block_size=16, num_blocks=32, - spec_block_sizes=spec_block_sizes or [16], layer_names=layer_names, + layer_block_sizes=layer_block_sizes or [16] * num_layers, group_indices=[0] * num_layers, - spec_indices=spec_indices or [0] * num_layers, block_strides=block_strides or [[128] for _ in range(num_layers)], block_lens=block_lens or [[128] for _ in range(num_layers)], block_shapes=block_shapes or [[(1, 16, 4)] for _ in range(num_layers)], @@ -109,6 +107,7 @@ def make_pp_metadata( metadata_by_tp_rank={ tp_rank: MooncakeTPTransferMetadata( te_rpc_port=9000 + tp_rank, + layer_indices=tp_layer_indices.get(tp_rank, list(range(num_layers))), kv_caches_base_addr=base_addrs, local_ip=f"10.0.0.{tp_rank + 1}", handshake_port=5000 + tp_rank, @@ -122,6 +121,7 @@ def make_metadata_groups( *, engine_id: str = "engine-p", tp_size: int = 1, + use_kv_pp: bool = False, pp_metadata: MooncakePPTransferMetadata | None = None, ) -> MooncakeTransferMetadataGroups: return MooncakeTransferMetadataGroups( @@ -132,6 +132,7 @@ def make_metadata_groups( pcp_size=1, dcp_size=1, tp_size=tp_size, + use_kv_pp=use_kv_pp, metadata_by_pp_rank={0: pp_metadata or make_pp_metadata()}, ) diff --git a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py index 3a80789666be..b94e5e06ae5a 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py @@ -75,6 +75,7 @@ def test_register_kv_caches_uses_config_order_and_publishes_tensor_metadata(monk metadata = worker.xfer_handshake_metadata assert metadata is not None assert metadata.layer_names == ["layer.0"] + assert metadata.layer_block_sizes == [spec.block_size] assert metadata.kv_caches_base_addr == [[k_cache.data_ptr(), v_cache.data_ptr()]] assert metadata.block_shapes == [[(1, 16, 8), (1, 16, 8)]] assert metadata.block_size_scales == [[2, 2]] diff --git a/tests/ut/kv_offload/mooncake_v2/test_metadata.py b/tests/ut/kv_offload/mooncake_v2/test_metadata.py index 892118eaa5ee..1f9334d461ff 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_metadata.py +++ b/tests/ut/kv_offload/mooncake_v2/test_metadata.py @@ -33,6 +33,14 @@ def test_transfer_metadata_rejects_misaligned_layer_count() -> None: ) +def test_transfer_metadata_rejects_misaligned_layer_block_sizes() -> None: + with pytest.raises(ValueError, match="layer_block_sizes.*1 layers, expected 2"): + make_transfer_metadata( + layer_names=["layer.0", "layer.1"], + layer_block_sizes=[16], + ) + + def test_transfer_metadata_rejects_misaligned_tensor_count() -> None: with pytest.raises(ValueError, match="layer 'layer.0'.*block_strides"): make_transfer_metadata( diff --git a/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py b/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py index 74cb3f4b5c42..cc4d971c4e80 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py @@ -33,6 +33,7 @@ def make_sending_thread( *, tp_size: int = 1, pp_size: int = 1, + dcp_size: int = 1, ) -> MooncakeSchedulerSendingThread: return MooncakeSchedulerSendingThread( host="127.0.0.1", @@ -42,7 +43,7 @@ def make_sending_thread( tp_size=tp_size, pp_size=pp_size, pcp_size=1, - dcp_size=1, + dcp_size=dcp_size, ready_event=threading.Event(), ) @@ -74,6 +75,7 @@ def test_sending_thread_merges_tp_private_and_pp_common_metadata() -> None: thread = make_sending_thread({0: first, 1: second}, tp_size=2) decoded = msgspec.msgpack.decode(thread.encoded_metadata, type=MooncakeTransferMetadataGroups) + assert decoded.use_kv_pp is False pp_metadata = decoded.metadata_by_pp_rank[0] assert pp_metadata.layer_names == first.layer_names assert pp_metadata.block_shapes == first.block_shapes @@ -82,6 +84,89 @@ def test_sending_thread_merges_tp_private_and_pp_common_metadata() -> None: assert pp_metadata.metadata_by_tp_rank[1].te_rpc_port == 9001 +def test_sending_thread_merges_layer_split_metadata_by_name() -> None: + tp0 = make_transfer_metadata( + layer_names=["layer.0", "layer.2"], + group_indices=[0, 0], + base_addrs=[[1000], [3000]], + block_strides=[[128], [128]], + block_lens=[[128], [128]], + block_shapes=[[(1, 16, 4)], [(1, 16, 4)]], + block_size_scales=[[1], [1]], + ) + tp1 = make_transfer_metadata( + te_rpc_port=9001, + local_ip="10.0.0.2", + layer_names=["layer.1", "layer.2"], + group_indices=[0, 0], + base_addrs=[[2000], [4000]], + block_strides=[[128], [128]], + block_lens=[[128], [128]], + block_shapes=[[(1, 16, 4)], [(1, 16, 4)]], + block_size_scales=[[1], [1]], + ) + + thread = make_sending_thread({0: tp0, 1: tp1}, tp_size=2) + decoded = msgspec.msgpack.decode(thread.encoded_metadata, type=MooncakeTransferMetadataGroups) + + assert decoded.use_kv_pp is True + pp_metadata = decoded.metadata_by_pp_rank[0] + assert pp_metadata.layer_names == ["layer.0", "layer.1", "layer.2"] + assert pp_metadata.layer_block_sizes == [16, 16, 16] + assert pp_metadata.metadata_by_tp_rank[0].layer_indices == [0, 2] + assert pp_metadata.metadata_by_tp_rank[0].kv_caches_base_addr == [[1000], [], [3000]] + assert pp_metadata.metadata_by_tp_rank[1].layer_indices == [1, 2] + assert pp_metadata.metadata_by_tp_rank[1].kv_caches_base_addr == [[], [2000], [4000]] + + +def test_sending_thread_rejects_layer_split_with_dcp() -> None: + tp0 = make_transfer_metadata(layer_names=["layer.0"], base_addrs=[[1000]]) + tp1 = make_transfer_metadata( + te_rpc_port=9001, + local_ip="10.0.0.2", + layer_names=["layer.1"], + base_addrs=[[2000]], + ) + + with pytest.raises(ValueError, match="LayerSplit cannot be combined with DCP"): + make_sending_thread({0: tp0, 1: tp1}, tp_size=2, dcp_size=2) + + +def test_sending_thread_aligns_different_tp_layer_orders() -> None: + tp0 = make_transfer_metadata( + layer_names=["layer.0", "layer.1"], + group_indices=[0, 0], + base_addrs=[[1000], [2000]], + block_strides=[[128], [128]], + block_lens=[[128], [128]], + block_shapes=[[(1, 16, 4)], [(1, 16, 4)]], + block_size_scales=[[1], [1]], + ) + tp1 = make_transfer_metadata( + te_rpc_port=9001, + local_ip="10.0.0.2", + layer_names=["layer.1", "layer.0"], + group_indices=[0, 0], + base_addrs=[[4000], [3000]], + block_strides=[[128], [128]], + block_lens=[[128], [128]], + block_shapes=[[(1, 16, 4)], [(1, 16, 4)]], + block_size_scales=[[1], [1]], + ) + + thread = make_sending_thread({0: tp0, 1: tp1}, tp_size=2) + decoded = msgspec.msgpack.decode(thread.encoded_metadata, type=MooncakeTransferMetadataGroups) + + pp_metadata = decoded.metadata_by_pp_rank[0] + assert pp_metadata.layer_names == ["layer.0", "layer.1"] + assert pp_metadata.metadata_by_tp_rank[0].layer_indices == [0, 1] + assert pp_metadata.metadata_by_tp_rank[1].layer_indices == [0, 1] + assert pp_metadata.metadata_by_tp_rank[1].kv_caches_base_addr == [ + [3000], + [4000], + ] + + def test_sending_thread_accepts_pp_aware_keys() -> None: pp0 = make_transfer_metadata(layer_names=["pp0.layer"]) pp1 = make_transfer_metadata(layer_names=["pp1.layer"], base_addrs=[[3000]]) 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 b6a9b3150520..15577aa9aa70 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py @@ -44,6 +44,7 @@ def make_thread(**overrides: object) -> MooncakePullRecvingThread: thread.layer_names = ["model.layers.0.self_attn"] thread.group_indices = [0] thread.spec_indices = [0] + thread.layer_block_sizes = [16] thread.kv_cache_specs = [make_full_spec()] thread.local_metadata = make_transfer_metadata() thread.kv_caches_base_addr = [[1000]] @@ -155,6 +156,7 @@ def test_build_remote_layout_matches_layers_across_pp_ranks() -> None: groups = SimpleNamespace( tp_size=1, dcp_size=1, + use_kv_pp=False, metadata_by_pp_rank={0: pp0, 1: pp1}, ) @@ -162,6 +164,40 @@ def test_build_remote_layout_matches_layers_across_pp_ranks() -> None: assert pairs == {0: [(0, 0)], 1: [(1, 0)]} assert rank_groups == {0: {(0, 0): [[0]]}, 1: {(1, 0): [[0]]}} + assert thread._get_layer_remote_tp_rank_groups.call_count == 1 + + +def test_build_remote_layout_filters_kv_parallel_owners_and_keeps_mtp_replicas() -> None: + thread = make_thread( + layer_names=["layer.0", "mtp.layer"], + spec_indices=[0, 0], + kv_cache_specs=[make_full_spec()], + ) + thread._get_layer_remote_tp_rank_groups = MagicMock(return_value=[[0, 1]]) # type: ignore[method-assign] + pp_metadata = make_pp_metadata( + layer_names=["layer.0", "mtp.layer"], + tp_base_addrs={ + 0: [[5000], [6000]], + 1: [[], [7000]], + }, + tp_layer_indices={0: [0, 1], 1: [1]}, + ) + groups = make_metadata_groups( + tp_size=2, + use_kv_pp=True, + pp_metadata=pp_metadata, + ) + + rank_groups, pairs = thread._build_remote_transfer_layout(groups) + + assert pairs == {0: [(0, 0), (1, 1)]} + assert rank_groups == { + 0: { + (0, 0): [[0]], + (1, 1): [[0, 1]], + } + } + assert thread._get_layer_remote_tp_rank_groups.call_count == 1 def test_build_remote_layout_rejects_missing_local_layer() -> None: @@ -308,6 +344,7 @@ def test_transfer_bucket_reuses_block_mapping_for_layers_of_same_spec() -> None: layer_names=["layer.0", "layer.1"], group_indices=[0, 0], spec_indices=[0, 0], + layer_block_sizes=[16, 16], block_shapes=[[(1, 16, 4)], [(1, 16, 4)]], block_strides=[[128], [128]], block_lens=[[128], [128]], @@ -316,14 +353,12 @@ def test_transfer_bucket_reuses_block_mapping_for_layers_of_same_spec() -> None: ) thread.local_metadata = make_transfer_metadata( layer_names=["layer.0", "layer.1"], - spec_indices=[0, 0], - spec_block_sizes=[16], + layer_block_sizes=[16, 16], ) thread._compute_group_block_ids = MagicMock(return_value=[(0, [10], [20])]) # type: ignore[method-assign] remote = make_pp_metadata( layer_names=["layer.0", "layer.1"], - spec_indices=[0, 0], - spec_block_sizes=[16], + layer_block_sizes=[16, 16], tp_base_addrs={0: [[5000], [6000]]}, ) @@ -341,6 +376,53 @@ def test_transfer_bucket_reuses_block_mapping_for_layers_of_same_spec() -> None: assert request_ids == {0: {"request"}} +def test_transfer_bucket_separates_same_spec_layers_with_different_tp_owners() -> None: + thread = make_thread( + layer_names=["layer.0", "layer.1"], + group_indices=[0, 0], + spec_indices=[0, 0], + layer_block_sizes=[16, 16], + block_shapes=[[(1, 16, 4)], [(1, 16, 4)]], + block_strides=[[128], [128]], + block_lens=[[128], [128]], + block_size_scales=[[1], [1]], + kv_caches_base_addr=[[1000], [2000]], + ) + thread.local_metadata = make_transfer_metadata( + layer_names=["layer.0", "layer.1"], + layer_block_sizes=[16, 16], + ) + + def compute_block_ids( + request_id: str, + remote_tp_rank_groups: list[list[int]], + *args: object, + ) -> list[tuple[int, list[int], list[int]]]: + del request_id, args + return [(remote_tp_rank_groups[0][0], [10], [20])] + + thread._compute_group_block_ids = MagicMock(side_effect=compute_block_ids) # type: ignore[method-assign] + remote = make_pp_metadata( + layer_names=["layer.0", "layer.1"], + layer_block_sizes=[16, 16], + tp_base_addrs={0: [[5000], []], 1: [[], [6000]]}, + tp_layer_indices={0: [0], 1: [1]}, + ) + + buckets, request_ids = thread._build_transfer_block_buckets( + remote, + [(0, 0), (1, 1)], + {(0, 0): [[0]], (1, 1): [[1]]}, + 1, + {"request": make_req_meta()}, + {}, + ) + + assert set(buckets) == {0, 1} + assert request_ids == {0: {"request"}, 1: {"request"}} + assert thread._compute_group_block_ids.call_count == 2 + + def test_attention_address_generation_handles_partial_head_overlap() -> None: thread = make_thread( tp_size=1, 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 1a58b1dd902a..23ce3d5ca8fb 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 @@ -155,11 +155,9 @@ def register_kv_caches( logger.info("num_blocks: %s", self.num_blocks) self.kv_caches = kv_caches self._build_kv_cache_spec_mappings() - spec_block_sizes = [spec.block_size for spec in self.kv_cache_specs] - layer_names: list[str] = [] + layer_block_sizes: list[int] = [] group_indices: list[int] = [] - spec_indices: list[int] = [] kv_caches_base_addr: list[list[int]] = [] block_strides_per_layer: list[list[int]] = [] block_lens_per_layer: list[list[int]] = [] @@ -196,8 +194,9 @@ def register_kv_caches( configured_layer_names.add(layer_name) layer_names.append(layer_name) + spec_index = self.layer_name_to_spec_index[layer_name] + layer_block_sizes.append(self.kv_cache_specs[spec_index].block_size) group_indices.append(self.layer_name_to_group_index[layer_name]) - spec_indices.append(self.layer_name_to_spec_index[layer_name]) kv_caches_base_addr.append(base_addrs) block_strides_per_layer.append(block_strides) block_lens_per_layer.append(block_lens) @@ -217,10 +216,9 @@ def register_kv_caches( te_rpc_port=self.te_rpc_port, block_size=self.block_size, num_blocks=self.num_blocks, - spec_block_sizes=spec_block_sizes, layer_names=layer_names, + layer_block_sizes=layer_block_sizes, group_indices=group_indices, - spec_indices=spec_indices, kv_caches_base_addr=kv_caches_base_addr, block_strides=block_strides_per_layer, block_lens=block_lens_per_layer, diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py index 0d974c0c6b4e..6e48c3025a33 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/metadata.py @@ -17,19 +17,17 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): """Worker transfer information exchanged during the P/D handshake. - Per-spec fields use the flattened KV-cache spec order. All per-layer - fields use layer_names order; each nested list preserves that layer's - cache tensor order. + All per-layer fields use layer_names order; each nested list preserves + that layer's cache tensor order. """ engine_id: str te_rpc_port: int block_size: int num_blocks: int - spec_block_sizes: list[int] layer_names: list[str] + layer_block_sizes: list[int] group_indices: list[int] - spec_indices: list[int] kv_caches_base_addr: list[list[int]] block_strides: list[list[int]] block_lens: list[list[int]] @@ -41,8 +39,8 @@ class MooncakeTransferMetadata(KVConnectorHandshakeMetadata): def __post_init__(self) -> None: num_layers = len(self.layer_names) per_layer_fields: tuple[tuple[str, Sequence[object]], ...] = ( + ("layer_block_sizes", self.layer_block_sizes), ("group_indices", self.group_indices), - ("spec_indices", self.spec_indices), ("kv_caches_base_addr", self.kv_caches_base_addr), ("block_strides", self.block_strides), ("block_lens", self.block_lens), @@ -77,6 +75,10 @@ class MooncakeTPTransferMetadata: """TP-private connection and KV-cache address information.""" te_rpc_port: int + # PP-union layer indices physically owned by this TP rank. The address + # table remains aligned with the PP-union layer order; unowned entries are + # empty lists. + layer_indices: list[int] kv_caches_base_addr: list[list[int]] local_ip: str handshake_port: int @@ -88,10 +90,9 @@ class MooncakePPTransferMetadata: block_size: int num_blocks: int - spec_block_sizes: list[int] layer_names: list[str] + layer_block_sizes: list[int] group_indices: list[int] - spec_indices: list[int] block_strides: list[list[int]] block_lens: list[list[int]] block_shapes: list[list[tuple[int, ...]]] @@ -110,6 +111,7 @@ class MooncakeTransferMetadataGroups: pcp_size: int dcp_size: int tp_size: int + use_kv_pp: bool metadata_by_pp_rank: dict[int, MooncakePPTransferMetadata] diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index 4caa8d6897da..f337bf501e1f 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -75,6 +75,7 @@ def __init__( self.pcp_size = pcp_size self.dcp_size = dcp_size self.engine_id = engine_id + self.use_kv_pp = False metadata_by_pp_rank = self._merge_metadata_by_pp_rank(metadata) self.encoded_metadata = encoder.encode( MooncakeTransferMetadataGroups( @@ -85,6 +86,7 @@ def __init__( pcp_size=pcp_size, dcp_size=dcp_size, tp_size=tp_size, + use_kv_pp=self.use_kv_pp, metadata_by_pp_rank=metadata_by_pp_rank, ) ) @@ -99,7 +101,14 @@ def _merge_metadata_by_pp_rank( self, metadata: Mapping[int | tuple[int, ...], KVConnectorHandshakeMetadata], ) -> dict[int, MooncakePPTransferMetadata]: - """Validate worker metadata and merge TP-invariant fields by PP.""" + """Merge worker metadata into one PP-wide layer table. + + LayerSplit may assign different layers to redundant TP ranks. Per-layer + fields are therefore aligned by layer name instead of requiring every + TP worker to expose the same layer list. TP-private address tables are + padded to the PP-union layer order, while ``layer_indices`` records the + entries physically owned by each TP rank. + """ workers_by_pp_rank: dict[int, dict[int, MooncakeTransferMetadata]] = {} for metadata_key, rank_metadata in metadata.items(): if isinstance(metadata_key, int): @@ -131,22 +140,9 @@ def _merge_metadata_by_pp_rank( if set(workers_by_pp_rank) != expected_pp_ranks: raise ValueError( "Mooncake worker metadata has incomplete PP ranks: expected " - f"{sorted(expected_pp_ranks)}, got " - f"{sorted(workers_by_pp_rank)}" + f"{sorted(expected_pp_ranks)}, got {sorted(workers_by_pp_rank)}" ) - common_fields = ( - "block_size", - "num_blocks", - "spec_block_sizes", - "layer_names", - "group_indices", - "spec_indices", - "block_strides", - "block_lens", - "block_shapes", - "block_size_scales", - ) expected_tp_ranks = set(range(self.tp_size)) merged: dict[int, MooncakePPTransferMetadata] = {} for pp_rank, workers_by_tp_rank in sorted(workers_by_pp_rank.items()): @@ -159,11 +155,10 @@ def _merge_metadata_by_pp_rank( reference_tp_rank = min(workers_by_tp_rank) reference = workers_by_tp_rank[reference_tp_rank] - for tp_rank, worker_metadata in workers_by_tp_rank.items(): mismatched_fields = [ field_name - for field_name in common_fields + for field_name in ("block_size", "num_blocks") if getattr(worker_metadata, field_name) != getattr(reference, field_name) ] if mismatched_fields: @@ -173,26 +168,97 @@ def _merge_metadata_by_pp_rank( f"TP {tp_rank} mismatch in {mismatched_fields}" ) + layer_source_by_name: dict[str, tuple[MooncakeTransferMetadata, int]] = {} + layer_signature_by_name: dict[str, tuple[object, ...]] = {} + tp_layer_names: list[set[str]] = [] + for tp_rank, worker_metadata in sorted(workers_by_tp_rank.items()): + if len(set(worker_metadata.layer_names)) != len(worker_metadata.layer_names): + raise ValueError( + f"Mooncake worker metadata for PP rank {pp_rank}, TP rank {tp_rank} contains duplicate layers" + ) + tp_layer_names.append(set(worker_metadata.layer_names)) + for local_layer_index, layer_name in enumerate(worker_metadata.layer_names): + layer_signature = ( + worker_metadata.group_indices[local_layer_index], + worker_metadata.layer_block_sizes[local_layer_index], + worker_metadata.block_strides[local_layer_index], + worker_metadata.block_lens[local_layer_index], + worker_metadata.block_shapes[local_layer_index], + worker_metadata.block_size_scales[local_layer_index], + ) + previous_signature = layer_signature_by_name.get(layer_name) + if previous_signature is not None and previous_signature != layer_signature: + layer_field_names = ( + "group_indices", + "layer_block_sizes", + "block_strides", + "block_lens", + "block_shapes", + "block_size_scales", + ) + mismatched_layer_fields = [ + field_name + for field_name, previous_value, value in zip( + layer_field_names, previous_signature, layer_signature + ) + if previous_value != value + ] + raise ValueError( + "Mooncake worker metadata differs across TP ranks for " + f"PP rank {pp_rank}, layer {layer_name!r}: mismatch in {mismatched_layer_fields}" + ) + if previous_signature is None: + layer_source_by_name[layer_name] = (worker_metadata, local_layer_index) + layer_signature_by_name[layer_name] = layer_signature + + layer_names = sorted(layer_source_by_name) + has_layer_split = any(layer_set != tp_layer_names[0] for layer_set in tp_layer_names[1:]) + self.use_kv_pp |= has_layer_split + if has_layer_split and self.dcp_size != 1: + raise ValueError(f"Mooncake LayerSplit cannot be combined with DCP, got dcp_size={self.dcp_size}") + + layer_block_sizes: list[int] = [] + group_indices: list[int] = [] + block_strides: list[list[int]] = [] + block_lens: list[list[int]] = [] + block_shapes: list[list[tuple[int, ...]]] = [] + block_size_scales: list[list[int]] = [] + for layer_name in layer_names: + worker_metadata, local_layer_index = layer_source_by_name[layer_name] + layer_block_sizes.append(worker_metadata.layer_block_sizes[local_layer_index]) + group_indices.append(worker_metadata.group_indices[local_layer_index]) + block_strides.append(worker_metadata.block_strides[local_layer_index]) + block_lens.append(worker_metadata.block_lens[local_layer_index]) + block_shapes.append(worker_metadata.block_shapes[local_layer_index]) + block_size_scales.append(worker_metadata.block_size_scales[local_layer_index]) + + layer_index_by_name = {layer_name: layer_index for layer_index, layer_name in enumerate(layer_names)} + metadata_by_tp_rank: dict[int, MooncakeTPTransferMetadata] = {} + for tp_rank, worker_metadata in sorted(workers_by_tp_rank.items()): + layer_indices = sorted(layer_index_by_name[layer_name] for layer_name in worker_metadata.layer_names) + aligned_base_addrs: list[list[int]] = [[] for _ in layer_names] + for local_layer_index, layer_name in enumerate(worker_metadata.layer_names): + layer_index = layer_index_by_name[layer_name] + aligned_base_addrs[layer_index] = worker_metadata.kv_caches_base_addr[local_layer_index] + metadata_by_tp_rank[tp_rank] = MooncakeTPTransferMetadata( + te_rpc_port=worker_metadata.te_rpc_port, + layer_indices=layer_indices, + kv_caches_base_addr=aligned_base_addrs, + local_ip=worker_metadata.local_ip, + handshake_port=worker_metadata.handshake_port, + ) + merged[pp_rank] = MooncakePPTransferMetadata( block_size=reference.block_size, num_blocks=reference.num_blocks, - spec_block_sizes=reference.spec_block_sizes, - layer_names=reference.layer_names, - group_indices=reference.group_indices, - spec_indices=reference.spec_indices, - block_strides=reference.block_strides, - block_lens=reference.block_lens, - block_shapes=reference.block_shapes, - block_size_scales=reference.block_size_scales, - metadata_by_tp_rank={ - tp_rank: MooncakeTPTransferMetadata( - te_rpc_port=worker_metadata.te_rpc_port, - kv_caches_base_addr=(worker_metadata.kv_caches_base_addr), - local_ip=worker_metadata.local_ip, - handshake_port=worker_metadata.handshake_port, - ) - for tp_rank, worker_metadata in sorted(workers_by_tp_rank.items()) - }, + layer_names=layer_names, + layer_block_sizes=layer_block_sizes, + group_indices=group_indices, + block_strides=block_strides, + block_lens=block_lens, + block_shapes=block_shapes, + block_size_scales=block_size_scales, + metadata_by_tp_rank=metadata_by_tp_rank, ) return merged 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 d6da079c5277..c0067d00dec6 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 @@ -84,10 +84,10 @@ def __init__( self.local_metadata = local_metadata self.block_size = local_metadata.block_size - self.spec_block_sizes = local_metadata.spec_block_sizes self.layer_names = local_metadata.layer_names + self.layer_block_sizes = local_metadata.layer_block_sizes self.group_indices = local_metadata.group_indices - self.spec_indices = local_metadata.spec_indices + self.spec_indices = [layer_name_to_spec_index[layer_name] for layer_name in self.layer_names] self.kv_caches_base_addr = local_metadata.kv_caches_base_addr self.block_strides = local_metadata.block_strides self.block_lens = local_metadata.block_lens @@ -215,9 +215,12 @@ def _build_remote_transfer_layout( dict[int, list[tuple[int, int]]], ]: """Build per-PP layer candidates and matching layer index pairs.""" + if remote_metadata.use_kv_pp and (self.dcp_size != 1 or remote_metadata.dcp_size != 1): + raise ValueError("Mooncake KV parallel cannot be combined with DCP") groups_by_pp_rank: dict[int, dict[tuple[int, int], list[list[int]]]] = {} layer_pairs_by_pp_rank: dict[int, list[tuple[int, int]]] = {} matched_local_layer_indices: set[int] = set() + raw_tp_rank_groups_by_spec: dict[int, list[list[int]]] = {} for remote_pp_rank, pp_metadata in sorted(remote_metadata.metadata_by_pp_rank.items()): layer_index_pairs: list[tuple[int, int]] = [] @@ -225,6 +228,12 @@ def _build_remote_transfer_layout( remote_layer_index_by_name = { layer_name: layer_index for layer_index, layer_name in enumerate(pp_metadata.layer_names) } + owner_tp_ranks_by_layer_index: list[set[int]] = [set() for _ in pp_metadata.layer_names] + if remote_metadata.use_kv_pp: + for remote_tp_rank, tp_metadata in pp_metadata.metadata_by_tp_rank.items(): + for remote_layer_index in tp_metadata.layer_indices: + owner_tp_ranks_by_layer_index[remote_layer_index].add(remote_tp_rank) + for local_layer_index, layer_name in enumerate(self.layer_names): remote_layer_index = remote_layer_index_by_name.get(layer_name) if remote_layer_index is None: @@ -234,14 +243,30 @@ def _build_remote_transfer_layout( layer_index_pairs.append(layer_pair) matched_local_layer_indices.add(local_layer_index) local_spec_index = self.spec_indices[local_layer_index] - groups_by_layer_pair[layer_pair] = self._get_layer_remote_tp_rank_groups( - local_layer_index, - remote_layer_index, - self.kv_cache_specs[local_spec_index], - pp_metadata, - remote_metadata.tp_size, - remote_metadata.dcp_size, - ) + raw_tp_rank_groups = raw_tp_rank_groups_by_spec.get(local_spec_index) + if raw_tp_rank_groups is None: + raw_tp_rank_groups = self._get_layer_remote_tp_rank_groups( + local_layer_index, + remote_layer_index, + self.kv_cache_specs[local_spec_index], + pp_metadata, + remote_metadata.tp_size, + remote_metadata.dcp_size, + ) + raw_tp_rank_groups_by_spec[local_spec_index] = raw_tp_rank_groups + + remote_tp_rank_groups = raw_tp_rank_groups + if remote_metadata.use_kv_pp: + owner_tp_ranks = owner_tp_ranks_by_layer_index[remote_layer_index] + remote_tp_rank_groups = [ + [tp_rank for tp_rank in group if tp_rank in owner_tp_ranks] for group in remote_tp_rank_groups + ] + if any(not group for group in remote_tp_rank_groups): + raise ValueError( + "Mooncake KV parallel metadata has no producer TP owning " + f"layer {layer_name!r} for every required TP group on PP rank {remote_pp_rank}" + ) + groups_by_layer_pair[layer_pair] = remote_tp_rank_groups layer_pairs_by_pp_rank[remote_pp_rank] = layer_index_pairs groups_by_pp_rank[remote_pp_rank] = groups_by_layer_pair @@ -416,7 +441,13 @@ def _handle_requests( remote_metadata = self._get_remote_metadata(remote_engine_id, remote_host, remote_port) tp_rank_groups_by_pp_rank = self.remote_tp_rank_groups[remote_engine_id] layer_pairs_by_pp_rank = self.remote_layer_index_pairs[remote_engine_id] - transfer_block_ids_by_spec: dict[str, dict[int, list[tuple[int, list[int], list[int]]]]] = {} + transfer_block_ids_by_spec: dict[ + str, + dict[ + tuple[int, tuple[tuple[int, ...], ...]], + list[tuple[int, list[int], list[int]]], + ], + ] = {} future_to_task: dict[Future[None], tuple[int, int, set[str]]] = {} submission_error: Exception | None = None try: @@ -480,7 +511,13 @@ def _build_transfer_block_buckets( tp_rank_groups_by_layer: dict[tuple[int, int], list[list[int]]], remote_dcp_size: int, requests: dict[str, ReqMeta], - transfer_block_ids_by_spec: dict[str, dict[int, list[tuple[int, list[int], list[int]]]]], + transfer_block_ids_by_spec: dict[ + str, + dict[ + tuple[int, tuple[tuple[int, ...], ...]], + list[tuple[int, list[int], list[int]]], + ], + ], ) -> tuple[ dict[int, dict[int, dict[tuple[int, int], list[tuple[str, list[int], list[int]]]]]], dict[int, set[str]], @@ -502,17 +539,20 @@ def _build_transfer_block_buckets( spec_index = self.spec_indices[local_layer_index] spec = self.kv_cache_specs[spec_index] layer_pair = (local_layer_index, remote_layer_index) - transfer_block_ids = request_block_ids_by_spec.get(spec_index) + remote_tp_rank_groups = tp_rank_groups_by_layer[layer_pair] + cache_key = ( + spec_index, + tuple(tuple(group) for group in remote_tp_rank_groups), + ) + transfer_block_ids = request_block_ids_by_spec.get(cache_key) if transfer_block_ids is None: group_index = self.group_indices[local_layer_index] - remote_spec_index = remote_metadata.spec_indices[remote_layer_index] if not isinstance(spec, MambaSpec): local_kernel_block_size = ( - self.local_metadata.spec_block_sizes[spec_index] - // self.block_size_scales[local_layer_index][0] + self.layer_block_sizes[local_layer_index] // self.block_size_scales[local_layer_index][0] ) remote_kernel_block_size = ( - remote_metadata.spec_block_sizes[remote_spec_index] + remote_metadata.layer_block_sizes[remote_layer_index] // remote_metadata.block_size_scales[remote_layer_index][0] ) assert local_kernel_block_size == remote_kernel_block_size, ( @@ -522,11 +562,11 @@ def _build_transfer_block_buckets( ) transfer_block_ids = self._compute_group_block_ids( request_id, - tp_rank_groups_by_layer[layer_pair], + remote_tp_rank_groups, remote_dcp_size, spec_index, - self.local_metadata.spec_block_sizes[spec_index], - remote_metadata.spec_block_sizes[remote_spec_index], + self.layer_block_sizes[local_layer_index], + remote_metadata.layer_block_sizes[remote_layer_index], request_metadata.local_block_ids[group_index], request_metadata.local_full_block_ids[group_index], request_metadata.remote_block_ids[group_index], @@ -538,7 +578,7 @@ def _build_transfer_block_buckets( spec, selection_index, ) - request_block_ids_by_spec[spec_index] = transfer_block_ids + request_block_ids_by_spec[cache_key] = transfer_block_ids for remote_tp_rank, local_block_ids, remote_block_ids in transfer_block_ids: request_ids_by_remote_tp_rank.setdefault(remote_tp_rank, set()).add(request_id) From da5aa121ba00afe3f85028a0b3647a446fc73cf5 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Sun, 23 Aug 2026 02:07:02 +0800 Subject: [PATCH 41/48] add comments & fix mypy Signed-off-by: nwpu-zxr --- .../kv_p2p/mooncake/pull_scheduler.py | 8 +++ .../kv_p2p/mooncake/pull_worker.py | 52 ++++++++++++++++--- 2 files changed, 52 insertions(+), 8 deletions(-) diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index f337bf501e1f..b6743a94bd0f 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -109,6 +109,7 @@ def _merge_metadata_by_pp_rank( padded to the PP-union layer order, while ``layer_indices`` records the entries physically owned by each TP rank. """ + # pp_rank -> tp_rank -> worker handshake metadata. workers_by_pp_rank: dict[int, dict[int, MooncakeTransferMetadata]] = {} for metadata_key, rank_metadata in metadata.items(): if isinstance(metadata_key, int): @@ -168,7 +169,12 @@ def _merge_metadata_by_pp_rank( f"TP {tp_rank} mismatch in {mismatched_fields}" ) + # layer_name -> (first owning worker metadata, its local layer index). + # The source supplies PP-shared fields after all duplicate owners + # have been checked against the corresponding structural signature. layer_source_by_name: dict[str, tuple[MooncakeTransferMetadata, int]] = {} + # layer_name -> (group index, block size, strides, lengths, shapes, + # block-size scales). layer_signature_by_name: dict[str, tuple[object, ...]] = {} tp_layer_names: list[set[str]] = [] for tp_rank, worker_metadata in sorted(workers_by_tp_rank.items()): @@ -236,6 +242,8 @@ def _merge_metadata_by_pp_rank( metadata_by_tp_rank: dict[int, MooncakeTPTransferMetadata] = {} for tp_rank, worker_metadata in sorted(workers_by_tp_rank.items()): layer_indices = sorted(layer_index_by_name[layer_name] for layer_name in worker_metadata.layer_names) + # PP-union layer index -> this TP's per-cache-tensor addresses; + # layers not owned by this TP retain an empty address list. aligned_base_addrs: list[list[int]] = [[] for _ in layer_names] for local_layer_index, layer_name in enumerate(worker_metadata.layer_names): layer_index = layer_index_by_name[layer_name] 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 c0067d00dec6..80e2fa72d2b2 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 @@ -122,10 +122,15 @@ def __init__( self.encoder = msgspec.msgpack.Encoder() self.decoder = msgspec.msgpack.Decoder(MooncakeTransferMetadataGroups) self.remote_metadata: SizedDict[str, MooncakeTransferMetadataGroups] = SizedDict() - # Candidate producer TP ranks, cached independently for each engine: - # engine_id -> remote_pp_rank -> (local_layer_index, remote_layer_index) -> candidate_groups + # Candidate producer TP ranks cached independently for each engine: + # engine_id -> remote_pp_rank -> (local_layer_index, remote_layer_index) + # -> candidate_groups. Each outer candidate group represents one + # remote head/TP piece required by this local worker. With remote DCP=1, + # ranks inside a group are interchangeable replicas and one is selected; + # with remote DCP>1, they are the DCP shards that jointly provide it. self.remote_tp_rank_groups: SizedDict[str, dict[int, dict[tuple[int, int], list[list[int]]]]] = SizedDict() - # engine_id -> remote_pp_rank -> [[local_layer_index, remote_layer_index], ...] + # engine_id -> remote_pp_rank + # -> [(local_layer_index, remote_layer_index), ...] self.remote_layer_index_pairs: SizedDict[str, dict[int, list[tuple[int, int]]]] = SizedDict() self.request_queue: queue.Queue[tuple[str, dict[str, ReqMeta]]] = queue.Queue() self.finished_requests: queue.SimpleQueue[str] = queue.SimpleQueue() @@ -214,12 +219,19 @@ def _build_remote_transfer_layout( dict[int, dict[tuple[int, int], list[list[int]]]], dict[int, list[tuple[int, int]]], ]: - """Build per-PP layer candidates and matching layer index pairs.""" + """Build per-PP layer matches and producer TP candidates. + + Candidate groups use the same layout as ``remote_tp_rank_groups``: + one outer entry per required remote head/TP piece and one inner list of + replica candidates (DCP=1) or participating DCP shards (DCP>1). + """ if remote_metadata.use_kv_pp and (self.dcp_size != 1 or remote_metadata.dcp_size != 1): raise ValueError("Mooncake KV parallel cannot be combined with DCP") groups_by_pp_rank: dict[int, dict[tuple[int, int], list[list[int]]]] = {} layer_pairs_by_pp_rank: dict[int, list[tuple[int, int]]] = {} matched_local_layer_indices: set[int] = set() + # local_spec_index -> topology-derived candidate groups before the + # per-layer KVPP owner filter. Layers sharing one spec reuse this value. raw_tp_rank_groups_by_spec: dict[int, list[list[int]]] = {} for remote_pp_rank, pp_metadata in sorted(remote_metadata.metadata_by_pp_rank.items()): @@ -228,11 +240,12 @@ def _build_remote_transfer_layout( remote_layer_index_by_name = { layer_name: layer_index for layer_index, layer_name in enumerate(pp_metadata.layer_names) } + # remote_layer_index -> TP ranks that physically own this layer. owner_tp_ranks_by_layer_index: list[set[int]] = [set() for _ in pp_metadata.layer_names] if remote_metadata.use_kv_pp: for remote_tp_rank, tp_metadata in pp_metadata.metadata_by_tp_rank.items(): - for remote_layer_index in tp_metadata.layer_indices: - owner_tp_ranks_by_layer_index[remote_layer_index].add(remote_tp_rank) + for owned_remote_layer_index in tp_metadata.layer_indices: + owner_tp_ranks_by_layer_index[owned_remote_layer_index].add(remote_tp_rank) for local_layer_index, layer_name in enumerate(self.layer_names): remote_layer_index = remote_layer_index_by_name.get(layer_name) @@ -389,7 +402,12 @@ def _get_attention_remote_tp_rank_groups( remote_dcp_size: int, total_num_kv_heads: int, ) -> list[list[int]]: - """Build remote TP groups from an already inferred head topology.""" + """Build remote TP groups from an already inferred head topology. + + Each result entry covers one intersection between this local rank's + head interval and a remote head shard. Its rank list contains either + interchangeable replicas or all DCP ranks for that remote head shard. + """ local_head_tp_size = self.tp_size // local_dcp_size remote_head_tp_size = remote_tp_size // remote_dcp_size local_head_tp_rank = self.tp_rank // local_dcp_size @@ -441,6 +459,13 @@ def _handle_requests( remote_metadata = self._get_remote_metadata(remote_engine_id, remote_host, remote_port) tp_rank_groups_by_pp_rank = self.remote_tp_rank_groups[remote_engine_id] layer_pairs_by_pp_rank = self.remote_layer_index_pairs[remote_engine_id] + # Cache block mappings for this producer-engine batch: + # request_id -> (local_spec_index, normalized_remote_tp_rank_groups) -> + # [(selected_remote_tp_rank, local_kernel_block_ids, + # remote_kernel_block_ids), ...]. + # Layers with the same spec and candidate topology share their block + # mapping across PP ranks. Candidate groups are part of the key because + # KVPP layers using the same spec can have different owner TP ranks. transfer_block_ids_by_spec: dict[ str, dict[ @@ -448,6 +473,8 @@ def _handle_requests( list[tuple[int, list[int], list[int]]], ], ] = {} + # future -> (remote_pp_rank, remote_tp_rank, affected_request_ids). + # A failed TP task marks only the requests represented in that bucket. future_to_task: dict[Future[None], tuple[int, int, set[str]]] = {} submission_error: Exception | None = None try: @@ -532,6 +559,8 @@ def _build_transfer_block_buckets( dict[tuple[int, int], list[tuple[str, list[int], list[int]]]], ], ] = {} + # remote_tp_rank -> requests with at least one transfer entry in its + # bucket; used to attribute a failed TP task to affected requests only. request_ids_by_remote_tp_rank: dict[int, set[str]] = {} for selection_index, (request_id, request_metadata) in enumerate(requests.items()): request_block_ids_by_spec = transfer_block_ids_by_spec.setdefault(request_id, {}) @@ -687,6 +716,9 @@ def _compute_group_block_ids( ) remote_tp_ranks = candidate_tp_ranks + # remote_tp_rank -> (local_kernel_block_ids, + # remote_kernel_block_ids). + # DCP remapping can split one logical request across ranks. block_ids_by_remote_tp_rank: dict[int, tuple[list[int], list[int]]] = {} if local_block_size != remote_block_size: for local_kernel_block_id, virtual_kernel_block_idx in local_virtual_kernel_blocks: @@ -1131,7 +1163,11 @@ def _execute_tp_transfer_bucket( dict[tuple[int, int], list[tuple[str, list[int], list[int]]]], ], ) -> None: - """Calculate addresses and execute one remote PP/TP transfer bucket.""" + """Calculate addresses and execute one remote PP/TP transfer bucket. + + ``transfer_entries_by_spec`` maps local spec index to layer-index pairs, + then to ``(request_id, local_block_ids, remote_block_ids)`` records. + """ src_list: list[int] = [] dst_list: list[int] = [] length_list: list[int] = [] From b9ba3e401909aee4bdf2d12f865d89446dbd5179 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Sun, 23 Aug 2026 02:39:34 +0800 Subject: [PATCH 42/48] support p-kvpp with d-dcp Signed-off-by: nwpu-zxr --- .../mooncake_v2/test_pull_worker.py | 23 +++++++++++++++++++ .../kv_p2p/mooncake/pull_worker.py | 4 ++-- 2 files changed, 25 insertions(+), 2 deletions(-) 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 15577aa9aa70..e933c8e61da8 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py @@ -200,6 +200,29 @@ def test_build_remote_layout_filters_kv_parallel_owners_and_keeps_mtp_replicas() assert thread._get_layer_remote_tp_rank_groups.call_count == 1 +def test_build_remote_layout_allows_local_dcp_with_remote_kv_parallel() -> None: + thread = make_thread(tp_size=2, dcp_size=2) + thread._get_layer_remote_tp_rank_groups = MagicMock(return_value=[[0, 1]]) # type: ignore[method-assign] + pp_metadata = make_pp_metadata( + tp_base_addrs={0: [[5000]], 1: [[]]}, + tp_layer_indices={0: [0], 1: []}, + ) + groups = make_metadata_groups(tp_size=2, use_kv_pp=True, pp_metadata=pp_metadata) + + rank_groups, pairs = thread._build_remote_transfer_layout(groups) + + assert pairs == {0: [(0, 0)]} + assert rank_groups == {0: {(0, 0): [[0]]}} + + +def test_build_remote_layout_rejects_remote_kv_parallel_with_remote_dcp() -> None: + thread = make_thread() + groups = replace(make_metadata_groups(use_kv_pp=True), dcp_size=2) + + with pytest.raises(ValueError, match="producer cannot enable KV parallel and DCP together"): + thread._build_remote_transfer_layout(groups) + + def test_build_remote_layout_rejects_missing_local_layer() -> None: thread = make_thread(layer_names=["layer.0", "missing"], spec_indices=[0, 0]) thread._get_layer_remote_tp_rank_groups = MagicMock(return_value=[[0]]) # type: ignore[method-assign] 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 80e2fa72d2b2..68f76faf79db 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 @@ -225,8 +225,8 @@ def _build_remote_transfer_layout( one outer entry per required remote head/TP piece and one inner list of replica candidates (DCP=1) or participating DCP shards (DCP>1). """ - if remote_metadata.use_kv_pp and (self.dcp_size != 1 or remote_metadata.dcp_size != 1): - raise ValueError("Mooncake KV parallel cannot be combined with DCP") + if remote_metadata.use_kv_pp and remote_metadata.dcp_size != 1: + raise ValueError("Mooncake producer cannot enable KV parallel and DCP together") groups_by_pp_rank: dict[int, dict[tuple[int, int], list[list[int]]]] = {} layer_pairs_by_pp_rank: dict[int, list[tuple[int, int]]] = {} matched_local_layer_indices: set[int] = set() From 30b6753ffa535549c1c638cdd0face09cc42f2fe Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Fri, 28 Aug 2026 10:16:04 +0800 Subject: [PATCH 43/48] fix indexerspec block size Signed-off-by: nwpu-zxr --- tests/ut/kv_offload/mooncake_v2/helpers.py | 14 +++++ .../mooncake_v2/test_base_worker.py | 36 ++++++++++- .../mooncake_v2/test_pull_worker.py | 63 +++++++++++++++++++ .../kv_p2p/mooncake/base_worker.py | 12 +++- 4 files changed, 123 insertions(+), 2 deletions(-) diff --git a/tests/ut/kv_offload/mooncake_v2/helpers.py b/tests/ut/kv_offload/mooncake_v2/helpers.py index 61721bd1adc4..fbdbfa2cb819 100644 --- a/tests/ut/kv_offload/mooncake_v2/helpers.py +++ b/tests/ut/kv_offload/mooncake_v2/helpers.py @@ -8,6 +8,7 @@ import torch from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, SlidingWindowSpec +from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( MooncakePPTransferMetadata, MooncakeTPTransferMetadata, @@ -44,6 +45,19 @@ def make_mamba_spec(block_size: int = 16) -> MambaSpec: ) +def make_sfa_indexer_spec( + block_size: int = 16, + replication_size: int = 2, +) -> AscendSFAIndexerCacheSpec: + return AscendSFAIndexerCacheSpec( + block_size=block_size, + num_kv_heads=1, + head_size=8, + dtype=torch.float16, + sfa_dcp_replicated_indexer_size=replication_size, + ) + + def make_transfer_metadata( *, engine_id: str = "engine-p", diff --git a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py index b94e5e06ae5a..707daf3098bc 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py @@ -17,7 +17,7 @@ ) from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.stats import MooncakeKVConnectorStats -from .helpers import make_full_spec, make_sliding_spec +from .helpers import make_full_spec, make_sfa_indexer_spec, make_sliding_spec def test_build_spec_mappings_expands_uniform_group_by_layer_spec() -> None: @@ -83,6 +83,40 @@ def test_register_kv_caches_uses_config_order_and_publishes_tensor_metadata(monk transfer_engine.register_buffer.assert_called_once() +def test_register_kv_caches_publishes_sfa_indexer_virtual_block_size(monkeypatch) -> None: + spec = make_sfa_indexer_spec(block_size=16, replication_size=2) + cache = torch.empty((4, 16, 1, 8), dtype=torch.float16) + config = KVCacheConfig( + num_blocks=2, + kv_cache_tensors=[KVCacheTensor(size=cache.nbytes, shared_by=["layer.0.indexer"])], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0.indexer"], kv_cache_spec=spec)], + ) + worker = MooncakeBaseConnectorWorker.__new__(MooncakeBaseConnectorWorker) + worker.kv_cache_config = config + worker.engine_id = "engine-d" + worker.te_rpc_port = 9000 + worker.block_size = 16 + worker.side_channel_host = "10.0.0.1" + worker.handshake_port = 5000 + transfer_engine = MagicMock() + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.global_te", + transfer_engine, + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.validate_register_region_count", + MagicMock(), + ) + + worker.register_kv_caches({"layer.0.indexer": cache}) + + metadata = worker.xfer_handshake_metadata + assert metadata is not None + assert metadata.layer_block_sizes == [32] + assert metadata.block_size_scales == [[2]] + assert metadata.layer_block_sizes[0] // metadata.block_size_scales[0][0] == spec.block_size + + def test_register_kv_caches_rejects_missing_and_unconfigured_layers() -> None: spec = make_full_spec() config = KVCacheConfig( 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 e933c8e61da8..eac6e9c7485d 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py @@ -29,6 +29,7 @@ make_mamba_spec, make_metadata_groups, make_pp_metadata, + make_sfa_indexer_spec, make_sliding_spec, make_transfer_metadata, ) @@ -262,6 +263,68 @@ def test_compute_full_attention_blocks_skips_remote_prefix_and_balances_replica( assert result == [(3, [10, 11], [21, 22])] +@pytest.mark.parametrize( + ("remote_block_size", "remote_scale", "expected_remote_blocks"), + [ + (32, 2, [42, 43, 44, 45]), + (64, 4, [82, 83, 84, 85]), + ], +) +def test_compute_sfa_indexer_blocks_uses_virtual_block_size_for_prefix( + remote_block_size: int, + remote_scale: int, + expected_remote_blocks: list[int], +) -> None: + thread = make_thread(dcp_size=2) + + result = thread._compute_group_block_ids( + request_id="request", + remote_tp_rank_groups=[[0]], + remote_dcp_size=remote_scale, + spec_index=0, + local_block_size=32, + remote_block_size=remote_block_size, + local_group_block_ids=[10, 11], + local_full_group_block_ids=[10, 11], + remote_group_block_ids=[20, 21, 22], + local_num_prompt_tokens=96, + remote_num_prompt_tokens=96, + num_computed_tokens=32, + local_block_size_scale=2, + remote_block_size_scale=remote_scale, + spec=make_sfa_indexer_spec(replication_size=2), + selection_index=0, + ) + + assert result == [(0, [20, 21, 22, 23], expected_remote_blocks)] + + +def test_transfer_bucket_accepts_sfa_indexer_virtual_block_sizes() -> None: + spec = make_sfa_indexer_spec(replication_size=2) + thread = make_thread( + dcp_size=2, + kv_cache_specs=[spec], + layer_block_sizes=[32], + block_size_scales=[[2]], + ) + remote = make_pp_metadata( + layer_block_sizes=[64], + block_size_scales=[[4]], + ) + + buckets, request_ids = thread._build_transfer_block_buckets( + remote_metadata=remote, + layer_pairs=[(0, 0)], + tp_rank_groups_by_layer={(0, 0): [[0]]}, + remote_dcp_size=4, + requests={"request": make_req_meta()}, + transfer_block_ids_by_spec={}, + ) + + assert buckets[0][0][(0, 0)] == [("request", [20, 21, 22, 23], [80, 81, 82, 83])] + assert request_ids == {0: {"request"}} + + def test_compute_sliding_window_blocks_uses_unhashed_suffix() -> None: thread = make_thread() 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 23ce3d5ca8fb..594aa1aa076c 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 @@ -24,6 +24,7 @@ from vllm.v1.kv_cache_interface import KVCacheSpec, UniformTypeKVCacheSpecs from vllm_ascend.ascend_config import get_ascend_config, init_ascend_config +from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( MooncakeConnectorMetadata, MooncakeTransferMetadata, @@ -195,7 +196,16 @@ def register_kv_caches( configured_layer_names.add(layer_name) layer_names.append(layer_name) spec_index = self.layer_name_to_spec_index[layer_name] - layer_block_sizes.append(self.kv_cache_specs[spec_index].block_size) + spec = self.kv_cache_specs[spec_index] + layer_block_size = spec.block_size + if isinstance(spec, AscendSFAIndexerCacheSpec): + # The cache manager treats one SFA indexer block as a DCP + # virtual block, while every worker physically stores all + # replicated indexer blocks. Publish the virtual token span + # so dividing it by the tensor block scale recovers the + # physical kernel block size. + layer_block_size *= spec.sfa_dcp_replicated_indexer_size + layer_block_sizes.append(layer_block_size) group_indices.append(self.layer_name_to_group_index[layer_name]) kv_caches_base_addr.append(base_addrs) block_strides_per_layer.append(block_strides) From a8af7c432dd91a21a934d07061bc1cb706fee4bf Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Mon, 7 Sep 2026 23:13:53 +0800 Subject: [PATCH 44/48] add block merge Signed-off-by: nwpu-zxr --- .../mooncake_v2/test_base_worker.py | 49 ++++++++ .../mooncake_v2/test_pull_worker.py | 54 +++++++- .../kv_p2p/mooncake/base_worker.py | 90 +++++++++++-- .../kv_p2p/mooncake/pull_worker.py | 119 ++++++++++++++---- 4 files changed, 274 insertions(+), 38 deletions(-) diff --git a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py index 707daf3098bc..1aaa461a8cd7 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py @@ -9,6 +9,7 @@ KVCacheConfig, KVCacheGroupSpec, KVCacheTensor, + MLAAttentionSpec, UniformTypeKVCacheSpecs, ) @@ -83,6 +84,54 @@ def test_register_kv_caches_uses_config_order_and_publishes_tensor_metadata(monk transfer_engine.register_buffer.assert_called_once() +def test_register_kv_caches_collapses_views_packed_in_one_page(monkeypatch) -> None: + spec = MLAAttentionSpec(block_size=16, num_kv_heads=1, head_size=32, dtype=torch.float16) + raw_cache = torch.empty(4 * 64, dtype=torch.float16) + k_cache = torch.as_strided(raw_cache, size=(4, 32), stride=(64, 1), storage_offset=0) + scale_cache = torch.as_strided(raw_cache, size=(4, 8), stride=(64, 1), storage_offset=32) + config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[KVCacheTensor(size=raw_cache.nbytes, shared_by=["layer.0"])], + kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=spec)], + ) + worker = MooncakeBaseConnectorWorker.__new__(MooncakeBaseConnectorWorker) + worker.kv_cache_config = config + worker.engine_id = "engine-d" + worker.te_rpc_port = 9000 + worker.block_size = 16 + worker.side_channel_host = "10.0.0.1" + worker.handshake_port = 5000 + transfer_engine = MagicMock() + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.global_te", + transfer_engine, + ) + monkeypatch.setattr( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.validate_register_region_count", + MagicMock(), + ) + + worker.register_kv_caches({"layer.0": [k_cache, scale_cache]}) + + metadata = worker.xfer_handshake_metadata + assert metadata is not None + assert metadata.kv_caches_base_addr == [[raw_cache.data_ptr()]] + assert metadata.block_strides == [[128]] + assert metadata.block_lens == [[128]] + assert metadata.block_shapes == [[(32,)]] + assert metadata.block_size_scales == [[1]] + + +def test_shared_storage_is_not_enough_to_identify_a_packed_page() -> None: + raw_cache = torch.empty(4 * 64, dtype=torch.float16) + k_cache = torch.as_strided(raw_cache, size=(4, 32), stride=(32, 1), storage_offset=0) + v_cache = torch.as_strided(raw_cache, size=(4, 32), stride=(32, 1), storage_offset=128) + worker = MooncakeBaseConnectorWorker.__new__(MooncakeBaseConnectorWorker) + worker.num_blocks = 4 + + assert worker._get_shared_page_metadata((k_cache, v_cache)) is None + + def test_register_kv_caches_publishes_sfa_indexer_virtual_block_size(monkeypatch) -> None: spec = make_sfa_indexer_spec(block_size=16, replication_size=2) cache = torch.empty((4, 16, 1, 8), dtype=torch.float16) 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 eac6e9c7485d..a70de26ee812 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py @@ -5,7 +5,7 @@ from concurrent.futures import Future from dataclasses import replace from types import SimpleNamespace -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import msgspec import pytest @@ -22,7 +22,10 @@ MooncakePullConnectorWorker, MooncakePullRecvingThread, ) -from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils import SizedDict +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils import ( + SizedDict, + group_concurrent_contiguous, +) from .helpers import ( make_full_spec, @@ -824,6 +827,53 @@ def test_whole_block_mla_address_generation_uses_independent_stride() -> None: assert lengths == [128, 128] +def test_whole_block_coalescing_is_prepared_per_request_and_reused_across_layers() -> None: + spec = MLAAttentionSpec(block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float16) + thread = make_thread( + layer_names=["layer.0", "layer.1"], + kv_cache_specs=[spec], + kv_caches_base_addr=[[1000], [2000]], + block_strides=[[128], [128]], + block_lens=[[128], [128]], + ) + remote = make_pp_metadata( + layer_names=["layer.0", "layer.1"], + block_strides=[[128], [128]], + block_lens=[[128], [128]], + tp_base_addrs={0: [[5000], [6000]]}, + ) + transfer_entries = [ + ("request-a", [1, 2], [3, 4]), + ("request-b", [3, 5], [5, 8]), + ] + src: list[int] = [] + dst: list[int] = [] + lengths: list[int] = [] + + with patch( + "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.group_concurrent_contiguous", + wraps=group_concurrent_contiguous, + ) as mock_group: + thread._append_spec_transfer_addresses( + 0, + 0, + 1, + 1, + {(0, 0): transfer_entries, (1, 1): transfer_entries}, + remote, + src, + dst, + lengths, + ) + + assert mock_group.call_count == 2 + mock_group.assert_any_call([1, 2], [3, 4]) + mock_group.assert_any_call([3, 5], [5, 8]) + assert src == [1128, 1384, 1640, 2128, 2384, 2640] + assert dst == [5384, 5640, 6024, 6384, 6640, 7024] + assert lengths == [256, 128, 128, 256, 128, 128] + + def test_mamba_unequal_tp_slices_conv_projections_and_state() -> None: spec = SimpleNamespace( dtypes=(torch.float16, torch.float16), 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 594aa1aa076c..328417724df9 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 @@ -21,7 +21,12 @@ ) from vllm.logger import logger from vllm.utils.network_utils import get_ip -from vllm.v1.kv_cache_interface import KVCacheSpec, UniformTypeKVCacheSpecs +from vllm.v1.kv_cache_interface import ( + KVCacheSpec, + MLAAttentionSpec, + SlidingWindowMLASpec, + UniformTypeKVCacheSpecs, +) from vllm_ascend.ascend_config import get_ascend_config, init_ascend_config from vllm_ascend.core.kv_cache_interface import AscendSFAIndexerCacheSpec @@ -41,6 +46,7 @@ ) from vllm_ascend.distributed.kv_transfer.utils.utils import ( get_transfer_timeout_value, + tensor_storage_key, validate_register_region_count, ) from vllm_ascend.distributed.utils import ( @@ -147,6 +153,54 @@ def _build_kv_cache_spec_mappings(self) -> None: self.layer_name_to_group_index[layer_name] = group_index self.layer_name_to_spec_index[layer_name] = spec_index + def _get_shared_page_metadata( + self, + caches: tuple[torch.Tensor, ...], + ) -> tuple[int, int, tuple[int, ...], int] | None: + """Describe views packed into one physical page allocation. + + Sharing a storage alone is insufficient: conventional planar K/V + views can share one allocation while their first blocks are many pages + apart. A packed layout must place every view's first block inside one + common page stride and use the same physical block count and stride. + """ + if not caches or len({tensor_storage_key(cache) for cache in caches}) != 1: + return None + + tensor_num_blocks = {cache.shape[0] for cache in caches} + page_strides = {cache.stride(0) * cache.element_size() for cache in caches} + if len(tensor_num_blocks) != 1 or len(page_strides) != 1: + return None + + 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: + return None + + page_base_addr = min(cache.data_ptr() for cache in caches) + block_lens = [math.prod(cache.shape[1:]) * cache.element_size() for cache in caches] + if any( + cache.data_ptr() - page_base_addr + block_len > page_stride for cache, block_len in zip(caches, block_lens) + ): + return None + + storage = caches[0].untyped_storage() + storage_base_addr = storage.data_ptr() + storage_end_addr = storage_base_addr + storage.nbytes() + if page_base_addr < storage_base_addr or page_base_addr + num_tensor_blocks * page_stride > storage_end_addr: + return None + + selected_cache = max( + zip(caches, block_lens), + key=lambda cache_and_len: cache_and_len[1], + )[0] + return ( + page_base_addr, + page_stride, + tuple(selected_cache.shape[1:]), + num_tensor_blocks // self.num_blocks, + ) + def register_kv_caches( self, kv_caches: dict[str, torch.Tensor | list[torch.Tensor]], @@ -183,20 +237,34 @@ def register_kv_caches( block_shapes: list[tuple[int, ...]] = [] block_size_scales: list[int] = [] - for cache in as_kv_cache_tensors(cache_or_caches): - tensor_num_blocks = cache.shape[0] - element_size = cache.element_size() - block_shape = tuple(cache.shape[1:]) - base_addrs.append(cache.data_ptr()) - block_strides.append(cache.stride(0) * element_size) - block_lens.append(math.prod(block_shape) * element_size) + spec_index = self.layer_name_to_spec_index[layer_name] + spec = self.kv_cache_specs[spec_index] + caches = as_kv_cache_tensors(cache_or_caches) + shared_page_metadata = ( + self._get_shared_page_metadata(caches) + if isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)) + else None + ) + if shared_page_metadata is not None: + page_base_addr, page_stride, block_shape, block_size_scale = shared_page_metadata + base_addrs.append(page_base_addr) + block_strides.append(page_stride) + block_lens.append(page_stride) block_shapes.append(block_shape) - block_size_scales.append(tensor_num_blocks // self.num_blocks) + block_size_scales.append(block_size_scale) + else: + for cache in caches: + tensor_num_blocks = cache.shape[0] + element_size = cache.element_size() + block_shape = tuple(cache.shape[1:]) + base_addrs.append(cache.data_ptr()) + 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) configured_layer_names.add(layer_name) layer_names.append(layer_name) - spec_index = self.layer_name_to_spec_index[layer_name] - spec = self.kv_cache_specs[spec_index] layer_block_size = spec.block_size if isinstance(spec, AscendSFAIndexerCacheSpec): # The cache manager treats one SFA indexer block as a DCP 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 68f76faf79db..1ca081812858 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 @@ -39,6 +39,7 @@ SizedDict, ensure_zmq_recv, ensure_zmq_send, + group_concurrent_contiguous, zmq_ctx, ) @@ -977,6 +978,27 @@ def _append_mamba_transfer_addresses( dst_list.append(remote_block_addr + remote_inner_offset) length_list.append(transfer_len) + @staticmethod + def _append_block_transfer_addresses( + local_block_groups: list[list[int]], + remote_block_groups: list[list[int]], + local_base_addr: int, + remote_base_addr: int, + local_block_stride: int, + remote_block_stride: int, + transfer_len: int, + src_list: list[int], + dst_list: list[int], + length_list: list[int], + local_inner_offset: int = 0, + remote_inner_offset: int = 0, + ) -> None: + """Append addresses from block groups prepared for one request.""" + for local_block_group, remote_block_group in zip(local_block_groups, remote_block_groups): + src_list.append(local_base_addr + local_block_group[0] * local_block_stride + local_inner_offset) + dst_list.append(remote_base_addr + remote_block_group[0] * remote_block_stride + remote_inner_offset) + length_list.append(transfer_len * len(local_block_group)) + def _append_spec_transfer_addresses( self, spec_index: int, @@ -990,21 +1012,6 @@ def _append_spec_transfer_addresses( length_list: list[int], ) -> None: """Append all layers and requests belonging to one local spec.""" - for (local_layer_index, _), transfer_entries in transfer_entries_by_layer.items(): - for request_id, local_block_ids, remote_block_ids in transfer_entries: - if len(local_block_ids) != len(remote_block_ids): - raise ValueError( - f"Mooncake block ID count mismatch for request {request_id!r}, " - f"layer {self.layer_names[local_layer_index]!r}: " - f"local={len(local_block_ids)}, remote={len(remote_block_ids)}" - ) - if not any( - local_block_ids - for transfer_entries in transfer_entries_by_layer.values() - for _, local_block_ids, _ in transfer_entries - ): - return - spec = self.kv_cache_specs[spec_index] if isinstance(spec, MambaSpec): self._append_mamba_transfer_addresses( @@ -1018,6 +1025,37 @@ def _append_spec_transfer_addresses( length_list, ) return + + block_ids_by_request: dict[str, tuple[list[int], list[int]]] = {} + for (local_layer_index, _), transfer_entries in transfer_entries_by_layer.items(): + for request_id, local_block_ids, remote_block_ids in transfer_entries: + if len(local_block_ids) != len(remote_block_ids): + raise ValueError( + f"Mooncake block ID count mismatch for request {request_id!r}, " + f"layer {self.layer_names[local_layer_index]!r}: " + f"local={len(local_block_ids)}, remote={len(remote_block_ids)}" + ) + block_ids = (local_block_ids, remote_block_ids) + previous_block_ids = block_ids_by_request.setdefault(request_id, block_ids) + if previous_block_ids != block_ids: + raise ValueError( + f"Mooncake block IDs differ within spec {spec_index} for request {request_id!r}" + ) + + if not any(local_block_ids for local_block_ids, _ in block_ids_by_request.values()): + return + + single_block_ids_by_request = { + request_id: ( + [[block_id] for block_id in local_block_ids], + [[block_id] for block_id in remote_block_ids], + ) + for request_id, (local_block_ids, remote_block_ids) in block_ids_by_request.items() + } + contiguous_block_ids_by_request = { + request_id: group_concurrent_contiguous(local_block_ids, remote_block_ids) + for request_id, (local_block_ids, remote_block_ids) in block_ids_by_request.items() + } remote_tp_metadata = remote_metadata.metadata_by_tp_rank[remote_tp_rank] transfer_whole_block = isinstance( spec, @@ -1045,11 +1083,25 @@ def _append_spec_transfer_addresses( f"{self.layer_names[local_layer_index]!r}, cache {cache_index}: " f"local={local_block_len}, remote={remote_block_len}" ) - for _, local_block_ids, remote_block_ids in transfer_entries: - for local_block_id, remote_block_id in zip(local_block_ids, remote_block_ids): - src_list.append(local_base_addr + local_block_id * local_block_stride) - dst_list.append(remote_base_addr + remote_block_id * remote_block_stride) - length_list.append(local_block_len) + block_ids_to_transfer = ( + contiguous_block_ids_by_request + if local_block_stride == remote_block_stride == local_block_len + else single_block_ids_by_request + ) + for request_id, _, _ in transfer_entries: + local_block_groups, remote_block_groups = block_ids_to_transfer[request_id] + self._append_block_transfer_addresses( + local_block_groups=local_block_groups, + remote_block_groups=remote_block_groups, + local_base_addr=local_base_addr, + remote_base_addr=remote_base_addr, + local_block_stride=local_block_stride, + remote_block_stride=remote_block_stride, + transfer_len=local_block_len, + src_list=src_list, + dst_list=dst_list, + length_list=length_list, + ) return if not isinstance(spec, (FullAttentionSpec, SlidingWindowSpec)): @@ -1145,11 +1197,28 @@ def _append_spec_transfer_addresses( transfer_len = num_transfer_heads * local_head_len local_inner_offset = local_head_offset * local_head_len remote_inner_offset = remote_head_offset * remote_head_len - for _, local_block_ids, remote_block_ids in transfer_entries: - for local_block_id, remote_block_id in zip(local_block_ids, remote_block_ids): - src_list.append(local_base_addr + local_block_id * local_block_stride + local_inner_offset) - dst_list.append(remote_base_addr + remote_block_id * remote_block_stride + remote_inner_offset) - length_list.append(transfer_len) + block_ids_to_transfer = ( + contiguous_block_ids_by_request + if local_block_stride == local_block_len == transfer_len + and remote_block_stride == remote_block_len == transfer_len + else single_block_ids_by_request + ) + for request_id, _, _ in transfer_entries: + local_block_groups, remote_block_groups = block_ids_to_transfer[request_id] + self._append_block_transfer_addresses( + local_block_groups=local_block_groups, + remote_block_groups=remote_block_groups, + local_base_addr=local_base_addr, + remote_base_addr=remote_base_addr, + local_block_stride=local_block_stride, + remote_block_stride=remote_block_stride, + transfer_len=transfer_len, + local_inner_offset=local_inner_offset, + remote_inner_offset=remote_inner_offset, + src_list=src_list, + dst_list=dst_list, + length_list=length_list, + ) def _execute_tp_transfer_bucket( self, From d07686af74a52bcc32cbffa13f4d18ebe35167bf Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Mon, 7 Sep 2026 23:23:57 +0800 Subject: [PATCH 45/48] fix ruff Signed-off-by: nwpu-zxr --- .../distributed/kv_transfer/kv_p2p/mooncake/pull_worker.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 1ca081812858..59d9b2a8dd02 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 @@ -1038,9 +1038,7 @@ def _append_spec_transfer_addresses( block_ids = (local_block_ids, remote_block_ids) previous_block_ids = block_ids_by_request.setdefault(request_id, block_ids) if previous_block_ids != block_ids: - raise ValueError( - f"Mooncake block IDs differ within spec {spec_index} for request {request_id!r}" - ) + raise ValueError(f"Mooncake block IDs differ within spec {spec_index} for request {request_id!r}") if not any(local_block_ids for local_block_ids, _ in block_ids_by_request.values()): return From 3bd67f32f01c58a881ffc16b2cde4b5485deeb27 Mon Sep 17 00:00:00 2001 From: wangxiaoteng Date: Mon, 7 Sep 2026 21:39:34 +0800 Subject: [PATCH 46/48] chore(mooncake): add connector V2 runtime logs Signed-off-by: wangxiaoteng --- .../kv_p2p/mooncake/base_worker.py | 10 +++++++++ .../kv_p2p/mooncake/pull_scheduler.py | 13 +++++++++++ .../kv_p2p/mooncake/pull_worker.py | 22 +++++++++++++++++++ 3 files changed, 45 insertions(+) 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 328417724df9..78d0b6d783c3 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 @@ -308,6 +308,16 @@ def register_kv_caches( self.transfer_metadata = transfer_metadata self.xfer_handshake_metadata = transfer_metadata + logger.info( + "Registered Mooncake KV caches: engine_id=%s, num_layers=%d, " + "num_blocks=%d, num_regions=%d, registered_bytes=%d", + self.engine_id, + len(layer_names), + self.num_blocks, + len(register_regions.ptrs), + sum(register_regions.lengths), + ) + logger.debug( "Mooncake KV cache transfer metadata: metadata=%s, register_ptrs=%s, register_lengths=%s", transfer_metadata, diff --git a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py index b6743a94bd0f..43b79ebf71a6 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/pull_scheduler.py @@ -90,6 +90,19 @@ def __init__( metadata_by_pp_rank=metadata_by_pp_rank, ) ) + logger.info( + "Initialized Mooncake scheduler handshake metadata: engine_id=%s, " + "worker_count=%d, pp_ranks=%s, tp_size=%d, pp_size=%d, " + "pcp_size=%d, dcp_size=%d, use_kv_pp=%s", + engine_id, + len(metadata), + sorted(metadata_by_pp_rank), + tp_size, + pp_size, + pcp_size, + dcp_size, + self.use_kv_pp, + ) self.ready_event = ready_event self.delayed_free_requests: OrderedDict[str, float] = OrderedDict() self.finished_requests: queue.SimpleQueue[str] = queue.SimpleQueue() 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 59d9b2a8dd02..e331698c8bd6 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 @@ -4,6 +4,7 @@ import queue import threading +import time from concurrent.futures import Future, ThreadPoolExecutor, as_completed from typing import Any @@ -1255,11 +1256,32 @@ def _execute_tp_transfer_bucket( return tp_metadata = remote_metadata.metadata_by_tp_rank[remote_tp_rank] session_id = f"{tp_metadata.local_ip}:{tp_metadata.te_rpc_port}" + request_ids = sorted( + { + request_id + for transfer_entries_by_layer in transfer_entries_by_spec.values() + for transfer_entries in transfer_entries_by_layer.values() + for request_id, _, _ in transfer_entries + } + ) + transfer_start_time = time.perf_counter() ret = self.engine.batch_transfer_sync_read(session_id, src_list, dst_list, length_list) if ret < 0: raise RuntimeError( f"Mooncake KV transfer failed for remote PP rank {remote_pp_rank}, TP rank {remote_tp_rank}, ret={ret}" ) + 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). " + "local_ip %s local_device_id %s remote_session_id %s", + request_ids, + transfer_elapsed_ms, + len(length_list), + sum(length_list), + self.local_metadata.local_ip, + self.tp_rank, + session_id, + ) @staticmethod def _validate_remote_metadata( From 5f12fbafef30592066432d9c22e28df03ea8ec15 Mon Sep 17 00:00:00 2001 From: nwpu-zxr Date: Wed, 9 Sep 2026 16:29:07 +0800 Subject: [PATCH 47/48] support new kv_cache_tensors Signed-off-by: nwpu-zxr --- .../mooncake_v2/test_base_worker.py | 31 ++++++++-- tests/ut/kv_offload/mooncake_v2/test_utils.py | 59 ++++++++++++++++++- tests/ut/worker/test_model_runner_v1.py | 2 + .../kv_p2p/mooncake/base_worker.py | 2 +- .../kv_transfer/kv_p2p/mooncake/utils.py | 52 ++++++++++++---- vllm_ascend/worker/model_runner_v1.py | 14 +++-- 6 files changed, 137 insertions(+), 23 deletions(-) diff --git a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py index 1aaa461a8cd7..ed7a44d37ae6 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py @@ -49,7 +49,9 @@ def test_register_kv_caches_uses_config_order_and_publishes_tensor_metadata(monk kv_cache_tensors=[ KVCacheTensor( size=k_cache.nbytes + v_cache.nbytes, - shared_by=["layer.0"], + layers=["layer.0"], + layer_stride=k_cache.nbytes + v_cache.nbytes, + block_stride=spec.page_size_bytes, ) ], kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=spec)], @@ -91,7 +93,14 @@ def test_register_kv_caches_collapses_views_packed_in_one_page(monkeypatch) -> N scale_cache = torch.as_strided(raw_cache, size=(4, 8), stride=(64, 1), storage_offset=32) config = KVCacheConfig( num_blocks=4, - kv_cache_tensors=[KVCacheTensor(size=raw_cache.nbytes, shared_by=["layer.0"])], + kv_cache_tensors=[ + KVCacheTensor( + size=raw_cache.nbytes, + layers=["layer.0"], + layer_stride=raw_cache.nbytes, + block_stride=128, + ) + ], kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=spec)], ) worker = MooncakeBaseConnectorWorker.__new__(MooncakeBaseConnectorWorker) @@ -137,7 +146,14 @@ def test_register_kv_caches_publishes_sfa_indexer_virtual_block_size(monkeypatch cache = torch.empty((4, 16, 1, 8), dtype=torch.float16) config = KVCacheConfig( num_blocks=2, - kv_cache_tensors=[KVCacheTensor(size=cache.nbytes, shared_by=["layer.0.indexer"])], + kv_cache_tensors=[ + KVCacheTensor( + size=cache.nbytes, + layers=["layer.0.indexer"], + layer_stride=cache.nbytes, + block_stride=cache.stride(0) * cache.element_size(), + ) + ], kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0.indexer"], kv_cache_spec=spec)], ) worker = MooncakeBaseConnectorWorker.__new__(MooncakeBaseConnectorWorker) @@ -170,7 +186,14 @@ def test_register_kv_caches_rejects_missing_and_unconfigured_layers() -> None: spec = make_full_spec() config = KVCacheConfig( num_blocks=2, - kv_cache_tensors=[KVCacheTensor(size=64, shared_by=["layer.0"])], + kv_cache_tensors=[ + KVCacheTensor( + size=64, + layers=["layer.0"], + layer_stride=64, + block_stride=spec.page_size_bytes, + ) + ], kv_cache_groups=[KVCacheGroupSpec(layer_names=["layer.0"], kv_cache_spec=spec)], ) worker = MooncakeBaseConnectorWorker.__new__(MooncakeBaseConnectorWorker) diff --git a/tests/ut/kv_offload/mooncake_v2/test_utils.py b/tests/ut/kv_offload/mooncake_v2/test_utils.py index 2bdb3940c07a..94acc3006289 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_utils.py +++ b/tests/ut/kv_offload/mooncake_v2/test_utils.py @@ -44,7 +44,22 @@ def test_sized_dict_evicts_oldest_entry_and_initializes_missing_value() -> None: def test_collect_register_regions_deduplicates_shared_storage() -> None: backing = torch.empty(128, dtype=torch.uint8) config = SimpleNamespace( - kv_cache_tensors=[SimpleNamespace(shared_by=["layer.0", "layer.1"], offset=0, size=backing.nbytes)] + kv_cache_tensors=[ + SimpleNamespace( + layers=["layer.0"], + layer_stride=64, + block_stride=64, + offset=0, + size=backing.nbytes, + ), + SimpleNamespace( + layers=["layer.1"], + layer_stride=64, + block_stride=64, + offset=64, + size=backing.nbytes, + ), + ] ) regions = collect_configured_register_regions( @@ -54,13 +69,23 @@ def test_collect_register_regions_deduplicates_shared_storage() -> None: assert regions.ptrs == [backing.data_ptr()] assert regions.lengths == [backing.nbytes] - assert regions.logical_tensor_count == 1 + assert regions.logical_tensor_count == 2 def test_collect_register_regions_handles_independent_storages() -> None: first = torch.empty(32, dtype=torch.uint8) second = torch.empty(48, dtype=torch.uint8) - config = SimpleNamespace(kv_cache_tensors=[SimpleNamespace(shared_by=["layer.0", "layer.1"], offset=0, size=80)]) + config = SimpleNamespace( + kv_cache_tensors=[ + SimpleNamespace( + layers=["layer.0", "layer.1"], + layer_stride=40, + block_stride=40, + offset=0, + size=80, + ) + ] + ) regions = collect_configured_register_regions(config, {"layer.0": first, "layer.1": second}) @@ -70,6 +95,34 @@ def test_collect_register_regions_handles_independent_storages() -> None: } +def test_collect_register_regions_recovers_aligned_backing_before_view() -> None: + alignment = 2 * 1024 * 1024 + backing_size = 2 * alignment + raw = torch.empty(backing_size + alignment, dtype=torch.uint8) + aligned_offset = (-raw.data_ptr()) % alignment + backing = raw[aligned_offset : aligned_offset + backing_size] + logical_view = backing[alignment + 64 :] + config = SimpleNamespace( + kv_cache_tensors=[ + SimpleNamespace( + layers=["layer.0"], + layer_stride=backing_size, + block_stride=alignment, + offset=alignment, + size=backing_size, + ) + ] + ) + + regions = collect_configured_register_regions( + config, + {"layer.0": logical_view}, + ) + + assert regions.ptrs == [backing.data_ptr()] + assert regions.lengths == [backing_size] + + def test_group_concurrent_contiguous_honors_byte_strides() -> None: src_groups, dst_groups = group_concurrent_contiguous( [1, 2, 4], diff --git a/tests/ut/worker/test_model_runner_v1.py b/tests/ut/worker/test_model_runner_v1.py index cb41bfa43c86..58dc34d8ee60 100644 --- a/tests/ut/worker/test_model_runner_v1.py +++ b/tests/ut/worker/test_model_runner_v1.py @@ -640,6 +640,8 @@ def test_hybrid_descriptors_share_standardized_backing_allocation(self): for kv_transfer_config in ( None, SimpleNamespace(kv_connector="ExampleHiddenStatesConnector"), + SimpleNamespace(kv_connector="MooncakeConnectorV2"), + SimpleNamespace(kv_connector="MooncakePullConnector"), ): with self.subTest(kv_transfer_config=kv_transfer_config): runner = self._build_runner() 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 78d0b6d783c3..fa1ebbea123b 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 @@ -221,7 +221,7 @@ def register_kv_caches( configured_layer_names: set[str] = set() for tensor_config in self.kv_cache_config.kv_cache_tensors: - for layer_name in tensor_config.shared_by: + for layer_name in tensor_config.layers: if layer_name in configured_layer_names: raise ValueError(f"Layer {layer_name!r} is referenced by more than one configured KV cache tensor.") if layer_name not in self.layer_name_to_group_index: 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 108952d59421..d945498e1859 100644 --- a/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py +++ b/vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake/utils.py @@ -27,6 +27,9 @@ from vllm.v1.kv_cache_interface import KVCacheConfig +KV_CACHE_BUFFER_ALIGNMENT = 2 * 1024 * 1024 + + @contextlib.contextmanager def zmq_ctx(socket_type: Any, addr: str) -> Iterator[Any]: """Create a Mooncake ROUTER or REQ socket and clean up its context.""" @@ -69,6 +72,15 @@ def _get_storage_nbytes(tensor: torch.Tensor) -> int: return tensor.nbytes +def _get_tensor_span_nbytes(tensor: torch.Tensor) -> int: + """Return the physical byte span touched by a strided tensor view.""" + if tensor.numel() == 0: + return 0 + return tensor.element_size() + sum( + (size - 1) * stride * tensor.element_size() for size, stride in zip(tensor.shape, tensor.stride()) + ) + + def collect_configured_register_regions( kv_cache_config: "KVCacheConfig", kv_caches: dict[str, Any], @@ -92,11 +104,11 @@ def merge_storage_range( ) for tensor_config in kv_cache_config.kv_cache_tensors: - if not tensor_config.shared_by: + if not tensor_config.layers: continue cache_tensors: list[torch.Tensor] = [] - for layer_name in tensor_config.shared_by: + for layer_name in tensor_config.layers: cache_tensors.extend(as_kv_cache_tensors(kv_caches.get(layer_name))) caches_by_storage: dict[int, list[torch.Tensor]] = {} @@ -105,16 +117,36 @@ def merge_storage_range( caches_by_storage.setdefault(storage_key, []).append(cache) if len(caches_by_storage) == 1: - # KVCacheTensor.size is authoritative when all layer views belong - # to its one configured allocation. Packed layouts expose each - # view at allocation_base + offset. + # The model runner over-allocates by one alignment unit and returns + # views into the aligned suffix of that raw storage. ``offset`` is + # relative to the standardized backing, not necessarily to the + # first logical view (for example, compressor scale follows K in a + # padded page), so recover the backing from the storage boundary. storage_key, storage_caches = next(iter(caches_by_storage.items())) - register_start = min(cache.data_ptr() for cache in storage_caches) - tensor_config.offset - merge_storage_range( - storage_key, - register_start, - register_start + tensor_config.size, + storage_end = storage_key + _get_storage_nbytes(storage_caches[0]) + aligned_start = ( + (storage_key + KV_CACHE_BUFFER_ALIGNMENT - 1) // KV_CACHE_BUFFER_ALIGNMENT * KV_CACHE_BUFFER_ALIGNMENT + ) + configured_end = aligned_start + tensor_config.size + views_fit_configured_range = all( + aligned_start <= cache.data_ptr() + and cache.data_ptr() + _get_tensor_span_nbytes(cache) <= configured_end + for cache in storage_caches ) + if configured_end <= storage_end and views_fit_configured_range: + merge_storage_range( + storage_key, + aligned_start, + configured_end, + ) + else: + # Cache types that allocate components independently do not + # expose the standardized shared backing. Register the actual + # allocation range instead. + register_start = min(cache.data_ptr() for cache in storage_caches) + if storage_end <= register_start: + raise ValueError(f"Invalid KV cache storage range: start={register_start}, end={storage_end}.") + merge_storage_range(storage_key, register_start, storage_end) else: # Some cache types initialize component tensors independently. # The config size cannot be applied to every storage, so register diff --git a/vllm_ascend/worker/model_runner_v1.py b/vllm_ascend/worker/model_runner_v1.py index 3b8b3f61035f..c678ce8955c5 100644 --- a/vllm_ascend/worker/model_runner_v1.py +++ b/vllm_ascend/worker/model_runner_v1.py @@ -4456,13 +4456,17 @@ def _allocate_kv_cache_tensors(self, kv_cache_config: KVCacheConfig) -> dict[str if kv_transfer_config is not None else None ) - # ExampleHiddenStatesConnector only consumes the dedicated cache-only - # layer and does not register or migrate the per-layer Attention/Mamba - # buffers. It therefore does not require the independent 2 MiB-aligned - # allocations needed by Mooncake/ADXL connectors. + # Mooncake V2 retains per-layer transfer metadata while registering the + # standardized Attention/Mamba backing allocation once. The example + # connector only consumes its dedicated cache-only layer. supports_shared_backing_with_kv_transfer = ( kv_transfer_config is None - or kv_connector == "ExampleHiddenStatesConnector" + or kv_connector + in { + "ExampleHiddenStatesConnector", + "MooncakeConnectorV2", + "MooncakePullConnector", + } ) # The restored DSV4 planner on main emits multiple descriptors into a From 7de770817a6c1925d4e3eb2fc5277b0f4a7be74f Mon Sep 17 00:00:00 2001 From: wangxiaoteng Date: Wed, 9 Sep 2026 17:58:03 +0800 Subject: [PATCH 48/48] test(mooncake): patch module objects to avoid CPU suite import failures Patch loaded Mooncake modules directly instead of resolving dotted paths through parent packages that other CPU tests can replace. Cover a detached kv_p2p parent attribute in the socket context test. Signed-off-by: wangxiaoteng --- .../mooncake_v2/test_base_scheduler.py | 19 ++++++--- .../mooncake_v2/test_base_worker.py | 22 +++++++---- .../mooncake_v2/test_pull_scheduler.py | 22 +++++++---- .../mooncake_v2/test_pull_worker.py | 39 ++++++++++++------- tests/ut/kv_offload/mooncake_v2/test_utils.py | 16 ++++++-- 5 files changed, 81 insertions(+), 37 deletions(-) diff --git a/tests/ut/kv_offload/mooncake_v2/test_base_scheduler.py b/tests/ut/kv_offload/mooncake_v2/test_base_scheduler.py index fa53aa842610..4c9c21daffc3 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_base_scheduler.py +++ b/tests/ut/kv_offload/mooncake_v2/test_base_scheduler.py @@ -6,6 +6,7 @@ import pytest from vllm.v1.kv_cache_interface import KVCacheGroupSpec, UniformTypeKVCacheSpecs +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake import base_scheduler from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler import ( MooncakeBaseConnectorScheduler, ) @@ -49,15 +50,18 @@ def test_base_scheduler_initializes_parallel_layout_and_control_port( ascend_config = object() init_ascend_config = MagicMock() monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.init_ascend_config", + base_scheduler, + "init_ascend_config", init_ascend_config, ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.get_ascend_config", + base_scheduler, + "get_ascend_config", MagicMock(return_value=ascend_config), ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.get_ip", + base_scheduler, + "get_ip", MagicMock(return_value="10.0.0.1"), ) @@ -90,15 +94,18 @@ def test_base_scheduler_rejects_invalid_transfer_roles( def test_base_scheduler_rejects_unsupported_pcp(monkeypatch: pytest.MonkeyPatch) -> None: config = make_scheduler_config(pcp_size=2) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.init_ascend_config", + base_scheduler, + "init_ascend_config", MagicMock(), ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.get_ascend_config", + base_scheduler, + "get_ascend_config", MagicMock(), ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler.get_ip", + base_scheduler, + "get_ip", MagicMock(return_value="10.0.0.1"), ) diff --git a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py index ed7a44d37ae6..e82159bcd0f6 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_base_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_base_worker.py @@ -13,6 +13,7 @@ UniformTypeKVCacheSpecs, ) +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake import base_worker from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker import ( MooncakeBaseConnectorWorker, ) @@ -65,11 +66,13 @@ def test_register_kv_caches_uses_config_order_and_publishes_tensor_metadata(monk worker.handshake_port = 5000 transfer_engine = MagicMock() monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.global_te", + base_worker, + "global_te", transfer_engine, ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.validate_register_region_count", + base_worker, + "validate_register_region_count", MagicMock(), ) @@ -112,11 +115,13 @@ def test_register_kv_caches_collapses_views_packed_in_one_page(monkeypatch) -> N worker.handshake_port = 5000 transfer_engine = MagicMock() monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.global_te", + base_worker, + "global_te", transfer_engine, ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.validate_register_region_count", + base_worker, + "validate_register_region_count", MagicMock(), ) @@ -165,11 +170,13 @@ def test_register_kv_caches_publishes_sfa_indexer_virtual_block_size(monkeypatch worker.handshake_port = 5000 transfer_engine = MagicMock() monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.global_te", + base_worker, + "global_te", transfer_engine, ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.validate_register_region_count", + base_worker, + "validate_register_region_count", MagicMock(), ) @@ -270,7 +277,8 @@ def patch_worker_runtime( } for name, value in patches.items(): monkeypatch.setattr( - f"vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.{name}", + base_worker, + name, value, ) monkeypatch.setattr(torch.npu, "current_device", MagicMock(return_value=7)) diff --git a/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py b/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py index cc4d971c4e80..7e98afb597c1 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_scheduler.py @@ -12,6 +12,7 @@ import zmq from vllm.v1.request import RequestStatus +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake import pull_scheduler from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_scheduler import ( MooncakeBaseConnectorScheduler, ) @@ -206,11 +207,13 @@ def test_sending_thread_force_frees_expired_request(monkeypatch: pytest.MonkeyPa thread = make_sending_thread() thread.add_delayed_request("expired", 10.0) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.envs.VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", + pull_scheduler.envs, + "VLLM_MOONCAKE_ABORT_REQUEST_TIMEOUT", 5, ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.time.time", + pull_scheduler.time, + "time", MagicMock(return_value=20.0), ) @@ -228,11 +231,13 @@ def test_recving_thread_reuses_socket_after_ack_and_discards_it_on_error( send = MagicMock() recv = MagicMock(return_value=ACK_MSG) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.ensure_zmq_send", + pull_scheduler, + "ensure_zmq_send", send, ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.ensure_zmq_recv", + pull_scheduler, + "ensure_zmq_recv", recv, ) @@ -260,11 +265,13 @@ def test_recving_thread_socket_pool_creates_once_and_reuses_by_endpoint( context_cls = MagicMock(return_value=context) make_socket = MagicMock(return_value=socket) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.zmq.Context", + pull_scheduler.zmq, + "Context", context_cls, ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.make_zmq_socket", + pull_scheduler, + "make_zmq_socket", make_socket, ) path = "tcp://10.0.0.1:6000" @@ -510,7 +517,8 @@ def make_fake_thread(*args, **_kwargs): thread_cls = MagicMock(side_effect=make_fake_thread) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_scheduler.MooncakeSchedulerSendingThread", + pull_scheduler, + "MooncakeSchedulerSendingThread", thread_cls, ) metadata = {0: make_transfer_metadata()} 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 a70de26ee812..64bc5ed63db4 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py +++ b/tests/ut/kv_offload/mooncake_v2/test_pull_worker.py @@ -13,6 +13,7 @@ from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.kv_cache_interface import MLAAttentionSpec +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake import base_worker, pull_worker from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.metadata import ( MooncakeConnectorMetadata, MooncakeTransferMetadataGroups, @@ -691,15 +692,18 @@ def test_get_remote_metadata_fetches_once_and_populates_layout_cache( return_value=({0: {(0, 0): [[0]]}}, {0: [(0, 0)]}) ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.zmq_ctx", + pull_worker, + "zmq_ctx", MagicMock(return_value=context), ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_send", + pull_worker, + "ensure_zmq_send", send, ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_recv", + pull_worker, + "ensure_zmq_recv", recv, ) @@ -723,15 +727,18 @@ def test_get_remote_metadata_rejects_empty_response(monkeypatch: pytest.MonkeyPa context = MagicMock() context.__enter__.return_value = socket monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.zmq_ctx", + pull_worker, + "zmq_ctx", MagicMock(return_value=context), ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_send", + pull_worker, + "ensure_zmq_send", MagicMock(), ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_recv", + pull_worker, + "ensure_zmq_recv", MagicMock(return_value=b""), ) @@ -850,8 +857,9 @@ def test_whole_block_coalescing_is_prepared_per_request_and_reused_across_layers dst: list[int] = [] lengths: list[int] = [] - with patch( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.group_concurrent_contiguous", + with patch.object( + pull_worker, + "group_concurrent_contiguous", wraps=group_concurrent_contiguous, ) as mock_group: thread._append_spec_transfer_addresses( @@ -1000,15 +1008,18 @@ def test_get_remote_metadata_propagates_network_error_without_caching( context = MagicMock() context.__enter__.return_value = MagicMock() monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.zmq_ctx", + pull_worker, + "zmq_ctx", MagicMock(return_value=context), ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_send", + pull_worker, + "ensure_zmq_send", MagicMock(), ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.ensure_zmq_recv", + pull_worker, + "ensure_zmq_recv", MagicMock(side_effect=RuntimeError("timed out")), ) @@ -1051,11 +1062,13 @@ def fake_thread(**kwargs: object) -> MagicMock: return thread monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.base_worker.MooncakeBaseConnectorWorker.register_kv_caches", + base_worker.MooncakeBaseConnectorWorker, + "register_kv_caches", fake_register, ) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.pull_worker.MooncakePullRecvingThread", + pull_worker, + "MooncakePullRecvingThread", fake_thread, ) monkeypatch.setattr(torch.npu, "current_device", MagicMock(return_value=0)) diff --git a/tests/ut/kv_offload/mooncake_v2/test_utils.py b/tests/ut/kv_offload/mooncake_v2/test_utils.py index 94acc3006289..8a2c60acd76a 100644 --- a/tests/ut/kv_offload/mooncake_v2/test_utils.py +++ b/tests/ut/kv_offload/mooncake_v2/test_utils.py @@ -1,5 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 +import importlib from types import SimpleNamespace from unittest.mock import MagicMock @@ -7,6 +8,7 @@ import torch import zmq +from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake import utils from vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils import ( SizedDict, as_kv_cache_tensors, @@ -156,12 +158,18 @@ def test_zmq_ctx_rejects_unknown_socket_type() -> None: pass -def test_zmq_ctx_destroys_context(monkeypatch: pytest.MonkeyPatch) -> None: +@pytest.mark.parametrize("detach_parent_package", [False, True]) +def test_zmq_ctx_destroys_context(monkeypatch: pytest.MonkeyPatch, detach_parent_package: bool) -> None: + if detach_parent_package: + # Other CPU tests can replace parent packages while child modules stay cached. + parent = importlib.import_module("vllm_ascend.distributed.kv_transfer") + monkeypatch.delattr(parent, "kv_p2p", raising=False) context = MagicMock() socket = MagicMock() monkeypatch.setattr(zmq, "Context", MagicMock(return_value=context)) monkeypatch.setattr( - "vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils.make_zmq_socket", + utils, + "make_zmq_socket", MagicMock(return_value=socket), ) @@ -172,7 +180,7 @@ def test_zmq_ctx_destroys_context(monkeypatch: pytest.MonkeyPatch) -> None: def test_zmq_retry_helpers_succeed_after_transient_error(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr("vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils.time.sleep", MagicMock()) + monkeypatch.setattr(utils.time, "sleep", MagicMock()) socket = MagicMock() socket.send.side_effect = [zmq.ZMQError("transient"), None] # type: ignore[attr-defined] socket.recv.side_effect = [zmq.ZMQError("transient"), b"response"] # type: ignore[attr-defined] @@ -184,7 +192,7 @@ def test_zmq_retry_helpers_succeed_after_transient_error(monkeypatch: pytest.Mon def test_zmq_retry_helpers_raise_after_limit(monkeypatch: pytest.MonkeyPatch) -> None: - monkeypatch.setattr("vllm_ascend.distributed.kv_transfer.kv_p2p.mooncake.utils.time.sleep", MagicMock()) + monkeypatch.setattr(utils.time, "sleep", MagicMock()) socket = MagicMock() socket.send.side_effect = zmq.ZMQError("failed") # type: ignore[attr-defined] socket.recv.side_effect = zmq.ZMQError("failed") # type: ignore[attr-defined]