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 ac308df63..b14211e17 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py @@ -342,7 +342,8 @@ def dispatch_custom_allreduce(): On AMD with 1-stage AR enabled, use sglang's CustomAllreduce. Otherwise use AiterCustomAllreduce if available. - Set SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=1 to use the JIT-compiled v2 implementation. + On CUDA, the JIT-compiled v2 implementation is used by default. + Set SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=0 to fall back to the legacy CustomAllreduce. """ if _is_cuda and envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get(): from .custom_all_reduce_v2 import CustomAllReduceV2 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 3167b8ee2..c87321c35 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -448,7 +448,7 @@ class Envs: # Set to 1: force enable (even without --enable-deterministic-inference) # Set to 0: force disable (use default Aiter AR even with --enable-deterministic-inference) SGLANG_USE_1STAGE_ALLREDUCE = EnvBool(False) - SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2 = EnvBool(False) + SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2 = EnvBool(True) SGLANG_FLASHINFER_PREFILL_SPLIT_TILE_SIZE = EnvInt(4096) SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE = EnvInt(2048) SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096)