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

Turn on JIT custom AR implementation by default (#24363)

Browse Source 
Co-authored-by: b8zhong <b8zhong@users.noreply.github.com>
This commit is contained in:
Brayden Zhong
2026-05-08 02:05:31 -07:00
committed by GitHub
co-authored by b8zhong
parent 73792629d4
commit 80d0226b68
2 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
+2 -1
View File
@@ -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
python/sglang/srt/environ.py
+1 -1
View File
@@ -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)