diff --git a/docs/advanced_features/pd_disaggregation.md b/docs/advanced_features/pd_disaggregation.md index 21ea7dffe..b40ab11b4 100644 --- a/docs/advanced_features/pd_disaggregation.md +++ b/docs/advanced_features/pd_disaggregation.md @@ -262,6 +262,26 @@ python -m sglang.launch_server \ --max-running-requests 128 ``` +### Advanced Configuration + +#### NIXL Backend Selection + +By default, NIXL uses the **UCX** backend for KV cache transfers. You can select a different NIXL plugin backend depending on your infrastructure using the environment variable `SGLANG_DISAGGREGATION_NIXL_BACKEND`. + +Example: `export SGLANG_DISAGGREGATION_NIXL_BACKEND=LIBFABRIC` + +**Available backends:** UCX (default), LIBFABRIC, or any installed NIXL plugin. + +Example usage: +```bash +export SGLANG_DISAGGREGATION_NIXL_BACKEND=LIBFABRIC +python -m sglang.launch_server \ + --model-path meta-llama/Llama-3.1-8B-Instruct \ + --disaggregation-mode prefill \ + --disaggregation-transfer-backend nixl \ + --port 30000 +``` + ## ASCEND ### Usage diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index d3e390c92..7905aae65 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/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: diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 5dbd90899..928ec998e 100644 --- a/python/sglang/srt/environ.py +++ b/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.