[NIXL] Add custom NIXL backend selection for KVManager (#17146)

Signed-off-by: Yoray Zack <yorayz@nvidia.com>

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

@@ -136,14 +136,30 @@ class NixlKVManager(CommonKVManager):
     ):
         super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
         try:
-            from nixl._api import nixl_agent
+            from nixl._api import nixl_agent, nixl_agent_config
         except ImportError as e:
             raise ImportError(
                 "Please install NIXL by following the instructions at "
                 "https://github.com/ai-dynamo/nixl/blob/main/README.md "
                 "to run SGLang with NixlTransferEngine."
             ) from e
-        self.agent = nixl_agent(str(uuid.uuid4()))
+
+        agent_config = nixl_agent_config(backends=[])
+        self.agent = nixl_agent(str(uuid.uuid4()), agent_config)
+
+        backend = envs.SGLANG_DISAGGREGATION_NIXL_BACKEND.get()
+
+        available_plugins = self.agent.get_plugin_list()
+        if backend not in available_plugins:
+            raise ValueError(
+                f"NIXL backend '{backend}' not found. Available: {available_plugins}. "
+                f"Please install the required NIXL plugin or choose from: {available_plugins}"
+            )
+
+        self.agent.create_backend(backend)
+        self.nixl_backend = backend
+        logger.info(f"NIXL KVManager initialized with backend: {backend}")
+
         self.register_buffer_to_engine()
 
         if self.disaggregation_mode == DisaggregationMode.PREFILL:

python/sglang/srt/environ.py

@@ -234,6 +234,7 @@ class Envs:
     SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL = EnvFloat(5.0)
     SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2)
     SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300)
+    SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX")
 
     # Scheduler: others:
     SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1)  # in seconds. Set if you observe high memory accumulation over a long serving period.