chore: import upstream snapshot with attribution

python/sglang/srt/disaggregation/ascend/__init__.py

@@ -0,0 +1,6 @@
+from sglang.srt.disaggregation.ascend.conn import (
+    AscendKVBootstrapServer,
+    AscendKVManager,
+    AscendKVReceiver,
+    AscendKVSender,
+)

python/sglang/srt/disaggregation/ascend/conn.py

@@ -0,0 +1,191 @@
+import concurrent.futures
+import logging
+from typing import List, Tuple
+
+import numpy as np
+import numpy.typing as npt
+
+from sglang.srt.disaggregation.ascend.transfer_engine import AscendTransferEngine
+from sglang.srt.disaggregation.common.utils import group_concurrent_contiguous
+from sglang.srt.disaggregation.mooncake.conn import (
+    MooncakeKVBootstrapServer,
+    MooncakeKVManager,
+    MooncakeKVReceiver,
+    MooncakeKVSender,
+)
+from sglang.srt.utils.network import get_local_ip_auto
+
+logger = logging.getLogger(__name__)
+
+
+class AscendKVManager(MooncakeKVManager):
+    def init_engine(self):
+        # TransferEngine initialized on ascend.
+        local_ip = get_local_ip_auto()
+        self.engine = AscendTransferEngine(
+            hostname=local_ip,
+            npu_id=self.kv_args.gpu_id,
+            disaggregation_mode=self.disaggregation_mode,
+        )
+
+    def register_buffer_to_engine(self):
+        self.engine.batch_register(self.kv_args.kv_data_ptrs, self.kv_args.kv_data_lens)
+        # The Ascend backend optimize batch registration for small memory blocks.
+        self.engine.batch_register(
+            self.kv_args.aux_data_ptrs, self.kv_args.aux_data_lens
+        )
+        # Batch register state/extra pool data buffers
+        for component_ptrs, component_lens in zip(
+            self.kv_args.state_data_ptrs or [],
+            self.kv_args.state_data_lens or [],
+        ):
+            self.engine.batch_register(component_ptrs, component_lens)
+
+    def get_mla_kv_ptrs_with_pp(
+        self, src_kv_ptrs: List[int], dst_kv_ptrs: List[int]
+    ) -> Tuple[List[int], List[int], int]:
+        # src_kv_ptrs: k_data, v_data, index_k_data(optional)
+        # dst_kv_ptrs: k_data, v_data, index_k_data(optional)
+        start_layer = self.kv_args.prefill_start_layer
+        kv_buf_groups = getattr(self.kv_args, "kv_buf_groups", 1)
+        total_kv_layers = getattr(self.kv_args, "total_kv_layers", 0)
+        src_layers = len(src_kv_ptrs) // kv_buf_groups
+        # When only speculative-algorithm is enabled for decode
+        # the KV has one more layer than prefill.
+        # The draft layer needs to be skipped.
+        dst_total_layers = (
+            min(len(dst_kv_ptrs) // kv_buf_groups, total_kv_layers)
+            if total_kv_layers
+            else len(dst_kv_ptrs) // kv_buf_groups
+        )
+        end_layer = start_layer + src_layers
+        if src_layers == dst_total_layers:
+            sliced_dst_kv_ptrs = dst_kv_ptrs
+        else:
+            sliced_dst_kv_ptrs = []
+            for i in range(kv_buf_groups):
+                layer_offset = i * dst_total_layers
+                sliced_dst_kv_ptrs.extend(
+                    dst_kv_ptrs[layer_offset + start_layer : layer_offset + end_layer]
+                )
+        layers_current_pp_stage = len(src_kv_ptrs)
+        return src_kv_ptrs, sliced_dst_kv_ptrs, layers_current_pp_stage
+
+    def send_kvcache(
+        self,
+        mooncake_session_id: str,
+        prefill_kv_indices: npt.NDArray[np.int32],
+        dst_kv_ptrs: list[int],
+        dst_kv_indices: npt.NDArray[np.int32],
+        executor: concurrent.futures.ThreadPoolExecutor,
+    ):
+        # Group by indices
+        prefill_kv_blocks, dst_kv_blocks = group_concurrent_contiguous(
+            prefill_kv_indices, dst_kv_indices
+        )
+
+        if self.pp_size > 1:
+            if self.is_mla_backend:
+                src_kv_ptrs, sliced_dst_kv_ptrs, layers_current_pp_stage = (
+                    self.get_mla_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
+                )
+                layers_params = [
+                    (
+                        src_kv_ptrs[layer_id],
+                        sliced_dst_kv_ptrs[layer_id],
+                        self.kv_args.kv_item_lens[layer_id],
+                    )
+                    for layer_id in range(layers_current_pp_stage)
+                ]
+            else:
+                (
+                    src_k_ptrs,
+                    src_v_ptrs,
+                    dst_k_ptrs,
+                    dst_v_ptrs,
+                    layers_current_pp_stage,
+                ) = self.get_mha_kv_ptrs_with_pp(self.kv_args.kv_data_ptrs, dst_kv_ptrs)
+
+                layers_params = [
+                    (
+                        src_k_ptrs[layer_id],
+                        dst_k_ptrs[layer_id],
+                        self.kv_args.kv_item_lens[layer_id],
+                    )
+                    for layer_id in range(layers_current_pp_stage)
+                ] + [
+                    (
+                        src_v_ptrs[layer_id],
+                        dst_v_ptrs[layer_id],
+                        self.kv_args.kv_item_lens[layers_current_pp_stage + layer_id],
+                    )
+                    for layer_id in range(layers_current_pp_stage)
+                ]
+        else:
+            num_layers = len(self.kv_args.kv_data_ptrs)
+            layers_params = [
+                (
+                    self.kv_args.kv_data_ptrs[layer_id],
+                    dst_kv_ptrs[layer_id],
+                    self.kv_args.kv_item_lens[layer_id],
+                )
+                for layer_id in range(num_layers)
+            ]
+
+        def set_transfer_blocks(
+            src_ptr: int, dst_ptr: int, item_len: int
+        ) -> List[Tuple[int, int, int]]:
+            transfer_blocks = []
+            for prefill_index, decode_index in zip(prefill_kv_blocks, dst_kv_blocks):
+                src_addr = src_ptr + int(prefill_index[0]) * item_len
+                dst_addr = dst_ptr + int(decode_index[0]) * item_len
+                length = item_len * len(prefill_index)
+                transfer_blocks.append((src_addr, dst_addr, length))
+            return transfer_blocks
+
+        # Worker function for processing a single layer
+        def process_layer(src_ptr: int, dst_ptr: int, item_len: int) -> int:
+            transfer_blocks = set_transfer_blocks(src_ptr, dst_ptr, item_len)
+            return self._transfer_data(mooncake_session_id, transfer_blocks)
+
+        # Worker function for processing all layers in a batch
+        def process_layers(layers_params: List[Tuple[int, int, int]]) -> int:
+            transfer_blocks = []
+            for src_ptr, dst_ptr, item_len in layers_params:
+                transfer_blocks.extend(set_transfer_blocks(src_ptr, dst_ptr, item_len))
+            return self._transfer_data(mooncake_session_id, transfer_blocks)
+
+        if self.enable_custom_mem_pool:
+            futures = [
+                executor.submit(
+                    process_layer,
+                    src_ptr,
+                    dst_ptr,
+                    item_len,
+                )
+                for (src_ptr, dst_ptr, item_len) in layers_params
+            ]
+            for future in concurrent.futures.as_completed(futures):
+                status = future.result()
+                if status != 0:
+                    for f in futures:
+                        f.cancel()
+                    return status
+        else:
+            # Combining all layers' params in one batch transfer is more efficient
+            # compared to using multiple threads
+            return process_layers(layers_params)
+
+        return 0
+
+
+class AscendKVSender(MooncakeKVSender):
+    pass
+
+
+class AscendKVReceiver(MooncakeKVReceiver):
+    pass
+
+
+class AscendKVBootstrapServer(MooncakeKVBootstrapServer):
+    pass

python/sglang/srt/disaggregation/ascend/transfer_engine.py

@@ -0,0 +1,103 @@
+import logging
+import os
+from typing import List
+
+import torch
+
+from sglang.srt.disaggregation.utils import DisaggregationMode
+from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
+    MooncakeTransferEngine,
+)
+from sglang.srt.utils.network import NetworkAddress
+
+try:
+    from memfabric_hybrid import TransferEngine
+
+    import_error = None
+except ImportError as e:
+    import_error = e
+    pass
+
+logger = logging.getLogger(__name__)
+
+
+class AscendTransferEngine(MooncakeTransferEngine):
+
+    def __init__(
+        self,
+        hostname: str,
+        npu_id: int,
+        disaggregation_mode: DisaggregationMode,
+    ):
+        if import_error is not None:
+            logger.warning(
+                "Please install memfabric_hybrid, for details, see docs/backend/pd_disaggregation.md"
+            )
+            raise import_error
+
+        self.engine = TransferEngine()
+        self.hostname = hostname
+        self.npu_id = npu_id
+
+        # Centralized storage address of the AscendTransferEngine
+        self.store_url = os.getenv("ASCEND_MF_STORE_URL")
+        if disaggregation_mode == DisaggregationMode.PREFILL:
+            self.role = "Prefill"
+        elif disaggregation_mode == DisaggregationMode.DECODE:
+            self.role = "Decode"
+        else:
+            logger.error(f"Unsupported DisaggregationMode: {disaggregation_mode}")
+            raise ValueError(f"Unsupported DisaggregationMode: {disaggregation_mode}")
+        self.session_id = NetworkAddress(
+            self.hostname, self.engine.get_rpc_port()
+        ).to_host_port_str()
+        self.initialize()
+
+    def initialize(self) -> None:
+        from sglang.srt.distributed.parallel_state import (
+            get_world_group,
+            get_world_size,
+        )
+
+        transfer_protocol = self._get_transfer_protocol()
+        if transfer_protocol is None or transfer_protocol == "sdma":
+            trans_op_type = TransferEngine.TransDataOpType.SDMA
+        else:
+            trans_op_type = TransferEngine.TransDataOpType.DEVICE_RDMA
+            """with device RDMA for PD transfer"""
+            tmp_tensor = torch.zeros(1, device="npu")
+            output_tensor_list = [
+                torch.empty_like(tmp_tensor) for _ in range(get_world_size())
+            ]
+            # Initialize hccl in advance through all_gather to avoid conflicts with rdma initialization.
+            torch.distributed.all_gather(
+                output_tensor_list, tmp_tensor, group=get_world_group().device_group
+            )
+        """Initialize the ascend transfer instance."""
+        ret_value = self.engine.initialize(
+            self.store_url, self.session_id, self.role, self.npu_id, trans_op_type
+        )
+        if ret_value != 0:
+            logger.error("Ascend Transfer Engine initialization failed.")
+            raise RuntimeError("Ascend Transfer Engine initialization failed.")
+
+    def batch_register(self, ptrs: List[int], lengths: List[int]):
+        try:
+            ret_value = self.engine.batch_register_memory(ptrs, lengths)
+        except Exception:
+            # Mark register as failed
+            ret_value = -1
+        if ret_value != 0:
+            logger.debug(f"Ascend memory registration for ptr {ptrs} failed.")
+
+    @staticmethod
+    def _get_transfer_protocol():
+        protocol = os.getenv("ASCEND_MF_TRANSFER_PROTOCOL")
+        allowed_protocols = {"device_rdma", "sdma"}
+        if protocol and protocol.lower() in allowed_protocols:
+            return protocol.lower()
+        else:
+            logger.warning(
+                "Invalid or no transfer protocol specified, using default protocol."
+            )
+            return None