Motivation

When using multi-GPU tensor parallelism (TP > 1), the first collective communication operation triggers NCCL/RCCL communicator initialization, causing severe P99 TTFT degradation (up to 1400ms) for the first 2-3 requests.

This PR implements NCCL/RCCL pre-warming during server startup to eliminate cold-start latency, inspired by InstantTensor's implementation.

Measured Impact on AMD MI355X:

• P99 TTFT improvement: 74.9% (1413ms → 357ms)
• Latency stability: 87.8% lower std dev (327ms → 40ms)
• Warmup overhead: 4.7s one-time cost

Default Behavior:

• Enabled by default for AMD/HIP (RCCL) - validated on MI355X
• Disabled by default for NVIDIA/CUDA (NCCL) - pending validation

Modifications

Server Arguments (server_args.py)

Added pre_warm_nccl field with platform-aware default:

pre_warm_nccl: bool = dataclasses.field(
    default_factory=lambda: is_hip()
)  # Default: True for AMD/HIP, False for NVIDIA/CUDA

Added CLI argument:

parser.add_argument(
    "--pre-warm-nccl",
    action="store_true",
    help="Pre-warm NCCL/RCCL communicators during startup to reduce P99 TTFT cold-start latency. Default: enabled for AMD/HIP (RCCL), disabled for NVIDIA/CUDA (NCCL).",
)

Model Runner (model_runner.py)

Added warmup logic during initialization:

if self.server_args.pre_warm_nccl and (self.tp_size > 1 or self.pp_size > 1 or self.moe_ep_size > 1):
    warmup_start = time.perf_counter()
    tp_group_handle = get_tp_group().device_group

    # Single warmup all_reduce to initialize NCCL/RCCL communicator
    warmup_tensor = torch.zeros(1, device=torch.cuda.current_device())
    dist.all_reduce(warmup_tensor, group=tp_group_handle)
    torch.cuda.synchronize()

    warmup_elapsed = time.perf_counter() - warmup_start
    logger.info(f"NCCL/RCCL warmup completed in {warmup_elapsed:.3f}s ...")

Accuracy Tests

No accuracy impact - latency optimization only, does not affect model outputs.

Validated with GSM8K (100 questions):

• Without pre-warm: 97.0%
• With pre-warm: 98.0%

Benchmarking and Profiling

Test Environment

• Platform: AMD MI355X (8 GPUs)
• Model: DeepSeek-R1-MXFP4-Preview
• Configuration: TP=8, 128-token prompts, 16 output tokens

Results

Key Findings:

• P99 TTFT: 74.9% improvement (1413ms → 357ms)
• Latency stability: 87.8% lower std dev (327ms → 40ms)
• Warmup overhead: 4.7s one-time cost
• ROI: Warmup pays for itself after 4-5 requests

Reproduction Commands

# Start server (disable via Python API: ServerArgs(pre_warm_nccl=False))
python3 -m sglang.launch_server \
  --model-path /data/DeepSeek-R1-MXFP4-Preview \
  --tp-size 8

# Send requests - first 2-3 will be slow (~1400ms)
for i in {1..20}; do
  time curl -X POST http://127.0.0.1:30000/generate \
    -H "Content-Type: application/json" \
    -d '{"text": "Hello", "sampling_params": {"max_new_tokens": 16}}'
done

# Start server (pre-warm enabled by default on AMD)
python3 -m sglang.launch_server \
  --model-path /data/DeepSeek-R1-MXFP4-Preview \
  --tp-size 8

# Expected log: "NCCL/RCCL warmup completed in 4.561s"

# Send requests - all fast (~300ms)
for i in {1..20}; do
  time curl -X POST http://127.0.0.1:30000/generate \
    -H "Content-Type: application/json" \
    -d '{"text": "Hello", "sampling_params": {"max_new_tokens": 16}}'
done

# Enable explicitly with --pre-warm-nccl
python3 -m sglang.launch_server \
  --model-path /data/model \
  --tp-size 8 \
  --pre-warm-nccl

Checklist

• Accuracy validation: No impact on model outputs (GSM8K: 97.0% vs 98.0%)
• Performance benchmarks: 74.9% P99 TTFT improvement on AMD MI355X
• Code style: Follows SGLang conventions
• Unit tests: TODO (warmup is tested via integration)
• Documentation: Inline comments added
• CI tests: Pending

Review Process

1. Ping Merge Oncalls to start PR flow
2. Get approvals from CODEOWNERS
3. Trigger CI tests: /tag-run-ci-label, /rerun-failed-ci
4. After green CI + approvals, merge