diff --git a/docs/references/environment_variables.md b/docs/references/environment_variables.md index 73a11be23..45e51b9ab 100644 --- a/docs/references/environment_variables.md +++ b/docs/references/environment_variables.md @@ -115,6 +115,7 @@ SGLang supports various environment variables that can be used to configure its | Environment Variable | Description | Default Value | | --- | --- | --- | | `SGLANG_USE_AITER` | Use AITER optimize implementation | `false` | +| `SGLANG_ROCM_USE_MULTI_STREAM` | Allocate alt CUDA/HIP stream on ROCm/AITER to overlap shared and routed experts in DeepseekV2 MoE. Requires the HIP env `GPU_MAX_HW_QUEUES>=5` (default `4`, the cap on HSA/ROCr HW queues HIP creates) so the alt stream gets its own queue instead of serializing with the main stream. Best paired with `--deepep-mode low_latency` so Mori's AsyncLL kernel offloads dispatch/combine to copy engines and frees CUs. | `false` | | `SGLANG_MOE_PADDING` | Enable MoE padding (sets padding size to 128 if value is `1`, often set to `1` in Docker builds) | `false` | | `SGLANG_CUTLASS_MOE` (deprecated) | Use Cutlass FP8 MoE kernel on Blackwell GPUs (deprecated, use --moe-runner-backend=cutlass) | `false` | diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d1d321e1f..55a0245eb 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -325,6 +325,9 @@ class Envs: SGLANG_ROCM_FUSED_DECODE_MLA = EnvBool(False) SGLANG_ROCM_DISABLE_LINEARQUANT = EnvBool(False) SGLANG_MORI_NUM_MAX_DISPATCH_TOKENS_PER_RANK = EnvInt(4096) + # Enable dual-stream MoE (shared experts vs routed experts) on the + # ROCm/AITER path. Requires GPU_MAX_HW_QUEUES>=5 to avoid HW-queue serialization. + SGLANG_ROCM_USE_MULTI_STREAM = EnvBool(False) # MPS (Apple Silicon) SGLANG_USE_MLX = EnvBool(False) diff --git a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py index 18c4eb45f..013ed92f0 100644 --- a/python/sglang/srt/layers/moe/token_dispatcher/moriep.py +++ b/python/sglang/srt/layers/moe/token_dispatcher/moriep.py @@ -932,6 +932,18 @@ class MoriEPDispatcher(BaseDispatcher): self.deepep_mode = deepep_mode + async_mode = self.deepep_mode.enable_low_latency() + if get_bool_env_var("SGLANG_ROCM_USE_MULTI_STREAM") and not async_mode: + logger.warning_once( + "SGLANG_ROCM_USE_MULTI_STREAM=1 is set but Mori AsyncLL is " + "not enabled (--deepep-mode=%s). The alt-stream overlap only " + "frees up CUs when dispatch/combine runs on the AsyncLL " + "copy-engine kernel; otherwise it stays on CUs and competes " + "with the alt-stream work. Pass --deepep-mode low_latency " + "(or auto) to enable the AsyncLL kernel.", + self.deepep_mode.value, + ) + common_kwargs = dict( group=group, router_topk=router_topk, diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index 595cb5a5c..38520b787 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -1921,7 +1921,12 @@ class DeepseekV2Model(nn.Module): self.alt_stream = ( torch.cuda.Stream() - if _is_cuda or _is_musa or envs.SGLANG_NPU_USE_MULTI_STREAM.get() + if ( + _is_cuda + or _is_musa + or envs.SGLANG_NPU_USE_MULTI_STREAM.get() + or envs.SGLANG_ROCM_USE_MULTI_STREAM.get() + ) else None )