From b7d3c3016d8cdc10087e65e6a5ea0ab6c697968f Mon Sep 17 00:00:00 2001 From: kk <43161300+kkHuang-amd@users.noreply.github.com> Date: Fri, 26 Jun 2026 04:45:18 +0800 Subject: [PATCH] [AMD] Feat/dsv4 aiter reduce scatter decode (#29103) Co-authored-by: wunhuang --- .../sglang/srt/distributed/parallel_state.py | 49 +++++++++++++++++++ python/sglang/srt/environ.py | 35 +++++++++++-- python/sglang/srt/layers/dp_attention.py | 12 ++++- python/sglang/srt/models/deepseek_v4.py | 37 ++++++++++++-- 4 files changed, 124 insertions(+), 9 deletions(-) diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 3d1a6b093..8a48047de 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -798,9 +798,58 @@ class GroupCoordinator: def reduce_scatter_tensor(self, output: torch.Tensor, input: torch.Tensor): if _is_npu: self._reduce_scatter_tensor(output, input) + elif self._maybe_aiter_reduce_scatter(output, input): + return else: reg_reduce_scatter_tensor(output, input, group_name=self.unique_name) + def _has_aiter_custom_reduce_scatter(self) -> bool: + ca_comm = self.ca_comm + return ( + ca_comm is not None + and not getattr(ca_comm, "disabled", True) + and hasattr(ca_comm, "should_custom_ar") + and hasattr(ca_comm, "reduce_scatter") + ) + + def _maybe_aiter_reduce_scatter( + self, output: torch.Tensor, input: torch.Tensor + ) -> bool: + # Aiter custom reduce-scatter (ROCm). Mirrors `_all_gather_into_tensor`'s + # custom all-gather path: an equal-chunk (no variable sizes) reduce-scatter + # using the registered symmetric-memory buffers, which is faster than the + # generic RCCL kernel for the small, latency-bound decode collective. + # Gated by SGLANG_DP_USE_REDUCE_SCATTER. Falls back (returns False) + # for non-ROCm / unsupported shape/size/topology so the caller uses RCCL. + if not ( + is_hip() + and envs.SGLANG_DP_USE_REDUCE_SCATTER.get() + and self._has_aiter_custom_reduce_scatter() + and input.is_contiguous() + and output.is_contiguous() + and input.dtype in (torch.float32, torch.float16, torch.bfloat16) + ): + return False + ca_comm = self.ca_comm + # input is the full (pre-reduce) buffer; should_custom_ar bounds its size. + if not ca_comm.should_custom_ar(input): + return False + # Equal-chunk only: input rows must split evenly into world_size chunks + # matching the per-rank output rows. + if input.shape[0] != output.shape[0] * self.world_size: + return False + if getattr(ca_comm, "_IS_CAPTURING", False): + if torch.cuda.is_current_stream_capturing(): + ca_comm.reduce_scatter(input, output, registered=True) + elif is_in_tc_piecewise_cuda_graph(): + ca_comm.reduce_scatter(input, output, registered=False) + else: + # True CUDA graph warmup: avoid a different host collective. + output.zero_() + return True + ca_comm.reduce_scatter(input, output, registered=False) + return True + def _all_to_all_single(self, output: torch.Tensor, input: torch.Tensor) -> None: torch.distributed.all_to_all_single(output, input, group=self.device_group) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index a1768ce05..eb9455dc2 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -1,3 +1,4 @@ +import functools import os import subprocess import warnings @@ -6,6 +7,23 @@ from enum import IntEnum from typing import Any, Optional +@functools.lru_cache(maxsize=1) +def _default_hip() -> bool: + """Lazy ROCm/HIP detection for platform-conditional env defaults. + + Avoids importing torch at environ import time (this module is intentionally + stdlib-only and loaded very early). Resolved on first EnvField.get() that uses + it as a default, by which point torch is already imported in any real run; + falls back to False if torch is unavailable. + """ + try: + import torch + + return torch.version.hip is not None + except Exception: + return False + + @contextmanager def temp_set_env(*, allow_sglang: bool = False, **env_vars: Any): """Temporarily set environment variables, restoring originals on exit. @@ -51,6 +69,11 @@ class EnvField: def parse(self, value: str) -> Any: raise NotImplementedError() + def _resolve_default(self) -> Any: + # Support a callable default for lazily/platform-computed defaults + # (e.g. EnvBool(_default_hip)); evaluated only when the env is unset. + return self.default() if callable(self.default) else self.default + def get(self) -> Any: value = os.getenv(self.name) @@ -61,15 +84,16 @@ class EnvField: # Not set, return default if value is None: - return self.default + return self._resolve_default() try: return self.parse(value) except ValueError as e: + default = self._resolve_default() warnings.warn( - f'Invalid value for {self.name}: {e}, using default "{self.default}"' + f'Invalid value for {self.name}: {e}, using default "{default}"' ) - return self.default + return default def is_set(self): return self.name in os.environ @@ -443,6 +467,11 @@ class Envs: # AMD & ROCm SGLANG_USE_AITER = EnvBool(False) SGLANG_USE_AITER_AG = EnvBool(True) + # Use reduce_scatter (instead of all_reduce + dp_scatter) for the equal-chunk + # MAX_LEN DP-MoE combine. Default ON for ROCm/HIP (uses the aiter custom + # symmetric-memory kernel), OFF elsewhere (would fall back to RCCL); override + # explicitly to force on/off on any platform. + SGLANG_DP_USE_REDUCE_SCATTER = EnvBool(_default_hip) SGLANG_USE_AITER_UNIFIED_ATTN = EnvBool(False) # Select the gate/up tile layout for AITER MoE: True -> interleave # (matches FlyDSL `gate_mode="interleave"` kernels), False -> separated diff --git a/python/sglang/srt/layers/dp_attention.py b/python/sglang/srt/layers/dp_attention.py index fc0d2d8b4..00495d84a 100644 --- a/python/sglang/srt/layers/dp_attention.py +++ b/python/sglang/srt/layers/dp_attention.py @@ -533,12 +533,20 @@ _USE_DP_GATHERV = get_bool_env_var("SGLANG_DP_USE_GATHERV") def is_dp_gatherv_active() -> bool: """Variable-length DP-MoE gather/scatter (all_gatherv + reduce_scatterv) is - enabled and the current parallel layout (attn_tp_size==1, tp_size==dp_size) - is supported. Env-gated by SGLANG_DP_USE_GATHERV; default off.""" + enabled and applicable to the CURRENT forward. Requires: + - env SGLANG_DP_USE_GATHERV (default off), + - supported layout (attn_tp_size==1, tp_size==dp_size), + - SUM_LEN padding mode. The gatherv pair (all_gatherv + reduce_scatterv) is + only valid under SUM_LEN; under MAX_LEN the buffer is equal-padded and the + gather/combine use all_gather / (aiter) reduce_scatter instead. Reading the + per-forward padding via _DpGatheredBufferWrapper.is_dp_max_padding() (set by + set_dp_buffer_len) keeps callers that lack a ForwardBatch (e.g. + dp_reduce_scatter_tensor) consistent.""" return ( _USE_DP_GATHERV and get_attention_tp_size() == 1 and get_tensor_model_parallel_world_size() == get_attention_dp_size() + and not _DpGatheredBufferWrapper.is_dp_max_padding() ) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 7d3265bd9..c0ce9c629 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -57,6 +57,7 @@ from sglang.srt.layers.dp_attention import ( attn_tp_all_reduce, dp_gather_partial, dp_gather_replicate, + dp_reduce_scatter_tensor, dp_scatter, get_dp_global_num_tokens, get_global_dp_buffer, @@ -1587,12 +1588,28 @@ class DeepseekV4DecoderLayer(nn.Module): # in ONE op, REPLACING the MoE-internal post-experts all_reduce — so we # MUST tell the MoE to skip it (use_reduce_scatter=True) or it # double-reduces. Env-gated via SGLANG_DP_USE_GATHERV, default OFF. - _use_gatherv_pair = ( + _use_reduce_scatterv = ( _use_tp_moe_gather and is_dp_gatherv_active() and forward_batch.dp_padding_mode is not None and not forward_batch.dp_padding_mode.is_max_len() ) + # SGLANG_DP_USE_REDUCE_SCATTER: in the MAX_LEN decode path (equal per-rank + # padding, gatherv inactive, no EP), replace the MoE-internal post-experts + # all_reduce + dp_scatter with an equal-chunk reduce_scatter. On ROCm this + # uses the aiter custom kernel (so BOTH gather and combine are aiter custom), + # elsewhere RCCL reduce_scatter; either way it cuts combine traffic ~2x vs + # all_reduce. tp_size==attn_dp_size required so the global buffer splits + # evenly into per-rank chunks. + _use_reduce_scatter = ( + envs.SGLANG_DP_USE_REDUCE_SCATTER.get() + and _use_tp_moe_gather + and not _use_reduce_scatterv + and not should_use_dp_reduce_scatterv() + and forward_batch.dp_padding_mode is not None + and forward_batch.dp_padding_mode.is_max_len() + and get_parallel().tp_size == get_parallel().attn_dp_size + ) # PoC (SGLANG_DP_SHARED_EXPERT_LOCAL): compute the replicated shared expert # on LOCAL hidden before the gather and add it back after the combine # (reduce_scatterv OR dp_scatter), instead of on the gathered global buffer. @@ -1638,8 +1655,9 @@ class DeepseekV4DecoderLayer(nn.Module): input_ids=input_ids, input_ids_global=input_ids_global, # Skip the MoE-internal post-experts all_reduce when we will do the - # reduce via reduce_scatterv at the combine below (else double-reduce). - use_reduce_scatter=_use_cp or _use_gatherv_pair, + # reduce via reduce_scatterv/reduce_scatter at the combine below + # (else double-reduce). + use_reduce_scatter=_use_cp or _use_reduce_scatterv or _use_reduce_scatter, skip_shared_experts=_do_shared_local, ) if _use_cp and get_moe_a2a_backend().is_none(): @@ -1649,7 +1667,7 @@ class DeepseekV4DecoderLayer(nn.Module): get_local_dp_buffer(get_tp_group()), hidden_states, ) - if should_use_dp_reduce_scatterv() or _use_gatherv_pair: + if should_use_dp_reduce_scatterv() or _use_reduce_scatterv: # SUM the TP-sharded per-rank partial expert outputs AND scatter # each rank its own token slice, in one op. Correct because the # MoE-internal all_reduce was skipped (use_reduce_scatter above). @@ -1659,6 +1677,17 @@ class DeepseekV4DecoderLayer(nn.Module): output=hidden_states, sizes=get_dp_global_num_tokens(), ) + elif _use_reduce_scatter: + # Equal-chunk reduce_scatter: SUM the TP-sharded per-rank partial + # expert outputs AND scatter each rank its own (MAX_LEN-padded) + # token chunk in one op (symmetric inverse of the MAX_LEN + # all_gather). Correct because the MoE-internal all_reduce was + # skipped (use_reduce_scatter above). dp_reduce_scatter_tensor + # routes to the equal-chunk reduce_scatter_tensor here (its + # variable-length reduce_scatterv branch is gated by + # is_dp_gatherv_active(), which is False under MAX_LEN), which in + # turn uses the aiter custom kernel when it fits (else RCCL). + dp_reduce_scatter_tensor(hidden_states, global_hidden_states) else: dp_scatter(hidden_states, global_hidden_states, forward_batch) # PoC: add the locally-computed shared-expert output to this rank's