[AMD][DSV4.1][CP] Support interleave prefill CP with breakable CUDA graphs on ROCm

python/sglang/srt/arg_groups/deepseek_v4_hook.py

@@ -138,15 +138,13 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None:
                 f"backend for both phases, got prefill={prefill_backend!r}, "
                 f"decode={decode_backend!r}."
             )
+        is_dsv41 = model_config_of(server_args).hf_config.model_type == "deepseek_v41"
         unsupported = (
             ("multiple nodes", cfg.nnodes > 1),
-            (
-                "DeepSeek-V4.1",
-                model_config_of(server_args).hf_config.model_type == "deepseek_v41",
-            ),
+            # V4.1 carries its own bounded-replay tail layout; V4 has none.
             (
                 "--enable-decoder-swa-bounded-replay",
-                cfg.enable_decoder_swa_bounded_replay,
+                cfg.enable_decoder_swa_bounded_replay and not is_dsv41,
             ),
             ("--enable-two-batch-overlap", cfg.enable_two_batch_overlap),
         )

python/sglang/srt/layers/attention/dsv4/low_ratio_backend_hip.py

@@ -345,10 +345,12 @@ def low_ratio_index_topk_hip_extend(
     q_lora,
     pos,
     forward_batch: ForwardBatch,
+    rows_per_request: Optional[List[int]] = None,
 ) -> None:
     """Ragged prefill: the FlyDSL prefill kernel scores every token's visible compressed positions,
     candidate masks are published or applied per request, one paged top-k selects every row.
-    Identity requests are written without scores; an all-identity batch skips the kernel."""
+    Identity requests are written without scores; an all-identity batch skips the kernel.
+    rows_per_request (prefill CP): this rank's rows of each request, request-major."""
     pool = backend.token_to_kv_pool
     metadata = backend.forward_metadata
     core = metadata.core_metadata
@@ -361,8 +363,10 @@ def low_ratio_index_topk_hip_extend(
         raw_indices.fill_(-1)
 
     seq_lens_cpu = _as_int_list(forward_batch.seq_lens_cpu)
-    # bounded replay counts cover only each request's tail
-    if metadata.late_layer_tail is not None:
+    if rows_per_request is not None:
+        extend_lens_cpu = list(rows_per_request)
+    elif metadata.late_layer_tail is not None:
+        # bounded replay counts cover only each request's tail
         extend_lens_cpu = metadata.late_layer_tail.extend_seq_lens_cpu
     else:
         extend_lens_cpu = _as_int_list(forward_batch.extend_seq_lens_cpu)

python/sglang/srt/layers/cp/bcg.py

@@ -31,15 +31,23 @@
     CPAttentionBackendKind,
     get_cp_strategy,
 )
-from sglang.srt.layers.cp.interleave import InterleaveCPStrategy
+from sglang.srt.layers.cp.interleave import (
+    _HIP_MIN_TOKENS_PER_RANK,
+    InterleaveCPStrategy,
+)
 from sglang.srt.layers.cp.padding import get_cp_padding_align_size
 from sglang.srt.layers.cp.utils import (
     cp_gather_after_forward,
+    cp_interleave_input_ids,
+    cp_shard_hidden_states,
     cp_split_before_forward,
+    drop_unconsumed_states_before_norm,
     prepare_cp_forward,
 )
 from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
+from sglang.srt.layers.moe.utils import get_moe_a2a_backend
 from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
+from sglang.srt.utils import is_hip
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -54,6 +62,12 @@
 ] = {
     (ContextParallelStrategyKind.ZIGZAG, CPAttentionBackendKind.TRTLLM_MHA),
 }
+if is_hip():
+    # The capture path shards input_ids with the hidden states for DSV4 hash
+    # routing; enabled on ROCm only until the CUDA path is validated as well.
+    PREFILL_CP_BCG_SUPPORTED_COMBINATIONS.add(
+        (ContextParallelStrategyKind.INTERLEAVE, CPAttentionBackendKind.DSV4)
+    )
 
 
 def supports_prefill_cp_bcg(server_args: ServerArgs) -> bool:
@@ -89,8 +103,13 @@ def filter_prefill_cp_bcg_capture_num_tokens(
 ) -> list[int]:
     """Keep only token buckets where the selected CP strategy can run."""
     min_num_tokens = resolved_view(server_args).attn_cp_size
-    if resolving_view(server_args).cp_strategy == "zigzag":
+    cp_strategy = resolving_view(server_args).cp_strategy
+    if cp_strategy == "zigzag":
         min_num_tokens *= 2
+    elif cp_strategy == "interleave" and is_hip():
+        # InterleaveCPStrategy.can_apply leaves smaller HIP batches unsharded,
+        # so a smaller bucket could never be captured or replayed with CP.
+        min_num_tokens *= _HIP_MIN_TOKENS_PER_RANK
     filtered = [size for size in capture_num_tokens if size >= min_num_tokens]
     if not filtered:
         raise ValueError(
@@ -120,9 +139,24 @@ class PrefillCPBCGInput:
     positions: torch.Tensor
     bucket_local_tokens: Dict[int, int] = field(default_factory=dict)
     live_local_tokens: int = 0
+    input_ids: Optional[torch.Tensor] = None
+    # Rank-major token ids for TP MoE, which routes the CP all-gathered rows.
+    global_input_ids: Optional[torch.Tensor] = None
 
     @classmethod
     def create(cls, runner: PrefillCudaGraphRunner) -> PrefillCPBCGInput:
+        strategy = get_cp_strategy()
+        global_id_capacity = 0
+        if (
+            isinstance(strategy, InterleaveCPStrategy)
+            and get_moe_a2a_backend().is_none()
+        ):
+            align = get_cp_padding_align_size()
+            local_capacity = (
+                runner.max_num_tokens + strategy.cp_size - 1
+            ) // strategy.cp_size
+            local_capacity = (local_capacity + align - 1) // align * align
+            global_id_capacity = strategy.cp_size * local_capacity
         with torch.device(runner.device):
             return cls(
                 input_embeds=torch.zeros(
@@ -136,6 +170,12 @@ def create(cls, runner: PrefillCudaGraphRunner) -> PrefillCPBCGInput:
                     (runner.max_num_tokens,),
                     dtype=torch.int64,
                 ),
+                input_ids=torch.zeros((runner.max_num_tokens,), dtype=torch.int64),
+                global_input_ids=(
+                    torch.zeros((global_id_capacity,), dtype=torch.int64)
+                    if global_id_capacity
+                    else None
+                ),
             )
 
     def required_local_tokens(self, extend_seq_lens: Any) -> Optional[int]:
@@ -248,6 +288,7 @@ def prepare(
         raw_tokens = int(forward_batch.extend_num_tokens)
         global_input_ids = forward_batch.input_ids[:raw_tokens]
         global_positions = forward_batch.positions[:raw_tokens]
+        local_input_ids = cp_shard_hidden_states(global_input_ids, forward_batch)
         global_input_embeds = runner.model_runner.model.get_input_embeddings()(
             global_input_ids
         )
@@ -278,12 +319,31 @@ def prepare(
 
         input_embeds = self.input_embeds[:captured_local_tokens]
         positions = self.positions[:captured_local_tokens]
+        assert self.input_ids is not None
+        input_ids = self.input_ids[:captured_local_tokens]
         input_embeds.zero_()
         positions.zero_()
+        input_ids.zero_()
         input_embeds[:live_local_tokens].copy_(local_input_embeds)
         positions[:live_local_tokens].copy_(local_positions)
+        input_ids[:live_local_tokens].copy_(local_input_ids)
         forward_batch.input_embeds = input_embeds
-        forward_batch.positions = positions
+        # The captured body takes CP-local positions and ids, while attention
+        # backends keep reading forward_batch.positions in global token order
+        # (e.g. the SWA ring store after the CP KV all-gather).
+        forward_batch._cp_positions = positions
+        forward_batch._cp_input_ids = input_ids
+        if self.global_input_ids is not None:
+            global_ids = cp_interleave_input_ids(global_input_ids, forward_batch)
+            if global_ids.numel() > self.global_input_ids.numel():
+                raise RuntimeError(
+                    f"CP global token ids need {global_ids.numel()} rows, but the "
+                    f"fixed buffer has capacity {self.global_input_ids.numel()}"
+                )
+            forward_batch.input_ids_global = self.global_input_ids[: global_ids.numel()]
+            forward_batch.input_ids_global.copy_(global_ids)
+        else:
+            forward_batch.input_ids_global = input_ids
         self.live_local_tokens = live_local_tokens
 
         if isinstance(get_cp_strategy(), InterleaveCPStrategy):
@@ -341,6 +401,11 @@ def execute_prefill_cp_bcg(
                 else hidden_states
             )
 
+        hidden_states = drop_unconsumed_states_before_norm(
+            hidden_states,
+            forward_batch=forward_batch,
+            has_aux_hidden_states=aux_hidden_states is not None,
+        )
         hidden_states = cp_gather_after_forward(
             hidden_states,
             static_forward_batch,

python/sglang/srt/layers/cp/cp_decode_attn_tp.py

@@ -157,7 +157,15 @@ def _get_linear_attrs(self, linear_instance) -> List[Tuple]:
             return []
 
         attrs = [(linear_instance, "weight", dim)]
-        for scale_name in ("weight_scale_inv", "weight_scale"):
+        # The gfx950 MXFP8 routes (fp8_hip) read derived copies, not the block scales;
+        # all keep (N, K) order with whole tiles, so a chunk is the rank's slice.
+        for scale_name in (
+            "weight_scale_inv",
+            "weight_scale",
+            "weight_scale_mx_e8m0",
+            "weight_scale_inv_mx",
+            "weight_bf16",
+        ):
             if getattr(linear_instance, scale_name, None) is not None:
                 attrs.append((linear_instance, scale_name, dim))
         return attrs

python/sglang/srt/layers/cp/interleave.py

@@ -48,6 +48,12 @@
     is_allocation_symmetric,
 )
 from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils import is_hip
+
+_is_hip = is_hip()
+# gfx950 aiter's small-M group32 MXFP8 GEMM tile returns NaN for DeepSeek-V4.1's
+# wqkv_a (N=1792, K=5120, M=5..8); shards that small run unsharded instead.
+_HIP_MIN_TOKENS_PER_RANK = 16
 
 
 def attn_cp_interleave_gather(hidden_states: torch.Tensor):
@@ -158,6 +164,8 @@ def can_apply(self, num_tokens: int, forward_batch) -> bool:
             return False
         cp_size = self.cp_size
         seq_len = sum(forward_batch.extend_seq_lens_cpu)
+        if _is_hip and seq_len < cp_size * _HIP_MIN_TOKENS_PER_RANK:
+            return False
         return seq_len > 0 and seq_len >= cp_size and cp_size > 1
 
     def build_metadata(

python/sglang/srt/layers/cp/utils.py

@@ -259,6 +259,19 @@ def cp_gather_after_forward(x: Any, forward_batch, stream: Optional[Any] = None)
     return strategy.gather_hidden_states(x, forward_batch, stream)
 
 
+def drop_unconsumed_states_before_norm(
+    hidden_states: Any, *, forward_batch, has_aux_hidden_states: bool
+) -> Any:
+    """Drop DSV4's pre-mHC-head states before the CP gather when logits won't read them."""
+    # The logits processor reads them only to store hidden states for MTP/EAGLE;
+    # gathering them otherwise moves ~4x the hidden-state bytes for nothing.
+    if not isinstance(hidden_states, tuple) or (
+        forward_batch.capture_hidden_mode.need_capture() and not has_aux_hidden_states
+    ):
+        return hidden_states
+    return hidden_states[0]
+
+
 def cp_materialize_global_token_order(
     x: Any, forward_batch, stream: Optional[Any] = None
 ):

python/sglang/srt/model_executor/runner/eager_runner.py

@@ -27,6 +27,7 @@
 from sglang.srt.layers.cp.utils import (
     cp_gather_after_forward,
     cp_shard_model_inputs,
+    drop_unconsumed_states_before_norm,
     get_cp_strategy,
     is_cp_active,
     prepare_cp_forward,
@@ -402,6 +403,11 @@ def _execute_extend_cp(
             )
 
         stream = torch.cuda.current_stream()
+        hidden_states = drop_unconsumed_states_before_norm(
+            hidden_states,
+            forward_batch=forward_batch,
+            has_aux_hidden_states=aux_hidden_states is not None,
+        )
         hidden_states = cp_gather_after_forward(hidden_states, forward_batch, stream)
         # DSpark aux tensors ride the same CP token split; gather them the same way.
         if aux_hidden_states is not None:

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -755,6 +755,9 @@ def _capture_num_token_non_padded(self, num_tokens: int) -> Optional[torch.Tenso
 
     def _get_layer_model_positions(self, forward_batch: ForwardBatch) -> torch.Tensor:
         """Mirror outer multimodal wrappers when BCG captures layer_model directly."""
+        cp_positions = getattr(forward_batch, "_cp_positions", None)
+        if cp_positions is not None:
+            return cp_positions
         if forward_batch.mrope_positions is None:
             return forward_batch.positions
 
@@ -837,7 +840,9 @@ def _run_forward(self, forward_batch: ForwardBatch, num_tokens: int):
             if self._uses_eager_prefill_tail():
                 # BCG / Full: capture the transformer body only.
                 positions = self._get_layer_model_positions(forward_batch)
-                input_ids = forward_batch.input_ids
+                input_ids = getattr(
+                    forward_batch, "_cp_input_ids", forward_batch.input_ids
+                )
                 kwargs = _build_layer_model_forward_kwargs(
                     self.layer_model, forward_batch, pp_proxy_tensors
                 )

python/sglang/srt/models/deepseek_v4.py

@@ -4127,9 +4127,15 @@ def _forward_layers_hc_pre_from_prev(
             if cp_extend:
                 # n-gram hashing needs each token's predecessors: hash the whole prompt
                 total = int(forward_batch.attn_cp_metadata.total_seq_lens)
-                hash_ids = self.engram_hasher(
-                    forward_batch.input_ids[:total], forward_batch
-                )
+                global_input_ids = forward_batch.input_ids[:total]
+                if is_in_breakable_cuda_graph():
+                    # The hasher reads per-request rows and commits the n-gram
+                    # history, so CP breakable graphs run it at an eager break.
+                    hash_ids = bcg_deepseek_v4_engram_hash_ids(
+                        self.engram_hasher, global_input_ids
+                    )
+                else:
+                    hash_ids = self.engram_hasher(global_input_ids, forward_batch)
                 parallel = get_parallel()
                 hash_ids = hash_ids[parallel.attn_cp_rank :: parallel.attn_cp_size]
                 pad_rows = hidden_states.shape[0] - hash_ids.shape[0]
@@ -4540,7 +4546,10 @@ def forward(
             )
         hidden_states = self.norm(hidden_states)
 
-        if tail is not None and not capture_dspark:
+        # HIP prefill CP: the runner gathers every output, DSpark captures included,
+        # by this rank's full shard layout, so the tail rows go back into it.
+        hip_cp_tail = _is_hip and tail is not None and is_cp_active(forward_batch)
+        if tail is not None and (not capture_dspark or hip_cp_tail):
             # The logits processor indexes rows by the full extend layout.
             num_tokens = input_ids.shape[0]
             hidden_states = _scatter_tail_rows(
@@ -4549,6 +4558,11 @@ def forward(
             pre_hc_head = _scatter_tail_rows(
                 tail=tail, rows=pre_hc_head, num_tokens=num_tokens
             )
+        if hip_cp_tail and capture_dspark:
+            dspark_aux_hidden_states = [
+                _scatter_tail_rows(tail=tail, rows=aux, num_tokens=input_ids.shape[0])
+                for aux in dspark_aux_hidden_states
+            ]
 
         if capture_dspark:
             return (hidden_states, pre_hc_head), dspark_aux_hidden_states

test/registered/amd/test_dsv4_hip_bcg_metadata.py

@@ -251,10 +251,14 @@ def test_eager_prefill_marks_host_proven_token_count_exact(self):
 
     def test_cp_prefill_attaches_before_reindex(self):
         for need_compress in (False, True):
-            with self.subTest(need_compress=need_compress):
-                self._check_cp_prefill_order(need_compress)
+            for use_prefill_cuda_graph in (False, True):
+                with self.subTest(
+                    need_compress=need_compress,
+                    use_prefill_cuda_graph=use_prefill_cuda_graph,
+                ):
+                    self._check_cp_prefill_order(need_compress, use_prefill_cuda_graph)
 
-    def _check_cp_prefill_order(self, need_compress):
+    def _check_cp_prefill_order(self, need_compress, use_prefill_cuda_graph):
         backend = object.__new__(DeepseekV4HipRadixBackend)
         backend.req_to_token = torch.zeros((2, 8), dtype=torch.int32)
         backend.has_c4 = False
@@ -273,8 +277,13 @@ def _check_cp_prefill_order(self, need_compress):
             side_effect=lambda *_args, **_kwargs: events.append("attach")
         )
         forward_batch = SimpleNamespace(
-            attn_cp_metadata=SimpleNamespace(per_rank_actual_token=[2, 2])
+            attn_cp_metadata=SimpleNamespace(
+                per_rank_actual_token=[2, 2], local_index=None
+            )
         )
+        # CP breakable prefill graphs pass the bucket-sized out_cache_loc while
+        # the live batch has 3 tokens; the global write fields follow the bucket.
+        num_write_tokens = 4 if use_prefill_cuda_graph else 3
 
         with (
             mock.patch(
@@ -294,19 +303,24 @@ def _check_cp_prefill_order(self, need_compress):
                 req_pool_indices=torch.tensor([7, 9], dtype=torch.int32),
                 seq_lens=torch.tensor([1, 3], dtype=torch.int32),
                 seq_lens_cpu=[1, 3],
-                out_cache_loc=torch.zeros(3, dtype=torch.int64),
+                out_cache_loc=torch.zeros(num_write_tokens, dtype=torch.int64),
                 num_tokens=3,
                 extend_seq_lens=torch.tensor([1, 2], dtype=torch.int32),
                 extend_seq_lens_cpu=[1, 2],
                 need_compress=need_compress,
+                use_prefill_cuda_graph=use_prefill_cuda_graph,
                 forward_batch=forward_batch,
             )
 
         self.assertEqual(expand_prefill.call_args.kwargs["padded_num_tokens"], 4)
+        self.assertEqual(expand_prefill.call_args.kwargs["num_tokens"], 3)
         self.assertEqual(
-            backend.make_core_attn_metadata.call_args.kwargs["num_tokens"], 3
+            backend.make_core_attn_metadata.call_args.kwargs["num_tokens"],
+            num_write_tokens,
+        )
+        core.apply_cp_reindex.assert_called_once_with(
+            num_tokens=num_write_tokens, local_index=None
         )
-        core.apply_cp_reindex.assert_called_once_with(num_tokens=3)
         # Target prefill rebuilds FlashMLA on the local rows; the draft skips it.
         self.assertEqual(
             events, ["attach", "reindex"] + (["flashmla"] if need_compress else [])