#27313 [2/n] [CP] Add context parallel strategy abstractions

原始 PR 作者 Fridge003 合并时间 2026-06-16 15:20 文件变更 9 提交数 3 评论 28 代码增减 +715 / -1

执行摘要

新增上下文并行策略抽象层

现有 Context Parallelism (CP) 的实现分散在模型文件和注意力后端中，server args 混乱、模型覆盖率不对称、通信代码散落。为系统化重构，需要引入 strategy-owned runtime abstraction，将 CP 分片逻辑统一到策略接口中。详见 Roadmap Issue #27252。

值得精读，尤其是 ContextParallelStrategy 抽象基类和包组织结构。该设计为后续模型无关的 CP 分片打下基础，是 CP 重构的核心抽象。关注后续的 #27314 等 PR 以了解逐步填充的实现。

讨论亮点

文件命名与组织：作者将 strategy.py 重命名为 utils.py，并将初始化逻辑、use_cp_v2 等函数移入 base.py，简化文件结构。同时明确了 shard_hidden_states/gather_hidden_states 等方法的命名和职责（在 base.py 中关联评论）。

CP v2 门控简化：最初 use_cp_v2() 函数被提出，但经 review 认为可直接通过 is_cp_enabled() 覆盖，故移除。SGLANG_ENABLE_CP_V2 改为直接使用 EnvBool 而非 EnvBoolWithAlias。

添加策略说明图：要求在 zigzag.py 和 interleave.py 文件头部添加 ASCII 图表说明分片方式，从旧 cp_utils.py 迁移。

测试适应性：测试中确保 get_cp_strategy() 在 CP v1 和 CP v2 下均返回非 None，且 init_cp_strategy 在两种情况下均被调用。

实现拆解

创建 CP 包和核心抽象：在 python/sglang/srt/layers/cp/ 下创建包，base.py 定义 ContextParallelStrategyKind 枚举（NONE/ZIGZAG/INTERLEAVE）、CPAttentionBackendKind 枚举、BaseContextParallelMetadata 数据类，以及 ContextParallelStrategy 抽象基类，声明 can_apply、build_metadata、shard_hidden_states、shard_position_ids、gather_hidden_states、gather_kv_cache 等抽象方法。
实现策略 Shell：zigzag.py 实现 ZigzagCPStrategy，interleave.py 实现 InterleaveCPStrategy。当前仅实现 can_apply（检查 cp_size 和 token 数）和 build_metadata（返回基本元数据），其余方法 raise NotImplementedError，留待后续 PR。
环境变量门控：在 environ.py 中添加 SGLANG_ENABLE_CP_V2（默认 False），并保留别名 SGLANG_ENABLE_REFACTOR_CP 作为兼容。新增 use_cp_v2() 函数（后因 review 删除，改用 is_cp_enabled 替代）。
集成 server args：在 server_args.py 的 _handle_context_parallelism 方法末尾调用 init_cp_strategy(self)，从规范化后的 server args 初始化策略单例。同时添加 cp_strategy 字段映射和验证。
公共导出与门面：在 __init__.py 和 utils.py 中集中导出所有公共符号，提供 is_cp_enabled、is_zigzag、is_interleave、get_cp_strategy、get_cp_strategy_kind 等便捷函数。
测试覆盖：新增 test/registered/cp/test_cp_strategy_unit.py，测试枚举解析、策略初始化、门控 SGLANG_ENABLE_CP_V2 兼容性；扩展 test_server_args.py 验证 context parallel handler 会初始化策略。

文件	模块	状态	重要度
`python/sglang/srt/layers/cp/base.py`	CP 抽象层	added	9.25
`python/sglang/srt/layers/cp/zigzag.py`	CP 抽象层	added	8.98
`python/sglang/srt/layers/cp/interleave.py`	CP 抽象层	added	8.85
`python/sglang/srt/layers/cp/utils.py`	CP 抽象层	added	7.18
`python/sglang/srt/layers/cp/__init__.py`	CP 抽象层	added	7.31
`python/sglang/srt/environ.py`	环境配置	modified	5.54
`python/sglang/srt/server_args.py`	服务配置	modified	5.11
`test/registered/cp/test_cp_strategy_unit.py`	单元测试	added	7.56
`test/registered/unit/server_args/test_server_args.py`	服务配置	modified	5.29

关键符号

ContextParallelStrategyKind.from_string ContextParallelStrategyKind.cli_value CPAttentionBackendKind.from_string ContextParallelStrategy.__init__ ContextParallelStrategy.cp_rank ContextParallelStrategy.can_apply ContextParallelStrategy.build_metadata ContextParallelStrategy.shard_hidden_states ContextParallelStrategy.shard_position_ids ContextParallelStrategy.gather_hidden_states ContextParallelStrategy.gather_kv_cache ZigzagCPStrategy.can_apply ZigzagCPStrategy.build_metadata InterleaveCPStrategy.can_apply InterleaveCPStrategy.build_metadata init_cp_strategy get_cp_strategy get_cp_strategy_kind is_cp_enabled is_zigzag is_interleave TestCPStrategyUnit.test_strategy_kind_maps_cli_values TestCPStrategyUnit.test_init_cp_strategy_binds_zigzag_strategy TestCPStrategyUnit.test_get_cp_strategy_is_initialized_under_cp_v1_and_cp_v2

关键源码片段

python/sglang/srt/layers/cp/zigzag.py dependency-wiring

Zigzag 策略的具体实现壳，包括 ZigzagContextParallelMetadata 数据类和 ZigzagCPStrategy 类，展示策略接口的初步填充。

# (ZigzagCPStrategy 部分实现 )
class ZigzagCPStrategy(ContextParallelStrategy):
    name = "zigzag"
    kind = ContextParallelStrategyKind.ZIGZAG

    def can_apply(self, num_tokens: int, forward_batch) -> bool:
        # 要求 cp_size > 1 且 token 数至少为 2 * cp_size
        if self.cp_size <= 1 or num_tokens < self.cp_size * 2:
            return False
        forward_mode = getattr(forward_batch, "forward_mode", None)
        # 仅允许在 context parallel extend 模式下应用
        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,
    ) -> ZigzagContextParallelMetadata:
        # 构建基础的元数据，后续 PR 会填充详细的分拆信息
        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]),
        )

    # 以下方法留待后续 PR 实现
    def shard_hidden_states(self, x: Any, forward_batch) -> Any:
        raise NotImplementedError("Zigzag hidden-state sharding will land in a follow-up PR")

    def shard_position_ids(self, positions: Any, forward_batch) -> Any:
        raise NotImplementedError("Zigzag 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("Zigzag 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("Zigzag 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("Zigzag attention dispatch will land in a follow-up PR")

python/sglang/srt/layers/cp/interleave.py dependency-wiring

Interleave 策略的具体实现壳，结构类似 zigzag.py，展示另一策略的骨架。

# (InterleaveCPStrategy 部分实现 )
class InterleaveCPStrategy(ContextParallelStrategy):
    name = "interleave"
    kind = ContextParallelStrategyKind.INTERLEAVE

    def can_apply(self, num_tokens: int, forward_batch) -> bool:
        # 要求 cp_size > 1 且 token 数至少为 cp_size
        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]),
        )

    # 以下方法留待后续 PR 实现
    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")

test/registered/cp/test_cp_strategy_unit.py test-coverage

单元测试文件，验证枚举解析、策略初始化、门控环境变量兼容性，保证基础功能正确。

import unittest
from types import SimpleNamespace
from unittest.mock import patch

from sglang.srt.layers.cp.base import (
    ContextParallelStrategyKind,
    get_cp_strategy,
    get_cp_strategy_kind,
    init_cp_strategy,
    is_cp_enabled,
    is_interleave,
    is_zigzag,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

register_cpu_ci(est_time=2, suite="base-a-test-cpu")


class TestCPStrategyUnit(CustomTestCase):
    def tearDown(self):
        # 每次测试后关闭 CP，避免影响其他测试
        init_cp_strategy(SimpleNamespace(enable_prefill_cp=False))

    def test_strategy_kind_maps_cli_values(self):
        # 验证枚举值与 CLI 字符串的映射
        self.assertEqual(ContextParallelStrategyKind.NONE.value, 0)
        self.assertEqual(
            ContextParallelStrategyKind.from_string("zigzag"),
            ContextParallelStrategyKind.ZIGZAG,
        )
        self.assertEqual(
            ContextParallelStrategyKind.from_string("interleave"),
            ContextParallelStrategyKind.INTERLEAVE,
        )
        self.assertEqual(ContextParallelStrategyKind.ZIGZAG.cli_value, "zigzag")
        self.assertEqual(ContextParallelStrategyKind.INTERLEAVE.cli_value, "interleave")

    def test_init_cp_strategy_binds_zigzag_strategy(self):
        # 模拟启用 CP 并选用 zigzag 策略，验证单例状态
        init_cp_strategy(
            SimpleNamespace(
                enable_prefill_cp=True,
                cp_strategy="zigzag",
                attn_cp_size=4,
            )
        )
        self.assertTrue(is_cp_enabled())
        self.assertTrue(is_zigzag())
        self.assertFalse(is_interleave())
        self.assertEqual(get_cp_strategy_kind(), ContextParallelStrategyKind.ZIGZAG)

    def test_get_cp_strategy_is_initialized_under_cp_v1_and_cp_v2(self):
        # 验证在 CP v1（SGLANG_ENABLE_CP_V2=False）和 CP v2 下策略均被初始化
        init_cp_strategy(
            SimpleNamespace(
                enable_prefill_cp=True,
                cp_strategy="interleave",
                attn_cp_size=4,
            )
        )
        with patch(
            "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=False
        ):
            self.assertIsNotNone(get_cp_strategy())
            self.assertTrue(is_cp_enabled())
            self.assertTrue(is_interleave())

        with patch(
            "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True
        ):
            self.assertIsNotNone(get_cp_strategy())


if __name__ == "__main__":
    unittest.main()

评论区精华

文件命名与组织 设计

作者在 review 中要求将 `strategy.py` 重命名为 `utils.py`，并将 `use_cp_v2` 等函数移入 `base.py`，同时要求将 copyright 统一为 2023-2026。

结论：已执行：文件重命名、函数移动、copyright 更新。 · 已解决

CP v2 门控简化 设计

作者认为 `use_cp_v2()` 函数不必要，因为 `is_cp_enabled()` 可覆盖此检测，且 `SGLANG_ENABLE_CP_V2` 无需使用 `EnvBoolWithAlias`。

结论：已移除 `use_cp_v2`，改用 `is_cp_enabled`；环境变量改为 `EnvBool`。 · 已解决

添加策略说明 ASCII 图 documentation

要求在 zigzag.py 和 interleave.py 文件头部添加分片方式的 ASCII 图表，从旧的 `cp_utils.py` 迁移。

结论：已添加图表说明，用户可直观了解分片方案。 · 已解决

抽象方法命名与职责明确 设计

作者对多个抽象方法名称提出修改建议：`request_to_cp_rank` → `shard_hidden_states`，`request_to_cp_position_ids` → `shard_position_ids`，并强调应在第一层调用分片、最后一层调用 gather。

结论：已按建议重命名并更新 docstring。 · 已解决

测试覆盖门控分支 测试

作者指出需要确保 `get_cp_strategy` 在 CP v1 和 CP v2 两种环境下均初始化为非 None 值，并且 `init_cp_strategy` 在两种情况下均被调用。

结论：测试已补充 `test_get_cp_strategy_is_initialized_under_cp_v1_and_cp_v2`，使用 patch 模拟环境变量。 · 已解决

风险与影响

兼容性风险：新抽象默认关闭，现有 CP 路径完全不受影响，风险低。
未实现的方法：shard_hidden_states 等方法抛出 NotImplementedError，若在 CP v2 路由启用前被意外调用将崩溃，但通过后续 PR 逐步填充可避免。
环境变量别名：SGLANG_ENABLE_REFACTOR_CP 作为已废弃别名保留，但删除旧环境变量的用户无感知。

用户：无用户可见变更，现有命令行参数继续生效。后续 CP v2 启用后将影响 CP 行为。
系统：增加少量内存开销（策略单例），无性能影响。
团队：为后续 CP 重构提供统一接口，降低模型接入成本。

新抽象默认关闭未实现方法可能误调环境变量别名兼容测试覆盖率有限 CP v2 尚未启用

关联 Issue

#27252 [Roadmap]Prefill Context Parallel Refactor

#27312 [1/n] [CP] Simplify prefill context parallel server args

完整报告

执行摘要

一句话：新增上下文并行策略抽象层
推荐动作：值得精读，尤其是 ContextParallelStrategy 抽象基类和包组织结构。该设计为后续模型无关的 CP 分片打下基础，是 CP 重构的核心抽象。关注后续的 #27314 等 PR 以了解逐步填充的实现。

功能与动机

实现拆解

创建 CP 包和核心抽象：在 python/sglang/srt/layers/cp/ 下创建包，base.py 定义 ContextParallelStrategyKind 枚举（NONE/ZIGZAG/INTERLEAVE）、CPAttentionBackendKind 枚举、BaseContextParallelMetadata 数据类，以及 ContextParallelStrategy 抽象基类，声明 can_apply、build_metadata、shard_hidden_states、shard_position_ids、gather_hidden_states、gather_kv_cache 等抽象方法。
实现策略 Shell：zigzag.py 实现 ZigzagCPStrategy，interleave.py 实现 InterleaveCPStrategy。当前仅实现 can_apply（检查 cp_size 和 token 数）和 build_metadata（返回基本元数据），其余方法 raise NotImplementedError，留待后续 PR。
环境变量门控：在 environ.py 中添加 SGLANG_ENABLE_CP_V2（默认 False），并保留别名 SGLANG_ENABLE_REFACTOR_CP 作为兼容。新增 use_cp_v2() 函数（后因 review 删除，改用 is_cp_enabled 替代）。
集成 server args：在 server_args.py 的 _handle_context_parallelism 方法末尾调用 init_cp_strategy(self)，从规范化后的 server args 初始化策略单例。同时添加 cp_strategy 字段映射和验证。
公共导出与门面：在 __init__.py 和 utils.py 中集中导出所有公共符号，提供 is_cp_enabled、is_zigzag、is_interleave、get_cp_strategy、get_cp_strategy_kind 等便捷函数。
测试覆盖：新增 test/registered/cp/test_cp_strategy_unit.py，测试枚举解析、策略初始化、门控 SGLANG_ENABLE_CP_V2 兼容性；扩展 test_server_args.py 验证 context parallel handler 会初始化策略。

关键文件：

python/sglang/srt/layers/cp/base.py（模块 CP 抽象层；类别 source；类型 dependency-wiring；符号 ContextParallelStrategyKind, from_string, cli_value, CPAttentionBackendKind）: 核心文件，定义策略枚举、注意力后端枚举、基础元数据类型和 ContextParallelStrategy 抽象基类，包含所有抽象方法签名和进程级单例辅助函数。
python/sglang/srt/layers/cp/zigzag.py（模块 CP 抽象层；类别 source；类型 dependency-wiring；符号 ZigzagContextParallelMetadata, ZigzagCPStrategy, can_apply, build_metadata）: Zigzag 策略的具体实现壳，包括 ZigzagContextParallelMetadata 数据类和 ZigzagCPStrategy 类，展示策略接口的初步填充。
python/sglang/srt/layers/cp/interleave.py（模块 CP 抽象层；类别 source；类型 dependency-wiring；符号 InterleaveContextParallelMetadata, InterleaveCPStrategy, can_apply, build_metadata）: Interleave 策略的具体实现壳，结构类似 zigzag.py，展示另一策略的骨架。
python/sglang/srt/layers/cp/utils.py（模块 CP 抽象层；类别 source；类型 dependency-wiring）: 公共导入门面，导出所有关键类型和辅助函数，便于外部模块引用。
python/sglang/srt/layers/cp/__init__.py（模块 CP 抽象层；类别 source；类型 dependency-wiring）: 包初始化文件，确保从 sglang.srt.layers.cp 直接导入主要符号。
python/sglang/srt/environ.py（模块环境配置；类别 source；类型 core-logic）: 添加 SGLANG_ENABLE_CP_V2 环境变量门控，控制是否启用新的 CP v2 策略路由。
python/sglang/srt/server_args.py（模块服务配置；类别 source；类型 dependency-wiring）: 在 _handle_context_parallelism 中调用 init_cp_strategy，确保服务器启动时根据规范化的参数初始化策略单例。
test/registered/cp/test_cp_strategy_unit.py（模块单元测试；类别 test；类型 test-coverage；符号 TestCPStrategyUnit, tearDown, test_strategy_kind_maps_cli_values, test_init_cp_strategy_binds_zigzag_strategy）: 单元测试文件，验证枚举解析、策略初始化、门控环境变量兼容性，保证基础功能正确。
test/registered/unit/server_args/test_server_args.py（模块服务配置；类别 test；类型 test-coverage；符号 test_context_parallel_handler_initializes_cp_strategy）: 扩展现有 server args 测试，验证 context parallel handler 会调用 init_cp_strategy。

关键符号：ContextParallelStrategyKind.from_string, ContextParallelStrategyKind.cli_value, CPAttentionBackendKind.from_string, ContextParallelStrategy.init, ContextParallelStrategy.cp_rank, ContextParallelStrategy.can_apply, ContextParallelStrategy.build_metadata, ContextParallelStrategy.shard_hidden_states, ContextParallelStrategy.shard_position_ids, ContextParallelStrategy.gather_hidden_states, ContextParallelStrategy.gather_kv_cache, ZigzagCPStrategy.can_apply, ZigzagCPStrategy.build_metadata, InterleaveCPStrategy.can_apply, InterleaveCPStrategy.build_metadata, init_cp_strategy, get_cp_strategy, get_cp_strategy_kind, is_cp_enabled, is_zigzag, is_interleave, TestCPStrategyUnit.test_strategy_kind_maps_cli_values, TestCPStrategyUnit.test_init_cp_strategy_binds_zigzag_strategy, TestCPStrategyUnit.test_get_cp_strategy_is_initialized_under_cp_v1_and_cp_v2

关键源码片段

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

Zigzag 策略的具体实现壳，包括 ZigzagContextParallelMetadata 数据类和 ZigzagCPStrategy 类，展示策略接口的初步填充。

# (ZigzagCPStrategy 部分实现 )
class ZigzagCPStrategy(ContextParallelStrategy):
    name = "zigzag"
    kind = ContextParallelStrategyKind.ZIGZAG

    def can_apply(self, num_tokens: int, forward_batch) -> bool:
        # 要求 cp_size > 1 且 token 数至少为 2 * cp_size
        if self.cp_size <= 1 or num_tokens < self.cp_size * 2:
            return False
        forward_mode = getattr(forward_batch, "forward_mode", None)
        # 仅允许在 context parallel extend 模式下应用
        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,
    ) -> ZigzagContextParallelMetadata:
        # 构建基础的元数据，后续 PR 会填充详细的分拆信息
        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]),
        )

    # 以下方法留待后续 PR 实现
    def shard_hidden_states(self, x: Any, forward_batch) -> Any:
        raise NotImplementedError("Zigzag hidden-state sharding will land in a follow-up PR")

    def shard_position_ids(self, positions: Any, forward_batch) -> Any:
        raise NotImplementedError("Zigzag 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("Zigzag 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("Zigzag 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("Zigzag attention dispatch will land in a follow-up PR")

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

Interleave 策略的具体实现壳，结构类似 zigzag.py，展示另一策略的骨架。

# (InterleaveCPStrategy 部分实现 )
class InterleaveCPStrategy(ContextParallelStrategy):
    name = "interleave"
    kind = ContextParallelStrategyKind.INTERLEAVE

    def can_apply(self, num_tokens: int, forward_batch) -> bool:
        # 要求 cp_size > 1 且 token 数至少为 cp_size
        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]),
        )

    # 以下方法留待后续 PR 实现
    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")

`test/registered/cp/test_cp_strategy_unit.py`

单元测试文件，验证枚举解析、策略初始化、门控环境变量兼容性，保证基础功能正确。

import unittest
from types import SimpleNamespace
from unittest.mock import patch

from sglang.srt.layers.cp.base import (
    ContextParallelStrategyKind,
    get_cp_strategy,
    get_cp_strategy_kind,
    init_cp_strategy,
    is_cp_enabled,
    is_interleave,
    is_zigzag,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase

register_cpu_ci(est_time=2, suite="base-a-test-cpu")


class TestCPStrategyUnit(CustomTestCase):
    def tearDown(self):
        # 每次测试后关闭 CP，避免影响其他测试
        init_cp_strategy(SimpleNamespace(enable_prefill_cp=False))

    def test_strategy_kind_maps_cli_values(self):
        # 验证枚举值与 CLI 字符串的映射
        self.assertEqual(ContextParallelStrategyKind.NONE.value, 0)
        self.assertEqual(
            ContextParallelStrategyKind.from_string("zigzag"),
            ContextParallelStrategyKind.ZIGZAG,
        )
        self.assertEqual(
            ContextParallelStrategyKind.from_string("interleave"),
            ContextParallelStrategyKind.INTERLEAVE,
        )
        self.assertEqual(ContextParallelStrategyKind.ZIGZAG.cli_value, "zigzag")
        self.assertEqual(ContextParallelStrategyKind.INTERLEAVE.cli_value, "interleave")

    def test_init_cp_strategy_binds_zigzag_strategy(self):
        # 模拟启用 CP 并选用 zigzag 策略，验证单例状态
        init_cp_strategy(
            SimpleNamespace(
                enable_prefill_cp=True,
                cp_strategy="zigzag",
                attn_cp_size=4,
            )
        )
        self.assertTrue(is_cp_enabled())
        self.assertTrue(is_zigzag())
        self.assertFalse(is_interleave())
        self.assertEqual(get_cp_strategy_kind(), ContextParallelStrategyKind.ZIGZAG)

    def test_get_cp_strategy_is_initialized_under_cp_v1_and_cp_v2(self):
        # 验证在 CP v1（SGLANG_ENABLE_CP_V2=False）和 CP v2 下策略均被初始化
        init_cp_strategy(
            SimpleNamespace(
                enable_prefill_cp=True,
                cp_strategy="interleave",
                attn_cp_size=4,
            )
        )
        with patch(
            "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=False
        ):
            self.assertIsNotNone(get_cp_strategy())
            self.assertTrue(is_cp_enabled())
            self.assertTrue(is_interleave())

        with patch(
            "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True
        ):
            self.assertIsNotNone(get_cp_strategy())


if __name__ == "__main__":
    unittest.main()

评论区精华

添加策略说明图：要求在 zigzag.py 和 interleave.py 文件头部添加 ASCII 图表说明分片方式，从旧 cp_utils.py 迁移。

测试适应性：测试中确保 get_cp_strategy() 在 CP v1 和 CP v2 下均返回非 None，且 init_cp_strategy 在两种情况下均被调用。

文件命名与组织 (design): 已执行：文件重命名、函数移动、copyright 更新。
CP v2 门控简化 (design): 已移除 use_cp_v2，改用 is_cp_enabled；环境变量改为 EnvBool。
添加策略说明 ASCII 图 (documentation): 已添加图表说明，用户可直观了解分片方案。
抽象方法命名与职责明确 (design): 已按建议重命名并更新 docstring。
测试覆盖门控分支 (testing): 测试已补充 test_get_cp_strategy_is_initialized_under_cp_v1_and_cp_v2，使用 patch 模拟环境变量。

风险与影响

风险：兼容性风险：新抽象默认关闭，现有 CP 路径完全不受影响，风险低。
未实现的方法：shard_hidden_states 等方法抛出 NotImplementedError，若在 CP v2 路由启用前被意外调用将崩溃，但通过后续 PR 逐步填充可避免。
环境变量别名：SGLANG_ENABLE_REFACTOR_CP 作为已废弃别名保留，但删除旧环境变量的用户无感知。
影响：用户：无用户可见变更，现有命令行参数继续生效。后续 CP v2 启用后将影响 CP 行为。
系统：增加少量内存开销（策略单例），无性能影响。
团队：为后续 CP 重构提供统一接口，降低模型接入成本。
风险标记：新抽象默认关闭, 未实现方法可能误调, 环境变量别名兼容, 测试覆盖率有限, CP v2 尚未启用

关联脉络

PR #27312 [1/n] [CP] Simplify prefill context parallel server args: 本 PR #27313 依赖于 #27312，后者简化了 server args 并规范化了 cp_strategy 字段。本 PR 在此基础上初始化策略单例。
PR #27252 [Roadmap] Prefill Context Parallel Refactor: 本 PR 是 Roadmap issue #27252 的第二项，整体设计遵循该 issue。
PR #27314 [3/n] [CP] Integrate CP-v2 strategy into model runner: 预期的后续 PR（未在历史中明确，但根据 roadmap 推测），将策略抽象集成到模型 runner。

#27313 [2/n] [CP] Add context parallel strategy abstractions

执行摘要

新增上下文并行策略抽象层

实现拆解

评论区精华

风险与影响

关联 Issue

完整报告

参与讨论