chore: import upstream snapshot with attribution

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

@@ -0,0 +1,55 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Context parallel strategy abstractions."""
+
+from sglang.srt.layers.cp.base import (
+    BaseContextParallelMetadata,
+    ContextParallelStrategy,
+    ContextParallelStrategyKind,
+    CPAttentionBackendKind,
+    get_cp_strategy,
+    get_cp_strategy_kind,
+    init_cp_strategy,
+    is_cp_enabled,
+    is_interleave,
+    is_zigzag,
+)
+from sglang.srt.layers.cp.interleave import (
+    InterleaveContextParallelMetadata,
+    InterleaveCPStrategy,
+)
+from sglang.srt.layers.cp.zigzag import (
+    ContextParallelMetadata,
+    ZigzagContextParallelMetadata,
+    ZigzagCPStrategy,
+)
+
+__all__ = [
+    "BaseContextParallelMetadata",
+    "CPAttentionBackendKind",
+    "ContextParallelMetadata",
+    "ContextParallelStrategy",
+    "ContextParallelStrategyKind",
+    "InterleaveCPStrategy",
+    "InterleaveContextParallelMetadata",
+    "ZigzagCPStrategy",
+    "ZigzagContextParallelMetadata",
+    "get_cp_strategy",
+    "get_cp_strategy_kind",
+    "init_cp_strategy",
+    "is_cp_enabled",
+    "is_interleave",
+    "is_zigzag",
+]

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

@@ -0,0 +1,277 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Base types and process-wide helpers for context parallel strategies.
+
+The strategy implementation is split across:
+
+* ``base.py``: base ABC, base metadata dataclass, enums, and singleton helpers.
+* ``zigzag.py``: former in-seq-split strategy and zigzag metadata.
+* ``interleave.py``: former round-robin-split strategy and interleave metadata.
+* ``utils.py``: public re-exports for import convenience.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import IntEnum
+from typing import TYPE_CHECKING, Any, Callable, List, Optional, Tuple
+
+from sglang.srt.runtime_context import get_parallel
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.server_args import ServerArgs
+
+
+class ContextParallelStrategyKind(IntEnum):
+    """Context parallel strategy identifiers."""
+
+    NONE = 0
+    ZIGZAG = 1
+    INTERLEAVE = 2
+
+    @classmethod
+    def from_string(cls, value: str) -> ContextParallelStrategyKind:
+        if value == "zigzag":
+            return cls.ZIGZAG
+        if value == "interleave":
+            return cls.INTERLEAVE
+        raise ValueError(
+            f"Unknown cp_strategy={value!r}; expected one of "
+            "{'zigzag', 'interleave'}"
+        )
+
+    @property
+    def cli_value(self) -> str:
+        return {
+            ContextParallelStrategyKind.NONE: "none",
+            ContextParallelStrategyKind.ZIGZAG: "zigzag",
+            ContextParallelStrategyKind.INTERLEAVE: "interleave",
+        }[self]
+
+
+class CPAttentionBackendKind(IntEnum):
+    """Attention backend calling convention used by CP strategy dispatch."""
+
+    FLASH_ATTENTION = 0
+
+    @classmethod
+    def from_string(cls, value: str) -> CPAttentionBackendKind:
+        if value in ("fa3", "flashinfer"):
+            return cls.FLASH_ATTENTION
+        raise ValueError(
+            f"Unsupported attention_backend={value!r} for CP strategy; expected one "
+            "of {'fa3', 'flashinfer'}"
+        )
+
+
+@dataclass
+class BaseContextParallelMetadata:
+    total_seq_lens: int = 0
+    bs: int = 1
+
+
+class ContextParallelStrategy(ABC):
+    """Owns process-wide policy for one context parallel layout."""
+
+    name: str
+    kind: ContextParallelStrategyKind
+
+    def __init__(self, cp_size: int):
+        self.cp_size = cp_size
+
+    @property
+    def cp_rank(self) -> int:
+
+        return get_parallel().attn_cp_rank
+
+    @property
+    def per_layer_attn_cp_comm(self) -> bool:
+        return _is_dsa_active()
+
+    @abstractmethod
+    def can_apply(self, num_tokens: int, forward_batch: ForwardBatch) -> bool:
+        """Return True if this strategy can shard the current forward."""
+
+    @abstractmethod
+    def build_metadata(
+        self,
+        num_tokens: int,
+        seqs_len: Optional[List[int]],
+        extend_seqs_len: Optional[List[int]] = None,
+    ) -> BaseContextParallelMetadata:
+        """Build per-forward metadata for this strategy."""
+
+    @abstractmethod
+    def shard_hidden_states(self, x: Any, forward_batch: ForwardBatch) -> Any:
+        """Shard hidden states to the current CP rank, usually at the first layer."""
+
+    @abstractmethod
+    def shard_position_ids(self, positions: Any, forward_batch: ForwardBatch) -> Any:
+        """Shard KV-cache slot position IDs for each token to the current CP rank."""
+
+    @abstractmethod
+    def gather_hidden_states(
+        self,
+        x: Any,
+        forward_batch: ForwardBatch,
+        stream: Optional[Any] = None,
+    ) -> Any:
+        """Gather rank-local hidden states, usually at the last layer."""
+
+    @abstractmethod
+    def gather_kv_cache(
+        self,
+        x: Any,
+        forward_batch: ForwardBatch,
+        stream: Optional[Any] = None,
+    ) -> Any:
+        """Gather rank-local KV payloads back to full token order."""
+
+    def shard_per_request(
+        self,
+        extend_seqs_cpu: List[int],
+        extend_seqs: Any,
+    ) -> Tuple[List[int], Any, List[int], Any]:
+        raise NotImplementedError(
+            f"{self.name} strategy does not support per-request sharding"
+        )
+
+    def split_before_forward(
+        self,
+        forward_batch: ForwardBatch,
+        input_ids: Optional[Any],
+        positions: Any,
+        input_embeds: Optional[Any] = None,
+    ) -> Optional[Any]:
+        """Shard model inputs before model.forward in CP-v2 paths."""
+        if input_ids is not None:
+            forward_batch.cp_v2_input_ids = self.shard_hidden_states(
+                input_ids, forward_batch
+            )
+        forward_batch.positions = self.shard_position_ids(positions, forward_batch)
+        if input_embeds is not None:
+            return self.shard_hidden_states(input_embeds, forward_batch)
+        return None
+
+    @abstractmethod
+    def run_attention(
+        self,
+        q: Any,
+        forward_batch: ForwardBatch,
+        device: Any,
+        attn_fn: Callable[[Any, Any, Any, int], Any],
+        attention_backend: CPAttentionBackendKind = CPAttentionBackendKind.FLASH_ATTENTION,
+    ) -> Any:
+        """Dispatch CP attention using the selected backend convention."""
+
+    @abstractmethod
+    def materialize_full_kv(
+        self,
+        forward_batch: ForwardBatch,
+        layer: Any,
+        k: Any,
+        v: Any,
+        swa_loc: Optional[Any] = None,
+    ) -> None:
+        """Write full-layout K/V to the backend cache if needed."""
+
+    def reindex_attn_metadata(self, core_attn_metadata: Any) -> None:
+        """Optional attention metadata rewrite for strategies that need it."""
+        return None
+
+
+def _is_dsa_active() -> bool:
+    from sglang.srt.runtime_context import get_server_args
+
+    sa = get_server_args()
+    return bool(
+        getattr(sa, "enable_prefill_cp", False)
+        and getattr(sa, "_is_dsa_model_arch", False)
+    )
+
+
+_STRATEGY: Optional[ContextParallelStrategy] = None
+
+
+def init_cp_strategy(server_args: ServerArgs) -> None:
+    """Bind the configured CP strategy for this process."""
+    global _STRATEGY
+
+    if not getattr(server_args, "enable_prefill_cp", False):
+        _STRATEGY = None
+        return
+
+    cp_size = getattr(server_args, "attn_cp_size", 1)
+    if cp_size <= 1:
+        _STRATEGY = None
+        return
+
+    kind = ContextParallelStrategyKind.from_string(server_args.cp_strategy)
+    if kind == ContextParallelStrategyKind.ZIGZAG:
+        from sglang.srt.layers.cp.zigzag import ZigzagCPStrategy
+
+        _STRATEGY = ZigzagCPStrategy(cp_size=cp_size)
+    elif kind == ContextParallelStrategyKind.INTERLEAVE:
+        from sglang.srt.layers.cp.interleave import InterleaveCPStrategy
+
+        _STRATEGY = InterleaveCPStrategy(cp_size=cp_size)
+    else:
+        raise ValueError(
+            f"Unsupported cp_strategy kind {kind} for "
+            f"cp_strategy={server_args.cp_strategy!r}"
+        )
+
+
+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
+    re-run ``ServerArgs.__post_init__`` because the pickled instance bypasses
+    ``__init__``. Lazy init lets worker processes recover the singleton from
+    global server args.
+    """
+    global _STRATEGY
+
+    if _STRATEGY is None:
+        from sglang.srt.runtime_context import get_server_args
+
+        try:
+            server_args = get_server_args()
+        except ValueError:
+            return None
+        if server_args is not None and getattr(server_args, "enable_prefill_cp", False):
+            init_cp_strategy(server_args)
+    return _STRATEGY
+
+
+def get_cp_strategy_kind() -> ContextParallelStrategyKind:
+    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
+
+
+def is_zigzag() -> bool:
+    return get_cp_strategy_kind() == ContextParallelStrategyKind.ZIGZAG
+
+
+def is_interleave() -> bool:
+    return get_cp_strategy_kind() == ContextParallelStrategyKind.INTERLEAVE

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

@@ -0,0 +1,107 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Interleave context parallel strategy shell.
+
+For ``cp_size = 4``, each rank owns every fourth token:
+
+    cp0: token0, token4, token8,  token12, token16, ...
+    cp1: token1, token5, token9,  token13, token17, ...
+    cp2: token2, token6, token10, token14, token18, ...
+    cp3: token3, token7, token11, token15, token19, ...
+
+After all-gather, tokens are restored to the original order:
+
+    token0, token1, token2, token3, token4, token5, token6, token7, ...
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, List, Optional
+
+from sglang.srt.layers.cp.base import (
+    BaseContextParallelMetadata,
+    ContextParallelStrategy,
+    ContextParallelStrategyKind,
+    CPAttentionBackendKind,
+)
+
+
+@dataclass
+class InterleaveContextParallelMetadata(BaseContextParallelMetadata):
+    """Interleave has no per-forward zigzag permutation payload."""
+
+
+class InterleaveCPStrategy(ContextParallelStrategy):
+    name = "interleave"
+    kind = ContextParallelStrategyKind.INTERLEAVE
+
+    def can_apply(self, num_tokens: int, forward_batch) -> bool:
+        if self.cp_size <= 1 or num_tokens < self.cp_size:
+            return False
+        forward_mode = getattr(forward_batch, "forward_mode", None)
+        return forward_mode is None or forward_mode.is_context_parallel_extend()
+
+    def build_metadata(
+        self,
+        num_tokens: int,
+        seqs_len: Optional[List[int]],
+        extend_seqs_len: Optional[List[int]] = None,
+    ) -> InterleaveContextParallelMetadata:
+        return InterleaveContextParallelMetadata(
+            total_seq_lens=sum(extend_seqs_len or seqs_len or [num_tokens]),
+            bs=len(extend_seqs_len or seqs_len or [num_tokens]),
+        )
+
+    def shard_hidden_states(self, x: Any, forward_batch) -> Any:
+        raise NotImplementedError(
+            "Interleave hidden-state sharding will land in a follow-up PR"
+        )
+
+    def shard_position_ids(self, positions: Any, forward_batch) -> Any:
+        raise NotImplementedError(
+            "Interleave position-id sharding will land in a follow-up PR"
+        )
+
+    def gather_hidden_states(
+        self, x: Any, forward_batch, stream: Optional[Any] = None
+    ) -> Any:
+        raise NotImplementedError(
+            "Interleave hidden-state gather will land in a follow-up PR"
+        )
+
+    def gather_kv_cache(
+        self, x: Any, forward_batch, stream: Optional[Any] = None
+    ) -> Any:
+        raise NotImplementedError("Interleave KV gather will land in a follow-up PR")
+
+    def run_attention(
+        self,
+        q: Any,
+        forward_batch,
+        device: Any,
+        attn_fn,
+        attention_backend: CPAttentionBackendKind = CPAttentionBackendKind.FLASH_ATTENTION,
+    ) -> Any:
+        raise NotImplementedError(
+            "Interleave attention dispatch will land in a follow-up PR"
+        )
+
+    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

@@ -0,0 +1,228 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Public import facade and runtime helpers for context parallel strategies."""
+
+from typing import TYPE_CHECKING, 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,
+    InterleaveCPStrategy,
+)
+from sglang.srt.layers.cp.zigzag import (
+    ContextParallelMetadata,
+    ZigzagContextParallelMetadata,
+    ZigzagCPStrategy,
+)
+from sglang.srt.runtime_context import get_parallel
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+CP_V2_DEFAULT_MODEL_CLASSES = frozenset(
+    {
+        "Qwen3MoeForCausalLM",
+    }
+)
+
+
+def is_glm_dsa_cache_layer_split_enabled(model_runner: "ModelRunner") -> bool:
+    """Whether DSA GPU KV/indexer cache layers are sharded across CP ranks.
+
+    Layer split is a prefill-CP-only optimization for DSA (DeepSeek Sparse
+    Attention) MLA models (e.g. GLM-5.2). Draft workers keep the full cache.
+    """
+    from sglang.srt.configs.model_config import is_deepseek_dsa
+
+    return (
+        not model_runner.is_draft_worker
+        and model_runner.server_args.enable_dsa_cache_layer_split
+        and model_runner.use_mla_backend
+        and is_deepseek_dsa(model_runner.model_config.hf_config)
+    )
+
+
+def get_glm_dsa_cp_layer_shard_info(
+    model_runner: "ModelRunner",
+) -> Tuple[Optional[int], int]:
+    """Return ``(layer_shard_rank, layer_shard_size)`` for the DSA KV pool.
+
+    ``(None, 1)`` disables sharding (feature off or only one CP rank).
+    """
+    if not is_glm_dsa_cache_layer_split_enabled(model_runner):
+        return None, 1
+    shard_size = get_parallel().attn_cp_size
+    if shard_size <= 1:
+        return None, 1
+    return get_parallel().attn_cp_rank, shard_size
+
+
+def get_glm_dsa_layer_split_effective_num_layers(
+    model_runner: "ModelRunner", num_layers: int
+) -> int:
+    """Per-rank owned layer count used when sizing the DSA KV cell.
+
+    Under layer split each CP rank only stores ``ceil(num_layers / shard_size)``
+    layers, plus one extra layer for the remote scratch buffer used when reading
+    a layer owned by another CP rank.
+    """
+    if not is_glm_dsa_cache_layer_split_enabled(model_runner):
+        return num_layers
+    shard_size = get_parallel().attn_cp_size
+    if shard_size <= 1:
+        return num_layers
+    owned_layers_upper_bound = (num_layers + shard_size - 1) // shard_size
+    return max(1, owned_layers_upper_bound + 1)
+
+
+def get_layer_shard_range(
+    rank: int, shard_size: int, total_layers: int
+) -> Tuple[int, int]:
+    """Contiguous ``[start, end)`` local-layer range owned by ``rank``.
+
+    Layers are split as evenly as possible; the first ``total_layers %
+    shard_size`` ranks own one extra layer.
+    """
+    base = total_layers // shard_size
+    rem = total_layers % shard_size
+    start = rank * base + min(rank, rem)
+    end = start + base + (1 if rank < rem else 0)
+    return start, end
+
+
+def get_layer_owner(local_layer_idx: int, shard_size: int, total_layers: int) -> int:
+    """CP rank that owns ``local_layer_idx`` under the contiguous split."""
+    for rank in range(shard_size):
+        start, end = get_layer_shard_range(rank, shard_size, total_layers)
+        if start <= local_layer_idx < end:
+            return rank
+    raise ValueError(
+        f"Invalid local_layer_idx={local_layer_idx} for "
+        f"shard_size={shard_size}, total_layers={total_layers}"
+    )
+
+
+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",
+    "ContextParallelMetadata",
+    "ContextParallelStrategy",
+    "ContextParallelStrategyKind",
+    "InterleaveCPStrategy",
+    "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",
+    "is_glm_dsa_cache_layer_split_enabled",
+    "get_glm_dsa_cp_layer_shard_info",
+    "get_glm_dsa_layer_split_effective_num_layers",
+    "get_layer_shard_range",
+    "get_layer_owner",
+]

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

@@ -0,0 +1,387 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Zigzag context parallel strategy shell.
+
+For ``cp_size = 4``, each sequence is split into ``2 * cp_size`` blocks. Each
+rank owns one early block and one late block:
+
+    cp0: block0, block7
+    cp1: block1, block6
+    cp2: block2, block5
+    cp3: block3, block4
+
+After all-gather, the blocks are reranged back to their original order:
+
+    block0 | block7 | block1 | block6 | block2 | block5 | block3 | block4
+      -> block0 | block1 | block2 | block3 | block4 | block5 | block6 | block7
+"""
+
+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 (
+    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
+from sglang.srt.runtime_context import get_parallel
+
+
+@dataclass
+class ZigzagContextParallelMetadata(BaseContextParallelMetadata):
+    # Layout lists have length bs * cp_segment_num (= bs * 2 * cp_size).
+    split_list: Optional[List[int]] = None
+    zigzag_index: Optional[List[int]] = None
+    cp_reverse_index: Optional[List[int]] = None
+    reverse_split_len: Optional[List[int]] = None
+
+    # Per-rank aggregate lists have length cp_size.
+    per_rank_actual_token: Optional[List[int]] = None
+    max_rank_len: Optional[List[int]] = None
+
+    # Per-sequence FlashAttention tensors (shape [bs] or [bs + 1]).
+    kv_len_prev_tensor: Optional[Any] = None
+    kv_len_next_tensor: Optional[Any] = None
+    actual_seq_q_prev_tensor: Optional[Any] = None
+    actual_seq_q_next_tensor: Optional[Any] = None
+    cu_seqlens_q_prev_tensor: Optional[Any] = None
+    cu_seqlens_q_next_tensor: Optional[Any] = None
+
+    # Scalars derived from the per-sequence lists above.
+    total_q_prev_tokens: int = 0
+    total_q_next_tokens: int = 0
+    max_seqlen_q_prev: int = 0
+    max_seqlen_q_next: int = 0
+
+    # Per-sequence CPU lists, useful for indexers and diagnostics.
+    kv_len_prev_list: Optional[List[int]] = None
+    kv_len_next_list: Optional[List[int]] = None
+    actual_seq_q_prev_list: Optional[List[int]] = None
+    actual_seq_q_next_list: Optional[List[int]] = None
+
+
+ContextParallelMetadata = ZigzagContextParallelMetadata
+
+
+class ZigzagCPStrategy(ContextParallelStrategy):
+    name = "zigzag"
+    kind = ContextParallelStrategyKind.ZIGZAG
+
+    def can_apply(self, num_tokens: int, forward_batch) -> bool:
+        if self.cp_size <= 1 or num_tokens < self.cp_size * 2:
+            return False
+        forward_mode = getattr(forward_batch, "forward_mode", None)
+        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,
+        num_tokens: int,
+        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.runtime_context import get_server_args
+
+        try:
+            device = torch.device(get_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(
+            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:
+        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:
+        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:
+        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:
+        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,
+        q: Any,
+        forward_batch,
+        device: Any,
+        attn_fn,
+        attention_backend: CPAttentionBackendKind = CPAttentionBackendKind.FLASH_ATTENTION,
+    ) -> Any:
+        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 _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_parallel().attn_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,
+        )