[3/N][CP] Implement zigzag CP strategy

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

@@ -13,6 +13,8 @@
     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 @@ def forward_extend(
                 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 @@ def _fa_cp_attn(
                         **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 @@ def materialize_full_kv(
         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 @@ def run_attention(
             "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,101 @@
     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."""
+    strategy = get_cp_strategy()
+    assert strategy is not None
+
+    input_ids = getattr(forward_batch, "input_ids", None)
+    assert input_ids is not None
+
+    num_tokens = len(input_ids)
+    assert strategy.can_apply(num_tokens, forward_batch)
+
+    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."""
+    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 maybe_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."""
+    if not is_cp_v2_active(forward_batch):
+        return x
+
+    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 +138,11 @@
     "InterleaveContextParallelMetadata",
     "ZigzagCPStrategy",
     "ZigzagContextParallelMetadata",
+    "CP_V2_DEFAULT_MODEL_CLASSES",
+    "enable_cp_v2",
+    "get_cp_strategy",
+    "is_cp_v2_active",
+    "maybe_cp_gather_after_forward",
+    "cp_split_before_forward",
+    "prepare_cp_forward",
 ]