[Cherry-pick to release/v0.5.16] Fix dynamo recompile limit in allreduce and bf16 gemm (#32239)

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

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
     )
 
 

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)
 

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():