[PP Prefill][NIXL] Fix PP mode transfer completion tracking to wait for all ranks

[YAMY1234]

Collaborated with @hlu1 on root cause analysis and fix

Motivation

Fix NIXL PP mode correctness bug: decode server prematurely considers KV transfer "complete" after receiving chunks from only one PP rank (instead of all ranks), causing accuracy drop.

Root cause: TransferStatus used Set[int] for chunk IDs without distinguishing PP ranks. Overlapping chunk IDs (0,1,2...) from different PP ranks got deduplicated.

Modifications

• Track chunks per PP rank: received_kvs_per_pp: Dict[int, Set[int]]
• Record expected count per PP rank: expected_kvs_per_pp: Dict[int, int]
• Update is_done(): wait for all PP ranks to complete all chunks
• Include pp_rank in notification format (backward compatible)
• Replace -1 sentinel with is_failure: bool

Accuracy Tests

GSM8K, PP=4 TP=4:

#!/bin/bash
export CUDA_VISIBLE_DEVICES=0,1,2,3
export MODEL=nvidia/DeepSeek-R1-0528-NVFP4-v2
export SERVED_NAME=dsr1
export HOST_IP=127.0.0.1

python3 -m sglang.launch_server \
  --model ${MODEL} \
  --served-model-name ${SERVED_NAME} \
  --host ${HOST_IP} \
  --port 12347 \
  --trust-remote-code \
  --disaggregation-mode prefill \
  --context-length 131072 \
  --attention-backend trtllm_mla \
  --moe-runner-backend flashinfer_trtllm \
  --quantization modelopt_fp4 \
  --kv-cache-dtype fp8_e4m3 \
  --page-size 64 \
  --decode-log-interval 1 \
  --disaggregation-transfer-backend nixl \
   --tensor-parallel-size 1 --pipeline-parallel-size 4 --expert-parallel-size 1 --chunked-prefill-size 1024 --cuda-graph-max-bs 32 --max-running-requests 36 --disable-radix-cache

#!/bin/bash
export CUDA_VISIBLE_DEVICES=4,5,6,7
export MODEL=nvidia/DeepSeek-R1-0528-NVFP4-v2
export SERVED_NAME=dsr1
export HOST_IP=127.0.0.1


python3 -m sglang.launch_server \
  --model ${MODEL} \
  --served-model-name ${SERVED_NAME} \
  --host ${HOST_IP} \
  --port 12346 \
  --trust-remote-code \
  --disaggregation-mode decode \
  --context-length 131072 \
  --attention-backend trtllm_mla \
  --moe-runner-backend flashinfer_trtllm \
  --quantization modelopt_fp4 \
  --kv-cache-dtype fp8_e4m3 \
  --page-size 64 \
  --decode-log-interval 1 \
  --disaggregation-transfer-backend nixl \
   --tensor-parallel-size 4 --pipeline-parallel-size 1 --expert-parallel-size 1 --chunked-prefill-size 1024 --cuda-graph-max-bs 32 --max-running-requests 36 --disable-radix-cache

Metric	Before	After
Accuracy	0.358	0.959

Checklist

• Format code with pre-commit
• Accuracy benchmark provided
• Follow SGLang code style

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[ShangmingCai]

Logic looks reasonable.
CC: @ishandhanani @shaharmor98

[ShangmingCai]

/tag-and-rerun-ci