[AMD][V4.1][*/N] Fix DP attention + TP MoE decode NaN and support bounded replay

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:

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)

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()
     )
 
 

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,

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

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,

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),