From 51c5ddbe65c1fe9790e8cba3c52c2b404686ceca Mon Sep 17 00:00:00 2001
From: Lucia Fang <116399278+luccafong@users.noreply.github.com>
Date: Fri, 10 Jul 2026 21:45:17 -0700
Subject: [PATCH] [eplb] chunk expert-weight P2P on CUDA to prevent NCCL
 rebalance hang (#30829)

---
 .../docs/references/environment_variables.mdx |  4 +-
 python/sglang/srt/environ.py                  |  5 ++
 .../srt/eplb/expert_location_updater.py       | 47 ++++++++-----------
 3 files changed, 26 insertions(+), 30 deletions(-)

diff --git a/docs_new/docs/references/environment_variables.mdx b/docs_new/docs/references/environment_variables.mdx
index 50802c1fb..d3268993d 100644
--- a/docs_new/docs/references/environment_variables.mdx
+++ b/docs_new/docs/references/environment_variables.mdx
@@ -59,8 +59,8 @@ SGLang supports various environment variables that can be used to configure its
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>0</code></td>
     </tr>
     <tr>
-      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_EPLB_ROCM_P2P_BATCH_CHUNK_SIZE</code></td>
-      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Number of logical expert IDs per batch when submitting P2P ops during EPLB rebalance on ROCm. Smaller values prevent RCCL GPU-side accumulation hangs but increase overhead.</td>
+      <td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}><code>SGLANG_EPLB_P2P_BATCH_CHUNK_SIZE</code></td>
+      <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Number of expert IDs per batch when submitting P2P ops during EPLB rebalance (CUDA and ROCm). Smaller values prevent NCCL/RCCL GPU-side accumulation hangs but increase overhead; set <code>&gt;= num_physical_experts</code> to submit a single batch. Deprecated alias: <code>SGLANG_EPLB_ROCM_P2P_BATCH_CHUNK_SIZE</code>.</td>
       <td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}><code>32</code></td>
     </tr>
     <tr>
diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py
index 78cd951b1..8303280d9 100644
--- a/python/sglang/srt/environ.py
+++ b/python/sglang/srt/environ.py
@@ -599,6 +599,11 @@ class Envs:
     SGLANG_EXPERT_DISTRIBUTION_RECORDER_DIR = EnvStr("/tmp")
     SGLANG_EPLB_HEATMAP_COLLECTION_INTERVAL = EnvInt(0)
     SGLANG_ENABLE_EPLB_BALANCEDNESS_METRIC = EnvBool(False)
+    # Chunk size for the rebalance expert-weight P2P exchange; set
+    # >= num_physical_experts to submit a single batch_isend_irecv.
+    SGLANG_EPLB_P2P_BATCH_CHUNK_SIZE = EnvIntWithAlias(
+        32, deprecated_name="SGLANG_EPLB_ROCM_P2P_BATCH_CHUNK_SIZE"
+    )
 
     # TBO
     SGLANG_TBO_DEBUG = EnvBool(False)
diff --git a/python/sglang/srt/eplb/expert_location_updater.py b/python/sglang/srt/eplb/expert_location_updater.py
index 694e62487..7873223f0 100644
--- a/python/sglang/srt/eplb/expert_location_updater.py
+++ b/python/sglang/srt/eplb/expert_location_updater.py
@@ -21,20 +21,19 @@ import torch.distributed
 from torch.distributed import P2POp
 
 from sglang.srt.elastic_ep.elastic_ep import ElasticEPStateManager
+from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_location import (
     ExpertLocationMetadata,
     get_global_expert_location_metadata,
 )
 from sglang.srt.runtime_context import get_server_args
-from sglang.srt.utils import get_bool_env_var, get_int_env_var, is_hip
+from sglang.srt.utils import get_bool_env_var
 
 logger = logging.getLogger(__name__)
 
 
 _LOG_INPUT = get_bool_env_var("SGLANG_EXPERT_LOCATION_UPDATER_LOG_INPUT")
 
-_is_hip = is_hip()
-
 
 class ExpertLocationUpdater:
     def __init__(self):
@@ -485,31 +484,23 @@ def update_expert_weights_single_layer(
         if len(p2p_ops) == 0:
             return
 
-        if _is_hip:
-            # Submit P2P ops in batches to prevent RCCL GPU-side
-            # accumulation hangs. All ranks use the same expert_id ranges
-            # (based on num_physical_experts) to ensure matching send/recv
-            # pairs land in the same batch. Setting batch_chunk_size >=
-            # num_physical_experts disables batching behavior.
-            batch_chunk_size = get_int_env_var(
-                "SGLANG_EPLB_ROCM_P2P_BATCH_CHUNK_SIZE", 32
-            )
-            ops_by_expert = {eid: ops for eid, ops in sorted_infos}
-            for start in range(0, num_physical_experts, batch_chunk_size):
-                batch_ops = []
-                for eid in range(
-                    start, min(start + batch_chunk_size, num_physical_experts)
-                ):
-                    if eid in ops_by_expert:
-                        batch_ops.extend(ops_by_expert[eid])
-                if batch_ops:
-                    reqs = torch.distributed.batch_isend_irecv(batch_ops)
-                    for req in reqs:
-                        req.wait()
-        else:
-            reqs = torch.distributed.batch_isend_irecv(p2p_ops)
-            for req in reqs:
-                req.wait()
+        # Submit P2P ops in batches to prevent NCCL/RCCL GPU-side accumulation
+        # hangs on large rebalances. All ranks use the same expert_id ranges
+        # (based on num_physical_experts) so matching send/recv pairs land in
+        # the same batch. Set batch_chunk_size >= num_physical_experts to disable.
+        batch_chunk_size = envs.SGLANG_EPLB_P2P_BATCH_CHUNK_SIZE.get()
+        ops_by_expert = {eid: ops for eid, ops in sorted_infos}
+        for start in range(0, num_physical_experts, batch_chunk_size):
+            batch_ops = []
+            for eid in range(
+                start, min(start + batch_chunk_size, num_physical_experts)
+            ):
+                if eid in ops_by_expert:
+                    batch_ops.extend(ops_by_expert[eid])
+            if batch_ops:
+                reqs = torch.distributed.batch_isend_irecv(batch_ops)
+                for req in reqs:
+                    req.wait()
 
     def _execute_buffer2weight_copies(buffer2weight_copy_infos):
         for (