minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Fix dynamo recompile limit in allreduce and bf16 gemm (#32239)

Browse Source
This commit is contained in:
Qiaolin Yu
2026-07-23 22:24:01 -07:00
committed by GitHub
parent 319055c191
commit 15d73f1e03
4 changed files with 141 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/parallel_state.py
+88 -26
View File
@@ -631,7 +631,10 @@ class GroupCoordinator:
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:
@@ -661,6 +664,32 @@ class GroupCoordinator:
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_)
@@ -675,31 +704,10 @@ class GroupCoordinator:
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_,
@@ -856,9 +864,63 @@ class GroupCoordinator:
except Exception:
return None
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
+5 -2
View File
@@ -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
+35 -6
View File
@@ -40,6 +40,7 @@ from sglang.srt.utils import (
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 (
@@ -109,6 +110,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:
@@ -214,15 +234,24 @@ class UnquantizedLinearMethod(LinearMethodBase):
and (bias is None or bias.dtype == torch.bfloat16)
and not layer.weight.requires_grad
and (bias is None or not bias.requires_grad)
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
+13
View File
@@ -4176,6 +4176,19 @@ class ServerArgs:
"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():