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

Rename SGLANG_USE_JIT_ALL_REDUCE to SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2 (#24297)

Browse Source 
Co-authored-by: DarkSharpness <2040703891@qq.com>
This commit is contained in:
Liangsheng Yin
2026-05-02 23:59:46 -07:00
committed by GitHub
co-authored by DarkSharpness
parent b7d8ceb444
commit fcc8b7b126
2 changed files with 6 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
+2 -4
View File
@@ -342,11 +342,9 @@ def dispatch_custom_allreduce():
On AMD with 1-stage AR enabled, use sglang's CustomAllreduce. On AMD with 1-stage AR enabled, use sglang's CustomAllreduce.
Otherwise use AiterCustomAllreduce if available. 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. if _is_cuda and envs.SGLANG_OPT_USE_CUSTOM_ALL_REDUCE_V2.get():
# Set SGLANG_USE_JIT_ALL_REDUCE=1 to enable.
if _is_cuda and get_bool_env_var("SGLANG_USE_JIT_ALL_REDUCE", default="false"):
from .custom_all_reduce_v2 import CustomAllReduceV2 from .custom_all_reduce_v2 import CustomAllReduceV2
logger.debug("[AR] Using CustomAllReduceV2 (JIT-compiled)") logger.debug("[AR] Using CustomAllReduceV2 (JIT-compiled)")
python/sglang/srt/environ.py
+4
View File
@@ -445,6 +445,7 @@ class Envs:
# Set to 1: force enable (even without --enable-deterministic-inference) # Set to 1: force enable (even without --enable-deterministic-inference)
# Set to 0: force disable (use default Aiter AR even with --enable-deterministic-inference) # Set to 0: force disable (use default Aiter AR even with --enable-deterministic-inference)
SGLANG_USE_1STAGE_ALLREDUCE = EnvBool(False) 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_PREFILL_SPLIT_TILE_SIZE = EnvInt(4096)
SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE = EnvInt(2048) SGLANG_FLASHINFER_DECODE_SPLIT_TILE_SIZE = EnvInt(2048)
SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096) SGLANG_TRITON_PREFILL_TRUNCATION_ALIGN_SIZE = EnvInt(4096)
@@ -604,6 +605,9 @@ def _convert_SGL_to_SGLANG():
"SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK", "SGLANG_ENABLE_TP_MEMORY_INBALANCE_CHECK",
) )
_print_deprecated_env("SGLANG_PER_TOKEN_GROUP_QUANT_8BIT_V2") _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 = { _deprecated_ms_to_s = {
"SGLANG_QUEUED_TIMEOUT_MS": "SGLANG_REQ_WAITING_TIMEOUT", "SGLANG_QUEUED_TIMEOUT_MS": "SGLANG_REQ_WAITING_TIMEOUT",
"SGLANG_FORWARD_TIMEOUT_MS": "SGLANG_REQ_RUNNING_TIMEOUT", "SGLANG_FORWARD_TIMEOUT_MS": "SGLANG_REQ_RUNNING_TIMEOUT",