[3/N][CP] Implement zigzag CP strategy (#28421)

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

@@ -13,6 +13,8 @@ from sglang.srt.layers.attention.triton_ops.metadata import (
     prepare_swa_spec_page_table_triton,
 )
 from sglang.srt.layers.attention.utils import assert_buffer_fits
+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.radix_attention import AttentionType
 from sglang.srt.layers.utils.cp_utils import (
     cp_allgather_and_save_kv_cache,
@@ -811,18 +813,26 @@ class FlashAttentionBackend(AttentionBackend):
                 elif is_cp_mode:
                     # Dense-MHA CP: k, v are still rank-local; backend
                     # all-gathers and writes to the per-rank pool.
-                    cp_allgather_and_save_kv_cache(
-                        forward_batch,
-                        layer,
-                        k,
-                        v,
-                        self.attn_cp_size,
-                        swa_loc=(
-                            self.forward_metadata.swa_out_cache_loc
-                            if self.use_sliding_window_kv_pool
-                            else None
-                        ),
+                    swa_loc = (
+                        self.forward_metadata.swa_out_cache_loc
+                        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,
+                        )
                 else:
                     self.token_to_kv_pool.set_kv_buffer(
                         layer,
@@ -967,12 +977,24 @@ class FlashAttentionBackend(AttentionBackend):
                         **kwargs,
                     )
 
-                result = cp_attn_forward_extend(
-                    forward_batch,
-                    q.contiguous().view(-1, layer.tp_q_head_num, layer.head_dim),
-                    self.device,
-                    _fa_cp_attn,
-                )
+                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,
+                    )
             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

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

@@ -185,6 +185,7 @@ class ContextParallelStrategy(ABC):
         layer: Any,
         k: Any,
         v: Any,
+        swa_loc: Optional[Any] = None,
     ) -> None:
         """Write full-layout K/V to the backend cache if needed."""
 
@@ -235,7 +236,7 @@ def init_cp_strategy(server_args: ServerArgs) -> None:
         )
 
 
-def _get_cp_strategy() -> Optional[ContextParallelStrategy]:
+def get_cp_strategy() -> Optional[ContextParallelStrategy]:
     """Return the configured strategy, initializing lazily on first call.
 
     Subprocesses re-import this module with ``_STRATEGY = None`` and never
@@ -257,20 +258,15 @@ def _get_cp_strategy() -> Optional[ContextParallelStrategy]:
     return _STRATEGY
 
 
-def get_cp_strategy() -> Optional[ContextParallelStrategy]:
-    """Return the configured CP strategy for runtime dispatch."""
-    return _get_cp_strategy()
-
-
 def get_cp_strategy_kind() -> ContextParallelStrategyKind:
-    strategy = _get_cp_strategy()
+    strategy = get_cp_strategy()
     if strategy is None:
         return ContextParallelStrategyKind.NONE
     return strategy.kind
 
 
 def is_cp_enabled() -> bool:
-    return _get_cp_strategy() is not None
+    return get_cp_strategy() is not None
 
 
 def is_zigzag() -> bool:

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

@@ -99,7 +99,9 @@ class InterleaveCPStrategy(ContextParallelStrategy):
             "Interleave attention dispatch will land in a follow-up PR"
         )
 
-    def materialize_full_kv(self, forward_batch, layer: Any, k: Any, v: Any) -> None:
+    def materialize_full_kv(
+        self, forward_batch, layer: Any, k: Any, v: Any, swa_loc: Optional[Any] = None
+    ) -> None:
         raise NotImplementedError(
             "Interleave KV materialization will land in a follow-up PR"
         )

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

@@ -12,13 +12,16 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Public import facade for context parallel strategy helpers."""
+"""Public import facade and runtime helpers for context parallel strategies."""
+
+from typing import Any, Optional, Tuple
 
 from sglang.srt.layers.cp.base import (
     BaseContextParallelMetadata,
     ContextParallelStrategy,
     ContextParallelStrategyKind,
     CPAttentionBackendKind,
+    get_cp_strategy,
 )
 from sglang.srt.layers.cp.interleave import (
     InterleaveContextParallelMetadata,
@@ -30,6 +33,96 @@ from sglang.srt.layers.cp.zigzag import (
     ZigzagCPStrategy,
 )
 
+CP_V2_DEFAULT_MODEL_CLASSES = frozenset(
+    {
+        "Qwen3MoeForCausalLM",
+    }
+)
+
+
+def enable_cp_v2() -> bool:
+    """Return whether the CP-v2 path is enabled for this process."""
+    from sglang.srt.environ import envs
+
+    return bool(envs.SGLANG_ENABLE_CP_V2.get())
+
+
+def is_cp_v2_active(forward_batch) -> bool:
+    """Return whether the current forward batch is running through CP-v2."""
+    if not enable_cp_v2():
+        return False
+    forward_mode = getattr(forward_batch, "forward_mode", None)
+    if forward_mode is None or not forward_mode.is_context_parallel_extend():
+        return False
+
+    strategy = get_cp_strategy()
+    if strategy is None:
+        return False
+
+    input_ids = getattr(forward_batch, "input_ids", None)
+    if input_ids is None:
+        return False
+
+    return strategy.can_apply(len(input_ids), forward_batch)
+
+
+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)
+    strategy = get_cp_strategy()
+    assert strategy is not None
+    num_tokens = len(forward_batch.input_ids)
+
+    seq_lens_cpu = _to_int_list(getattr(forward_batch, "seq_lens_cpu", None))
+    extend_lens_cpu = _to_int_list(getattr(forward_batch, "extend_seq_lens_cpu", None))
+    forward_batch.attn_cp_metadata = strategy.build_metadata(
+        num_tokens=num_tokens,
+        seqs_len=seq_lens_cpu,
+        extend_seqs_len=extend_lens_cpu,
+    )
+
+
+def cp_split_before_forward(
+    complete_hidden_states: Any,
+    complete_position_ids: Any,
+    forward_batch,
+) -> Tuple[Optional[Any], Optional[Any]]:
+    """Shard embeddings and positions for CP-v2 model-runner forwarding."""
+    assert is_cp_v2_active(forward_batch)
+    strategy = get_cp_strategy()
+    assert strategy is not None
+    assert complete_hidden_states is not None
+    assert getattr(forward_batch, "attn_cp_metadata", None) is not None
+    return (
+        strategy.shard_hidden_states(complete_hidden_states, forward_batch),
+        strategy.shard_position_ids(complete_position_ids, forward_batch),
+    )
+
+
+def cp_gather_after_forward(x: Any, forward_batch, stream: Optional[Any] = None):
+    """Gather CP-v2 hidden states at the model boundary when this batch is active."""
+    assert is_cp_v2_active(forward_batch)
+    strategy = get_cp_strategy()
+    assert strategy is not None
+
+    if isinstance(x, tuple):
+        hidden_states, *rest = x
+        hidden_states = strategy.gather_hidden_states(
+            hidden_states, forward_batch, stream
+        )
+        return (hidden_states, *rest)
+
+    return strategy.gather_hidden_states(x, forward_batch, stream)
+
+
+def _to_int_list(values) -> Optional[list[int]]:
+    if values is None:
+        return None
+    if hasattr(values, "tolist"):
+        values = values.tolist()
+    return [int(x) for x in values]
+
+
 __all__ = [
     "BaseContextParallelMetadata",
     "CPAttentionBackendKind",
@@ -40,4 +133,11 @@ __all__ = [
     "InterleaveContextParallelMetadata",
     "ZigzagCPStrategy",
     "ZigzagContextParallelMetadata",
+    "CP_V2_DEFAULT_MODEL_CLASSES",
+    "enable_cp_v2",
+    "get_cp_strategy",
+    "is_cp_v2_active",
+    "cp_gather_after_forward",
+    "cp_split_before_forward",
+    "prepare_cp_forward",
 ]

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

@@ -30,15 +30,29 @@ After all-gather, the blocks are reranged back to their original order:
 
 from __future__ import annotations
 
+from contextlib import nullcontext
 from dataclasses import dataclass
+from itertools import accumulate
 from typing import Any, List, Optional
 
+import torch
+import torch.nn.functional as F
+
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    use_symmetric_memory,
+)
 from sglang.srt.layers.cp.base import (
     BaseContextParallelMetadata,
     ContextParallelStrategy,
     ContextParallelStrategyKind,
     CPAttentionBackendKind,
 )
+from sglang.srt.layers.dp_attention import (
+    get_attention_cp_group,
+    is_allocation_symmetric,
+)
+from sglang.srt.mem_cache.memory_pool import KVWriteLoc
+from sglang.srt.model_executor.forward_context import get_token_to_kv_pool
 
 
 @dataclass
@@ -85,7 +99,13 @@ class ZigzagCPStrategy(ContextParallelStrategy):
         if self.cp_size <= 1 or num_tokens < self.cp_size * 2:
             return False
         forward_mode = getattr(forward_batch, "forward_mode", None)
-        return forward_mode is None or forward_mode.is_context_parallel_extend()
+        if forward_mode is not None and not forward_mode.is_context_parallel_extend():
+            return False
+
+        extend_lens = getattr(forward_batch, "extend_seq_lens_cpu", None)
+        if extend_lens is None:
+            return True
+        return all(int(length) >= self.cp_size * 2 for length in extend_lens)
 
     def build_metadata(
         self,
@@ -93,32 +113,193 @@ class ZigzagCPStrategy(ContextParallelStrategy):
         seqs_len: Optional[List[int]],
         extend_seqs_len: Optional[List[int]] = None,
     ) -> ZigzagContextParallelMetadata:
+        if extend_seqs_len is None:
+            extend_seqs_len = seqs_len or [num_tokens]
+        extend_seqs_len = [int(x) for x in extend_seqs_len]
+
+        pad_len = int(num_tokens) - sum(extend_seqs_len)
+        if pad_len > 0:
+            extend_seqs_len[-1] += pad_len
+            if seqs_len is not None and len(seqs_len) == len(extend_seqs_len):
+                seqs_len = list(seqs_len)
+                seqs_len[-1] += pad_len
+
+        bs = len(extend_seqs_len)
+        cp_segment_num = self.cp_size * 2
+        if seqs_len is not None and len(seqs_len) == bs:
+            prefix_offsets = [
+                max(int(seqs_len[i]) - extend_seqs_len[i], 0) for i in range(bs)
+            ]
+        else:
+            prefix_offsets = [0] * bs
+
+        # TODO: move these per-request layout/index computations to a Triton
+        # kernel if Python-side metadata construction becomes a bottleneck.
+        per_seq_block_sizes: List[List[int]] = []
+        split_list: List[int] = []
+        for length in extend_seqs_len:
+            base = length // cp_segment_num
+            rem = length % cp_segment_num
+            block_sizes = [
+                base + 1 if block_id < rem else base
+                for block_id in range(cp_segment_num)
+            ]
+            per_seq_block_sizes.append(block_sizes)
+            split_list.extend(block_sizes)
+
+        per_rank_actual_token = []
+        for rank in range(self.cp_size):
+            per_rank_actual_token.append(
+                sum(
+                    block_sizes[rank] + block_sizes[cp_segment_num - 1 - rank]
+                    for block_sizes in per_seq_block_sizes
+                )
+            )
+        max_rank_len = [max(per_rank_actual_token)] * self.cp_size
+
+        cp_rank = self.cp_rank
+        zigzag_index = list(
+            range(cp_rank, cp_rank + bs * cp_segment_num, cp_segment_num)
+        ) + list(
+            range(
+                cp_segment_num - cp_rank - 1,
+                bs * cp_segment_num,
+                cp_segment_num,
+            )
+        )
+
+        cp_reverse_index: List[int] = []
+        for batch_id in range(bs):
+            cp_reverse_index.extend(
+                list(range(batch_id, cp_segment_num * bs, 2 * bs))
+                + list(
+                    range(
+                        (cp_segment_num - 1) * bs + batch_id,
+                        0,
+                        -2 * bs,
+                    )
+                )
+            )
+
+        reverse_split_len: List[int] = []
+        for rank in range(self.cp_size):
+            for batch_id in range(bs):
+                reverse_split_len.append(per_seq_block_sizes[batch_id][rank])
+            for batch_id in range(bs):
+                reverse_split_len.append(
+                    per_seq_block_sizes[batch_id][cp_segment_num - 1 - rank]
+                )
+
+        kv_len_prev_list: List[int] = []
+        kv_len_next_list: List[int] = []
+        actual_seq_q_prev_list: List[int] = []
+        actual_seq_q_next_list: List[int] = []
+        for batch_id, block_sizes in enumerate(per_seq_block_sizes):
+            kv_len_prev_list.append(
+                prefix_offsets[batch_id] + sum(block_sizes[: cp_rank + 1])
+            )
+            kv_len_next_list.append(
+                prefix_offsets[batch_id] + sum(block_sizes[: cp_segment_num - cp_rank])
+            )
+            actual_seq_q_prev_list.append(block_sizes[cp_rank])
+            actual_seq_q_next_list.append(block_sizes[cp_segment_num - cp_rank - 1])
+
+        from sglang.srt.server_args import get_global_server_args
+
+        try:
+            device = torch.device(get_global_server_args().device)
+        except Exception:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        cu_prev = [0] + list(accumulate(actual_seq_q_prev_list))
+        cu_next = [0] + list(accumulate(actual_seq_q_next_list))
+
+        total_seq_lens = sum(extend_seqs_len)
+        assert len(split_list) == bs * cp_segment_num
+        assert sum(split_list) == total_seq_lens
+        assert len(zigzag_index) == 2 * bs
+        assert len(cp_reverse_index) == bs * cp_segment_num
+        assert sorted(cp_reverse_index) == list(range(bs * cp_segment_num))
+        assert sum(per_rank_actual_token) == total_seq_lens
+
         return ZigzagContextParallelMetadata(
-            total_seq_lens=sum(extend_seqs_len or seqs_len or [num_tokens]),
-            bs=len(extend_seqs_len or seqs_len or [num_tokens]),
+            split_list=split_list,
+            zigzag_index=zigzag_index,
+            cp_reverse_index=cp_reverse_index,
+            reverse_split_len=reverse_split_len,
+            per_rank_actual_token=per_rank_actual_token,
+            max_rank_len=max_rank_len,
+            kv_len_prev_tensor=torch.tensor(
+                kv_len_prev_list, device=device, dtype=torch.int32
+            ),
+            kv_len_next_tensor=torch.tensor(
+                kv_len_next_list, device=device, dtype=torch.int32
+            ),
+            actual_seq_q_prev_tensor=torch.tensor(
+                actual_seq_q_prev_list, device=device, dtype=torch.int32
+            ),
+            actual_seq_q_next_tensor=torch.tensor(
+                actual_seq_q_next_list, device=device, dtype=torch.int32
+            ),
+            cu_seqlens_q_prev_tensor=torch.tensor(
+                cu_prev, device=device, dtype=torch.int32
+            ),
+            cu_seqlens_q_next_tensor=torch.tensor(
+                cu_next, device=device, dtype=torch.int32
+            ),
+            total_q_prev_tokens=cu_prev[-1],
+            total_q_next_tokens=cu_next[-1],
+            max_seqlen_q_prev=(
+                max(actual_seq_q_prev_list) if actual_seq_q_prev_list else 0
+            ),
+            max_seqlen_q_next=(
+                max(actual_seq_q_next_list) if actual_seq_q_next_list else 0
+            ),
+            kv_len_prev_list=kv_len_prev_list,
+            kv_len_next_list=kv_len_next_list,
+            actual_seq_q_prev_list=actual_seq_q_prev_list,
+            actual_seq_q_next_list=actual_seq_q_next_list,
+            total_seq_lens=total_seq_lens,
+            bs=bs,
         )
 
     def shard_hidden_states(self, x: Any, forward_batch) -> Any:
-        raise NotImplementedError(
-            "Zigzag hidden-state sharding will land in a follow-up PR"
+        chunks = torch.split(x, forward_batch.attn_cp_metadata.split_list, dim=0)
+        return torch.cat(
+            [chunks[i] for i in forward_batch.attn_cp_metadata.zigzag_index], dim=0
         )
 
     def shard_position_ids(self, positions: Any, forward_batch) -> Any:
-        raise NotImplementedError(
-            "Zigzag position-id sharding will land in a follow-up PR"
+        chunks = torch.split(
+            positions, forward_batch.attn_cp_metadata.split_list, dim=-1
+        )
+        return torch.cat(
+            [chunks[i] for i in forward_batch.attn_cp_metadata.zigzag_index], dim=-1
         )
 
     def gather_hidden_states(
         self, x: Any, forward_batch, stream: Optional[Any] = None
     ) -> Any:
-        raise NotImplementedError(
-            "Zigzag hidden-state gather will land in a follow-up PR"
+        gathered = self._all_gather_reorganized(x, forward_batch, stream)
+        chunks = torch.split(
+            gathered, forward_batch.attn_cp_metadata.reverse_split_len, dim=0
+        )
+        return torch.cat(
+            [chunks[i] for i in forward_batch.attn_cp_metadata.cp_reverse_index], dim=0
         )
 
     def gather_kv_cache(
         self, x: Any, forward_batch, stream: Optional[Any] = None
     ) -> Any:
-        raise NotImplementedError("Zigzag KV gather will land in a follow-up PR")
+        gathered = self._all_gather_reorganized(x, forward_batch, stream)
+        chunks = torch.split(
+            gathered, forward_batch.attn_cp_metadata.reverse_split_len, dim=0
+        )
+        return torch.cat(
+            [chunks[i] for i in forward_batch.attn_cp_metadata.cp_reverse_index], dim=0
+        )
+
+    def get_supported_attention_backend(self):
+        return [CPAttentionBackendKind.FLASH_ATTENTION]
 
     def run_attention(
         self,
@@ -128,11 +309,79 @@ class ZigzagCPStrategy(ContextParallelStrategy):
         attn_fn,
         attention_backend: CPAttentionBackendKind = CPAttentionBackendKind.FLASH_ATTENTION,
     ) -> Any:
-        raise NotImplementedError(
-            "Zigzag attention dispatch will land in a follow-up PR"
+        assert (
+            attention_backend in self.get_supported_attention_backend()
+        ), f"{self.name} CP does not support {attention_backend=}"
+
+        meta = forward_batch.attn_cp_metadata
+        q_prev = q[: meta.total_q_prev_tokens]
+        q_next = q[meta.total_q_prev_tokens :]
+
+        result_prev = attn_fn(
+            q_prev,
+            meta.cu_seqlens_q_prev_tensor,
+            meta.kv_len_prev_tensor,
+            meta.max_seqlen_q_prev,
+        )
+        result_next = attn_fn(
+            q_next,
+            meta.cu_seqlens_q_next_tensor,
+            meta.kv_len_next_tensor,
+            meta.max_seqlen_q_next,
+        )
+        return torch.cat([result_prev, result_next], dim=0)
+
+    def materialize_full_kv(
+        self, forward_batch, layer: Any, k: Any, v: Any, swa_loc: Optional[Any] = None
+    ) -> None:
+        cache_loc = (
+            forward_batch.out_cache_loc
+            if not layer.is_cross_attention
+            else forward_batch.encoder_out_cache_loc
+        )
+        key_cache_full = self.gather_kv_cache(
+            k.contiguous(), forward_batch, torch.cuda.current_stream()
+        )
+        value_cache_full = self.gather_kv_cache(
+            v.contiguous(), forward_batch, torch.cuda.current_stream()
+        )
+        get_token_to_kv_pool().set_kv_buffer(
+            layer,
+            KVWriteLoc(cache_loc, swa_loc),
+            key_cache_full,
+            value_cache_full,
+            layer.k_scale,
+            layer.v_scale,
         )
 
-    def materialize_full_kv(self, forward_batch, layer: Any, k: Any, v: Any) -> None:
-        raise NotImplementedError(
-            "Zigzag KV materialization will land in a follow-up PR"
+    def _all_gather_reorganized(self, x: torch.Tensor, forward_batch, stream):
+        meta = forward_batch.attn_cp_metadata
+        max_len = meta.max_rank_len[0]
+        pad_size = max_len - x.shape[0]
+        if pad_size > 0:
+            padding = [0, 0] * (x.ndim - 1) + [0, pad_size]
+            x = F.pad(x, padding, mode="constant", value=0)
+
+        group = get_attention_cp_group()
+        ctx = (
+            use_symmetric_memory(group, disabled=not is_allocation_symmetric())
+            if x.is_cuda
+            else nullcontext()
+        )
+        with ctx:
+            gathered = torch.empty(
+                max_len * self.cp_size,
+                *x.shape[1:],
+                device=x.device,
+                dtype=x.dtype,
+            )
+        group.cp_all_gather_into_tensor_async(gathered, x, stream)
+
+        chunks = torch.split(gathered, meta.max_rank_len, dim=0)
+        return torch.cat(
+            [
+                chunks[rank][:per_rank_len]
+                for rank, per_rank_len in enumerate(meta.per_rank_actual_token)
+            ],
+            dim=0,
         )

python/sglang/srt/model_executor/model_runner.py

@@ -123,6 +123,13 @@ from sglang.srt.layers.attention.attention_registry import (
 )
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.attention.tbo_backend import TboAttnBackend
+from sglang.srt.layers.cp.utils import (
+    cp_gather_after_forward,
+    cp_split_before_forward,
+    get_cp_strategy,
+    is_cp_v2_active,
+    prepare_cp_forward,
+)
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
     get_attention_tp_group,
@@ -3411,6 +3418,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             self.prefill_cuda_graph_runner is not None
             and self.prefill_cuda_graph_runner.can_run(forward_batch)
         )
+        if get_cp_strategy() is not None:
+            can_run_graph = False
         if can_run_graph:
             # TODO: device_timer.wrap is too broad here — it also includes
             # replay_prepare time. Move timing into the prefill cuda graph
@@ -3434,6 +3443,21 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 # e.g. Moss-VL's prefill cross-attention custom mask.
                 self.model.prepare_forward_batch(forward_batch)
             self.attn_backend.init_forward_metadata(forward_batch)
+        cp_v2_active = is_cp_v2_active(forward_batch)
+        forward_positions = forward_batch.positions
+        if cp_v2_active:
+            prepare_cp_forward(forward_batch)
+            complete_hidden_states = kwargs.get("input_embeds")
+            if complete_hidden_states is None:
+                embed_layer = self.model.get_input_embeddings()
+                complete_hidden_states = embed_layer(forward_batch.input_ids)
+            sharded_hidden_states, sharded_positions = cp_split_before_forward(
+                complete_hidden_states,
+                forward_batch.positions,
+                forward_batch,
+            )
+            kwargs["input_embeds"] = sharded_hidden_states
+            forward_positions = sharded_positions
 
         ctx = (
             self.device_timer.wrap(metadata={"category": "extend"})
@@ -3441,7 +3465,11 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             else contextlib.nullcontext()
         )
         with ctx:
-            if _is_hip and self.prefill_cuda_graph_runner is not None:
+            if (
+                _is_hip
+                and self.prefill_cuda_graph_runner is not None
+                and not cp_v2_active
+            ):
                 # AMD/HIP: when PCG is enabled but the batch exceeds max captured
                 # size, run eagerly under enable_tc_piecewise_cuda_graph() and
                 # set_tc_piecewise_forward_context() so that (a) Dynamo guards on
@@ -3461,14 +3489,47 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 ):
                     ret = self.model.forward(
                         forward_batch.input_ids,
-                        forward_batch.positions,
+                        forward_positions,
                         forward_batch,
                         **kwargs,
                     )
+            elif cp_v2_active:
+                hidden_states = self.model.model(
+                    forward_batch.input_ids,
+                    forward_positions,
+                    forward_batch,
+                    input_embeds=kwargs.get("input_embeds"),
+                    pp_proxy_tensors=kwargs.get("pp_proxy_tensors"),
+                )
+
+                aux_hidden_states = None
+                capture_aux_hidden_states = getattr(
+                    self.model, "capture_aux_hidden_states", False
+                )
+                if capture_aux_hidden_states:
+                    hidden_states, aux_hidden_states = hidden_states
+
+                if self.model.pp_group.is_last_rank:
+                    hidden_states = cp_gather_after_forward(
+                        hidden_states,
+                        forward_batch,
+                        torch.cuda.current_stream(),
+                    )
+                    ret = self.model.logits_processor(
+                        forward_batch.input_ids,
+                        hidden_states,
+                        self.model.lm_head,
+                        forward_batch,
+                        aux_hidden_states,
+                    )
+                elif capture_aux_hidden_states:
+                    ret = hidden_states, aux_hidden_states
+                else:
+                    ret = hidden_states
             else:
                 ret = self.model.forward(
                     forward_batch.input_ids,
-                    forward_batch.positions,
+                    forward_positions,
                     forward_batch,
                     **kwargs,
                 )

python/sglang/srt/models/qwen2_moe.py

@@ -42,6 +42,7 @@ from sglang.srt.layers.communicator import (
     LayerScatterModes,
     ScatterMode,
 )
+from sglang.srt.layers.cp.utils import is_cp_v2_active
 from sglang.srt.layers.dp_attention import (
     is_dp_attention_enabled,
 )
@@ -889,6 +890,7 @@ class Qwen2MoeModel(nn.Module):
 
         if (
             is_prefill_context_parallel_enabled()
+            and not is_cp_v2_active(forward_batch)
             and forward_batch.forward_mode.is_context_parallel_extend()
             and forward_batch.attn_cp_metadata is not None
         ):
@@ -944,6 +946,7 @@ class Qwen2MoeModel(nn.Module):
 
         if (
             self.pp_group.is_last_rank
+            and not is_cp_v2_active(forward_batch)
             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,6 +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.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
     QKVParallelLinear,
@@ -988,7 +989,7 @@ class Qwen3MoeForCausalLM(nn.Module):
         input_embeds: torch.Tensor = None,
         pp_proxy_tensors: Optional[PPProxyTensors] = None,
     ) -> torch.Tensor:
-        if is_prefill_context_parallel_enabled():
+        if is_prefill_context_parallel_enabled() and not is_cp_v2_active(forward_batch):
             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/server_args.py

@@ -3436,6 +3436,17 @@ class ServerArgs:
         self.prefill_cp_mode = mode
 
     def _handle_context_parallelism(self):
+        if parse_connector_type(self.model_path) != ConnectorType.INSTANCE:
+            from sglang.srt.layers.cp.utils import CP_V2_DEFAULT_MODEL_CLASSES
+
+            model_config = self.get_model_config()
+            model_arch = model_config.hf_config.architectures[0]
+            if (
+                model_arch in CP_V2_DEFAULT_MODEL_CLASSES
+                and not envs.SGLANG_ENABLE_CP_V2.is_set()
+            ):
+                envs.SGLANG_ENABLE_CP_V2.set(True)
+
         if self.enable_prefill_cp and self.cp_strategy is None:
             raise ValueError(
                 "--cp-strategy must be set when --enable-prefill-cp is enabled."