From afebb7ab7893869ec8cf7ed6dd6f06981bc8ccef Mon Sep 17 00:00:00 2001 From: Yuan Luo Date: Sun, 1 Feb 2026 18:59:31 +0800 Subject: [PATCH] Optimize custom-all-reduce (#17674) Co-authored-by: luoyuan.luo --- docs/references/environment_variables.md | 1 + .../csrc/allreduce/custom_all_reduce.cuh | 47 ++++++++++++++----- 2 files changed, 36 insertions(+), 12 deletions(-) diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 4a9690d11..037210668 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -44,6 +44,7 @@ SGLang supports various environment variables that can be used to configure its | `SGLANG_MM_PRECOMPUTE_HASH` | Enable precomputing of hash values for MultimodalDataItem | `false` | | `SGLANG_NCCL_ALL_GATHER_IN_OVERLAP_SCHEDULER_SYNC_BATCH` | Enable NCCL for gathering when preparing mlp sync batch under overlap scheduler (without this flag gloo is used for gathering) | `false` | | `SGLANG_SYMM_MEM_PREALLOC_GB_SIZE` | Size of preallocated GPU buffer (in GB) for NCCL symmetric memory pool to limit memory fragmentation. Only have an effect when server arg `--enable-symm-mem` is set. | `4` | +| `SGLANG_CUSTOM_ALLREDUCE_ALGO` | The algorithom of custom all-reduce. Set to `oneshot` or `1stage` to force use one-shot. Set to `twoshot` or `2stage` to force use two-shot. | `` | ## DeepGEMM Configuration (Advanced Optimization) diff --git a/sgl-kernel/csrc/allreduce/custom_all_reduce.cuh b/sgl-kernel/csrc/allreduce/custom_all_reduce.cuh index ec223bdeb..4a2d62b28 100644 --- a/sgl-kernel/csrc/allreduce/custom_all_reduce.cuh +++ b/sgl-kernel/csrc/allreduce/custom_all_reduce.cuh @@ -442,21 +442,44 @@ class CustomAllreduce { size /= d; auto bytes = size * sizeof(typename packed_t::P); int blocks = std::min(block_limit, (size + threads - 1) / threads); + + // Check environment variable once + const char* env_algo = std::getenv("SGLANG_CUSTOM_ALLREDUCE_ALGO"); + bool force_1stage = false; + bool force_2stage = false; + if (env_algo != nullptr) { + if (std::strcmp(env_algo, "1stage") == 0 || std::strcmp(env_algo, "oneshot") == 0) { + force_1stage = true; + } else if (std::strcmp(env_algo, "2stage") == 0 || std::strcmp(env_algo, "twoshot") == 0) { + force_2stage = true; + } else { + throw std::runtime_error( + "Invalid SGLANG_CUSTOM_ALLREDUCE_ALGO: " + std::string(env_algo) + + ". Valid values: 1stage, oneshot, 2stage, twoshot"); + } + } + #define KL(ngpus, name) name<<>>(ptrs, sg_, self_sg_, output, rank_, size); // TODO(hanzhi713): Threshold is different for A100 and H100. // Add per device threshold. -#define REDUCE_CASE(ngpus) \ - case ngpus: { \ - if (world_size_ == 2) { \ - KL(ngpus, cross_device_reduce_1stage); \ - } else if (full_nvlink_) { \ - if ((world_size_ <= 4 && bytes < 512 * 1024) || (world_size_ <= 8 && bytes < 256 * 1024)) { \ - KL(ngpus, cross_device_reduce_1stage); \ - } else { \ - KL(ngpus, cross_device_reduce_2stage); \ - } \ - } \ - break; \ +#define REDUCE_CASE(ngpus) \ + case ngpus: { \ + if (force_1stage) { \ + KL(ngpus, cross_device_reduce_1stage); \ + } else if (force_2stage) { \ + KL(ngpus, cross_device_reduce_2stage); \ + } else { \ + if (world_size_ == 2) { \ + KL(ngpus, cross_device_reduce_1stage); \ + } else if (full_nvlink_) { \ + if ((world_size_ <= 4 && bytes < 512 * 1024) || (world_size_ <= 8 && bytes < 256 * 1024)) { \ + KL(ngpus, cross_device_reduce_1stage); \ + } else { \ + KL(ngpus, cross_device_reduce_2stage); \ + } \ + } \ + } \ + break; \ } switch (world_size_) {