[Feat] [AMD] DeepSeek-v4.1: Support prefill context parallel

python/sglang/srt/arg_groups/deepseek_v4_hook.py

@@ -267,15 +267,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.
             (
                 "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),
         )
@@ -284,6 +282,11 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None:
                 raise ValueError(
                     f"DeepSeekV4 prefill CP on HIP does not support {feature} yet."
                 )
+        if is_dsv41 and not envs.SGLANG_CP_MIN_TOKENS_PER_RANK.is_set():
+            # TODO(At1a8): temporary work-around for the gfx950 aiter MXFP8 GEMM NaN
+            # (its small-M group32 tile returns NaN for wqkv_a: N=1792, K=5120, M=5..8),
+            # it can be safely removed later, once the aiter kernel is fixed
+            envs.SGLANG_CP_MIN_TOKENS_PER_RANK.set(16)
 
     attn_dp_size = context_parallel_attn_dp_size(cfg, "DeepSeek-V4 context parallelism")
     declare_resolution(

python/sglang/srt/environ.py

@@ -834,6 +834,9 @@ class Envs:
     )
     SGLANG_ENABLE_COLOCATED_BATCH_GEN = EnvBool(False)
     SGLANG_SHARED_EXPERT_TP1 = EnvBool(False)
+    # Interleave prefill CP runs a batch unsharded when its smallest shard would hold
+    # fewer rows than this.
+    SGLANG_CP_MIN_TOKENS_PER_RANK = EnvInt(0)
     # Replicate the input embedding across TP ranks instead of sharding it
     # along the vocab dimension (saves an all-reduce/all-gather in the embed
     # lookup at the cost of replicated embedding weights). Drives both the

python/sglang/srt/layers/attention/deepseek_v4_backend_hip_radix.py

@@ -66,6 +66,7 @@
     DeepseekV4AttnBackend,
     LateLayerTail,
     _as_int_list,
+    _first_rows,
     _tail_rows,
     late_layer_num_tokens,
     late_layer_whole_extend,
@@ -102,7 +103,15 @@
     resolve_hip_flashmla_backend,
 )
 from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
-from sglang.srt.layers.cp.utils import is_cp_active
+from sglang.srt.layers.cp.interleave import (
+    InterleaveContextParallelMetadata,
+    interleave_rows_per_request,
+)
+from sglang.srt.layers.cp.utils import cp_materialize_global_token_order, is_cp_active
+from sglang.srt.layers.dp_attention import (
+    get_local_dp_buffer_len,
+    set_local_dp_buffer_len,
+)
 from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool
 from sglang.srt.mem_cache.dsv41_request_window import window_layout
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -888,9 +897,10 @@ class DeepseekV4HipRadixBackend(
     forward = DeepseekV4AttnBackend.forward
 
     # the ratio-1/2 compressor and indexer orchestration is the CUDA backend's, taken
-    # unbound; HIP differs only in the paged top-k (_low_ratio_index_topk below)
+    # unbound; HIP differs only in the paged top-k (_low_ratio_index_topk below) and
+    # in the prefill-CP sources (_forward_low_ratio_sources_cp below)
     forward_low_ratio_sources = DeepseekV4AttnBackend.forward_low_ratio_sources
-    _forward_low_ratio_sources_cp = DeepseekV4AttnBackend._forward_low_ratio_sources_cp
+    _late_layer_tail_cp_layout = DeepseekV4AttnBackend._late_layer_tail_cp_layout
     write_tail_source_kv = DeepseekV4AttnBackend.write_tail_source_kv
     # No two-level candidates on HIP: the ratio-1/2 sources stay one eager break.
     low_ratio_prefill_graph = False
@@ -1180,13 +1190,23 @@ def init_forward_metadata_prefill(
         # token map skip an implicit D2H.
         exact_num_tokens: bool = True,
         forward_batch: Optional[ForwardBatch] = None,
+        # the late-layer tail passes its own layout; the full extend reads the batch's
+        cp_metadata: Optional[InterleaveContextParallelMetadata] = None,
     ) -> DSV4Metadata:
         padded_num_tokens = out_cache_loc.shape[0]
         cp_active = forward_batch is not None and is_cp_active(forward_batch)
         if cp_active:
-            cp_metadata = forward_batch.attn_cp_metadata
+            if cp_metadata is None:
+                cp_metadata = forward_batch.attn_cp_metadata
             assert cp_metadata is not None
             padded_num_tokens = sum(cp_metadata.per_rank_actual_token)
+            if (
+                swa_replay_start is not None
+                and swa_replay_start.shape[0] < padded_num_tokens
+            ):
+                swa_replay_start = F.pad(
+                    swa_replay_start, (0, padded_num_tokens - swa_replay_start.shape[0])
+                )
 
         from sglang.kernels.ops.attention.dsv4.attn_metadata_kernels import (
             ExpandPrefillCausally,
@@ -1238,7 +1258,9 @@ def init_forward_metadata_prefill(
             exact_num_tokens=exact_num_tokens or host_proves_exact_num_tokens,
         )
         if cp_active:
-            core_attn_metadata.apply_cp_reindex(num_tokens=num_tokens)
+            core_attn_metadata.apply_cp_reindex(
+                num_tokens=num_tokens, local_index=cp_metadata.local_index
+            )
             if need_compress:
                 core_attn_metadata.init_flashmla_related()
         if attach_decode_streams:
@@ -2265,6 +2287,13 @@ def _build_late_layer_tail_metadata(
         )
         tail_lens = torch.tensor(tail_lens_cpu, dtype=torch.int32, device=device)
         num_tokens = sum(tail_lens_cpu)
+        # prefill CP: every rank runs the tail rows it already holds, padded to the
+        # largest share; the tail's compact metadata is sliced by its own local_index
+        cp_tail = (
+            self._late_layer_tail_cp_layout(forward_batch, token_indices, tail_lens)
+            if is_cp_active(forward_batch)
+            else None
+        )
         metadata = self.init_forward_metadata_prefill(
             max_seq_len=int(seq_lens_cpu.max().item()),
             req_pool_indices=forward_batch.req_pool_indices,
@@ -2276,6 +2305,8 @@ def _build_late_layer_tail_metadata(
             extend_seq_lens_cpu=tail_lens_cpu,
             extend_start_loc=torch.cumsum(tail_lens, dim=0) - tail_lens,
             swa_replay_start=swa_replay_start,
+            forward_batch=forward_batch if cp_tail is not None else None,
+            cp_metadata=cp_tail["cp_metadata"] if cp_tail is not None else None,
         )
         request_layout = metadata.core_attn_metadata.request_window_layout
         swa_out_cache_loc = (
@@ -2297,14 +2328,30 @@ def _build_late_layer_tail_metadata(
             contiguous_start=contiguous_start,
         )
         metadata.low_ratio_pos_i64 = positions.to(torch.int64)
-        metadata.late_layer_tail = LateLayerTail(
-            token_indices=token_indices,
-            positions=positions,
-            extend_seq_lens=tail_lens,
-            extend_seq_lens_cpu=tail_lens_cpu,
-            swa_out_cache_loc=swa_out_cache_loc,
-            contiguous_start=contiguous_start,
-        )
+        if cp_tail is None:
+            metadata.late_layer_tail = LateLayerTail(
+                token_indices=token_indices,
+                positions=positions,
+                extend_seq_lens=tail_lens,
+                extend_seq_lens_cpu=tail_lens_cpu,
+                swa_out_cache_loc=swa_out_cache_loc,
+                contiguous_start=contiguous_start,
+            )
+        else:
+            # rows index this rank's shard; the SWA store and the compressor read
+            # the gathered tail, so they keep the global rows
+            metadata.late_layer_tail = LateLayerTail(
+                token_indices=cp_tail["local_token_indices"],
+                positions=cp_tail["local_positions"],
+                extend_seq_lens=tail_lens,
+                extend_seq_lens_cpu=tail_lens_cpu,
+                swa_out_cache_loc=swa_out_cache_loc,
+                pad_rows=cp_tail["pad_rows"],
+                cp_metadata=cp_tail["cp_metadata"],
+                local_lens_cpu=cp_tail["local_lens_cpu"],
+                req_global=metadata.low_ratio_req_indices,
+                pos_global=metadata.low_ratio_pos_i64,
+            )
         self._refresh_fp4_prefill_workspace(forward_batch, metadata)
         return metadata
 
@@ -2317,15 +2364,30 @@ def enter_late_layer_tail(self, forward_batch: ForwardBatch) -> tuple:
         )
         tail_metadata = self.tail_forward_metadata
         assert tail_metadata is not None, "no tail metadata for this forward"
-        saved = (self.forward_metadata, self.candidate_masks)
+        saved = (
+            self.forward_metadata,
+            self.candidate_masks,
+            forward_batch.attn_cp_metadata,
+            get_local_dp_buffer_len(),
+        )
+        tail = tail_metadata.late_layer_tail
         self.forward_metadata = tail_metadata
         window = self.token_to_kv_pool.request_window
         if window is not None:
             window.activate(tail_metadata.core_attn_metadata.request_window_layout)
+        if tail.cp_metadata is not None:
+            forward_batch.attn_cp_metadata = tail.cp_metadata
+            set_local_dp_buffer_len(sum(tail.cp_metadata.per_rank_actual_token))
         return saved
 
     def exit_late_layer_tail(self, saved: tuple, forward_batch: ForwardBatch) -> None:
-        self.forward_metadata, self.candidate_masks = saved
+        (
+            self.forward_metadata,
+            self.candidate_masks,
+            forward_batch.attn_cp_metadata,
+            local_dp_buffer_len,
+        ) = saved
+        set_local_dp_buffer_len(local_dp_buffer_len)
         self.tail_source_kv_layer_id = None
         window = self.token_to_kv_pool.request_window
         if window is not None:
@@ -3496,18 +3558,67 @@ def _low_ratio_index_topk(
                 self, layer, x, q_lora, pos, forward_batch=forward_batch
             )
         elif forward_batch.forward_mode.is_extend():
-            # HIP rejects V4.1 prefill CP, the only caller passing rows_per_request
-            assert rows_per_request is None, "V4.1 prefill CP is not supported on HIP"
             assert (
                 forward_batch.seq_lens_cpu is not None
                 and forward_batch.extend_seq_lens_cpu is not None
             ), "the HIP low-ratio prefill indexer needs the batch's CPU lengths"
-            low_ratio_index_topk_hip_extend(self, layer, x, q_lora, pos, forward_batch)
+            low_ratio_index_topk_hip_extend(
+                self,
+                layer,
+                x,
+                q_lora,
+                pos,
+                forward_batch,
+                rows_per_request=rows_per_request,
+            )
         else:
             raise NotImplementedError(
                 f"low-ratio indexer for {forward_batch.forward_mode} on HIP"
             )
 
+    def _forward_low_ratio_sources_cp(
+        self, *, layer, x, q_lora, positions, forward_batch, run_compressor, run_indexer
+    ) -> None:
+        """Every rank writes the whole prompt's compressed state and scores only its own
+        rows, which are request-major like the extend's (interleave keeps the order)."""
+        cp_meta = forward_batch.attn_cp_metadata
+        total = int(cp_meta.total_seq_lens)
+        tail = self.forward_metadata.late_layer_tail
+        if tail is not None:
+            q_lens_cpu = tail.local_lens_cpu
+            req_global, pos_global = tail.req_global, tail.pos_global
+        else:
+            q_lens_cpu = interleave_rows_per_request(
+                _as_int_list(forward_batch.extend_seq_lens_cpu),
+                get_parallel().attn_cp_rank,
+                get_parallel().attn_cp_size,
+            )
+            req_global = token_req_indices(forward_batch, num_tokens=total)
+            pos_global = forward_batch.positions[:total].to(torch.int64)
+        num_local = sum(q_lens_cpu)
+        if run_compressor and layer.compressor is not None:
+            x_global = cp_materialize_global_token_order(
+                x.contiguous(), forward_batch, torch.cuda.current_stream()
+            )[:total]
+            if tail is None:
+                # x_global is the extend itself, so the fused prefill store of the
+                # non-CP path applies with the batch's own extend offsets
+                self._low_ratio_compress(
+                    layer, x_global, req_global, pos_global, forward_batch
+                )
+            else:
+                self._low_ratio_compress_torch(layer, x_global, req_global, pos_global)
+        if run_indexer and layer.indexer is not None:
+            self._low_ratio_index_topk(
+                layer,
+                x[:num_local],
+                _first_rows(q_lora, num_local),
+                None,
+                positions[:num_local].to(torch.int64),
+                forward_batch,
+                rows_per_request=q_lens_cpu,
+            )
+
     def expand_extend_with_same_length(
         self,
         bs: int,

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

@@ -495,10 +495,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
@@ -511,8 +513,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/interleave.py

@@ -36,6 +36,7 @@
 from sglang.srt.distributed.device_communicators.pynccl_allocator import (
     use_symmetric_memory,
 )
+from sglang.srt.environ import envs
 from sglang.srt.layers.cp.base import (
     BaseContextParallelMetadata,
     ContextParallelStrategy,
@@ -172,6 +173,9 @@ 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)
+        # tokens are dealt by global index, so the smallest shard holds seq_len // cp_size
+        if seq_len // cp_size < envs.SGLANG_CP_MIN_TOKENS_PER_RANK.get():
+            return False
         return seq_len > 0 and seq_len >= cp_size and cp_size > 1
 
     def build_metadata(

python/sglang/srt/models/deepseek_v4.py

@@ -4655,7 +4655,10 @@ def _forward_impl(
             )
         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(
@@ -4664,6 +4667,13 @@ def _forward_impl(
             pre_hc_head = _scatter_tail_rows(
                 tail=tail, rows=pre_hc_head, num_tokens=num_tokens
             )
+        if hip_cp_tail and capture_dspark:
+            aux = _scatter_tail_rows(
+                tail=tail,
+                rows=dspark_aux_hidden_states.finalize(),
+                num_tokens=input_ids.shape[0],
+            )
+            return (hidden_states, pre_hc_head), aux
 
         if capture_dspark:
             return (hidden_states, pre_hc_head), dspark_aux_hidden_states.finalize()