From 0be30d4b0d5468daf4ea34b01df1e142463a08aa Mon Sep 17 00:00:00 2001 From: Yuwei An Date: Thu, 19 Feb 2026 00:48:06 -0800 Subject: [PATCH] Fix PCG MoE Error (#17739) --- python/sglang/srt/compilation/compile.py | 4 +-- .../compilation/piecewise_context_manager.py | 7 +++-- .../device_communicators/custom_all_reduce.py | 17 +++++++++-- .../device_communicators/pynccl.py | 29 +++++++++++++++++++ .../sglang/srt/distributed/parallel_state.py | 14 +++++++-- .../piecewise_cuda_graph_runner.py | 8 +++++ .../test_piecewise_cuda_graph_2_gpu.py | 4 +-- 7 files changed, 71 insertions(+), 12 deletions(-) diff --git a/python/sglang/srt/compilation/compile.py b/python/sglang/srt/compilation/compile.py index b9ff7f6bd..b3feace52 100644 --- a/python/sglang/srt/compilation/compile.py +++ b/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) diff --git a/python/sglang/srt/compilation/piecewise_context_manager.py b/python/sglang/srt/compilation/piecewise_context_manager.py index a87d31f3e..e588b6335 100644 --- a/python/sglang/srt/compilation/piecewise_context_manager.py +++ b/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 diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py index 6e7ea07e7..b761bb69b 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/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, diff --git a/python/sglang/srt/distributed/device_communicators/pynccl.py b/python/sglang/srt/distributed/device_communicators/pynccl.py index 86c53f26b..660582ad3 100644 --- a/python/sglang/srt/distributed/device_communicators/pynccl.py +++ b/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, diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 2c4e99c38..8f1069c00 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/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 diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 0cb778c7b..a8ea3af57 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/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 diff --git a/test/registered/cuda_graph/test_piecewise_cuda_graph_2_gpu.py b/test/registered/cuda_graph/test_piecewise_cuda_graph_2_gpu.py index dfaca8270..ac8176210 100644 --- a/test/registered/cuda_graph/test_piecewise_cuda_graph_2_gpu.py +++ b/test/registered/cuda_graph/test_piecewise_cuda_graph_2_gpu.py @@ -12,9 +12,7 @@ from sglang.test.test_utils import ( ) # CI Registration - 2-GPU tests (80GB GPUs required) -register_cuda_ci( - est_time=160, suite="stage-b-test-large-2-gpu", disabled="see issue #16691" -) +register_cuda_ci(est_time=160, suite="stage-b-test-large-2-gpu") class TestPiecewiseCudaGraphTP(CustomTestCase):