Fix PCG MoE Error (#17739)

python/sglang/srt/compilation/compile.py

@@ -88,12 +88,12 @@ class _MaybeIntermediateTensors:
 
 def _mark_dynamic_on_value(val, dims):
     if isinstance(val, torch.Tensor):
-        torch._dynamo.mark_dynamic(val, _normalize_dims(dims, val.ndim))
+        torch._dynamo.maybe_mark_dynamic(val, _normalize_dims(dims, val.ndim))
     else:
         mit = _MaybeIntermediateTensors(val)
         if mit.is_intermediate:
             for t in mit.obj.tensors.values():
-                torch._dynamo.mark_dynamic(t, _normalize_dims(dims, t.ndim))
+                torch._dynamo.maybe_mark_dynamic(t, _normalize_dims(dims, t.ndim))
         # else: ignore (None or non-tensor)
 
 

python/sglang/srt/compilation/piecewise_context_manager.py

@@ -1,10 +1,13 @@
+from __future__ import annotations
+
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any, List, Optional
+from typing import TYPE_CHECKING, Any, List, Optional
 
 import torch
 
-from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
 _in_piecewise_cuda_graph = False
 _in_pcg_torch_compile = False

python/sglang/srt/distributed/device_communicators/custom_all_reduce.py

@@ -11,6 +11,7 @@ import torch.distributed as dist
 from torch.distributed import ProcessGroup
 
 import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops
+from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
 from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
     gpu_p2p_access_check,
@@ -414,9 +415,19 @@ class CustomAllreduce:
                 else:
                     return self.all_reduce(input, registered=not self.tms_cudagraph)
             else:
-                # If warm up, mimic the allocation pattern since custom
-                # allreduce is out-of-place.
-                return torch.zeros_like(input)
+                # Could be warmup OR piecewise cuda graph split op execution.
+                # In piecewise cuda graph, split ops run eagerly outside the graph
+                # but _IS_CAPTURING is still True. We need to do real all-reduce.
+                if is_in_piecewise_cuda_graph():
+                    # Split op execution - do real all-reduce
+                    if _is_hip:
+                        return self.all_reduce_unreg(input)
+                    else:
+                        return self.all_reduce(input, registered=False)
+                else:
+                    # True warmup - mimic the allocation pattern since custom
+                    # allreduce is out-of-place.
+                    return torch.zeros_like(input)
         else:
             if _is_hip:
                 # note: outside of cuda graph context,

python/sglang/srt/distributed/device_communicators/pynccl.py

@@ -164,6 +164,35 @@ class PyNcclCommunicator:
             cudaStream_t(stream.cuda_stream),
         )
 
+    def outplace_all_reduce(
+        self,
+        in_tensor: torch.Tensor,
+        out_tensor: Optional[torch.Tensor] = None,
+        op: ReduceOp = ReduceOp.SUM,
+        stream=None,
+    ) -> Optional[torch.Tensor]:
+        if self.disabled:
+            return None
+        assert in_tensor.device == self.device, (
+            f"this nccl communicator is created to work on {self.device}, "
+            f"but the input tensor is on {in_tensor.device}"
+        )
+
+        if out_tensor is None:
+            out_tensor = torch.empty_like(in_tensor)
+
+        stream = self._resolve_stream(stream)
+        self.nccl.ncclAllReduce(
+            buffer_type(in_tensor.data_ptr()),  # sendbuff
+            buffer_type(out_tensor.data_ptr()),  # recvbuff - DIFFERENT pointer
+            in_tensor.numel(),
+            ncclDataTypeEnum.from_torch(in_tensor.dtype),
+            ncclRedOpTypeEnum.from_torch(op),
+            self.comm,
+            cudaStream_t(stream.cuda_stream),
+        )
+        return out_tensor
+
     def all_gather(
         self,
         output_tensor: torch.Tensor,

python/sglang/srt/distributed/parallel_state.py

@@ -41,6 +41,7 @@ import torch.distributed
 from torch.distributed import Backend, ProcessGroup
 
 from sglang.srt.compilation.compilation_config import register_split_op
+from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.environ import envs
 from sglang.srt.utils import (
     get_bool_env_var,
@@ -612,6 +613,9 @@ class GroupCoordinator:
             and self.torch_symm_mem_comm.should_torch_symm_mem_allreduce(input_)
         ):
             outplace_all_reduce_method = "torch_symm_mem"
+        elif is_in_piecewise_cuda_graph():
+            # For piecewise cuda graph, we use pynccl outplace allreduce
+            outplace_all_reduce_method = "pynccl"
         if outplace_all_reduce_method is not None:
             return outplace_all_reduce(
                 input_,
@@ -629,7 +633,8 @@ class GroupCoordinator:
         qr_comm = self.qr_comm
         pymscclpp_comm = self.pymscclpp_comm
         torch_symm_mem_comm = self.torch_symm_mem_comm
-        assert any([qr_comm, ca_comm, pymscclpp_comm, torch_symm_mem_comm])
+        pynccl_comm = self.pynccl_comm
+        assert any([qr_comm, ca_comm, pymscclpp_comm, torch_symm_mem_comm, pynccl_comm])
         if outplace_all_reduce_method == "ca":
             assert not ca_comm.disabled
             out = ca_comm.custom_all_reduce(input_)
@@ -639,9 +644,14 @@ class GroupCoordinator:
         elif outplace_all_reduce_method == "torch_symm_mem":
             assert not torch_symm_mem_comm.disabled
             out = torch_symm_mem_comm.all_reduce(input_)
-        else:
+        elif outplace_all_reduce_method == "pymscclpp":
             assert not pymscclpp_comm.disabled
             out = pymscclpp_comm.all_reduce(input_)
+        elif outplace_all_reduce_method == "pynccl":
+            with pynccl_comm.change_state(
+                enable=True, stream=get_current_device_stream_fast()
+            ):
+                out = pynccl_comm.outplace_all_reduce(input_)
         assert out is not None
         return out
 

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -561,6 +561,13 @@ class PiecewiseCudaGraphRunner:
             self.out_cache_loc.zero_()
             if self.out_cache_loc_swa is not None:
                 self.out_cache_loc_swa.zero_()
+            self.input_ids[num_tokens:static_num_tokens].zero_()
+            self.positions[num_tokens:static_num_tokens].zero_()
+            if self.is_multimodal:
+                self.input_embeds[:, num_tokens:static_num_tokens].zero_()
+            if forward_batch.mrope_positions is not None:
+                self.mrope_positions[:, num_tokens:static_num_tokens].zero_()
+
         bs = forward_batch.batch_size
 
         self.input_ids[:num_tokens].copy_(forward_batch.input_ids)
@@ -685,6 +692,7 @@ class PiecewiseCudaGraphRunner:
         **kwargs,
     ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
         with enable_piecewise_cuda_graph():
+            # Due to the dispatch kernel for MLA model, we init the metadata with original forward_batch
             self.model_runner.attn_backend.init_forward_metadata(forward_batch)
             static_forward_batch = self.replay_prepare(forward_batch, **kwargs)
             # Replay