chore: import upstream snapshot with attribution

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