[NPU]support mf device urma and host rdma trans type (#38174)

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

@@ -8,8 +8,6 @@ from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import (
     MooncakeTransferEngine,
 )
-from sglang.srt.environ import envs
-from sglang.srt.utils.common import run_with_deadline
 from sglang.srt.utils.network import NetworkAddress
 
 try:
@@ -22,6 +20,8 @@ except ImportError as e:
 
 logger = logging.getLogger(__name__)
 
+_DEFAULT_PROTOCOL = "sdma"
+
 
 class AscendTransferEngine(MooncakeTransferEngine):
     def __init__(
@@ -63,30 +63,21 @@ class AscendTransferEngine(MooncakeTransferEngine):
         )
 
         transfer_protocol = self._get_transfer_protocol()
-        if transfer_protocol is None or transfer_protocol == "sdma":
+        if transfer_protocol == "device_rdma":
-            trans_op_type = TransferEngine.TransDataOpType.SDMA
+            # with device RDMA for PD transfer: initialize hccl in advance
-        elif transfer_protocol == "device_urma":
+            # through all_gather to avoid conflicts with rdma initialization.
-            trans_op_type = TransferEngine.TransDataOpType.DEVICE_URMA
-        elif transfer_protocol == "device_uboe":
-            trans_op_type = TransferEngine.TransDataOpType.DEVICE_UBOE
-        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
             )
+
+        trans_op_type = self._resolve_trans_op_type(transfer_protocol)
         """Initialize the ascend transfer instance."""
-        ret_value = run_with_deadline(
+        ret_value = self.engine.initialize(
-            lambda: self.engine.initialize(
+            self.store_url, self.session_id, self.role, self.npu_id, trans_op_type
-                self.store_url, self.session_id, self.role, self.npu_id, trans_op_type
-            ),
-            timeout_s=envs.SGLANG_DISAGGREGATION_ENGINE_INIT_TIMEOUT.get(),
-            what=f"Ascend TransferEngine.initialize({self.store_url!r}, {self.session_id!r})",
         )
         if ret_value != 0:
             logger.error("Ascend Transfer Engine initialization failed.")
@@ -102,13 +93,19 @@ class AscendTransferEngine(MooncakeTransferEngine):
             logger.debug(f"Ascend memory registration for ptr {ptrs} failed.")
 
     @staticmethod
-    def _get_transfer_protocol():
+    def _get_transfer_protocol() -> str:
         protocol = os.getenv("ASCEND_MF_TRANSFER_PROTOCOL")
-        allowed_protocols = {"device_rdma", "sdma", "device_urma", "device_uboe"}
+        return protocol.strip().lower() if protocol else _DEFAULT_PROTOCOL
-        if protocol and protocol.lower() in allowed_protocols:
+
-            return protocol.lower()
+    @staticmethod
-        else:
+    def _resolve_trans_op_type(protocol: str):
+        op_type = getattr(TransferEngine.TransDataOpType, protocol.upper(), None)
+        if op_type is None:
             logger.warning(
-                "Invalid or no transfer protocol specified, using default protocol."
+                "Transfer protocol %r is not supported by the installed "
+                "memfabric_hybrid, falling back to %r.",
+                protocol,
+                _DEFAULT_PROTOCOL,
             )
-            return None
+            op_type = TransferEngine.TransDataOpType.SDMA
+        return op_type