Enable TRT AllReduce Fusion by default (#14764)
Co-authored-by: Brayden Zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
@@ -41,6 +41,7 @@ from sglang.srt.utils.common import (
|
|||||||
get_bool_env_var,
|
get_bool_env_var,
|
||||||
get_device,
|
get_device,
|
||||||
get_device_memory_capacity,
|
get_device_memory_capacity,
|
||||||
|
get_device_name,
|
||||||
get_device_sm,
|
get_device_sm,
|
||||||
is_blackwell_supported,
|
is_blackwell_supported,
|
||||||
is_cuda,
|
is_cuda,
|
||||||
@@ -1109,12 +1110,6 @@ class ServerArgs:
|
|||||||
|
|
||||||
# common to all Deepseek MoE models
|
# common to all Deepseek MoE models
|
||||||
if is_cuda() and is_sm100_supported():
|
if is_cuda() and is_sm100_supported():
|
||||||
# workaround for https://github.com/flashinfer-ai/flashinfer/issues/2006
|
|
||||||
if not self.enable_dp_attention and self.nnodes == 1:
|
|
||||||
self.enable_flashinfer_allreduce_fusion = True
|
|
||||||
logger.info(
|
|
||||||
"Enable FlashInfer AllReduce Fusion on sm100 for DeepseekV3ForCausalLM"
|
|
||||||
)
|
|
||||||
quantization_config = getattr(hf_config, "quantization_config", None)
|
quantization_config = getattr(hf_config, "quantization_config", None)
|
||||||
quant_method = (
|
quant_method = (
|
||||||
quantization_config.get("quant_method")
|
quantization_config.get("quant_method")
|
||||||
@@ -1166,13 +1161,6 @@ class ServerArgs:
|
|||||||
f"- Decode: {decode_attn_backend}\n"
|
f"- Decode: {decode_attn_backend}\n"
|
||||||
)
|
)
|
||||||
|
|
||||||
if is_blackwell_supported():
|
|
||||||
# workaround for https://github.com/flashinfer-ai/flashinfer/issues/2006
|
|
||||||
if not self.enable_dp_attention and self.nnodes == 1:
|
|
||||||
self.enable_flashinfer_allreduce_fusion = True
|
|
||||||
logger.info(
|
|
||||||
"Enable FlashInfer AllReduce Fusion on sm100 for GptOssForCausalLM"
|
|
||||||
)
|
|
||||||
quantization_config = getattr(hf_config, "quantization_config", None)
|
quantization_config = getattr(hf_config, "quantization_config", None)
|
||||||
is_mxfp4_quant_format = (
|
is_mxfp4_quant_format = (
|
||||||
quantization_config is not None
|
quantization_config is not None
|
||||||
@@ -1437,6 +1425,32 @@ class ServerArgs:
|
|||||||
self.disable_radix_cache = True
|
self.disable_radix_cache = True
|
||||||
self.disable_overlap_schedule = False
|
self.disable_overlap_schedule = False
|
||||||
|
|
||||||
|
# TRTLLM AllReduce Fusion supports SM90/100/120, enable it by default
|
||||||
|
# for models with explicit support (DeepseekV3, GptOss, Glm4Moe, Qwen3Moe)
|
||||||
|
# TODO: currently, it is only supported in the single node scenario. https://github.com/flashinfer-ai/flashinfer/issues/2006
|
||||||
|
# TODO: there is currently a bug on H20 device specifically, https://github.com/flashinfer-ai/flashinfer/issues/2204
|
||||||
|
device_name = get_device_name()
|
||||||
|
is_h20_device = "H20" in device_name and "H200" not in device_name
|
||||||
|
if (
|
||||||
|
not self.enable_flashinfer_allreduce_fusion
|
||||||
|
and model_arch
|
||||||
|
in [
|
||||||
|
"DeepseekV3ForCausalLM",
|
||||||
|
"GptOssForCausalLM",
|
||||||
|
"Glm4MoeForCausalLM",
|
||||||
|
"Qwen3MoeForCausalLM",
|
||||||
|
]
|
||||||
|
and (is_sm90_supported() or is_blackwell_supported())
|
||||||
|
and not self.enable_dp_attention
|
||||||
|
and self.nnodes == 1
|
||||||
|
and not is_h20_device
|
||||||
|
and self.moe_a2a_backend == "none"
|
||||||
|
):
|
||||||
|
self.enable_flashinfer_allreduce_fusion = True
|
||||||
|
logger.info(
|
||||||
|
f"Enable FlashInfer AllReduce Fusion by default for {model_arch}"
|
||||||
|
)
|
||||||
|
|
||||||
def _handle_sampling_backend(self):
|
def _handle_sampling_backend(self):
|
||||||
if self.sampling_backend is None:
|
if self.sampling_backend is None:
|
||||||
self.sampling_backend = (
|
self.sampling_backend = (
|
||||||
|
|||||||
@@ -226,14 +226,14 @@ def is_blackwell():
|
|||||||
|
|
||||||
@lru_cache(maxsize=1)
|
@lru_cache(maxsize=1)
|
||||||
def is_blackwell_supported(device=None) -> bool:
|
def is_blackwell_supported(device=None) -> bool:
|
||||||
if not is_cuda_alike():
|
if not is_cuda():
|
||||||
return False
|
return False
|
||||||
return is_sm100_supported(device) or is_sm120_supported(device)
|
return is_sm100_supported(device) or is_sm120_supported(device)
|
||||||
|
|
||||||
|
|
||||||
@lru_cache(maxsize=1)
|
@lru_cache(maxsize=1)
|
||||||
def is_sm120_supported(device=None) -> bool:
|
def is_sm120_supported(device=None) -> bool:
|
||||||
if not is_cuda_alike():
|
if not is_cuda():
|
||||||
return False
|
return False
|
||||||
return (torch.cuda.get_device_capability(device)[0] == 12) and (
|
return (torch.cuda.get_device_capability(device)[0] == 12) and (
|
||||||
torch.version.cuda >= "12.8"
|
torch.version.cuda >= "12.8"
|
||||||
|
|||||||
Reference in New Issue
Block a user