From 943f34f6426f5801e483bd68cfd02d23271692b3 Mon Sep 17 00:00:00 2001
From: Hubert Lu <55214931+hubertlu-tw@users.noreply.github.com>
Date: Mon, 16 Mar 2026 20:23:14 -0700
Subject: [PATCH] Add NCCL/RCCL pre-warming to reduce P99 TTFT cold-start
 latency (#20477)

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
---
 .../sglang/srt/model_executor/model_runner.py | 19 +++++++++++++++++++
 python/sglang/srt/server_args.py              | 18 ++++++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py
index e73ea7bac..02b856c6f 100644
--- a/python/sglang/srt/model_executor/model_runner.py
+++ b/python/sglang/srt/model_executor/model_runner.py
@@ -829,6 +829,25 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             if is_npu():
                 register_sgl_tp_rank(self.gpu_id)
 
+            # Pre-warm NCCL/RCCL to eliminate cold-start latency in first request
+            # Controlled by --pre-warm-nccl flag (default: enabled on AMD GPUs)
+            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 "
+                    f"(tp_size={self.tp_size}, pp_size={self.pp_size}, ep_size={self.moe_ep_size})"
+                )
+
         pre_model_load_memory = get_available_gpu_memory(
             self.device,
             self.gpu_id,
diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py
index 16e21b73c..2b1a11dfe 100644
--- a/python/sglang/srt/server_args.py
+++ b/python/sglang/srt/server_args.py
@@ -614,6 +614,9 @@ class ServerArgs:
     disable_custom_all_reduce: bool = False
     enable_mscclpp: bool = False
     enable_torch_symm_mem: bool = False
+    pre_warm_nccl: bool = dataclasses.field(
+        default_factory=lambda: is_hip()
+    )  # Pre-warm NCCL/RCCL to reduce P99 TTFT cold-start latency (default: True for AMD/HIP, False for others)
     disable_overlap_schedule: bool = False
     enable_mixed_chunk: bool = False
     enable_dp_attention: bool = False
@@ -771,6 +774,7 @@ class ServerArgs:
         self._handle_kv4_compatibility()
         self._handle_page_size()
         self._handle_amd_specifics()
+        self._handle_nccl_pre_warm()
         self._handle_grammar_backend()
 
         # Handle Hicache settings.
@@ -2411,6 +2415,15 @@ class ServerArgs:
         if is_hip():
             self.triton_attention_num_kv_splits = 16
 
+    def _handle_nccl_pre_warm(self):
+        # pre_warm_nccl is only used with CUDA or HIP hardware
+        if self.pre_warm_nccl and not (is_cuda() or is_hip()):
+            logger.warning(
+                "pre_warm_nccl is only applicable for CUDA or HIP hardware. "
+                "Ignoring pre_warm_nccl setting on current hardware."
+            )
+            self.pre_warm_nccl = False
+
     def _handle_grammar_backend(self):
         if self.grammar_backend is None:
             self.grammar_backend = "xgrammar"
@@ -5093,6 +5106,11 @@ class ServerArgs:
             action="store_true",
             help="Enable using torch symm mem for all-reduce kernel and fall back to NCCL. Only supports CUDA device SM90 and above. SM90 supports world size 4, 6, 8. SM100 supports world size 6, 8.",
         )
+        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).",
+        )
         parser.add_argument(
             "--disable-overlap-schedule",
             action="store_true",