diff --git a/python/sglang/srt/arg_groups/deepseek_v4_hook.py b/python/sglang/srt/arg_groups/deepseek_v4_hook.py index 14b9c29553a5..21995ae36fe5 100644 --- a/python/sglang/srt/arg_groups/deepseek_v4_hook.py +++ b/python/sglang/srt/arg_groups/deepseek_v4_hook.py @@ -18,7 +18,7 @@ from sglang.srt.environ import envs from sglang.srt.model_executor.cuda_graph_config import Backend, Phase, with_phase from sglang.srt.runtime_context import attn_dp_enabled_of, get_platform -from sglang.srt.utils.common import is_gfx95_supported, is_npu +from sglang.srt.utils.common import is_gfx95_supported, is_hip, is_npu if TYPE_CHECKING: from sglang.srt.server_args import ServerArgs @@ -325,8 +325,17 @@ def validate_deepseek_v41_features(server_args: ServerArgs) -> None: "the prefill CUDA graph", cfg.cuda_graph_config.prefill.backend != Backend.DISABLED, ), - # input_ids_global is a DP-wide gather, so the tail slice cannot apply. - ("DP attention", attn_dp_enabled_of(cfg)), + # DP attention + TP MoE re-sizes the MoE gather to every rank's tail + # (deepseek_v4 _late_layer_dp_counts, ROCm only); other DP layouts do not. + ( + "DP attention other than TP MoE with attention TP 1 on ROCm", + attn_dp_enabled_of(cfg) + and not ( + is_hip() + and cfg.moe_a2a_backend == "none" + and cfg.attn_dp_size == cfg.tp_size + ), + ), ) for feature, enabled in incompatible: if enabled: diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 7f403b85dc9c..383d496c7dfa 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -84,6 +84,11 @@ _is_cpu = is_cpu() _is_xpu = is_xpu() _is_musa = is_musa() +# Under CUDA-graph capture, stage aiter custom all-gather / reduce-scatter inputs +# through the pre-registered IPC pool instead of registering graph-pool tensors: +# the registered path makes peers read stale data over IPC on replay (garbage +# rows / NaN under DP attention + TP MoE on MI355X). Set to 0 to restore it. +_AITER_CAPTURE_COPY_IN = os.environ.get("SGLANG_AITER_CAPTURE_COPY_IN", "0") == "1" TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"]) @@ -1190,7 +1195,7 @@ def _maybe_aiter_reduce_scatter( return False if getattr(ca_comm, "_IS_CAPTURING", False): if torch.cuda.is_current_stream_capturing(): - if envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get(): + if envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get() or _AITER_CAPTURE_COPY_IN: ca_comm.reduce_scatter(input, output, registered=False) else: ca_comm.reduce_scatter(input, output, registered=True) @@ -1309,7 +1314,10 @@ def _all_gather_into_tensor(self, output: torch.Tensor, input: torch.Tensor): ): if getattr(ca_comm, "_IS_CAPTURING", False): if torch.cuda.is_current_stream_capturing(): - if envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get(): + if ( + envs.SGLANG_MEMORY_SAVER_CUDA_GRAPH.get() + or _AITER_CAPTURE_COPY_IN + ): ca_comm.all_gather_unreg(input, out=output, dim=0) else: ca_comm.all_gather_reg(input, out=output, dim=0) diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 4ba4a6347db7..3e6ac6df7abe 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -31,7 +31,7 @@ import warnings from dataclasses import dataclass from enum import IntEnum, auto -from functools import total_ordering +from functools import lru_cache, total_ordering from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union import torch @@ -62,6 +62,7 @@ ) from sglang.srt.speculative.spec_info import SpecInputType from sglang.srt.utils import ( + get_bool_env_var, is_cpu, is_hip, is_npu, @@ -2115,10 +2116,27 @@ def enable_num_token_non_padded(): # Elastic joiners also need graph padding masked after joining WORLD; a # customized A2A backend consumes the count whatever the EP size. - return ( + if ( get_parallel().moe_ep_size > 1 or world_dp_gather_enabled() or get_moe_a2a_backend().is_customized() + ): + return True + # DP attention + TP MoE gathers each rank's whole padded chunk into the MoE, + # so the model needs the real count to zero the pad rows (see + # deepseek_v4 _run_moe_ffn_dp_sync). + return _dp_tp_moe_masks_pad_rows() + + +@lru_cache(maxsize=1) +def _dp_tp_moe_masks_pad_rows() -> bool: + from sglang.srt.layers.moe.utils import get_moe_a2a_backend + + return ( + _is_hip + and get_bool_env_var("SGLANG_DP_MASK_PAD_ROWS", default="true") + and get_parallel().attn_dp_size > 1 + and get_moe_a2a_backend().is_none() ) diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index aa8e4c1a2549..321379c5a443 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -70,6 +70,7 @@ ForwardBatch, ForwardMode, PPProxyTensors, + _dp_tp_moe_masks_pad_rows, compute_local_num_token_non_padded, enable_num_token_non_padded, get_required_capture_hidden_mode, @@ -1358,6 +1359,14 @@ def run_once(): post_warmup_hook=post_warmup_hook, ) + def _set_dp_real_num_tokens(self, raw_num_token: int) -> None: + # An idle DP rank replays a fabricated batch whose count may not reach + # the slot, leaving the previous replay's value; the model masks DP pad + # rows by this count, so write the real one on every replay. + buf = self.buffers.num_token_non_padded + if buf is not None and _dp_tp_moe_masks_pad_rows(): + buf.fill_(raw_num_token) + def _validate_capture_hidden_mode(self, forward_batch: ForwardBatch) -> None: if self.capture_hidden_mode < forward_batch.capture_hidden_mode: raise RuntimeError( @@ -1410,6 +1419,7 @@ def load_batch( f"{ragged_layout.graph_num_tokens}" ) self._stage_ragged_verify_layout(ragged_layout, graph_size_key) + self._set_dp_real_num_tokens(self.raw_num_token) self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) if ( @@ -1501,6 +1511,7 @@ def load_batch( buffers.input_embeds[:raw_num_token].copy_(forward_batch.input_embeds) # Padded tokens aren't read, so skip zeroing. Ragged input_ids arrive # from the planner already padded to the tier, invalid slots zeroed. + self._set_dp_real_num_tokens(raw_num_token) if self.enable_two_batch_overlap: self.tbo_plugin.replay_prepare( forward_mode=self.capture_forward_mode, diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 635a15666fba..939cdb8adaac 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/python/sglang/srt/models/deepseek_v4.py @@ -76,6 +76,7 @@ ) from sglang.srt.layers.deep_gemm_wrapper.configurer import DEEPGEMM_SCALE_UE8M0 from sglang.srt.layers.dp_attention import ( + DpPaddingMode, _tbo_event, attn_tp_all_gather, attn_tp_all_reduce, @@ -94,6 +95,7 @@ is_allocation_symmetric, is_dp_attention_enabled, is_dp_gatherv_active, + set_dp_buffer_len_from_batch, ) from sglang.srt.layers.engram import Engram, EngramHasher, EngramLayout from sglang.srt.layers.layer_boundary import get_attn_tp_context @@ -3470,6 +3472,11 @@ def _run_moe_ffn_dp_sync( ) if _do_shared_local and local_hidden_states.shape[0] > 0: _shared_local = self.mlp._forward_shared_experts(local_hidden_states) + # The gather ships the whole padded chunk and the MoE runs on it, but + # graph replay never refreshes pad rows, so they hold stale pool + # memory (bf16-max scale) that overflows the MoE and the combine. + if _is_hip: + _mask_dp_pad_rows(local_hidden_states, forward_batch) # self_attn has already reduced across attention TP, so these hidden # states are replicated and must not be summed by a partial gather. dp_gather_replicate(hidden_states, local_hidden_states, forward_batch) @@ -3809,6 +3816,79 @@ def op_combine_b(self, state): state.hidden_states_mlp_output = hidden +def _mask_dp_pad_rows(local_hidden_states, forward_batch) -> None: + n_real = forward_batch.num_token_non_padded + num_rows = local_hidden_states.shape[0] + if n_real is None or num_rows == 0: + return + pad = torch.arange(num_rows, device=local_hidden_states.device) >= n_real + local_hidden_states.masked_fill_( + pad.view(-1, *([1] * (local_hidden_states.dim() - 1))), 0 + ) + + +def _late_layer_dp_counts( + forward_batch: ForwardBatch, local_late_rows: int +) -> Optional[List[int]]: + # Bounded replay shrinks an extending rank's late layers to its tail, but the + # TP-MoE gather / combine there spans every DP rank, so each rank needs all + # ranks' late-layer row counts. Eager forwards only: every rank takes this + # branch together, and graph replays never reach the late-layer switch. + early = forward_batch.global_num_tokens_padded_cpu + if early is None: + early = forward_batch.global_num_tokens_cpu + if early is None or torch.cuda.is_current_stream_capturing(): + return None + device = forward_batch.global_num_tokens_gpu.device + local = torch.tensor([local_late_rows], dtype=torch.int64, device=device) + counts = torch.empty(len(early), dtype=torch.int64, device=device) + get_parallel().tp_group.all_gather_into_tensor(counts, local) + counts = counts.tolist() + return None if counts == list(early) else counts + + +def _enter_late_layer_dp(forward_batch: ForwardBatch, counts: List[int]) -> tuple: + fb = forward_batch + saved = ( + fb.global_num_tokens_cpu, + fb.global_num_tokens_gpu, + fb.global_num_tokens_padded_cpu, + fb.global_num_tokens_for_logprob_cpu, + fb.global_dp_buffer_len, + fb.dp_padding_mode, + fb.dp_local_start_pos, + fb.dp_local_num_tokens, + ) + fb.global_num_tokens_cpu = counts + fb.global_num_tokens_gpu = torch.tensor( + counts, dtype=torch.int64, device=fb.global_num_tokens_gpu.device + ) + fb.global_num_tokens_padded_cpu = counts + # The MoE gatherv reads its sizes from the logprob counts when present. + fb.global_num_tokens_for_logprob_cpu = counts + fb.global_dp_buffer_len = sum(counts) + # Tail rows are real rows only, so the late layers gather variable lengths. + fb.dp_padding_mode = DpPaddingMode.SUM_LEN + fb.dp_local_start_pos = fb.dp_local_num_tokens = None + set_dp_buffer_len_from_batch(fb) + return saved + + +def _exit_late_layer_dp(forward_batch: ForwardBatch, saved: tuple) -> None: + fb = forward_batch + ( + fb.global_num_tokens_cpu, + fb.global_num_tokens_gpu, + fb.global_num_tokens_padded_cpu, + fb.global_num_tokens_for_logprob_cpu, + fb.global_dp_buffer_len, + fb.dp_padding_mode, + fb.dp_local_start_pos, + fb.dp_local_num_tokens, + ) = saved + set_dp_buffer_len_from_batch(fb) + + def _scatter_tail_rows( tail: LateLayerTail, rows: torch.Tensor, num_tokens: int ) -> torch.Tensor: @@ -4059,6 +4139,22 @@ def _forward_layers_hc_pre_from_prev( self._check_late_layer_tail_readers(forward_batch) attn_backend = get_attn_backend() tail = attn_backend.tail_forward_metadata.late_layer_tail + late_dp_counts = None + saved_dp = None + if ( + _is_hip + and self.late_layer_start is not None + and get_parallel().attn_dp_size > 1 + and get_moe_a2a_backend().is_none() + ): + late_dp_counts = _late_layer_dp_counts( + forward_batch, + ( + tail.token_indices.shape[0] + if tail is not None + else hidden_states.shape[0] + ), + ) saved_full = None # A pending post never meets a residual reader: the HIP boundary's defer_post # gate excludes Engram/DSpark-capture layers and the model end. @@ -4077,6 +4173,11 @@ def _forward_layers_hc_pre_from_prev( positions = tail.positions if hash_ids is not None: hash_ids = tail.rows(hash_ids) + # late_dp_counts is only not None on hip + if late_dp_counts is not None and i == self.late_layer_start: + saved_dp = _enter_late_layer_dp(forward_batch, late_dp_counts) + # Gathered over the full extend; no late layer is hash-routed. + input_ids_global = None engram = self.layers[i].engram if engram is not None: before_engram = state.residual @@ -4133,6 +4234,9 @@ def _forward_layers_hc_pre_from_prev( mega_mhc_seam_open=i + 1 != self.late_layer_start, ) state = state.materialized(self.layers[self.end_layer - 1].hc_cfg) + # saved_dp is only not None on hip + if saved_dp is not None: + _exit_late_layer_dp(forward_batch, saved_dp) if saved_full is not None: attn_backend.exit_late_layer_tail(saved_full, forward_batch) return state.residual, state.pre, tail diff --git a/test/registered/unit/dllm/test_gemma4_cuda_graph.py b/test/registered/unit/dllm/test_gemma4_cuda_graph.py index 41506a6a00c1..f42282311fa7 100644 --- a/test/registered/unit/dllm/test_gemma4_cuda_graph.py +++ b/test/registered/unit/dllm/test_gemma4_cuda_graph.py @@ -165,6 +165,7 @@ def test_preplanned_replay_refreshes_self_conditioning(self): input_ids=torch.zeros(16, dtype=torch.long), positions=torch.zeros(16, dtype=torch.long), input_embeds=torch.full((16, 7), -1.0), + num_token_non_padded=None, ) batch = SimpleNamespace( needs_forward_metadata_init=lambda: False, diff --git a/test/registered/unit/layer_boundary/test_moe_output_reduction.py b/test/registered/unit/layer_boundary/test_moe_output_reduction.py index 53048f05115b..b515105b537c 100644 --- a/test/registered/unit/layer_boundary/test_moe_output_reduction.py +++ b/test/registered/unit/layer_boundary/test_moe_output_reduction.py @@ -207,6 +207,7 @@ def scatter(output, value, batch): ) ), _SHARED_EXPERT_LOCAL=False, + _is_hip=False, nullcontext=contextlib.nullcontext, get_global_dp_buffer=lambda g: torch.empty(4, 3), get_local_dp_buffer=lambda g: torch.empty(2, 3),