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 c9e055bb2..ac308df63 100644 --- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py +++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py @@ -342,11 +342,9 @@ def dispatch_custom_allreduce(): On AMD with 1-stage AR enabled, use sglang's CustomAllreduce. Otherwise use AiterCustomAllreduce if available. - Set SGLANG_USE_JIT_ALL_REDUCE=1 to use the JIT-compiled v2 implementation. + Set SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2=1 to use the JIT-compiled v2 implementation. """ - # HARDCODED: opt-in flag for v2 JIT all-reduce. - # Set SGLANG_USE_JIT_ALL_REDUCE=1 to enable. - if _is_cuda and get_bool_env_var("SGLANG_USE_JIT_ALL_REDUCE", default="false"): + if _is_cuda and envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get(): from .custom_all_reduce_v2 import CustomAllReduceV2 logger.debug("[AR] Using CustomAllReduceV2 (JIT-compiled)") diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index c7247261d..d1d321e1f 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -445,6 +445,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_FLASHINFER_PREFILL_SPLIT_TILE_SIZE = EnvInt(4096) SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE = EnvInt(2048) SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096) @@ -604,6 +605,9 @@ def _convert_SGL_to_SGLANG(): "SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK", ) _print_deprecated_env("SGLANG_PER_TOKEN_GROUP_QUANT_8BIT_V2") + _print_deprecated_env( + "SGLANG_USE_JIT_ALL_REDUCE", "SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2" + ) _deprecated_ms_to_s = { "SGLANG_QUEUED_TIMEOUT_MS": "SGLANG_REQ_WAITING_TIMEOUT", "SGLANG_FORWARD_TIMEOUT_MS": "SGLANG_REQ_RUNNING_TIMEOUT",