[CP V1 Deprecation 3/5] Remove generic prefill CP v1 runtime (#36228)

python/sglang/srt/arg_groups/deepseek_v4_hook.py

@@ -168,7 +168,7 @@ def validate_deepseek_v4_cp(server_args: ServerArgs) -> None:
             f"DeepSeekV4 only supports interleave CP strategy, got {cfg.cp_strategy}"
         )
 
-    if get_platform().is_hip or get_platform().is_npu:
+    if get_platform().is_hip or get_platform().is_npu or get_platform().is_musa:
         # Protected platform implementations still consume the legacy runtime
         # fields. Generic backends use enable_prefill_cp/cp_strategy directly.
         declare_resolution(

python/sglang/srt/arg_groups/kv_cache_hook.py

@@ -487,9 +487,9 @@ def validate_prefill_only_disable_kv_cache_args(server_args: Any):
             "radix cache indexes KV pool slots that no longer hold real data."
         )
 
-    # Context-parallel prefill stages K/V through cp_allgather_and_save_kv_cache,
-    # which writes to the pool via set_kv_buffer. NoOpMHATokenToKVPool intentionally
-    # raises on writes, so the engine would boot fine but fail on the first request.
+    # Context-parallel prefill writes K/V to the pool via set_kv_buffer.
+    # NoOpMHATokenToKVPool intentionally raises on writes, so the engine would
+    # boot fine but fail on the first request.
     if resolved_view(server_args).attn_cp_size > 1:
         raise ValueError(
             "--prefill-only-disable-kv-cache is incompatible with --attn-cp-size > 1: "

python/sglang/srt/arg_groups/parallel_hook.py

@@ -39,7 +39,7 @@ def handle_context_parallelism(server_args: Any):
             cfg.enable_prefill_cp
             and model_arch == "DeepseekV32ForCausalLM"
             and cfg.cp_strategy == "zigzag"
-            and not (platform.is_hip or platform.is_npu)
+            and not (platform.is_hip or platform.is_npu or platform.is_musa)
         ):
             raise ValueError(
                 "DeepSeek V3.2 prefill CP does not support --cp-strategy "
@@ -563,7 +563,7 @@ def handle_eplb_and_dispatch(server_args: Any):
 def handle_platform_cp_compatibility(server_args: Any):
     cfg = resolving_view(server_args)
     platform = get_platform()
-    is_protected_platform = platform.is_hip or platform.is_npu
+    is_protected_platform = platform.is_hip or platform.is_npu or platform.is_musa
     if not is_protected_platform:
         if (
             cfg.enable_prefill_context_parallel
@@ -571,7 +571,7 @@ def handle_platform_cp_compatibility(server_args: Any):
         ):
             raise ValueError(
                 "Legacy prefill context-parallel options are supported only "
-                "by protected HIP or Ascend NPU paths. Use "
+                "by protected HIP, Ascend NPU, or MUSA paths. Use "
                 "--enable-prefill-cp with --cp-strategy."
             )
         return
@@ -603,7 +603,10 @@ def handle_platform_cp_compatibility(server_args: Any):
 
 
 def handle_legacy_cp_runtime_compatibility(server_args: Any):
-    """Project canonical CP settings for runtime consumers removed by PR3."""
+    """Project canonical CP settings only for protected platform runtimes."""
+    platform = get_platform()
+    if not (platform.is_hip or platform.is_npu or platform.is_musa):
+        return
     cfg = resolving_view(server_args)
 
     if cfg.enable_prefill_context_parallel and cfg.enable_dsa_prefill_context_parallel:

python/sglang/srt/arg_groups/pipeline.py

@@ -287,8 +287,7 @@ def run_resolution_pipeline(server_args: Any) -> None:
     # Normalize load balancing defaults.
     handle_load_balance_method(server_args)
 
-    # The old runtime distinguishes DSA from other CP paths through legacy
-    # fields, so project only after attention_backend has been resolved.
+    # Protected runtimes still consume platform CP fields after backend selection.
     handle_legacy_cp_runtime_compatibility(server_args)
 
     # Handle context parallelism.

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import contextlib
 import logging
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, Union
 
 import torch
 from einops import rearrange
@@ -29,7 +29,6 @@ from sglang.srt.layers.attention.dsa.paged_mqa_logits_backend import (
 from sglang.srt.layers.attention.dsa.utils import (
     aiter_can_use_preshuffle_paged_mqa,
     is_dsa_enable_prefill_cp,
-    is_dsa_prefill_cp_in_seq_split,
     is_graph_dsa_split_op_surface,
 )
 from sglang.srt.layers.layernorm import LayerNorm, RMSNorm
@@ -1418,167 +1417,6 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
             return None
         return raw_topk_result
 
-    def _get_topk_ragged_with_cp(
-        self,
-        forward_batch: ForwardBatch,
-        layer_id: int,
-        q_fp8: torch.Tensor,
-        weights: torch.Tensor,
-        metadata: BaseIndexerMetadata,
-        kv_len: int,
-        actual_seq_q: int,
-        cp_index: List[Tuple[int, int, int]] = None,
-    ) -> torch.Tensor:
-        assert not _is_in_piecewise_or_breakable_cuda_graph(), (
-            "DSA context parallel (_get_topk_ragged_with_cp) not supported under "
-            "piecewise/breakable CUDA graph"
-        )
-        if TYPE_CHECKING:
-            assert isinstance(get_token_to_kv_pool(), DSATokenToKVPool)
-
-        page_size = get_token_to_kv_pool().page_size
-        if _is_xpu:
-            assert page_size in (
-                64,
-                128,
-            ), f"XPU DSA requires page_size 64 or 128, got {page_size}"
-        else:
-            assert page_size == 64, "only support page size 64"
-        assert len(weights.shape) == 3
-        weights = weights.squeeze(-1)
-        k_fp8_list = []
-        k_scale_list = []
-        ks_list = []
-        ke_offset_list = []
-        offset = 0
-        actual_seq_q_list = []
-        batch_idx_list = []
-
-        block_tables = metadata.get_page_table_64()
-
-        assert (
-            forward_batch.seq_lens_cpu is not None
-            and forward_batch.extend_seq_lens_cpu is not None
-        )
-        if cp_index is not None:
-            # TODO Multi-batch support has accuracy issues
-            for batch_idx, start_seq_position, end_seq_position in cp_index:
-                pre_chunk_offset = (
-                    forward_batch.seq_lens_cpu[batch_idx].item()
-                    - forward_batch.extend_seq_lens_cpu[batch_idx]
-                )
-                start_seq_position += pre_chunk_offset
-                end_seq_position += pre_chunk_offset
-                if offset == 0 and batch_idx != 0:
-                    offset += forward_batch.extend_seq_lens_cpu[batch_idx - 1]
-                k_fp8 = get_token_to_kv_pool().get_index_k_continuous(
-                    layer_id,
-                    end_seq_position,
-                    block_tables[batch_idx],
-                )
-                k_scale = get_token_to_kv_pool().get_index_k_scale_continuous(
-                    layer_id,
-                    end_seq_position,
-                    block_tables[batch_idx],
-                )
-
-                extend_seq_len = end_seq_position - start_seq_position
-                ks = torch.full(
-                    (extend_seq_len,), offset, dtype=torch.int32, device="cuda"
-                )
-                k_fp8_list.append(k_fp8)
-                k_scale_list.append(k_scale)
-                ks_list.append(ks)
-                ke_offset = torch.arange(
-                    start_seq_position + 1,
-                    end_seq_position + 1,
-                    dtype=torch.int32,
-                    device="cuda",
-                )
-                ke_offset_list.append(ke_offset)
-                actual_seq_q = torch.tensor(
-                    [extend_seq_len], dtype=torch.int32, device="cuda"
-                )
-                actual_seq_q_list.append(actual_seq_q)
-                batch_idx_list.append(batch_idx)
-
-            k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn)
-            k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1)
-            kv_fp8 = (k_fp8, k_scale)
-            ks = torch.cat(ks_list, dim=0)
-            ke_offset = torch.cat(ke_offset_list, dim=0)
-            ke = ks + ke_offset
-            actual_seq_q = torch.cat(actual_seq_q_list, dim=0)
-            with self._with_real_sm_count():
-                q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(q_fp8, weights)
-                logits = deep_gemm.fp8_mqa_logits(
-                    q_padded,
-                    kv_fp8,
-                    w_padded,
-                    ks,
-                    ke,
-                    clean_logits=False,
-                )
-            topk_result = metadata.topk_transform(
-                logits,
-                self.index_topk,
-                ks=ks,
-                cu_seqlens_q=actual_seq_q,
-                ke_offset=ke_offset,
-                batch_idx_list=batch_idx_list,
-            )
-        else:
-            kv_len = (
-                forward_batch.seq_lens_cpu[0].item()
-                - forward_batch.extend_seq_lens_cpu[0]
-                + kv_len
-            )
-            k_fp8 = get_token_to_kv_pool().get_index_k_continuous(
-                layer_id,
-                kv_len,
-                block_tables[0],
-            )
-            k_scale = get_token_to_kv_pool().get_index_k_scale_continuous(
-                layer_id,
-                kv_len,
-                block_tables[0],
-            )
-
-            k_fp8 = k_fp8.view(torch.float8_e4m3fn)
-            k_scale = k_scale.view(torch.float32).squeeze(-1)
-            kv_fp8 = (k_fp8, k_scale)
-            ks = torch.full((actual_seq_q,), offset, dtype=torch.int32, device="cuda")
-            ke_offset = torch.arange(
-                (kv_len - actual_seq_q) + 1,
-                kv_len + 1,
-                dtype=torch.int32,
-                device="cuda",
-            )
-            ke = ks + ke_offset
-
-            with self._with_real_sm_count():
-                q_padded, w_padded, _ = self._pad_heads_for_deep_gemm(q_fp8, weights)
-                logits = deep_gemm.fp8_mqa_logits(
-                    q_padded,
-                    kv_fp8,
-                    w_padded,
-                    ks,
-                    ke,
-                    clean_logits=False,
-                )
-            actual_seq_q = torch.tensor([actual_seq_q], dtype=torch.int32).to(
-                device="cuda", non_blocking=True
-            )
-            topk_result = metadata.topk_transform(
-                logits,
-                self.index_topk,
-                ks=ks,
-                cu_seqlens_q=actual_seq_q,
-                ke_offset=ke_offset,
-            )
-
-        return topk_result
-
     def _store_index_k_cache(
         self,
         forward_batch: ForwardBatch,
@@ -1955,67 +1793,21 @@ class Indexer(DSANPUIndexerMixin, BaseFusedOp):
                     forward_batch, layer_id, q_fp8, weights, metadata
                 )
             else:
-                if (
-                    forward_batch.attn_cp_metadata is not None
-                    and is_dsa_prefill_cp_in_seq_split()
-                ):
-                    kv_len_prev = forward_batch.attn_cp_metadata.kv_len_prev_list[0]
-                    kv_len_next = forward_batch.attn_cp_metadata.kv_len_next_list[0]
-                    actual_seq_q_prev = (
-                        forward_batch.attn_cp_metadata.actual_seq_q_prev_list[0]
-                    )
-                    actual_seq_q_next = (
-                        forward_batch.attn_cp_metadata.actual_seq_q_next_list[0]
-                    )
-
-                    # TODO support mutil-batch
-                    # cp_batch_seq_index_prev = forward_batch.attn_cp_metadata["cp_batch_seq_index_prev"]
-                    # cp_batch_seq_index_next = forward_batch.attn_cp_metadata["cp_batch_seq_index_next"]
-                    # TODO prev, next, combined into a single call
-                    q_fp8_prev, q_fp8_next = torch.split(
-                        q_fp8, (q_fp8.shape[0] + 1) // 2, dim=0
-                    )
-                    weights_prev, weights_next = torch.split(
-                        weights, (weights.shape[0] + 1) // 2, dim=0
-                    )
-                    topk_result_prev = self._get_topk_ragged_with_cp(
-                        forward_batch,
-                        layer_id,
-                        q_fp8_prev,
-                        weights_prev,
-                        metadata,
-                        kv_len_prev,
-                        actual_seq_q_prev,
-                    )
-
-                    topk_result_next = self._get_topk_ragged_with_cp(
-                        forward_batch,
-                        layer_id,
-                        q_fp8_next,
-                        weights_next,
-                        metadata,
-                        kv_len_next,
-                        actual_seq_q_next,
-                    )
-                    topk_result = torch.cat([topk_result_prev, topk_result_next], dim=0)
-                    topk_result = _broadcast_indexer_topk_from_rank0(topk_result)
-                    return maybe_capture_indexer_topk(layer_id, topk_result)
-                else:
-                    # In-graph (PCG/BCG) non-CP prefill is handled earlier by the
-                    # graph DSA split-op dispatch, so only the eager path reaches
-                    # here.
-                    assert not in_piecewise_or_breakable_cuda_graph, (
-                        "Internal error: in-graph DSA prefill must go through the "
-                        "graph DSA split-op dispatch"
-                    )
-                    topk_result = self._get_topk_ragged(
-                        enable_dual_stream,
-                        forward_batch,
-                        layer_id,
-                        q_fp8,
-                        weights,
-                        metadata,
-                    )
+                # In-graph (PCG/BCG) non-CP prefill is handled earlier by the
+                # graph DSA split-op dispatch, so only the eager path reaches
+                # here.
+                assert not in_piecewise_or_breakable_cuda_graph, (
+                    "Internal error: in-graph DSA prefill must go through the "
+                    "graph DSA split-op dispatch"
+                )
+                topk_result = self._get_topk_ragged(
+                    enable_dual_stream,
+                    forward_batch,
+                    layer_id,
+                    q_fp8,
+                    weights,
+                    metadata,
+                )
         else:
             raise NotImplementedError("DSA indexer only supports CUDA, HIP, and NPU")
         topk_result = _broadcast_indexer_topk_from_rank0(topk_result)

python/sglang/srt/layers/attention/dsa/utils.py

@@ -18,7 +18,7 @@ from sglang.srt.runtime_context import (
     get_parallel,
     process_model_config,
 )
-from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, is_npu
+from sglang.srt.utils import get_bool_env_var, is_cuda, is_hip, is_musa, is_npu
 from sglang.srt.utils.common import ceil_align, ceil_div
 
 
@@ -115,7 +115,7 @@ def should_use_dsa_fused_topk(seed_dsa_topk_from_draft_extend: bool) -> bool:
 
 
 def is_dsa_enable_prefill_cp():
-    if is_hip() or is_npu():
+    if is_hip() or is_npu() or is_musa():
         return get_parallel().enable_dsa_prefill_context_parallel
 
     # Generic prefill CP derives activation from the runtime topology and model
@@ -128,13 +128,6 @@ def is_dsa_enable_prefill_cp():
     return is_deepseek_dsa(hf_config) or is_deepseek_v4(hf_config)
 
 
-def is_dsa_prefill_cp_in_seq_split():
-    return (
-        is_dsa_enable_prefill_cp()
-        and get_parallel().dsa_prefill_cp_mode == "in-seq-split"
-    )
-
-
 def is_dsa_prefill_cp_round_robin_split():
     return (
         is_dsa_enable_prefill_cp()

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

@@ -80,7 +80,6 @@ from sglang.srt.layers.attention.dsa.utils import (
     dsa_cp_round_robin_split_q_seqs,
     dsa_use_prefill_cp,
     is_dsa_enable_prefill_cp,
-    is_dsa_prefill_cp_in_seq_split,
     pad_dsa_cache_seqlens,
     should_use_dsa_fused_topk,
 )
@@ -141,6 +140,31 @@ def _all_gather_dsa_trtllm_fp8_kv(
     return kv.split((kv_lora_rank, qk_rope_head_dim), dim=-1)
 
 
+def prepare_kv_for_attention(
+    attn_mla,
+    forward_batch: ForwardBatch,
+    k_nope: torch.Tensor,
+    k_pe: torch.Tensor,
+    *,
+    defer_materialization: bool,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Materialize KV needed before attention for the active layout."""
+    if (
+        defer_materialization
+        or not dsa_use_prefill_cp(forward_batch)
+        or not is_cp_v2_active(forward_batch)
+    ):
+        return k_nope, k_pe
+    strategy = get_cp_strategy()
+    assert strategy is not None
+    return strategy.materialize_full_mla_kv(
+        forward_batch,
+        attn_mla.attn_mqa,
+        k_nope,
+        k_pe,
+    )
+
+
 def materialize_full_kv_cp(
     attn_mla,
     forward_batch: ForwardBatch,
@@ -148,13 +172,18 @@ def materialize_full_kv_cp(
     k_nope: torch.Tensor,
     k_pe: torch.Tensor,
 ) -> tuple[torch.Tensor, torch.Tensor]:
+    """Materialize generic CP KV, retaining the ROCm DSA fallback."""
     if is_cp_v2_active(forward_batch):
-        return get_cp_strategy().materialize_full_mla_kv(
+        strategy = get_cp_strategy()
+        assert strategy is not None
+        return strategy.materialize_full_mla_kv(
             forward_batch,
             attn_mla.attn_mqa,
             k_nope,
             k_pe,
         )
+
+    assert is_hip(), "Legacy DSA KV materialization is HIP-only"
     return attn_mla.rebuild_cp_kv_cache(latent_cache, forward_batch, k_nope, k_pe)
 
 
@@ -3583,15 +3612,6 @@ class DeepseekSparseAttnBackend(
         block_tables = page_table_1.unsqueeze(1)
         seq_lens = metadata.cache_seqlens_int32 if seq_lens is None else seq_lens
 
-        if (
-            dsa_use_prefill_cp(forward_batch)
-            and is_dsa_prefill_cp_in_seq_split()
-            and forward_batch.attn_cp_metadata is not None
-        ):
-            cp_meta = forward_batch.attn_cp_metadata
-            seq_chunks = list(torch.split(seq_lens, cp_meta.split_list, dim=0))
-            seq_lens = torch.cat([seq_chunks[i] for i in cp_meta.zigzag_index], dim=0)
-
         out = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
             query=q,
             kv_cache=kv,

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

@@ -24,12 +24,8 @@ from sglang.srt.configs.model_config import AttentionArch
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_verify_mask
 from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
-from sglang.srt.layers.cp.utils import is_cp_v2_active
+from sglang.srt.layers.cp.utils import enable_cp_v2, is_cp_v2_active
 from sglang.srt.layers.radix_attention import AttentionType
-from sglang.srt.layers.utils.cp_utils import (
-    cp_allgather_and_save_kv_cache,
-    cp_attn_forward_extend,
-)
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
@@ -1012,7 +1008,7 @@ class FlashAttentionBackend(AttentionBackend):
             # (req_to_token is zero-init) and outputs for padding queries are
             # discarded downstream.
             if (
-                not is_cp_v2_active(forward_batch)
+                not enable_cp_v2()
                 and self.attn_cp_size > 1
                 and forward_batch.global_num_tokens_cpu is not None
                 and forward_batch.extend_num_tokens is not None
@@ -1266,11 +1262,7 @@ class FlashAttentionBackend(AttentionBackend):
     ):
         if score_mod is not None and self.fa_impl_ver != 4:
             raise RuntimeError("score_mod is only supported by the FA4 backend.")
-        is_cp_mode = (
-            forward_batch.forward_mode.is_context_parallel_extend()
-            and forward_batch.attn_cp_metadata is not None
-            and self.attn_cp_size > 1
-        )
+        cp_active = is_cp_v2_active(forward_batch)
 
         if k is not None:
             assert v is not None
@@ -1282,25 +1274,20 @@ class FlashAttentionBackend(AttentionBackend):
                     else forward_batch.encoder_out_cache_loc
                 )
                 if self.use_mla:
-                    if is_cp_v2_active(forward_batch):
-                        # CP-v2: k/k_rope are rank-local; the strategy gathers
-                        # the latent to full sequence and writes it.
+                    if cp_active:
                         cp_strategy = get_cp_strategy()
                         assert cp_strategy is not None
                         cp_strategy.materialize_full_mla_kv(
                             forward_batch, layer, k, k_rope
                         )
                     else:
-                        # CP-v1: k/k_rope arrive full-sequence (rebuild_cp_kv_cache
-                        # ran upstream); rank-local when CP is off. out_cache_loc is
-                        # never zigzag-split, so the write lands in the right slots.
                         self.token_to_kv_pool.set_mla_kv_buffer(
                             layer,
                             cache_loc,
                             k,
                             k_rope,
                         )
-                elif is_cp_mode:
+                elif cp_active:
                     # Dense-MHA CP: k, v are still rank-local; backend
                     # all-gathers and writes to the per-rank pool.
                     swa_loc = (
@@ -1308,21 +1295,11 @@ class FlashAttentionBackend(AttentionBackend):
                         if self.use_sliding_window_kv_pool
                         else None
                     )
-                    if is_cp_v2_active(forward_batch):
-                        cp_strategy = get_cp_strategy()
-                        assert cp_strategy is not None
-                        cp_strategy.materialize_full_kv(
-                            forward_batch, layer, k, v, swa_loc=swa_loc
-                        )
-                    else:
-                        cp_allgather_and_save_kv_cache(
-                            forward_batch,
-                            layer,
-                            k,
-                            v,
-                            self.attn_cp_size,
-                            swa_loc=swa_loc,
-                        )
+                    cp_strategy = get_cp_strategy()
+                    assert cp_strategy is not None
+                    cp_strategy.materialize_full_kv(
+                        forward_batch, layer, k, v, swa_loc=swa_loc
+                    )
                 else:
                     k_scale = k_descale if self.kv_cache_is_mxfp8 else layer.k_scale
                     v_scale = v_descale if self.kv_cache_is_mxfp8 else layer.v_scale
@@ -1459,11 +1436,7 @@ class FlashAttentionBackend(AttentionBackend):
                 cu_seqlens_k = metadata.encoder_cu_seqlens_k
                 window_size = (-1, -1)
 
-            if (
-                forward_batch.forward_mode.is_context_parallel_extend()
-                and forward_batch.attn_cp_metadata is not None
-                and self.attn_cp_size > 1
-            ):
+            if cp_active:
 
                 def _fa_cp_attn(
                     q_chunk, cu_seqlens_q_cp, cache_seqlens_cp, max_seqlen_q_cp
@@ -1488,23 +1461,15 @@ class FlashAttentionBackend(AttentionBackend):
                     )
 
                 q_cp = q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim)
-                if is_cp_v2_active(forward_batch):
-                    cp_strategy = get_cp_strategy()
-                    assert cp_strategy is not None
-                    result = cp_strategy.run_attention(
-                        q_cp,
-                        forward_batch,
-                        self.device,
-                        _fa_cp_attn,
-                        attention_backend=CPAttentionBackendKind.FLASH_ATTENTION,
-                    )
-                else:
-                    result = cp_attn_forward_extend(
-                        forward_batch,
-                        q_cp,
-                        self.device,
-                        _fa_cp_attn,
-                    )
+                cp_strategy = get_cp_strategy()
+                assert cp_strategy is not None
+                result = cp_strategy.run_attention(
+                    q_cp,
+                    forward_batch,
+                    self.device,
+                    _fa_cp_attn,
+                    attention_backend=CPAttentionBackendKind.FLASH_ATTENTION,
+                )
             elif self.fa_skip_kv_cache:
                 # Embedding mode: skip KV cache read and use raw K/V tensors
                 # directly via flash_attn_varlen_func. The KV cache write is
@@ -1731,16 +1696,15 @@ class FlashAttentionBackend(AttentionBackend):
                     q_nope = q_all[:, :, : layer.v_head_dim]
                     q_rope = q_all[:, :, layer.v_head_dim :]
 
-                if is_cp_mode:
+                if cp_active:
                     # MLA CP: q is rank-local zigzag-split; run the
                     # absorbed-MLA kernel twice (prev/next halves) against
-                    # the full latent KV pool (which rebuild_cp_kv_cache
-                    # populated upstream) via cp_attn_forward_extend.
+                    # the full latent KV pool through the selected strategy.
                     # Concat q_nope + q_rope along dim=-1 so the wrapper's
                     # chunk(2, dim=0) keeps their alignment; split back
                     # inside the closure.
                     assert not use_cascade_attn, (
-                        "Cascade attention under MLA CP is not supported in v1."
+                        "Cascade attention under MLA CP is not supported."
                     )
                     q_fused = torch.cat([q_nope, q_rope], dim=-1)
 
@@ -1773,20 +1737,15 @@ class FlashAttentionBackend(AttentionBackend):
                             ver=self.fa_impl_ver,
                         )
 
-                    if is_cp_v2_active(forward_batch):
-                        cp_strategy = get_cp_strategy()
-                        assert cp_strategy is not None
-                        o = cp_strategy.run_attention(
-                            q_fused,
-                            forward_batch,
-                            self.device,
-                            _mla_cp_attn,
-                            attention_backend=CPAttentionBackendKind.FLASH_ATTENTION,
-                        )
-                    else:
-                        o = cp_attn_forward_extend(
-                            forward_batch, q_fused, self.device, _mla_cp_attn
-                        )
+                    cp_strategy = get_cp_strategy()
+                    assert cp_strategy is not None
+                    o = cp_strategy.run_attention(
+                        q_fused,
+                        forward_batch,
+                        self.device,
+                        _mla_cp_attn,
+                        attention_backend=CPAttentionBackendKind.FLASH_ATTENTION,
+                    )
                 else:
                     result = flash_attn_with_kvcache(
                         q=q_rope,

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

@@ -3,6 +3,8 @@ from __future__ import annotations
 from contextlib import contextmanager
 from typing import TYPE_CHECKING, Iterator, Optional
 
+from sglang.srt.layers.cp.utils import cp_gather_after_forward, is_cp_v2_active
+
 if TYPE_CHECKING:
     import torch
 
@@ -33,8 +35,11 @@ class IndexTopKShareState:
 
     @property
     def _seed_buf(self) -> Optional[torch.Tensor]:
-        if self._forward_batch.forward_mode.is_extend(include_draft_extend_v2=True):
-            return self._forward_batch.spec_info.dsa_seed_topk_capture
+        spec_info = self._forward_batch.spec_info
+        if spec_info is not None and self._forward_batch.forward_mode.is_extend(
+            include_draft_extend_v2=True
+        ):
+            return spec_info.dsa_seed_topk_capture
         return None
 
     @property
@@ -46,6 +51,12 @@ class IndexTopKShareState:
         return self._topk_indices
 
     def update(self, topk_indices: Optional[torch.Tensor]) -> None:
+        if (
+            topk_indices is not None
+            and self.should_publish
+            and is_cp_v2_active(self._forward_batch)
+        ):
+            topk_indices = cp_gather_after_forward(topk_indices, self._forward_batch)
         self._topk_indices = topk_indices
 
     def publish(self) -> None:

python/sglang/srt/layers/communicator.py

@@ -37,6 +37,10 @@ from sglang.srt.layers.attention.dsa.utils import (
     is_dsa_enable_prefill_cp,
 )
 from sglang.srt.layers.aux_hidden_states import AuxHiddenStateAccumulator
+from sglang.srt.layers.cp.utils import (
+    is_mla_prefill_cp_enabled,
+    mla_use_prefill_cp,
+)
 from sglang.srt.layers.dp_attention import (
     attn_tp_all_gather_into_tensor,
     attn_tp_reduce_scatter_tensor,
@@ -64,10 +68,6 @@ from sglang.srt.layers.quantization.fp8_utils import (
     _use_aiter_bpreshuffle_gfx95,
     materialize_bpreshuffle_fp8_scale_tuple,
 )
-from sglang.srt.layers.utils.cp_utils import (
-    is_mla_prefill_cp_enabled,
-    mla_use_prefill_cp,
-)
 from sglang.srt.model_executor.cuda_graph_config import (
     Backend,
     Phase,

python/sglang/srt/layers/communicator_dsa_cp.py

@@ -20,7 +20,6 @@ import torch
 
 from sglang.srt.layers.attention.dsa.utils import (
     dsa_use_prefill_cp,
-    is_dsa_enable_prefill_cp,
 )
 from sglang.srt.layers.communicator import (
     CommunicateContext,
@@ -31,24 +30,17 @@ from sglang.srt.layers.communicator import (
     LayerScatterModes,
     ScatterMode,
 )
+from sglang.srt.layers.cp.utils import mla_use_prefill_cp
 from sglang.srt.layers.dp_attention import (
     attn_cp_all_gather_into_tensor,
     attn_cp_reduce_scatter_tensor,
     get_local_dp_buffer,
 )
-from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
 from sglang.srt.runtime_context import get_parallel
 
 
-def dsa_enable_prefill_cp():
-    # After using cp, the communication mode of this part changes.
-    # The three parts of prepare_attn, prepare_mlp, and postprocess_layer
-    # no longer require additional communication for reduce, scatter, etc.
-    return is_dsa_enable_prefill_cp()
-
-
 def maybe_prefetch_next_full_attention_kv(
     forward_batch: ForwardBatch,
     next_full_attention_layer_id: Optional[int],

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

@@ -23,11 +23,10 @@ from sglang.srt.runtime_context import get_parallel
 
 def get_cp_padding_align_size() -> int:
     """Return the token-count alignment required by the active CP strategy."""
-    from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split
-    from sglang.srt.layers.utils.cp_utils import is_prefill_cp_in_seq_split
+    from sglang.srt.layers.cp.base import is_zigzag
 
     attn_cp_size = get_parallel().attn_cp_size
-    if is_prefill_cp_in_seq_split() or is_dsa_prefill_cp_in_seq_split():
+    if is_zigzag():
         return attn_cp_size * 2
     return attn_cp_size
 

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

@@ -23,6 +23,7 @@ from sglang.srt.layers.cp.base import (
     ContextParallelStrategyKind,
     CPAttentionBackendKind,
     get_cp_strategy,
+    is_cp_enabled,
 )
 from sglang.srt.layers.cp.interleave import (
     InterleaveContextParallelMetadata,
@@ -35,7 +36,7 @@ from sglang.srt.layers.cp.zigzag import (
     ZigzagCPStrategy,
 )
 from sglang.srt.layers.moe.utils import get_moe_a2a_backend
-from sglang.srt.runtime_context import get_parallel
+from sglang.srt.runtime_context import get_parallel, uses_mla_backend
 
 if TYPE_CHECKING:
     from sglang.srt.model_executor.model_runner import ModelRunner
@@ -119,9 +120,9 @@ def get_layer_owner(local_layer_idx: int, shard_size: int, total_layers: int) ->
 
 def enable_cp_v2() -> bool:
     """Return whether the strategy-based generic prefill CP path is available."""
-    from sglang.srt.utils import is_hip, is_npu
+    from sglang.srt.utils import is_hip, is_musa, is_npu
 
-    return not (is_hip() or is_npu())
+    return not (is_hip() or is_npu() or is_musa())
 
 
 def is_cp_v2_active(forward_batch) -> bool:
@@ -143,6 +144,24 @@ def is_cp_v2_active(forward_batch) -> bool:
     return strategy.can_apply(len(input_ids), forward_batch)
 
 
+def is_mla_prefill_cp_enabled() -> bool:
+    """Return whether prefill CP is configured for an MLA attention backend."""
+    if enable_cp_v2():
+        return is_cp_enabled() and uses_mla_backend()
+    return get_parallel().enable_prefill_context_parallel and uses_mla_backend()
+
+
+def mla_use_prefill_cp(forward_batch) -> bool:
+    """Return whether this MLA forward batch is using prefill CP."""
+    if enable_cp_v2():
+        return is_mla_prefill_cp_enabled() and is_cp_v2_active(forward_batch)
+    return (
+        getattr(forward_batch, "attn_cp_metadata", None) is not None
+        and is_mla_prefill_cp_enabled()
+        and forward_batch.forward_mode.is_context_parallel_extend()
+    )
+
+
 def prepare_cp_forward(forward_batch) -> None:
     """Build CP-v2 metadata for an active context-parallel prefill batch."""
     assert is_cp_v2_active(forward_batch)
@@ -251,8 +270,8 @@ def cp_materialize_global_token_order(
         assert strategy is not None
         return strategy.gather_kv_cache(x, forward_batch, stream)
 
-    # TODO(hzh0425): Keep the legacy gather temporarily for CP-v1 compatibility. Remove it
-    # with the follow-up CP-v1 cleanup.
+    # HIP/NPU still materialize their protected platform layout through the
+    # legacy collective until those backends migrate independently.
     from sglang.srt.layers.utils.cp_utils import cp_all_gather_rerange_output
 
     return cp_all_gather_rerange_output(
@@ -265,6 +284,7 @@ def cp_shard_model_inputs(
     complete_hidden_states: Any,
     complete_position_ids: Any,
     forward_batch,
+    complete_input_ids: Optional[Any] = None,
 ):
     """Restore the shared batch so logits processing keeps full-batch metadata."""
     assert is_cp_v2_active(forward_batch)
@@ -272,6 +292,18 @@ def cp_shard_model_inputs(
         complete_hidden_states, forward_batch
     )
     sharded_positions = cp_shard_position_ids(complete_position_ids, forward_batch)
+    model_input_ids = (
+        cp_shard_hidden_states(complete_input_ids, forward_batch)
+        if complete_input_ids is not None
+        else None
+    )
+
+    had_input_ids_global = hasattr(forward_batch, "input_ids_global")
+    input_ids_global_backup = getattr(forward_batch, "input_ids_global", None)
+    if complete_input_ids is not None:
+        forward_batch.input_ids_global = cp_round_robin_input_ids_v2(
+            complete_input_ids, forward_batch
+        )
 
     spec_info = getattr(forward_batch, "spec_info", None)
     spec_hidden_states = getattr(spec_info, "hidden_states", None)
@@ -286,10 +318,14 @@ def cp_shard_model_inputs(
         )
 
     try:
-        yield sharded_hidden_states, sharded_positions
+        yield sharded_hidden_states, sharded_positions, model_input_ids
     finally:
         if spec_hidden_states_backup is not None:
             spec_info.hidden_states = spec_hidden_states_backup
+        if had_input_ids_global:
+            forward_batch.input_ids_global = input_ids_global_backup
+        elif hasattr(forward_batch, "input_ids_global"):
+            delattr(forward_batch, "input_ids_global")
 
 
 def _to_int_list(values) -> Optional[list[int]]:
@@ -313,6 +349,8 @@ __all__ = [
     "enable_cp_v2",
     "get_cp_strategy",
     "is_cp_v2_active",
+    "is_mla_prefill_cp_enabled",
+    "mla_use_prefill_cp",
     "cp_gather_after_forward",
     "cp_materialize_global_token_order",
     "cp_round_robin_input_ids_v2",

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

@@ -1,6 +1,8 @@
+"""Legacy prefill CP helpers retained for HIP, NPU, and MUSA callers."""
+
 from dataclasses import dataclass
 from itertools import accumulate
-from typing import Callable, List
+from typing import List
 
 import torch
 import torch.nn.functional as F
@@ -66,13 +68,6 @@ def is_prefill_context_parallel_enabled():
     return get_parallel().enable_prefill_context_parallel
 
 
-def is_prefill_cp_in_seq_split():
-    return (
-        is_prefill_context_parallel_enabled()
-        and get_parallel().prefill_cp_mode == "in-seq-split"
-    )
-
-
 def is_mla_prefill_cp_enabled() -> bool:
     return get_parallel().enable_prefill_context_parallel and uses_mla_backend()
 
@@ -477,48 +472,6 @@ def cp_allgather_and_save_kv_cache(forward_batch, layer, k, v, cp_size, swa_loc=
     )
 
 
-def cp_attn_forward_extend(
-    forward_batch,
-    q: torch.Tensor,
-    device: torch.device,
-    attn_fn: Callable[[torch.Tensor, torch.Tensor, torch.Tensor, int], torch.Tensor],
-) -> torch.Tensor:
-    """
-    Split q into prev/next zigzag halves based on CP metadata, call the
-    backend-specific attention function twice with appropriate per-half
-    metadata, and concatenate the results.
-
-    For bs > 1, q is laid out as [all_prev_tokens_across_seqs,
-    all_next_tokens_across_seqs]; the split point is total_q_prev_tokens.
-    cu_seqlens_q_prev/next tensors have shape [bs+1] and carry the
-    per-sequence boundaries through FlashAttention's variable-length API.
-
-    attn_fn signature:
-        attn_fn(q, cu_seqlens_q, cache_seqlens, max_seqlen_q) -> result
-    where only these four CP-varying parameters differ between halves.
-    All other backend-specific args should be captured in the closure.
-    """
-    cp_meta = forward_batch.attn_cp_metadata
-
-    q_prev = q[: cp_meta.total_q_prev_tokens]
-    q_next = q[cp_meta.total_q_prev_tokens :]
-
-    result_prev = attn_fn(
-        q_prev,
-        cp_meta.cu_seqlens_q_prev_tensor,
-        cp_meta.kv_len_prev_tensor,
-        cp_meta.max_seqlen_q_prev,
-    )
-    result_next = attn_fn(
-        q_next,
-        cp_meta.cu_seqlens_q_next_tensor,
-        cp_meta.kv_len_next_tensor,
-        cp_meta.max_seqlen_q_next,
-    )
-
-    return torch.concat([result_prev, result_next], dim=0)
-
-
 def prepare_context_parallel_metadata(
     kv_len,
     cp_rank,

python/sglang/srt/managers/schedule_policy.py

@@ -39,8 +39,6 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Set, Union
 import torch
 
 from sglang.srt.dllm.config import DllmConfig
-from sglang.srt.layers.attention.dsa.utils import is_dsa_prefill_cp_in_seq_split
-from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled
 from sglang.srt.managers.schedule_batch import (
     Req,
     ScheduleBatch,
@@ -582,9 +580,7 @@ class PrefillAdder:
         self.priority_scheduling_preemption_threshold = (
             priority_scheduling_preemption_threshold
         )
-        self.dsa_prefill_cp_in_seq_split = is_dsa_prefill_cp_in_seq_split()
         self.max_running_requests = max_running_requests
-        self.prefill_context_parallel_enabled = is_prefill_context_parallel_enabled()
         self.prefill_max_requests = prefill_max_requests
         self.prefill_delayer_single_pass = prefill_delayer_single_pass
         self.max_prefill_bs = max_prefill_bs
@@ -1200,12 +1196,6 @@ class PrefillAdder:
     def add_one_req(
         self, req: Req, has_chunked_req: bool, truncation_align_size: Optional[int]
     ):
-        # TODO support cp with multiple requests
-        # Enabling context parallelism currently presents precision issues;
-        # therefore, the prefill-batch setting is temporarily set to 1.
-        if (self.dsa_prefill_cp_in_seq_split) and len(self.can_run_list) >= 1:
-            return AddReqResult.OTHER
-
         if (x := self.prefill_max_requests) is not None and len(self.can_run_list) >= x:
             return AddReqResult.OTHER
 

python/sglang/srt/model_executor/forward_batch_info.py

@@ -67,9 +67,9 @@ from sglang.srt.utils import (
 from sglang.srt.utils.common import ceil_align, is_pin_memory_available
 
 if TYPE_CHECKING:
+    from sglang.srt.layers.cp.base import BaseContextParallelMetadata
     from sglang.srt.layers.dcp.metadata import DecodeContextParallelMetadata
     from sglang.srt.layers.logits_processor import LogitsProcessorOutput
-    from sglang.srt.layers.utils.cp_utils import ContextParallelMetadata
     from sglang.srt.managers.schedule_batch import MultimodalInputs, ScheduleBatch
     from sglang.srt.model_executor.model_runner import ModelRunner
     from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
@@ -304,8 +304,8 @@ def compute_local_num_token_non_padded_cpu(
 def prefill_graph_tolerates_sum_len() -> bool:
     """Whether MegaMoE may replay prefill graphs with local shapes."""
     from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
+    from sglang.srt.layers.cp.utils import is_mla_prefill_cp_enabled
     from sglang.srt.layers.moe.utils import get_moe_a2a_backend
-    from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
 
     if not get_moe_a2a_backend().is_megamoe():
         return False
@@ -598,7 +598,7 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     tbo_padded_len: Optional[int] = None
     tbo_children: Optional[List[ForwardBatch]] = None
 
-    attn_cp_metadata: Optional[ContextParallelMetadata] = None
+    attn_cp_metadata: Optional[BaseContextParallelMetadata] = None
 
     # For decode context parallel.
     # NOTE: DecodeContextParallelMetadata is imported under TYPE_CHECKING only (see the

python/sglang/srt/model_executor/model_runner.py

@@ -77,10 +77,10 @@ from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.cp.utils import (
     get_cp_strategy,
     is_cp_v2_active,
+    is_mla_prefill_cp_enabled,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.sampler import create_sampler
-from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
 from sglang.srt.lora.lora_manager import LoRAManager, init_lora_cuda_graph_moe_buffers
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.managers.schedule_batch import sanity_check_mm_pad_shift_value
@@ -264,7 +264,7 @@ class SamplingPrewarmResult:
 def _prefill_cuda_graph_allows_context_parallel(
     prefill_runner, forward_batch: ForwardBatch
 ) -> bool:
-    """Allow CP only through a runner that captured the validated CP-v2 body."""
+    """Allow CP only through a runner that captured the validated CP body."""
     return get_cp_strategy() is None or (
         bool(getattr(prefill_runner, "enable_cp_v2_bcg_capture", False))
         and is_cp_v2_active(forward_batch)

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

@@ -49,13 +49,13 @@ from sglang.srt.layers.attention.base_attn_backend import (
     SharedReadEnds,
 )
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
+from sglang.srt.layers.cp.utils import is_mla_prefill_cp_enabled
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
     set_dp_buffer_len,
     set_is_extend_in_batch,
 )
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
-from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
 from sglang.srt.model_executor.cuda_graph_buffer_registry import (
     CudaGraphBufferRegistry,
     build_decode_registry,

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

@@ -380,7 +380,7 @@ class EagerRunner(BaseRunner):
     def _execute_extend_cp_v2(
         self, forward_batch: ForwardBatch, kwargs: dict
     ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
-        """CP-v2 extend: shard inputs at the model boundary, run the body on the
+        """CP extend: shard inputs at the model boundary, run the body on the
         rank-local slice, then gather hidden states before the logits step.
         """
         model = self.model_runner.model
@@ -389,13 +389,16 @@ class EagerRunner(BaseRunner):
         if input_embeds is None:
             input_embeds = model.get_input_embeddings()(forward_batch.input_ids)
         with cp_shard_model_inputs(
-            input_embeds, forward_batch.positions, forward_batch
-        ) as (sharded_input_embeds, sharded_positions):
+            input_embeds,
+            forward_batch.positions,
+            forward_batch,
+            forward_batch.input_ids,
+        ) as (sharded_input_embeds, sharded_positions, model_input_ids):
             model_kwargs = {"input_embeds": sharded_input_embeds}
             if (pp_proxy_tensors := kwargs.get("pp_proxy_tensors")) is not None:
                 model_kwargs["pp_proxy_tensors"] = pp_proxy_tensors
             hidden_states = model.model(
-                forward_batch.input_ids,
+                model_input_ids,
                 sharded_positions,
                 forward_batch,
                 **model_kwargs,

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

@@ -64,7 +64,10 @@ from sglang.srt.layers.cp.bcg import (
     execute_prefill_cp_bcg,
     filter_prefill_cp_bcg_capture_num_tokens,
 )
-from sglang.srt.layers.cp.utils import is_cp_v2_active
+from sglang.srt.layers.cp.utils import (
+    is_cp_v2_active,
+    is_mla_prefill_cp_enabled,
+)
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
     set_dp_buffer_len,
@@ -72,7 +75,6 @@ from sglang.srt.layers.dp_attention import (
 )
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.pooler import EmbeddingPoolerOutput
-from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
 from sglang.srt.model_executor.cuda_graph_buffer_registry import (
     CudaGraphBufferRegistry,
     build_prefill_registry,

python/sglang/srt/models/deepseek_common/attention_backend_handler.py

@@ -1,5 +1,6 @@
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
+from sglang.srt.layers.cp.utils import enable_cp_v2, is_cp_v2_active
 from sglang.srt.layers.utils.cp_utils import mla_use_prefill_cp
 from sglang.srt.model_executor.forward_context import get_attn_backend
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
@@ -111,10 +112,12 @@ def _handle_attention_backend(attn, forward_batch, backend_name):
     if is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph():
         return AttnForwardMethod.MLA
 
-    # MLA prefill CP forces absorbed MLA regardless of prefix length: the
-    # CP path gathers latent KV via rebuild_cp_kv_cache and feeds the
-    # backend's absorbed-MLA kernel.
-    if mla_use_prefill_cp(forward_batch):
+    # Strategy CP gathers latent KV in the backend's absorbed MLA path;
+    # normal MHA would write rank-local KV against full cache locations.
+    # Protected platform CP retains its model-side materialization path.
+    if is_cp_v2_active(forward_batch) or (
+        not enable_cp_v2() and mla_use_prefill_cp(forward_batch)
+    ):
         return _dispatch_mla_subtype(attn, forward_batch)
 
     sum_extend_prefix_lens = _get_sum_extend_prefix_lens(forward_batch)

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -14,12 +14,10 @@ from sglang.kernels.ops.quantization.fp8_kernel import (
 from sglang.srt.compilation.compilation_config import register_split_op
 from sglang.srt.environ import envs
 from sglang.srt.layers import deep_gemm_wrapper
-from sglang.srt.layers.attention.dsa.utils import (
-    dsa_use_prefill_cp,
-    is_graph_dsa_split_op_surface,
-)
+from sglang.srt.layers.attention.dsa.utils import is_graph_dsa_split_op_surface
+from sglang.srt.layers.attention.dsa_backend import prepare_kv_for_attention
 from sglang.srt.layers.communicator import get_attn_tp_context
-from sglang.srt.layers.cp.utils import is_cp_v2_active
+from sglang.srt.layers.cp.utils import enable_cp_v2
 from sglang.srt.layers.dcp import (
     all_gather_kv_cache_for_mla_extend,
     all_gather_q_for_mla_decode,
@@ -115,6 +113,7 @@ def should_defer_dsa_cp_kv_gather(
     dsa_prefill_cp: bool,
     fuse_rope_for_trtllm_mla: bool,
 ) -> bool:
+    """Compatibility predicate imported by the unchanged ROCm MLA path."""
     return dsa_prefill_cp and fuse_rope_for_trtllm_mla
 
 
@@ -266,10 +265,6 @@ class DeepseekMLAForwardMixin:
             return None
         if get_parallel().dcp_enabled:
             return None
-        # Context-parallel prefill reshuffles the KV side; keep the handshake
-        # out of those paths.
-        if dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch):
-            return None
         # Kernel shape constraints (tl.arange / tl.dot / block tiling).  K
         # (qk_nope_head_dim) needs only K % 16 == 0 and K <= 256: power-of-2
         # K (DeepSeek 128) takes the kernel's preload-once path, other K
@@ -623,30 +618,17 @@ class DeepseekMLAForwardMixin:
                 num_tokens, self.num_local_heads, self.kv_lora_rank, q_nope.device
             )
 
-        dsa_prefill_cp = dsa_use_prefill_cp(forward_batch)
-        mla_prefill_cp = mla_use_prefill_cp(forward_batch)
-        defer_kv_gather_until_after_rope = should_defer_dsa_cp_kv_gather(
-            dsa_prefill_cp=dsa_prefill_cp,
-            fuse_rope_for_trtllm_mla=fuse_rope_for_trtllm_mla,
+        k_nope, k_pe = prepare_kv_for_attention(
+            self,
+            forward_batch,
+            k_nope,
+            k_pe,
+            defer_materialization=fuse_rope_for_trtllm_mla,
         )
-        if dsa_prefill_cp and not defer_kv_gather_until_after_rope:
-            from sglang.srt.layers.attention.dsa_backend import materialize_full_kv_cp
 
-            k_nope, k_pe = materialize_full_kv_cp(
-                self,
-                forward_batch,
-                latent_cache,
-                k_nope,
-                k_pe,
-            )
-        elif mla_prefill_cp and not is_cp_v2_active(forward_batch):
-            # CP-v1 gathers the latent here; CP-v2 gathers it in the attention
-            # backend via the strategy (materialize_full_mla_kv).
+        if not enable_cp_v2() and mla_use_prefill_cp(forward_batch):
             k_nope, k_pe = self.rebuild_cp_kv_cache(
-                latent_cache,
-                forward_batch,
-                k_nope,
-                k_pe,
+                latent_cache, forward_batch, k_nope, k_pe
             )
 
         # all_gather q_pe, q_nope_out,take tp8 as an example， q_pe [B, H, ROPE_DIM], q_nope_out [B, H, NOPE_DIM] gathered to [B, H * dcp_world_size, ROPE_DIM] [B, H * dcp_world_size, NOPE_DIM] for decode batch, and all gather k_pe, k_nope for extend batch.

python/sglang/srt/models/deepseek_nextn.py

@@ -25,7 +25,6 @@ from torch import nn
 from transformers import PretrainedConfig
 
 from sglang.kernels.ops.layernorm.fused_eh_norm import fused_eh_norm
-from sglang.srt.configs.model_config import is_deepseek_dsa
 from sglang.srt.distributed import get_pp_group
 from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
@@ -36,7 +35,7 @@ from sglang.srt.layers.attention.dsa.utils import (
     is_dsa_prefill_cp_round_robin_split,
 )
 from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState
-from sglang.srt.layers.cp.utils import cp_gather_after_forward, is_cp_v2_active
+from sglang.srt.layers.cp.utils import enable_cp_v2
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import ReplicatedLinear
 from sglang.srt.layers.logits_processor import LogitsProcessor
@@ -166,14 +165,6 @@ class DeepseekModelNextN(nn.Module):
             layer_name = "layers." + str(config.num_hidden_layers)
 
         self.quant_config = quant_config
-        self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
-        self.mla_enable_prefill_cp = (
-            is_mla_prefill_cp_enabled() and not is_deepseek_dsa(config)
-        )
-        if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp:
-            self.cp_size = get_parallel().attn_cp_size
-        else:
-            self.cp_size = None
         self.decoder = DeepseekV2DecoderLayer(
             config,
             0,
@@ -183,8 +174,6 @@ class DeepseekModelNextN(nn.Module):
             prefix=add_prefix(layer_name, prefix),
             alt_stream=self.alt_stream,
             skip_rope=config.qk_rope_head_dim == 0,
-            dsa_enable_prefill_cp=self.dsa_enable_prefill_cp,
-            mla_enable_prefill_cp=self.mla_enable_prefill_cp,
         )
 
         self.shared_head = nn.Module()
@@ -271,13 +260,11 @@ class DeepseekModelNextN(nn.Module):
                 else:
                     hidden_states = self.eh_proj(eh_input)
 
-            # CP-v2 shards/gathers hidden states at the eager-runner boundary.
-            cp_v2_active = is_cp_v2_active(forward_batch)
-            use_cp_v1 = (
-                dsa_use_prefill_cp(forward_batch, self.dsa_enable_prefill_cp)
-                or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp)
-            ) and not cp_v2_active
-            if use_cp_v1:
+            # Protected platforms retain their model-side token split.
+            use_platform_cp = not enable_cp_v2() and (
+                dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch)
+            )
+            if use_platform_cp:
                 hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
                 positions = cp_split_and_rebuild_position(forward_batch, positions)
             residual = None
@@ -297,11 +284,11 @@ class DeepseekModelNextN(nn.Module):
                 else:
                     hidden_states = self.shared_head.norm(hidden_states)
 
-                if use_cp_v1:
+                if use_platform_cp:
                     local_num_tokens = hidden_states.shape[0]
                     hidden_states = cp_all_gather_rerange_output(
                         hidden_states,
-                        self.cp_size,
+                        get_parallel().attn_cp_size,
                         forward_batch,
                         torch.cuda.current_stream(),
                     )
@@ -309,16 +296,10 @@ class DeepseekModelNextN(nn.Module):
                         topk_indices = _gather_dsa_topk_indices_for_cp(
                             topk_indices,
                             local_num_tokens,
-                            self.cp_size,
+                            get_parallel().attn_cp_size,
                             forward_batch,
                             torch.cuda.current_stream(),
                         )
-                elif (
-                    cp_v2_active
-                    and index_topk_share.should_publish
-                    and topk_indices is not None
-                ):
-                    topk_indices = cp_gather_after_forward(topk_indices, forward_batch)
             index_topk_share.update(topk_indices)
             index_topk_share.publish()
         finally:
@@ -369,16 +350,6 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
         # if not set, model load will be broken in DeepseekV3ForCausalLM load_weights()
         self.pp_group = get_pp_group()
         self.determine_num_fused_shared_experts()
-        self.use_dsa = is_deepseek_dsa(config)
-        self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
-        self.mla_enable_prefill_cp = is_mla_prefill_cp_enabled() and not self.use_dsa
-        if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp:
-            self.cp_rank = get_parallel().attn_cp_rank
-            self.cp_size = get_parallel().attn_cp_size
-        else:
-            self.cp_rank = None
-            self.cp_size = None
-
         nextn_quant_config = self._resolve_nextn_quant_config(config, quant_config)
 
         self.model = DeepseekModelNextN(
@@ -400,25 +371,31 @@ class DeepseekV3ForCausalLMNextN(DeepseekV3ForCausalLM):
         positions: torch.Tensor,
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
-        # TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
-        if not is_cp_v2_active(forward_batch):
-            if self.dsa_enable_prefill_cp:
+        if not enable_cp_v2():
+            if is_dsa_enable_prefill_cp():
                 if can_dsa_cp_split(
-                    len(input_ids), self.cp_size, self.use_dsa, forward_batch
+                    len(input_ids),
+                    get_parallel().attn_cp_size,
+                    self.model.decoder.self_attn.use_dsa,
+                    forward_batch,
                 ):
                     forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
                         len(input_ids),
-                        self.cp_rank,
-                        self.cp_size,
+                        get_parallel().attn_cp_rank,
+                        get_parallel().attn_cp_size,
                         forward_batch.seq_lens_cpu.tolist(),
                         extend_seqs_len=forward_batch.extend_seq_lens_cpu,
                     )
-            elif self.mla_enable_prefill_cp:
-                if can_cp_split(len(input_ids), self.cp_size, forward_batch):
+            elif (
+                is_mla_prefill_cp_enabled() and not self.model.decoder.self_attn.use_dsa
+            ):
+                if can_cp_split(
+                    len(input_ids), get_parallel().attn_cp_size, forward_batch
+                ):
                     forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
                         len(input_ids),
-                        self.cp_rank,
-                        self.cp_size,
+                        get_parallel().attn_cp_rank,
+                        get_parallel().attn_cp_size,
                         forward_batch.seq_lens_cpu.tolist(),
                         extend_seqs_len=forward_batch.extend_seq_lens_cpu,
                     )

python/sglang/srt/models/deepseek_v2.py

@@ -85,7 +85,7 @@ from sglang.srt.layers.communicator_dsa_cp import (
     maybe_prefetch_next_full_attention_kv,
 )
 from sglang.srt.layers.cp.cp_decode_attn_tp import get_cp_decode_attn_tp_ctx
-from sglang.srt.layers.cp.utils import is_cp_v2_active
+from sglang.srt.layers.cp.utils import enable_cp_v2
 from sglang.srt.layers.dcp.planner import (
     prepare_decode_context_parallel_metadata,
 )
@@ -470,14 +470,10 @@ class MoEGate(nn.Module):
         config,
         quant_config,
         prefix: str = "",
-        is_nextn: bool = False,
         is_hash_moe: bool = False,
         is_deepseek_v4: bool = False,
-        dsa_enable_prefill_cp: bool = False,
-        mla_enable_prefill_cp: bool = False,
     ):
         super().__init__()
-        self.is_nextn = is_nextn
         self.is_deepseek_v4 = is_deepseek_v4
         self.weight = nn.Parameter(
             torch.empty(
@@ -509,9 +505,6 @@ class MoEGate(nn.Module):
             self.e_score_correction_bias = None
         if _is_cpu and _is_cpu_amx_available:
             self.quant_method = PackWeightMethod(weight_names=["weight"])
-        self.use_dsa = is_deepseek_dsa(config)
-        self.dsa_enable_prefill_cp = dsa_enable_prefill_cp
-        self.mla_enable_prefill_cp = mla_enable_prefill_cp
         self.tiny_router_gemm_max_tokens = tiny_router_gemm_max_tokens(
             num_experts=config.n_routed_experts,
             hidden_size=config.hidden_size,
@@ -539,36 +532,29 @@ class MoEGate(nn.Module):
             return F.linear(hidden_states, self.weight, None)
 
         if (
-            not self.is_deepseek_v4
+            not enable_cp_v2()
+            and not self.is_deepseek_v4
             and forward_batch is not None
-            and (
-                dsa_use_prefill_cp(forward_batch, self.dsa_enable_prefill_cp)
-                or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp)
-            )
+            and (dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch))
         ):
-            if _is_cuda:
-                from sglang.kernels.ops.attention.dsv4 import linear_bf16_fp32
-
-                return linear_bf16_fp32(hidden_states, self.weight)
             return F.linear(hidden_states, self.weight, None)
+
+        if hidden_states.shape[0] <= self.tiny_router_gemm_max_tokens:
+            logits = tiny_gemm_bf16(
+                hidden_states,
+                self.weight,
+                out_dtype=torch.float32,
+                max_m=self.tiny_router_gemm_max_tokens,
+            )
+        elif _use_aiter:
+            logits = aiter_dsv3_router_gemm(hidden_states, self.weight)
+        elif not _is_cuda:
+            logits = F.linear(hidden_states, self.weight, None)
         else:
-            if hidden_states.shape[0] <= self.tiny_router_gemm_max_tokens:
-                logits = tiny_gemm_bf16(
-                    hidden_states,
-                    self.weight,
-                    out_dtype=torch.float32,
-                    max_m=self.tiny_router_gemm_max_tokens,
-                )
+            # cuBLAS bf16 x bf16 -> fp32 GEMM (torch.mm's out_dtype kwarg is CUDA-only)
+            from sglang.kernels.ops.attention.dsv4 import linear_bf16_fp32
 
-            elif _use_aiter:
-                logits = aiter_dsv3_router_gemm(hidden_states, self.weight)
-            elif not _is_cuda:
-                logits = F.linear(hidden_states, self.weight, None)
-            else:
-                # cuBLAS bf16 x bf16 -> fp32 GEMM (torch.mm's out_dtype kwarg is CUDA-only)
-                from sglang.kernels.ops.attention.dsv4 import linear_bf16_fp32
-
-                logits = linear_bf16_fp32(hidden_states, self.weight)
+            logits = linear_bf16_fp32(hidden_states, self.weight)
 
         return logits
 
@@ -583,8 +569,6 @@ class DeepseekV2MoE(nn.Module):
         alt_stream: Optional[torch.cuda.Stream] = None,
         is_nextn: bool = False,
         is_deepseek_v4: bool = False,
-        dsa_enable_prefill_cp: bool = False,
-        mla_enable_prefill_cp: bool = False,
     ):
         super().__init__()
         self.tp_size = get_parallel().tp_size
@@ -643,11 +627,8 @@ class DeepseekV2MoE(nn.Module):
             config=config,
             quant_config=quant_config,
             prefix=add_prefix("gate", prefix),
-            is_nextn=is_nextn,
             is_hash_moe=self.is_hash,
             is_deepseek_v4=is_deepseek_v4,
-            dsa_enable_prefill_cp=dsa_enable_prefill_cp,
-            mla_enable_prefill_cp=mla_enable_prefill_cp,
         )
 
         # scaling factor for fused shared experts on AMD-platform.
@@ -1771,8 +1752,6 @@ class DeepseekV2AttentionMLA(
         alt_stream: Optional[torch.cuda.Stream] = None,
         skip_rope: bool = False,
         is_nextn: bool = False,
-        dsa_enable_prefill_cp: bool = False,
-        mla_enable_prefill_cp: bool = False,
     ) -> None:
         super().__init__()
         self.layer_id = layer_id
@@ -1788,15 +1767,6 @@ class DeepseekV2AttentionMLA(
         attn_tp_rank = get_parallel().attn_tp_rank
         attn_tp_size = get_parallel().attn_tp_size
         self.use_dsa = is_deepseek_dsa(config)
-        self.dsa_enable_prefill_cp = dsa_enable_prefill_cp
-        self.mla_enable_prefill_cp = mla_enable_prefill_cp
-        if self.dsa_enable_prefill_cp:
-            assert self.use_dsa, "CP currently only supports deepseek v3.2 model"
-        # cp reuses the attn_tp comm group but needs to duplicate the weights;
-        # store cp_size whenever either CP flavor is active so rebuild_cp_kv_cache
-        # and the FA3 MLA wrapper can reach it on the dense MLA path too.
-        if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp:
-            self.cp_size = get_parallel().attn_cp_size
         self.num_heads = num_heads
         assert num_heads % attn_tp_size == 0
         self.num_local_heads = num_heads // attn_tp_size
@@ -2291,12 +2261,12 @@ class DeepseekV2AttentionMLA(
         return q.view(-1, self.num_local_heads, self.qk_head_dim)
 
     def rebuild_cp_kv_cache(self, latent_cache, forward_batch, k_nope, k_pe):
-        # support allgather+rerrange
+        # Retained for the platform MLA paths.
         latent_cache[..., : self.kv_lora_rank] = k_nope.squeeze(1)
         latent_cache[..., self.kv_lora_rank :] = k_pe.squeeze(1)
         latent_cache_output = cp_all_gather_rerange_output(
             latent_cache.contiguous(),
-            self.cp_size,
+            get_parallel().attn_cp_size,
             forward_batch,
             torch.cuda.current_stream(),
         )
@@ -2327,8 +2297,6 @@ class DeepseekV2DecoderLayer(nn.Module):
         prefix: str = "",
         alt_stream: Optional[torch.cuda.Stream] = None,
         skip_rope: bool = False,
-        dsa_enable_prefill_cp: bool = False,
-        mla_enable_prefill_cp: bool = False,
     ) -> None:
         super().__init__()
         self.hidden_size = config.hidden_size
@@ -2345,8 +2313,6 @@ class DeepseekV2DecoderLayer(nn.Module):
         self.speculative_algorithm = SpeculativeAlgorithm.from_string(
             get_spec().speculative_algorithm
         )
-        self.dsa_enable_prefill_cp = dsa_enable_prefill_cp
-        self.mla_enable_prefill_cp = mla_enable_prefill_cp
         self.layer_id = layer_id
         self.is_nextn = is_nextn
         if is_nextn and getattr(config, "mla_nope", False):
@@ -2374,8 +2340,6 @@ class DeepseekV2DecoderLayer(nn.Module):
             alt_stream=alt_stream,
             skip_rope=skip_rope,
             is_nextn=is_nextn,
-            dsa_enable_prefill_cp=dsa_enable_prefill_cp,
-            mla_enable_prefill_cp=mla_enable_prefill_cp,
         )
         if not hasattr(config, "q_lora_rank") and envs.SGLANG_USE_AG_AFTER_QLORA.get():
             raise ValueError(
@@ -2402,8 +2366,6 @@ class DeepseekV2DecoderLayer(nn.Module):
                 layer_id=self.layer_id,
                 alt_stream=alt_stream,
                 is_nextn=is_nextn,
-                dsa_enable_prefill_cp=dsa_enable_prefill_cp,
-                mla_enable_prefill_cp=mla_enable_prefill_cp,
             )
         else:
             if enable_moe_dense_fully_dp():
@@ -2428,31 +2390,21 @@ class DeepseekV2DecoderLayer(nn.Module):
 
         self._gfx95_quant_format = self._detect_gfx95_quant_format()
 
-        if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp:
-            # DSACPLayerCommunicator is flavor-agnostic; its internal gates
-            # read both dsa_use_prefill_cp and mla_use_prefill_cp. The rename
-            # to CPLayerCommunicator is deferred to a cleanup PR.
-            self.layer_communicator = DSACPLayerCommunicator(
-                layer_scatter_modes=self.layer_scatter_modes,
-                input_layernorm=self.input_layernorm,
-                post_attention_layernorm=self.post_attention_layernorm,
-                allow_reduce_scatter=True,
-                is_last_layer=(
-                    is_nextn or (self.layer_id == self.config.num_hidden_layers - 1)
-                ),
-                qkv_latent_func=self.self_attn.prepare_qkv_latent,
-            )
-        else:
-            self.layer_communicator = LayerCommunicator(
-                layer_scatter_modes=self.layer_scatter_modes,
-                input_layernorm=self.input_layernorm,
-                post_attention_layernorm=self.post_attention_layernorm,
-                allow_reduce_scatter=True,
-                is_last_layer=(
-                    is_nextn or (self.layer_id == self.config.num_hidden_layers - 1)
-                ),
-                qkv_latent_func=self.self_attn.prepare_qkv_latent,
-            )
+        communicator_cls = (
+            DSACPLayerCommunicator
+            if get_parallel().enable_prefill_cp
+            else LayerCommunicator
+        )
+        self.layer_communicator = communicator_cls(
+            layer_scatter_modes=self.layer_scatter_modes,
+            input_layernorm=self.input_layernorm,
+            post_attention_layernorm=self.post_attention_layernorm,
+            allow_reduce_scatter=True,
+            is_last_layer=(
+                is_nextn or (self.layer_id == self.config.num_hidden_layers - 1)
+            ),
+            qkv_latent_func=self.self_attn.prepare_qkv_latent,
+        )
 
     def _detect_gfx95_quant_format(self) -> str:
         if not _is_gfx95_supported:
@@ -2577,10 +2529,7 @@ class DeepseekV2DecoderLayer(nn.Module):
                     gemm_output_zero_allocator,
                 )
 
-        if (
-            not (self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp)
-            and fuse_mlp_allreduce
-        ):
+        if fuse_mlp_allreduce:
             hidden_states._sglang_needs_allreduce_fusion = True
 
         if not fuse_mlp_allreduce:
@@ -2666,14 +2615,6 @@ class DeepseekV2Model(nn.Module):
         self.vocab_size = config.vocab_size
         self.first_k_dense_replace = config.first_k_dense_replace
         self.pp_group = get_pp_group()
-        self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
-        self.mla_enable_prefill_cp = (
-            is_prefill_context_parallel_enabled() and not self.use_dsa
-        )
-        if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp:
-            self.cp_size = get_parallel().attn_cp_size
-        else:
-            self.cp_size = None
 
         if self.pp_group.is_first_rank:
             self.embed_tokens = VocabParallelEmbedding(
@@ -2704,8 +2645,6 @@ class DeepseekV2Model(nn.Module):
                 prefix=prefix,
                 alt_stream=self.alt_stream,
                 skip_rope=config.qk_rope_head_dim == 0,
-                dsa_enable_prefill_cp=self.dsa_enable_prefill_cp,
-                mla_enable_prefill_cp=self.mla_enable_prefill_cp,
             ),
             pp_rank=self.pp_group.rank_in_group,
             pp_size=self.pp_group.world_size,
@@ -2864,13 +2803,11 @@ class DeepseekV2Model(nn.Module):
             else None
         )
 
-        # CP-v2 shards/gathers at the eager-runner boundary instead.
-        use_cp_v1 = (
-            dsa_use_prefill_cp(forward_batch, self.dsa_enable_prefill_cp)
-            or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp)
-        ) and not is_cp_v2_active(forward_batch)
-
-        if use_cp_v1:
+        # HIP/NPU/MUSA retain their model-side CP boundary.
+        use_platform_cp = not enable_cp_v2() and (
+            dsa_use_prefill_cp(forward_batch) or mla_use_prefill_cp(forward_batch)
+        )
+        if use_platform_cp:
             if self.pp_group.is_first_rank:
                 hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
             positions = cp_split_and_rebuild_position(forward_batch, positions)
@@ -2974,11 +2911,11 @@ class DeepseekV2Model(nn.Module):
                 else:
                     hidden_states, _ = self.norm(hidden_states, residual)
 
-        if self.pp_group.is_last_rank and use_cp_v1:
+        if self.pp_group.is_last_rank and use_platform_cp:
             # allgather + rerrange
             hidden_states = cp_all_gather_rerange_output(
                 hidden_states,
-                self.cp_size,
+                get_parallel().attn_cp_size,
                 forward_batch,
                 torch.cuda.current_stream(),
             )
@@ -3051,16 +2988,6 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
         )
         self.capture_aux_hidden_states = False
 
-        self.dsa_enable_prefill_cp = is_dsa_enable_prefill_cp()
-        self.mla_enable_prefill_cp = (
-            is_prefill_context_parallel_enabled() and not is_deepseek_dsa(config)
-        )
-        if self.dsa_enable_prefill_cp or self.mla_enable_prefill_cp:
-            self.cp_rank = get_parallel().attn_cp_rank
-            self.cp_size = get_parallel().attn_cp_size
-        else:
-            self.cp_rank = self.cp_size = None
-
         q_lora_rank = config.q_lora_rank if hasattr(config, "q_lora_rank") else None
         get_attn_tp_context().init_context(q_lora_rank, is_deepseek_dsa(config))
 
@@ -3161,24 +3088,29 @@ class DeepseekV2ForCausalLM(nn.Module, DeepseekV2WeightLoaderMixin):
             len_input_ids = input_embeds.shape[0]
         else:
             len_input_ids = pp_proxy_tensors["hidden_states"].shape[0]
-        if not is_cp_v2_active(forward_batch):
-            if self.dsa_enable_prefill_cp:
+        if not enable_cp_v2():
+            if is_dsa_enable_prefill_cp():
                 if can_dsa_cp_split(
-                    len_input_ids, self.cp_size, self.use_dsa, forward_batch
+                    len_input_ids,
+                    get_parallel().attn_cp_size,
+                    self.use_dsa,
+                    forward_batch,
                 ):
                     forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
                         len_input_ids,
-                        self.cp_rank,
-                        self.cp_size,
+                        get_parallel().attn_cp_rank,
+                        get_parallel().attn_cp_size,
                         forward_batch.seq_lens_cpu.tolist(),
                         extend_seqs_len=forward_batch.extend_seq_lens_cpu,
                     )
-            elif self.mla_enable_prefill_cp:
-                if can_cp_split(len_input_ids, self.cp_size, forward_batch):
+            elif is_prefill_context_parallel_enabled() and not self.use_dsa:
+                if can_cp_split(
+                    len_input_ids, get_parallel().attn_cp_size, forward_batch
+                ):
                     forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
                         len_input_ids,
-                        self.cp_rank,
-                        self.cp_size,
+                        get_parallel().attn_cp_rank,
+                        get_parallel().attn_cp_size,
                         forward_batch.seq_lens_cpu.tolist(),
                         extend_seqs_len=forward_batch.extend_seq_lens_cpu,
                     )

python/sglang/srt/models/deepseek_v4.py

@@ -62,7 +62,7 @@ from sglang.srt.layers.communicator_dsa_cp import (
 from sglang.srt.layers.cp.cp_decode_attn_tp import get_cp_decode_attn_tp_ctx
 from sglang.srt.layers.cp.utils import (
     cp_materialize_global_token_order,
-    cp_round_robin_input_ids_v2,
+    enable_cp_v2,
     is_cp_v2_active,
 )
 from sglang.srt.layers.dp_attention import (
@@ -3223,8 +3223,6 @@ class DeepseekV4Model(nn.Module):
         input_embeds: Optional[torch.Tensor],
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> Union[torch.Tensor, PPProxyTensors]:
-        cp_v2_active = is_cp_v2_active(forward_batch)
-        use_prefill_cp = dsa_use_prefill_cp(forward_batch)
         if self.pp_group.is_first_rank:
             if input_embeds is None:
                 hidden_states = self.embed_tokens(input_ids)
@@ -3254,7 +3252,7 @@ class DeepseekV4Model(nn.Module):
             )
             input_ids_global = input_ids_global.squeeze(-1)
         else:
-            input_ids_global = input_ids
+            input_ids_global = getattr(forward_batch, "input_ids_global", input_ids)
 
         capture_dspark = self.dspark_layers_to_capture is not None
         dspark_aux_hidden_states: List[torch.Tensor] = []
@@ -3262,16 +3260,12 @@ class DeepseekV4Model(nn.Module):
         # execution cannot expose per-layer completed hidden states), so skip
         # TBO when capturing -- a perf-only downgrade, not a correctness one.
         run_tbo = self._can_run_tbo(forward_batch) and not capture_dspark
-        if use_prefill_cp and not run_tbo:
-            if cp_v2_active:
-                input_ids = cp_round_robin_input_ids_v2(input_ids, forward_batch)
-            else:
-                if self.pp_group.is_first_rank:
-                    hidden_states = cp_split_and_rebuild_data(
-                        forward_batch, hidden_states
-                    )
-                positions = cp_split_and_rebuild_position(forward_batch, positions)
-                input_ids = cp_round_robin_input_ids(input_ids)
+        use_platform_cp = not enable_cp_v2() and dsa_use_prefill_cp(forward_batch)
+        if use_platform_cp and not run_tbo:
+            if self.pp_group.is_first_rank:
+                hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
+            positions = cp_split_and_rebuild_position(forward_batch, positions)
+            input_ids = cp_round_robin_input_ids(input_ids)
             input_ids_global = input_ids
 
         # Reset Compressor's per-step freqs_cis cache from any previous step.
@@ -3323,12 +3317,7 @@ class DeepseekV4Model(nn.Module):
                 )
 
         # CP all-gather only on the last PP rank; PP IPC carries CP-split tensors.
-        if (
-            self.pp_group.is_last_rank
-            and use_prefill_cp
-            and not cp_v2_active
-            and not run_tbo
-        ):
+        if self.pp_group.is_last_rank and use_platform_cp and not run_tbo:
             stream = torch.cuda.current_stream()
             hidden_states = cp_all_gather_rerange_output(
                 hidden_states,
@@ -3489,7 +3478,7 @@ class DeepseekV4ForCausalLM(nn.Module):
         input_embeds: Optional[torch.Tensor] = None,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> torch.Tensor:
-        if self.dsa_enable_prefill_cp:
+        if not enable_cp_v2() and self.dsa_enable_prefill_cp:
             if can_dsa_cp_split(len(input_ids), self.cp_size, True, forward_batch):
                 forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
                     len(input_ids),

python/sglang/srt/models/deepseek_v4_nextn.py

@@ -14,8 +14,7 @@ from sglang.srt.layers.attention.dsa.utils import (
     is_dsa_prefill_cp_round_robin_split,
 )
 from sglang.srt.layers.cp.utils import (
-    cp_round_robin_input_ids_v2,
-    is_cp_v2_active,
+    enable_cp_v2,
 )
 from sglang.srt.layers.dp_attention import (
     dp_gather_replicate,
@@ -144,8 +143,7 @@ class DeepseekV4ModelNextN(nn.Module):
         forward_batch: ForwardBatch,
         input_embeds: torch.Tensor = None,
     ) -> torch.Tensor:
-        cp_v2_active = is_cp_v2_active(forward_batch)
-        use_prefill_cp = dsa_use_prefill_cp(forward_batch)
+        use_platform_cp = not enable_cp_v2() and dsa_use_prefill_cp(forward_batch)
         if input_embeds is None:
             hidden_states = self.embed_tokens(input_ids)
         else:
@@ -179,15 +177,12 @@ class DeepseekV4ModelNextN(nn.Module):
             )
             input_ids_global = input_ids_global.squeeze(-1)
         else:
-            input_ids_global = input_ids
+            input_ids_global = getattr(forward_batch, "input_ids_global", input_ids)
 
-        if use_prefill_cp:
-            if cp_v2_active:
-                input_ids = cp_round_robin_input_ids_v2(input_ids, forward_batch)
-            else:
-                hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
-                positions = cp_split_and_rebuild_position(forward_batch, positions)
-                input_ids = cp_round_robin_input_ids(input_ids)
+        if use_platform_cp:
+            hidden_states = cp_split_and_rebuild_data(forward_batch, hidden_states)
+            positions = cp_split_and_rebuild_position(forward_batch, positions)
+            input_ids = cp_round_robin_input_ids(input_ids)
             input_ids_global = input_ids
 
         hidden_states, residual, post, comb = self.decoder(
@@ -202,7 +197,7 @@ class DeepseekV4ModelNextN(nn.Module):
             # deferred fused hc_post state.
             hidden_states = self.decoder.hc_post(hidden_states, residual, post, comb)
 
-        if use_prefill_cp and not cp_v2_active:
+        if use_platform_cp:
             hidden_states = cp_all_gather_rerange_output(
                 hidden_states,
                 self.cp_size,
@@ -260,7 +255,7 @@ class DeepseekV4ForCausalLMNextN(DeepseekV4ForCausalLM):
         positions: torch.Tensor,
         forward_batch: ForwardBatch,
     ) -> torch.Tensor:
-        if self.dsa_enable_prefill_cp and not is_cp_v2_active(forward_batch):
+        if self.dsa_enable_prefill_cp and not enable_cp_v2():
             if can_dsa_cp_split(len(input_ids), self.cp_size, True, forward_batch):
                 forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
                     len(input_ids),

python/sglang/srt/models/qwen2_moe.py

@@ -48,7 +48,7 @@ from sglang.srt.layers.communicator import (
     LayerScatterModes,
     ScatterMode,
 )
-from sglang.srt.layers.cp.utils import is_cp_v2_active
+from sglang.srt.layers.cp.utils import enable_cp_v2
 from sglang.srt.layers.dp_attention import (
     is_dp_attention_enabled,
 )
@@ -1132,7 +1132,7 @@ class Qwen2MoeModel(nn.Module):
 
         if (
             is_prefill_context_parallel_enabled()
-            and not is_cp_v2_active(forward_batch)
+            and not enable_cp_v2()
             and forward_batch.forward_mode.is_context_parallel_extend()
             and forward_batch.attn_cp_metadata is not None
         ):
@@ -1198,7 +1198,7 @@ class Qwen2MoeModel(nn.Module):
 
         if (
             self.pp_group.is_last_rank
-            and not is_cp_v2_active(forward_batch)
+            and not enable_cp_v2()
             and is_prefill_context_parallel_enabled()
             and forward_batch.forward_mode.is_context_parallel_extend()
             and forward_batch.attn_cp_metadata is not None

python/sglang/srt/models/qwen3_moe.py

@@ -34,7 +34,7 @@ from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_r
 from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
 from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
 from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
-from sglang.srt.layers.cp.utils import is_cp_v2_active
+from sglang.srt.layers.cp.utils import enable_cp_v2
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
     QKVParallelLinear,
@@ -995,7 +995,7 @@ class Qwen3MoeForCausalLM(nn.Module):
         input_embeds: torch.Tensor = None,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> torch.Tensor:
-        if is_prefill_context_parallel_enabled() and not is_cp_v2_active(forward_batch):
+        if is_prefill_context_parallel_enabled() and not enable_cp_v2():
             if can_cp_split(len(input_ids), self.attn_cp_size, forward_batch):
                 forward_batch.attn_cp_metadata = prepare_context_parallel_metadata(
                     len(input_ids),

python/sglang/srt/models/utils.py

@@ -30,14 +30,14 @@ from sglang.kernels.ops.layernorm.norm import (
     fused_inplace_qknorm,
 )
 from sglang.srt.environ import envs
+from sglang.srt.layers.cp.utils import is_cp_v2_active
 from sglang.srt.layers.radix_attention import RadixAttention
-from sglang.srt.layers.utils.cp_utils import is_prefill_context_parallel_enabled
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
 from sglang.srt.model_executor.runner import get_is_capture_mode
 from sglang.srt.model_loader.weight_utils import default_weight_loader
-from sglang.srt.runtime_context import get_exec
+from sglang.srt.runtime_context import get_exec, get_parallel
 from sglang.srt.utils import get_current_device_stream_fast, is_cuda, is_hip
 from sglang.srt.utils.custom_op import register_custom_op
 
@@ -296,11 +296,11 @@ def enable_fused_set_kv_buffer(forward_batch: ForwardBatch):
         _is_cuda
         and pool.dtype == torch.bfloat16
         and not isinstance(pool, SWAKVPool)
-        and not is_prefill_context_parallel_enabled()
+        and not is_cp_v2_active(forward_batch)
         and getattr(forward_batch, "dcp_kv_mask", None) is None
     ) or (
         _is_hip
-        and not is_prefill_context_parallel_enabled()
+        and not get_parallel().enable_prefill_context_parallel
         and getattr(forward_batch, "dcp_kv_mask", None) is None
     )
 

python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py

@@ -342,12 +342,11 @@ class DSAMockModelRunner(ModelRunner):
             dllm_algorithm_config=None,
             dp_size=1,
             dsa_decode_backend=dsa_decode_backend,
-            dsa_prefill_cp_mode="round-robin-split",
             dsa_prefill_backend=dsa_prefill_backend,
             device=device,
             enable_deterministic_inference=False,
             enable_dp_attention=False,
-            enable_dsa_prefill_context_parallel=False,
+            enable_prefill_cp=False,
             enable_mis=False,
             is_embedding=False,
             kv_cache_dtype="auto",