[PD] Support KV transfer with MORI-IO (#14626)

Co-authored-by: cwortman-amd <cwortman@amd.com>

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

@@ -0,0 +1,6 @@
+from sglang.srt.disaggregation.mori.conn import (
+    MoriKVBootstrapServer,
+    MoriKVManager,
+    MoriKVReceiver,
+    MoriKVSender,
+)

python/sglang/srt/disaggregation/utils.py

@@ -231,6 +231,7 @@ class MetadataBuffers:
 
 class TransferBackend(Enum):
     MOONCAKE = "mooncake"
+    MORI = "mori"
     NIXL = "nixl"
     ASCEND = "ascend"
     FAKE = "fake"
@@ -266,6 +267,23 @@ def get_kv_class(
             KVClassType.BOOTSTRAP_SERVER: MooncakeKVBootstrapServer,
         }
         return class_mapping.get(class_type)
+    elif transfer_backend == TransferBackend.MORI:
+        from sglang.srt.disaggregation.base import KVArgs
+        from sglang.srt.disaggregation.mori import (
+            MoriKVBootstrapServer,
+            MoriKVManager,
+            MoriKVReceiver,
+            MoriKVSender,
+        )
+
+        class_mapping = {
+            KVClassType.KVARGS: KVArgs,
+            KVClassType.MANAGER: MoriKVManager,
+            KVClassType.SENDER: MoriKVSender,
+            KVClassType.RECEIVER: (MoriKVReceiver),
+            KVClassType.BOOTSTRAP_SERVER: MoriKVBootstrapServer,
+        }
+        return class_mapping.get(class_type)
     elif transfer_backend == TransferBackend.ASCEND:
         from sglang.srt.disaggregation.ascend import (
             AscendKVBootstrapServer,

python/sglang/srt/server_args.py

@@ -143,7 +143,7 @@ ATTENTION_BACKEND_CHOICES = [
 
 LORA_BACKEND_CHOICES = ["triton", "csgmv", "ascend", "torch_native"]
 
-DISAGG_TRANSFER_BACKEND_CHOICES = ["mooncake", "nixl", "ascend", "fake"]
+DISAGG_TRANSFER_BACKEND_CHOICES = ["mooncake", "nixl", "ascend", "fake", "mori"]
 
 ENCODER_TRANSFER_BACKEND_CHOICES = ["zmq_to_scheduler", "zmq_to_tokenizer", "mooncake"]