[NIXL] Add custom NIXL backend selection for KVManager (#17146)
Signed-off-by: Yoray Zack <yorayz@nvidia.com>
This commit is contained in:
@@ -262,6 +262,26 @@ python -m sglang.launch_server \
|
|||||||
--max-running-requests 128
|
--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
|
## ASCEND
|
||||||
|
|
||||||
### Usage
|
### Usage
|
||||||
|
|||||||
@@ -136,14 +136,30 @@ class NixlKVManager(CommonKVManager):
|
|||||||
):
|
):
|
||||||
super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
|
super().__init__(args, disaggregation_mode, server_args, is_mla_backend)
|
||||||
try:
|
try:
|
||||||
from nixl._api import nixl_agent
|
from nixl._api import nixl_agent, nixl_agent_config
|
||||||
except ImportError as e:
|
except ImportError as e:
|
||||||
raise ImportError(
|
raise ImportError(
|
||||||
"Please install NIXL by following the instructions at "
|
"Please install NIXL by following the instructions at "
|
||||||
"https://github.com/ai-dynamo/nixl/blob/main/README.md "
|
"https://github.com/ai-dynamo/nixl/blob/main/README.md "
|
||||||
"to run SGLang with NixlTransferEngine."
|
"to run SGLang with NixlTransferEngine."
|
||||||
) from e
|
) 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()
|
self.register_buffer_to_engine()
|
||||||
|
|
||||||
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
if self.disaggregation_mode == DisaggregationMode.PREFILL:
|
||||||
|
|||||||
@@ -234,6 +234,7 @@ class Envs:
|
|||||||
SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL = EnvFloat(5.0)
|
SGLANG_DISAGGREGATION_HEARTBEAT_INTERVAL = EnvFloat(5.0)
|
||||||
SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2)
|
SGLANG_DISAGGREGATION_HEARTBEAT_MAX_FAILURE = EnvInt(2)
|
||||||
SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300)
|
SGLANG_DISAGGREGATION_WAITING_TIMEOUT = EnvInt(300)
|
||||||
|
SGLANG_DISAGGREGATION_NIXL_BACKEND = EnvStr("UCX")
|
||||||
|
|
||||||
# Scheduler: others:
|
# Scheduler: others:
|
||||||
SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period.
|
SGLANG_EMPTY_CACHE_INTERVAL = EnvFloat(-1) # in seconds. Set if you observe high memory accumulation over a long serving period.
|
||||||
|
|||||||
Reference in New Issue
Block a user