From 4a7945713924b964d6840d6ae55780b417397590 Mon Sep 17 00:00:00 2001 From: Qiaolin Yu Date: Thu, 23 Jul 2026 22:24:01 -0700 Subject: [PATCH] Fix dynamo recompile limit in allreduce and bf16 gemm (#32239) --- .../sglang/srt/distributed/parallel_state.py | 114 ++++++++++++++---- python/sglang/srt/layers/communicator.py | 7 +- .../sglang/srt/layers/quantization/unquant.py | 41 ++++++- python/sglang/srt/server_args.py | 13 ++ 4 files changed, 141 insertions(+), 34 deletions(-) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index b1d26d32f723..990d1b98ccfb 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -626,7 +626,10 @@ def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: In addition, PyTorch custom ops do not support mutation or returning a new tensor in the same op. So we need to figure out if the op is - in-place or out-of-place ahead of time. + in-place or out-of-place ahead of time — except under Dynamo tracing, + where the method selection would guard on the symbolic shape; there we + always emit the out-of-place op with method "auto" and resolve the + method at runtime inside the op. """ # Bypass the function if we are using only 1 GPU. if self.world_size == 1: @@ -656,6 +659,32 @@ def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: if self.npu_communicator is not None and not self.npu_communicator.disabled: return self.npu_communicator.all_reduce(input_) + if torch.compiler.is_compiling(): + # Byte-size thresholds in method selection (e.g. `_pick_algo` or + # `should_mscclpp_allreduce`) would guard on the symbolic token dim + # and recompile per shape; defer the selection to runtime inside + # the opaque custom op. Groups without any accelerated + # communicator keep the inplace split op so their collective + # stays outside captured graphs. The symmetric-memory in-place + # path below is deliberately bypassed under compile: its raw + # pynccl call is untraceable (hard error with fullgraph, graph + # break otherwise) and its in-place contract does not fit the + # outplace custom op. + if ( + self.ca_comm is None + and self.qr_comm is None + and self.pymscclpp_comm is None + and self.torch_symm_mem_comm is None + and self.pynccl_comm is None + ): + inplace_all_reduce(input_, group_name=self.unique_name) + return input_ + return outplace_all_reduce( + input_, + group_name=self.unique_name, + outplace_all_reduce_method="auto", + ) + should_use_pymscclpp_allreduce = ( self.pymscclpp_comm is not None and self.pymscclpp_comm.should_mscclpp_allreduce(input_) @@ -670,31 +699,10 @@ def all_reduce(self, input_: torch.Tensor) -> torch.Tensor: self.pynccl_comm.all_reduce(input_) return input_ - outplace_all_reduce_method = None - if ( - self.ca_comm is not None - and not self.ca_comm.disabled - and not should_use_pymscclpp_allreduce - and self.ca_comm.should_custom_ar(input_) - ): - outplace_all_reduce_method = "ca" - elif ( - self.qr_comm is not None - and not self.qr_comm.disabled - and self.qr_comm.should_quick_allreduce(input_) - ): - outplace_all_reduce_method = "qr" - elif self.pymscclpp_comm is not None and should_use_pymscclpp_allreduce: - outplace_all_reduce_method = "pymscclpp" - elif ( - self.torch_symm_mem_comm is not None - and not self.torch_symm_mem_comm.disabled - and self.torch_symm_mem_comm.should_torch_symm_mem_allreduce(input_) - ): - outplace_all_reduce_method = "torch_symm_mem" - elif is_in_tc_piecewise_cuda_graph() and self.pynccl_comm is not None: - # For piecewise cuda graph, we use pynccl outplace allreduce - outplace_all_reduce_method = "pynccl" + outplace_all_reduce_method = self._resolve_outplace_all_reduce_method( + input_=input_, + should_use_pymscclpp_allreduce=should_use_pymscclpp_allreduce, + ) if outplace_all_reduce_method is not None: return outplace_all_reduce( input_, @@ -780,9 +788,63 @@ def fused_allreduce_rmsnorm( ) return fused_outputs + def _resolve_outplace_all_reduce_method( + self, + input_: torch.Tensor, + should_use_pymscclpp_allreduce: Optional[bool] = None, + ) -> Optional[str]: + if should_use_pymscclpp_allreduce is None: + should_use_pymscclpp_allreduce = ( + self.pymscclpp_comm is not None + and self.pymscclpp_comm.should_mscclpp_allreduce(input_) + ) + if ( + self.ca_comm is not None + and not self.ca_comm.disabled + and not should_use_pymscclpp_allreduce + and self.ca_comm.should_custom_ar(input_) + ): + return "ca" + if ( + self.qr_comm is not None + and not self.qr_comm.disabled + and self.qr_comm.should_quick_allreduce(input_) + ): + return "qr" + if self.pymscclpp_comm is not None and should_use_pymscclpp_allreduce: + return "pymscclpp" + if ( + self.torch_symm_mem_comm is not None + and not self.torch_symm_mem_comm.disabled + and self.torch_symm_mem_comm.should_torch_symm_mem_allreduce(input_) + ): + return "torch_symm_mem" + if is_in_tc_piecewise_cuda_graph() and self.pynccl_comm is not None: + # For piecewise cuda graph, we use pynccl outplace allreduce + return "pynccl" + return None + def _all_reduce_out_place( self, input_: torch.Tensor, outplace_all_reduce_method: str ) -> torch.Tensor: + if outplace_all_reduce_method == "auto": + outplace_all_reduce_method = self._resolve_outplace_all_reduce_method( + input_ + ) + if outplace_all_reduce_method == "pymscclpp": + # pymscclpp reduces in place and returns its input; feed it a + # clone to honor the op's no-mutation contract. + input_ = input_.clone() + elif outplace_all_reduce_method is None: + # Force pynccl over the in-place fallback: it is graph-capture + # safe and NCCL is natively out-of-place, avoiding the clone + # the in-place fallback needs. + if self.pynccl_comm is not None: + outplace_all_reduce_method = "pynccl" + else: + out = input_.clone() + self._all_reduce_in_place(out) + return out ca_comm = self.ca_comm qr_comm = self.qr_comm pymscclpp_comm = self.pymscclpp_comm diff --git a/python/sglang/srt/layers/communicator.py b/python/sglang/srt/layers/communicator.py index 3271b7c845c7..a23e6e4f17fc 100644 --- a/python/sglang/srt/layers/communicator.py +++ b/python/sglang/srt/layers/communicator.py @@ -167,11 +167,14 @@ def apply_flashinfer_allreduce_fusion(batch_size: int): # Ref: https://github.com/sgl-project/sglang/issues/17237 (_is_sm90_supported or _is_sm100_supported) and _is_flashinfer_available - and batch_size > 0 - and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE and not is_dp_attention_enabled() and get_server_args().flashinfer_allreduce_fusion_backend is not None and not is_flashinfer_allreduce_unavailable() + # Symbolic size checks stay last: under Dynamo tracing they guard on + # the dynamic token dim, so statically-off configs must short-circuit + # before reaching them. + and batch_size > 0 + and batch_size <= FUSE_ALLREDUCE_MAX_BATCH_SIZE ) diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 5c77f22a7b61..d2774728c1dc 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -40,6 +40,7 @@ use_intel_amx_backend, use_intel_xpu_backend, ) +from sglang.srt.utils.custom_op import register_custom_op if TYPE_CHECKING: from sglang.srt.layers.moe.token_dispatcher import ( @@ -104,6 +105,25 @@ def initialize_bf16_gemm_config(server_args: ServerArgs) -> None: _BF16_GEMM_BACKEND = backend +def _bf16_gemm_dispatch_fake( + x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] +) -> torch.Tensor: + return x.new_empty((*x.shape[:-1], weight.shape[0])) + + +@register_custom_op(fake_impl=_bf16_gemm_dispatch_fake) +def bf16_gemm_dispatch( + x: torch.Tensor, weight: torch.Tensor, bias: Optional[torch.Tensor] +) -> torch.Tensor: + if _use_cutedsl_bf16_gemm is not None and _use_cutedsl_bf16_gemm( + x.numel() // x.shape[-1], weight.shape[0], weight.shape[1] + ): + return _cutedsl_bf16_gemm(x.view(-1, x.shape[-1]), weight, bias).view( + *x.shape[:-1], -1 + ) + return F.linear(x, weight, bias) + + def get_bf16_gemm_backend() -> Bf16GemmBackend: global _BF16_GEMM_BACKEND if _BF16_GEMM_BACKEND is None: @@ -207,15 +227,24 @@ def apply( and x.dtype == torch.bfloat16 and layer.weight.dtype == torch.bfloat16 and (bias is None or bias.dtype == torch.bfloat16) - and _use_cutedsl_bf16_gemm( + ): + if torch.compiler.is_compiling(): + # The m-dependent kernel heuristic would guard on the symbolic + # token dim under Dynamo and recompile per shape bucket; the + # opaque op resolves it at runtime with concrete shapes, + # keeping the per-shape kernel choice. + return bf16_gemm_dispatch(x, layer.weight, bias) + if _use_cutedsl_bf16_gemm( x.numel() // x.shape[-1], layer.weight.shape[0], layer.weight.shape[1], - ) - ): - x_shapes = x.shape - output = _cutedsl_bf16_gemm(x.view(-1, x_shapes[-1]), layer.weight, bias) - return output.view(*x_shapes[:-1], -1) + ): + x_shapes = x.shape + output = _cutedsl_bf16_gemm( + x.view(-1, x_shapes[-1]), layer.weight, bias + ) + return output.view(*x_shapes[:-1], -1) + return F.linear(x, layer.weight, bias) return F.linear(x, layer.weight, bias) diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 838b925a34db..66b7f770f291 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3755,6 +3755,19 @@ def _disable_tc_piecewise_cudagraph_if_incompatible(self): "decode context parallel (dcp_size > 1)", lambda: self.dcp_size > 1, ), + # TcPiecewise makes the trtllm_mla prefill fall back to the + # flashinfer-MLA implementation, which faults (illegal address) + # on an FP8 KV cache. + ( + "MLA attention with FP8 KV cache", + lambda: self.kv_cache_dtype.startswith("fp8") + and ( + _resolved_view(self).attention_backend + in ("trtllm_mla", "flashinfer_mla") + or _resolved_view(self).prefill_attention_backend + in ("trtllm_mla", "flashinfer_mla") + ), + ), ] for _name, predicate in rules: if predicate():