[PD] Add NIXL transfer backend

[trevor-m]

Motivation

This PR adds NIXL as a disaggregation transfer backend via --disaggregation-transfer-backend nixl.

This PR replaces my previous draft #5006

Installing NIXL

The easiest way to use NIXL is to start with an NVIDIA pytorch container which already has dependencies like UCX installed. Otherwise you'll need to follow the instructions in the NIXL readme to install those dependencies.

# Inside docker image nvcr.io/nvidia/pytorch:25.03-py3
git clone https://github.com/ai-dynamo/nixl.git
cd nixl
pip install . --config-settings=setup-args="-Ducx_path=/opt/hpcx/ucx"

Example usage:

python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode prefill --port 30001 --disaggregation-transfer-backend nixl

python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode decode --port 30002 --base-gpu-id 1 --disaggregation-transfer-backend nixl

python3 -m sglang.srt.disaggregation.mini_lb --prefill http://0.0.0.0:30001/ --decode http://0.0.0.0:30002/ --host 0.0.0.0 --port 30000

curl -X POST http://127.0.0.1:30000/generate -H "Content-Type: application/json" -d '{
  "text": "Let me tell you a lonnng story ",
  "sampling_params": {
    "temperature": 0
  }
}'

Example usage with tp-size 2:

python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode prefill --port 30001 --disaggregation-transfer-backend nixl --tp-size 2

python -m sglang.launch_server --model-path meta-llama/Llama-3.1-8B-Instruct --disaggregation-mode decode --port 30002 --base-gpu-id 2 --disaggregation-transfer-backend nixl --tp-size 2

python3 -m sglang.srt.disaggregation.mini_lb --prefill http://0.0.0.0:30001/ --decode http://0.0.0.0:30002/ --host 0.0.0.0 --port 30000

curl -X POST http://127.0.0.1:30000/generate -H "Content-Type: application/json" -d '{
  "text": "Let me tell you a lonnng story ",
  "sampling_params": {
    "temperature": 0
  }
}'

Modifications

I reused the MooncakeKVBootstrapServer with some modifications to allow the prefill and decode servers to find eachother.
Unlike the mooncake transfer engine, there is no multithreading for the sender or receiver since the transfer is done asynchronously. We also don't need to use sockets to notify the receiver when the transfer is done since the receiver can check the status through nixl notifications.

Chunked prefill is supported.

Benchmark

python3 -m sglang.bench_serving --backend sglang --dataset-name random --num-prompts 50 --random-input 8192 --random-output 512 --random-range-ratio 0.5 --port 30000 --dataset-name "random" --model meta-llama/Llama-3.1-8B-Instruct --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json --request-rate 1.0

============ Serving Benchmark Result ============
Backend:                                 sglang
Traffic request rate:                    1.0
Max reqeuest concurrency:                not set
Successful requests:                     50
Benchmark duration (s):                  107.70
Total input tokens:                      304483
Total generated tokens:                  19429
Total generated tokens (retokenized):    19429
Request throughput (req/s):              0.46
Input token throughput (tok/s):          2827.22
Output token throughput (tok/s):         180.40
Total token throughput (tok/s):          3007.62
Concurrency:                             18.39
----------------End-to-End Latency----------------
Mean E2E Latency (ms):                   39603.46
Median E2E Latency (ms):                 41336.43
---------------Time to First Token----------------
Mean TTFT (ms):                          37582.10
Median TTFT (ms):                        39873.36
P99 TTFT (ms):                           68202.43
---------------Inter-Token Latency----------------
Mean ITL (ms):                           5.22
Median ITL (ms):                         5.16
P95 ITL (ms):                            5.77
P99 ITL (ms):                            5.99
Max ITL (ms):                            67.70
==================================================

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[trevor-m]

@ByronHsu @hnyls2002 @jokerwyt Could you please take a look?

[jokerwyt]

Hi, engine_rank is shared across this TP worker. Why do we need to put it into bootstrap_key to distinguish different remote prefill TP workers?

[trevor-m]

Thank you @jokerwyt for the review!

Each TP worker has a separate process, so the engine rank is used to distinguish between those.

[jokerwyt]

Hi, why do we need this lock since the bootstrap server will run in a single thread?

[trevor-m]

This is needed since the boostrap server is using python coroutines.
The bootstrap server code is the same as mooncake's:

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

Line 479 in 06d0a3d

class MooncakeKVBootstrapServer(BaseKVBootstrapServer):

[jokerwyt]

How do we make sure the coming TransferInfo is for the request we want to transfer?
Is it only for max_running_request=1 right now?

[trevor-m]

Hmm, good point. I've been using the default max_running_request and didn't run into any issue but it may just be by chance. Let me fix this

[trevor-m]

Just pushed a fix for this, thanks.

[jokerwyt]

Cool. I think we are ready to merge now! @ByronHsu @hnyls2002

[jokerwyt]

I found a concurrency bug here during testing.
We have many TP workers, and they call check_transfer_done separately. If we do the clean up here, requests will have no chance to get KVPoll.Success from all TP workers.
We may need a flag in KVReceiver to check whether check_transfer_done returns True in previous poll.

[trevor-m]

Thanks @jokerwyt, let me fix that.

[trevor-m]

Hi @jokerwyt I removed the cleanup for now to fix the bug.
I think we can add proper cleanup later (also needed for kvsender) but that might require more hooks in decode.py and prefill.py.

[jokerwyt]

Okay, it can be a temporary fix for testing.

decode.py and prefill.py will manage the lifecycle of KVReceiver and KVSender. Can we put the related flags and data structure in them so that clean up can happen naturally?

[msharmavikram]

@trevor-m nixl supports pip install nixl (https://pypi.org/project/nixl/)

[ayrnb]

@trevor-m Does the NIXL transfer backend support DeepEP?

[CSEEduanyu]

If there is cross-machine transmission, is GPU direct or NVSHMEM used in the current scheme? @trevor-m

[msharmavikram]

As of today we are using whatever supported in UCX.

(So not using nvshmem currently but still goes through nvlink. If you are interested in gpu initiated transfers, can you provide details on the use case. )

[CSEEduanyu]

As of today we are using whatever supported in UCX.

(So not using nvshmem currently but still goes through nvlink. If you are interested in gpu initiated transfers, can you provide details on the use case. )

So if kv_cache is transferred between two different machines, does it need to go through the CPU memory as an intermediary?Will the performance be relatively poor in this case? @msharmavikram

[msharmavikram]

It should not. NIXL picks the fastest path to transfer between the worker instances.