sglang 0.5.3rc0__py3-none-any.whl → 0.5.3rc2__py3-none-any.whl

sglang/srt/disaggregation/ascend/conn.py

@@ -13,7 +13,7 @@ from sglang.srt.disaggregation.mooncake.conn import (
     MooncakeKVReceiver,
     MooncakeKVSender,
 )
-from sglang.srt.utils import get_local_ip_by_remote
+from sglang.srt.utils import get_local_ip_auto
 
 logger = logging.getLogger(__name__)
 
@@ -21,7 +21,7 @@ logger = logging.getLogger(__name__)
 class AscendKVManager(MooncakeKVManager):
     def init_engine(self):
         # TransferEngine initialized on ascend.
-        local_ip = get_local_ip_by_remote()
+        local_ip = get_local_ip_auto()
         self.engine = AscendTransferEngine(
             hostname=local_ip,
             npu_id=self.kv_args.gpu_id,

sglang/srt/disaggregation/ascend/transfer_engine.py

@@ -2,9 +2,19 @@ import logging
 import os
 from typing import List, Optional
 
+import torch
+
 from sglang.srt.disaggregation.mooncake.transfer_engine import MooncakeTransferEngine
 from sglang.srt.disaggregation.utils import DisaggregationMode
 
+try:
+    from mf_adapter import TransferEngine
+
+    import_error = None
+except ImportError as e:
+    import_error = e
+    pass
+
 logger = logging.getLogger(__name__)
 
 
@@ -13,12 +23,11 @@ class AscendTransferEngine(MooncakeTransferEngine):
     def __init__(
         self, hostname: str, npu_id: int, disaggregation_mode: DisaggregationMode
     ):
-        try:
-            from mf_adapter import TransferEngine
-        except ImportError as e:
-            raise ImportError(
+        if import_error is not None:
+            logger.warning(
                 "Please install mf_adapter, for details, see docs/backend/pd_disaggregation.md"
-            ) from e
+            )
+            raise import_error
 
         self.engine = TransferEngine()
         self.hostname = hostname
@@ -37,12 +46,29 @@ class AscendTransferEngine(MooncakeTransferEngine):
         self.initialize()
 
     def initialize(self) -> None:
+        from sglang.srt.layers.dp_attention import (
+            get_tensor_model_parallel_world_size,
+            get_tp_group,
+        )
+
+        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_tensor_model_parallel_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_tp_group().device_group
+            )
         """Initialize the ascend transfer instance."""
         ret_value = self.engine.initialize(
-            self.store_url,
-            self.session_id,
-            self.role,
-            self.npu_id,
+            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.")
@@ -56,3 +82,15 @@ class AscendTransferEngine(MooncakeTransferEngine):
             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