[eplb] chunk expert-weight P2P on CUDA to prevent NCCL rebalance hang (#30829)

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)

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 (