[CP V1 Deprecation 1/5] Migrate tests to strategy-based prefill CP (#36222)

test/manual/dsv4/test_b200_flash.py

@@ -87,9 +87,9 @@ class TestB200FlashCP(DSV4FlashAime25TestBase):
         "4",
         "--moe-a2a-backend",
         "deepep",
-        "--enable-dsa-prefill-context-parallel",
+        "--enable-prefill-cp",
-        "--dsa-prefill-cp-mode",
+        "--cp-strategy",
-        "round-robin-split",
+        "interleave",
         "--chunked-prefill-size",
         "16384",
         "--mem-fraction-static",

test/manual/dsv4/test_b200_pro.py

@@ -101,9 +101,9 @@ class TestB200ProCP(DSV4ProAime25TestBase):
         "8",
         "--moe-a2a-backend",
         "deepep",
-        "--enable-dsa-prefill-context-parallel",
+        "--enable-prefill-cp",
-        "--dsa-prefill-cp-mode",
+        "--cp-strategy",
-        "round-robin-split",
+        "interleave",
         "--chunked-prefill-size",
         "16384",
         "--mem-fraction-static",

test/manual/dsv4/test_b300_flash.py

@@ -89,9 +89,9 @@ class TestB300FlashCP(DSV4FlashAime25TestBase):
         "4",
         "--moe-a2a-backend",
         "deepep",
-        "--enable-dsa-prefill-context-parallel",
+        "--enable-prefill-cp",
-        "--dsa-prefill-cp-mode",
+        "--cp-strategy",
-        "round-robin-split",
+        "interleave",
         "--chunked-prefill-size",
         "16384",
         "--mem-fraction-static",

test/manual/dsv4/test_b300_pro.py

@@ -103,9 +103,9 @@ class TestB300ProCP(DSV4ProAime25TestBase):
         "8",
         "--moe-a2a-backend",
         "deepep",
-        "--enable-dsa-prefill-context-parallel",
+        "--enable-prefill-cp",
-        "--dsa-prefill-cp-mode",
+        "--cp-strategy",
-        "round-robin-split",
+        "interleave",
         "--chunked-prefill-size",
         "16384",
         "--mem-fraction-static",

test/manual/dsv4/test_gb300_flash.py

@@ -87,9 +87,9 @@ class TestGB300FlashCP(DSV4FlashAime25TestBase):
         "4",
         "--moe-a2a-backend",
         "deepep",
-        "--enable-dsa-prefill-context-parallel",
+        "--enable-prefill-cp",
-        "--dsa-prefill-cp-mode",
+        "--cp-strategy",
-        "round-robin-split",
+        "interleave",
         "--chunked-prefill-size",
         "16384",
         "--mem-fraction-static",

test/manual/dsv4/test_gb300_pro.py

@@ -103,9 +103,9 @@ class TestGB300ProCP(DSV4ProAime25TestBase):
         "4",
         "--moe-a2a-backend",
         "deepep",
-        "--enable-dsa-prefill-context-parallel",
+        "--enable-prefill-cp",
-        "--dsa-prefill-cp-mode",
+        "--cp-strategy",
-        "round-robin-split",
+        "interleave",
         "--chunked-prefill-size",
         "16384",
         "--mem-fraction-static",

test/manual/dsv4/test_h200_fp8_flash.py

@@ -98,9 +98,9 @@ class TestH200Fp8FlashCP(DSV4FlashAime25TestBase):
         "4",
         "--moe-a2a-backend",
         "deepep",
-        "--enable-dsa-prefill-context-parallel",
+        "--enable-prefill-cp",
-        "--dsa-prefill-cp-mode",
+        "--cp-strategy",
-        "round-robin-split",
+        "interleave",
         "--chunked-prefill-size",
         "16384",
         "--mem-fraction-static",

test/manual/test_kv_events.py

@@ -329,7 +329,9 @@ class TestKvEvents(CustomTestCase):
                 2,
                 "--moe-dp-size",
                 2,
-                "--enable-prefill-context-parallel",
+                "--enable-prefill-cp",
+                "--cp-strategy",
+                "zigzag",
                 "--trust-remote-code",
                 "--max-total-tokens",
                 4096,

test/manual/test_qwen3_235b.py

@@ -80,7 +80,8 @@ class TestQwen3235BFP8(unittest.TestCase):
             "--moe-dp-size=2",
             "--attn-cp-size=2",
             "--ep-size=4",
-            "--enable-prefill-context-parallel",
+            "--enable-prefill-cp",
+            "--cp-strategy=zigzag",
         ]
 
         MTP_ARGS = [

test/registered/cp/test_cp_strategy_unit.py

@@ -89,7 +89,7 @@ class TestCPStrategyUnit(CustomTestCase):
         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):
+    def test_get_cp_strategy_is_initialized_under_cp_v2(self):
         init_cp_strategy(
             SimpleNamespace(
                 enable_prefill_cp=True,
@@ -99,17 +99,12 @@ class TestCPStrategyUnit(CustomTestCase):
         )
 
         with patch(
-            "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=False
+            "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True
         ):
             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())
-
 
 class TestPrefillCPBCGReplay(CustomTestCase):
     def tearDown(self):
@@ -312,12 +307,6 @@ class TestCPZigzagStrategy(CustomTestCase):
             extend_seq_lens_cpu=[7],
         )
 
-        with patch(
-            "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=False
-        ):
-            self.assertFalse(enable_cp_v2())
-            self.assertFalse(is_cp_v2_active(active_batch))
-
         with patch(
             "sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True
         ):

test/registered/cp/test_deepseek_v3_cp_single_node.py

@@ -37,7 +37,9 @@ class TestDeepseekV3CPInSeqSplit(CustomTestCase):
             "8",
             "--dp",
             "2",
-            "--enable-prefill-context-parallel",
+            "--enable-prefill-cp",
+            "--cp-strategy",
+            "zigzag",
             "--attention-backend",
             "fa3",
             "--mem-frac",
@@ -79,7 +81,7 @@ class TestDeepseekV3CPInSeqSplit(CustomTestCase):
 
         if is_in_ci():
             write_github_step_summary(
-                f"### test_a_gsm8k (deepseek-v3-mla-cp-in-seq-split)\n"
+                f"### test_a_gsm8k (deepseek-v3-mla-cp-zigzag)\n"
                 f'{metrics["score"]=:.3f}\n'
             )
             self.assertGreater(metrics["score"], GSM8K_ACCURACY_THRESHOLD)

test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py

@@ -29,7 +29,6 @@ SERVER_LAUNCH_TIMEOUT = 3600
 DEEPEP_CONFIG = '{"normal_dispatch":{"num_sms":96},"normal_combine":{"num_sms":96}}'
 
 _DEEPEP_ENV = {
-    "SGLANG_ENABLE_CP_V2": "1",
     "SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK": "1024",
     # The draft-extend graph pool costs ~4.5 GB here (DeepEP MoE workspace is
     # captured at full dispatch capacity), which starves the eager prefill
@@ -39,7 +38,6 @@ _DEEPEP_ENV = {
 }
 
 _MEGAMOE_ENV = {
-    "SGLANG_ENABLE_CP_V2": "1",
     "SGLANG_OPT_DEEPGEMM_MEGA_MOE_NUM_MAX_TOKENS_PER_RANK": "8320",
 }
 
@@ -188,7 +186,6 @@ class TestDSV4FlashFP4B200Balanced_CP_NonDeepEP(
                 "--moe-runner-backend",  # for fp4 checkpoint
                 "flashinfer_mxfp4",
             ],
-            env={"SGLANG_ENABLE_CP_V2": "1"},
         )
 
     @classmethod
@@ -232,7 +229,6 @@ class TestDSV4FlashFP4B200_CP_DSpark(
                 "--moe-runner-backend",  # for fp4 checkpoint
                 "flashinfer_mxfp4",
             ],
-            env={"SGLANG_ENABLE_CP_V2": "1"},
         )
 
     @classmethod

test/registered/cp/test_dsa_prefill_cp.py

@@ -54,7 +54,6 @@ class TestDSACPV2Interleave(CustomTestCase):
             cls.base_url,
             timeout=SERVER_LAUNCH_TIMEOUT,
             other_args=other_args,
-            env={"SGLANG_ENABLE_CP_V2": "1"},
         )
 
     @classmethod

test/registered/cp/test_dsa_prefill_cp_legacy.py

@@ -1,157 +0,0 @@
-import unittest
-from types import SimpleNamespace
-
-from sglang.srt.utils import kill_process_tree
-from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.run_eval import run_eval
-from sglang.test.test_utils import (
-    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-    DEFAULT_URL_FOR_TEST,
-    CustomTestCase,
-    is_in_ci,
-    popen_launch_server,
-    write_github_step_summary,
-)
-
-register_cuda_ci(est_time=616, stage="extra-b", runner_config="8-gpu-h200")
-DEEPSEEK_V32_MODEL_PATH = "deepseek-ai/DeepSeek-V3.2"
-
-
-class TestDeepseekV32CPInSeqSplit(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEEPSEEK_V32_MODEL_PATH
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        other_args = [
-            "--trust-remote-code",
-            "--tp",
-            "8",
-            "--enable-dp-attention",
-            "--dp",
-            "2",
-            "--attn-cp-size",
-            "4",
-            "--enable-dsa-prefill-context-parallel",
-            "--dsa-prefill-cp-mode",
-            "in-seq-split",
-            "--speculative-algorithm",
-            "EAGLE",
-            "--speculative-num-steps",
-            "3",
-            "--speculative-eagle-topk",
-            "1",
-            "--speculative-num-draft-tokens",
-            "4",
-            "--mem-frac",
-            "0.85",
-            "--cuda-graph-max-bs-decode",
-            "32",
-            "--max-running-requests",
-            "32",
-            "--model-loader-extra-config",
-            '{"enable_multithread_load": true, "num_threads": 64}',
-        ]
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=other_args,
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def test_a_gsm8k(
-        self,
-    ):  # Append an "a" to make this test run first (alphabetically) to warm up the server
-        args = SimpleNamespace(
-            base_url=self.base_url,
-            model=self.model,
-            eval_name="gsm8k",
-            api="completion",
-            max_tokens=512,
-            num_examples=500,
-            num_threads=32,
-            num_shots=20,
-        )
-        metrics = run_eval(args)
-        print(f"{metrics=}")
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_a_gsm8k (deepseek-v32-cp-in-seq-split)\n"
-                f'{metrics["score"]=:.3f}\n'
-            )
-            self.assertGreater(metrics["score"], 0.935)
-
-
-class TestDeepseekV32CPRoundRobinSplit(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        cls.model = DEEPSEEK_V32_MODEL_PATH
-        cls.base_url = DEFAULT_URL_FOR_TEST
-        other_args = [
-            "--trust-remote-code",
-            "--tp",
-            "8",
-            "--attn-cp-size",
-            "8",
-            "--enable-dsa-prefill-context-parallel",
-            "--dsa-prefill-cp-mode",
-            "round-robin-split",
-            "--speculative-algorithm",
-            "EAGLE",
-            "--speculative-num-steps",
-            "3",
-            "--speculative-eagle-topk",
-            "1",
-            "--speculative-num-draft-tokens",
-            "4",
-            "--mem-frac",
-            "0.85",
-            "--cuda-graph-max-bs-decode",
-            "32",
-            "--max-running-requests",
-            "32",
-            "--model-loader-extra-config",
-            '{"enable_multithread_load": true, "num_threads": 64}',
-        ]
-        cls.process = popen_launch_server(
-            cls.model,
-            cls.base_url,
-            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
-            other_args=other_args,
-            env={"SGLANG_ENABLE_CP_V2": "0"},
-        )
-
-    @classmethod
-    def tearDownClass(cls):
-        kill_process_tree(cls.process.pid)
-
-    def test_a_gsm8k(
-        self,
-    ):  # Append an "a" to make this test run first (alphabetically) to warm up the server
-        args = SimpleNamespace(
-            base_url=self.base_url,
-            model=self.model,
-            eval_name="gsm8k",
-            api="completion",
-            max_tokens=512,
-            num_examples=500,
-            num_threads=32,
-            num_shots=20,
-        )
-        metrics = run_eval(args)
-        print(f"{metrics=}")
-
-        if is_in_ci():
-            write_github_step_summary(
-                f"### test_a_gsm8k (deepseek-v32-cp-in-seq-split)\n"
-                f'{metrics["score"]=:.3f}\n'
-            )
-            self.assertGreater(metrics["score"], 0.935)
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/cp/test_gqa_prefill_cp.py

@@ -48,7 +48,6 @@ class TestGQACP2TP2EP2(CustomTestCase):
                 "--model-loader-extra-config",
                 '{"enable_multithread_load": true, "num_threads": 64}',
             ],
-            env={"SGLANG_ENABLE_CP_V2": "1"},
         )
 
     @classmethod
@@ -107,7 +106,6 @@ class TestGQACPTP2CP2EP4(CustomTestCase):
                 "--model-loader-extra-config",
                 '{"enable_multithread_load": true, "num_threads": 64}',
             ],
-            env={"SGLANG_ENABLE_CP_V2": "1"},
         )
 
     @classmethod

test/registered/hicache/test_hicache_storage_mooncake_backend.py

@@ -279,7 +279,8 @@ class TestMooncakeBackendQwen330BCP2(
                 "--tp-size": 2,
                 "--moe-dp-size": 2,
                 "--attn-cp-size": 2,
-                "--enable-prefill-context-parallel": True,
+                "--enable-prefill-cp": True,
+                "--cp-strategy": "zigzag",
                 "--trust-remote-code": True,
                 "--cuda-graph-max-bs-decode": 32,
                 "--max-running-requests": 32,

test/registered/kernels/ops/attention/test_cp_prefix_len_fa3_parity.py

@@ -1,143 +0,0 @@
-"""
-FA3 parity test for `prepare_context_parallel_metadata`.
-
-Drives the real function and feeds its `kv_len_prev/next_tensor` into FA3
-via `flash_attn_with_kvcache`. Compares per-rank CP output against a
-full-sequence FA3 reference computed over the unpadded `(prefix + extend)`
-KV. Any discrepancy indicates the metadata function emitted wrong
-`cache_seqlens` for at least one rank.
-"""
-
-import unittest
-from unittest.mock import patch
-
-import torch
-
-from sglang.srt.layers.utils.cp_utils import prepare_context_parallel_metadata
-from sglang.srt.utils.common import ceil_align
-from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.test_utils import CustomTestCase
-
-register_cuda_ci(est_time=5, stage="extra-a", runner_config="1-gpu-large")
-
-_DSA_UTILS = "sglang.srt.layers.attention.dsa.utils"
-_DEVICE = "cuda"
-_DTYPE = torch.bfloat16
-_HEAD_NUM = 8
-_HEAD_DIM = 128
-_SCALE = _HEAD_DIM**-0.5
-
-
-class TestCPPrefixLenFA3Parity(CustomTestCase):
-    """Per-rank FA3 output under CP must match a full-sequence reference."""
-
-    def _run_parity(self, prefix_len: int, extend_len: int, cp_size: int):
-        from sgl_kernel.flash_attn import flash_attn_with_kvcache
-
-        torch.manual_seed(extend_len * 1_000_003 + prefix_len * 101 + cp_size)
-
-        padded_extend = ceil_align(extend_len, cp_size)
-        pad = padded_extend - extend_len
-        self.assertGreaterEqual(
-            padded_extend,
-            2 * cp_size,
-            "runtime `can_cp_split` would skip this case; pick a larger extend",
-        )
-
-        # Reference: one full-sequence FA3 call over the unpadded KV.
-        q_full = torch.randn(
-            extend_len, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
-        )
-        k_full = torch.randn(
-            prefix_len + extend_len, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
-        )
-        v_full = torch.randn(
-            prefix_len + extend_len, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
-        )
-        ref = flash_attn_with_kvcache(
-            q=q_full.unsqueeze(0),
-            k_cache=k_full.unsqueeze(0),
-            v_cache=v_full.unsqueeze(0),
-            cache_seqlens=torch.tensor(
-                [k_full.shape[0]], dtype=torch.int32, device=_DEVICE
-            ),
-            softmax_scale=_SCALE,
-            causal=True,
-        ).squeeze(0)
-
-        # CP path sees tensors padded to `ceil_align(extend, cp_size)`,
-        # matching what `prepare_mlp_sync_batch` does in production.
-        zeros = torch.zeros(pad, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE)
-        q_padded = torch.cat([q_full, zeros], dim=0)
-        k_padded = torch.cat([k_full, zeros], dim=0)
-        v_padded = torch.cat([v_full, zeros], dim=0)
-
-        seqs_len = [prefix_len + extend_len]
-        extend_lens = [extend_len]
-
-        def _call_meta(rank: int):
-            return prepare_context_parallel_metadata(
-                padded_extend, rank, cp_size, seqs_len, extend_seqs_len=extend_lens
-            )
-
-        # Exercise the non-DSA branch; the DSA branch uses a separate
-        # `prefix_len` pathway re-added by `_get_topk_ragged_with_cp`.
-        with (
-            patch(f"{_DSA_UTILS}.is_dsa_enable_prefill_cp", return_value=False),
-            patch(
-                f"{_DSA_UTILS}.is_dsa_prefill_cp_round_robin_split",
-                return_value=False,
-            ),
-        ):
-            meta0 = _call_meta(0)
-            cp_segment_num = 2 * cp_size
-            blocks_q = list(torch.split(q_padded, meta0.split_list, dim=0))
-            outs = [None] * cp_segment_num
-
-            for rank in range(cp_size):
-                meta = meta0 if rank == 0 else _call_meta(rank)
-                for idx, cs_tensor in (
-                    (rank, meta.kv_len_prev_tensor),
-                    (cp_size * 2 - rank - 1, meta.kv_len_next_tensor),
-                ):
-                    if meta0.split_list[idx] == 0:
-                        outs[idx] = torch.empty(
-                            0, _HEAD_NUM, _HEAD_DIM, device=_DEVICE, dtype=_DTYPE
-                        )
-                        continue
-                    outs[idx] = flash_attn_with_kvcache(
-                        q=blocks_q[idx].unsqueeze(0),
-                        k_cache=k_padded.unsqueeze(0),
-                        v_cache=v_padded.unsqueeze(0),
-                        cache_seqlens=cs_tensor,
-                        softmax_scale=_SCALE,
-                        causal=True,
-                    ).squeeze(0)
-
-        cp_out = torch.cat(outs, dim=0)
-        err = (cp_out[:extend_len].float() - ref.float()).abs().max().item()
-
-        self.assertLess(
-            err,
-            1e-2,
-            f"CP output diverges from full-sequence FA3 reference by "
-            f"max_err={err:.5f} "
-            f"(prefix_len={prefix_len}, extend_len={extend_len}, "
-            f"cp_size={cp_size}, pad={pad})",
-        )
-
-    def test_cp2_prefix1_extend3(self):
-        """cp_size=2, prefix_len=1, extend_len=3 (pad=1)."""
-        self._run_parity(prefix_len=1, extend_len=3, cp_size=2)
-
-    def test_cp4_prefix1_extend7(self):
-        """cp_size=4, prefix_len=1, extend_len=7 (pad=1)."""
-        self._run_parity(prefix_len=1, extend_len=7, cp_size=4)
-
-    def test_cp8_prefix1_extend17(self):
-        """cp_size=8, prefix_len=1, extend_len=17 (pad=7)."""
-        self._run_parity(prefix_len=1, extend_len=17, cp_size=8)
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/kernels/ops/attention/test_mla_cp_fa3_parity.py

@@ -1,210 +0,0 @@
-"""FA3 numerical parity for MLA prefill CP.
-
-Verifies the rank-local zigzag-split FA3 path (``_mla_cp_attn`` +
-``cp_attn_forward_extend`` in ``flashattention_backend.py``) matches a
-single non-CP ``flash_attn_with_kvcache`` over the full sequence.
-
-Single-process, single-layer, pre-populated paged KV cache. Requires
-FA3 ver=3 (Hopper+).
-"""
-
-import math
-import sys
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-from sglang.srt.layers.utils.cp_utils import (
-    ContextParallelMetadata,
-    cp_attn_forward_extend,
-)
-from sglang.test.ci.ci_register import register_cuda_ci
-
-register_cuda_ci(est_time=30, stage="extra-a", runner_config="1-gpu-large")
-
-if not torch.cuda.is_available():
-    pytest.skip(reason="CUDA required for FA3", allow_module_level=True)
-
-_cap = torch.cuda.get_device_capability(0)
-if _cap[0] < 9:
-    pytest.skip(
-        reason=f"FA3 ver=3 requires Hopper (sm90+); got sm{_cap[0]}{_cap[1]}",
-        allow_module_level=True,
-    )
-
-try:
-    from sgl_kernel.flash_attn import flash_attn_with_kvcache
-except ImportError as e:
-    pytest.skip(
-        reason=f"sgl_kernel.flash_attn unavailable: {e}",
-        allow_module_level=True,
-    )
-
-DEVICE = torch.device("cuda")
-DTYPE = torch.bfloat16
-
-# Default shape is DeepSeek V3/R1 TP=8 MLA: 16 heads, v=512, rope=64.
-NUM_HEADS = 16
-V_HEAD_DIM = 512
-QK_ROPE_HEAD_DIM = 64
-PAGE_SIZE = 1
-
-
-def _build_cache_and_q(seq_len):
-    """Pre-populated paged KV cache + full-sequence q.
-
-    Pre-population mirrors upstream ``rebuild_cp_kv_cache``, which all-gathers
-    rank-local KV into the global pool before the attention call, so each
-    rank's FA3 invocation sees the same fully-populated cache.
-    """
-    num_pages = (seq_len + PAGE_SIZE - 1) // PAGE_SIZE
-    c_kv_cache = torch.randn(
-        num_pages, PAGE_SIZE, 1, V_HEAD_DIM, dtype=DTYPE, device=DEVICE
-    )
-    k_rope_cache = torch.randn(
-        num_pages, PAGE_SIZE, 1, QK_ROPE_HEAD_DIM, dtype=DTYPE, device=DEVICE
-    )
-    q_nope = torch.randn(seq_len, NUM_HEADS, V_HEAD_DIM, dtype=DTYPE, device=DEVICE)
-    q_rope = torch.randn(
-        seq_len, NUM_HEADS, QK_ROPE_HEAD_DIM, dtype=DTYPE, device=DEVICE
-    )
-    page_table = torch.arange(num_pages, dtype=torch.int32, device=DEVICE).unsqueeze(0)
-    return c_kv_cache, k_rope_cache, q_nope, q_rope, page_table
-
-
-def _full_seq_attn(
-    seq_len, q_nope, q_rope, c_kv_cache, k_rope_cache, page_table, softmax_scale
-):
-    """Non-CP reference: single flash_attn_with_kvcache over the full seq."""
-    return flash_attn_with_kvcache(
-        q=q_rope,
-        qv=q_nope,
-        k_cache=k_rope_cache,
-        v_cache=c_kv_cache,
-        page_table=page_table,
-        cache_seqlens=torch.tensor([seq_len], dtype=torch.int32, device=DEVICE),
-        cu_seqlens_q=torch.tensor([0, seq_len], dtype=torch.int32, device=DEVICE),
-        cu_seqlens_k_new=None,
-        max_seqlen_q=seq_len,
-        softmax_scale=softmax_scale,
-        causal=True,
-        ver=3,
-    )
-
-
-def _cp_attn_for_rank(
-    rank,
-    cp_size,
-    block_size,
-    q_nope,
-    q_rope,
-    c_kv_cache,
-    k_rope_cache,
-    page_table,
-    softmax_scale,
-):
-    """Run the rank-local CP closure from ``flashattention_backend.py``.
-
-    Zigzag layout: rank r gets blocks [r, num_blocks - 1 - r] where
-    num_blocks = cp_size * 2. kv_len for each half is the cumulative KV
-    extent through the end of that block.
-    """
-    num_blocks = cp_size * 2
-    b_prev, b_next = rank, num_blocks - 1 - rank
-    prev_slice = slice(b_prev * block_size, (b_prev + 1) * block_size)
-    next_slice = slice(b_next * block_size, (b_next + 1) * block_size)
-
-    q_nope_local = torch.cat([q_nope[prev_slice], q_nope[next_slice]], dim=0)
-    q_rope_local = torch.cat([q_rope[prev_slice], q_rope[next_slice]], dim=0)
-    q_fused = torch.cat([q_nope_local, q_rope_local], dim=-1)
-
-    cp_meta = ContextParallelMetadata(
-        kv_len_prev_tensor=torch.tensor(
-            [(b_prev + 1) * block_size], dtype=torch.int32, device=DEVICE
-        ),
-        kv_len_next_tensor=torch.tensor(
-            [(b_next + 1) * block_size], dtype=torch.int32, device=DEVICE
-        ),
-        cu_seqlens_q_prev_tensor=torch.tensor(
-            [0, block_size], dtype=torch.int32, device=DEVICE
-        ),
-        cu_seqlens_q_next_tensor=torch.tensor(
-            [0, block_size], dtype=torch.int32, device=DEVICE
-        ),
-        max_seqlen_q_prev=block_size,
-        max_seqlen_q_next=block_size,
-        total_q_prev_tokens=block_size,
-    )
-    fb = SimpleNamespace(attn_cp_metadata=cp_meta)
-
-    def _mla_cp_attn(q_chunk, cu_seqlens_q_cp, cache_seqlens_cp, max_seqlen_q_cp):
-        q_nope_chunk = q_chunk[..., :V_HEAD_DIM]
-        q_rope_chunk = q_chunk[..., V_HEAD_DIM:]
-        return flash_attn_with_kvcache(
-            q=q_rope_chunk,
-            qv=q_nope_chunk,
-            k_cache=k_rope_cache,
-            v_cache=c_kv_cache,
-            page_table=page_table,
-            cache_seqlens=cache_seqlens_cp,
-            cu_seqlens_q=cu_seqlens_q_cp,
-            cu_seqlens_k_new=None,
-            max_seqlen_q=max_seqlen_q_cp,
-            softmax_scale=softmax_scale,
-            causal=True,
-            ver=3,
-        )
-
-    local_out = cp_attn_forward_extend(fb, q_fused, DEVICE, _mla_cp_attn)
-    return local_out, prev_slice, next_slice
-
-
-@pytest.mark.parametrize(
-    "cp_size, block_size",
-    [
-        (2, 64),  # DSv3 TP=8 baseline
-        (2, 128),  # longer per-block seq
-        (4, 32),  # multi-rank zigzag: rank r gets blocks [r, 7-r]
-    ],
-)
-def test_cp_parity(cp_size, block_size):
-    torch.manual_seed(0)
-    seq_len = block_size * cp_size * 2
-    softmax_scale = 1.0 / math.sqrt(V_HEAD_DIM + QK_ROPE_HEAD_DIM)
-
-    c_kv_cache, k_rope_cache, q_nope, q_rope, page_table = _build_cache_and_q(seq_len)
-    ref_out = _full_seq_attn(
-        seq_len, q_nope, q_rope, c_kv_cache, k_rope_cache, page_table, softmax_scale
-    )
-
-    for rank in range(cp_size):
-        local_out, prev_slice, next_slice = _cp_attn_for_rank(
-            rank,
-            cp_size,
-            block_size,
-            q_nope,
-            q_rope,
-            c_kv_cache,
-            k_rope_cache,
-            page_table,
-            softmax_scale,
-        )
-        torch.testing.assert_close(
-            local_out[:block_size],
-            ref_out[prev_slice],
-            rtol=1e-3,
-            atol=5e-3,
-            msg=f"rank={rank} prev-half mismatch",
-        )
-        torch.testing.assert_close(
-            local_out[block_size:],
-            ref_out[next_slice],
-            rtol=1e-3,
-            atol=5e-3,
-            msg=f"rank={rank} next-half mismatch",
-        )
-
-
-if __name__ == "__main__":
-    sys.exit(pytest.main([__file__, "-v"]))

test/registered/pp/test_pp_parallel_compat.py

@@ -90,7 +90,6 @@ class TestQwen3MoePPxCP(_Qwen3MoePPCompatMixin, CustomTestCase):
         "--cp-strategy",
         "zigzag",
     ]
-    server_env = {"SGLANG_ENABLE_CP_V2": "1"}
 
 
 class TestQwen3MoePPxDP(_Qwen3MoePPCompatMixin, CustomTestCase):

test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py

@@ -39,7 +39,9 @@ class TestUnifiedQwen3HiCacheCP(UnifiedRadixTreeTestMixin, CustomTestCase):
                 "4",
                 "--attn-cp-size",
                 "2",
-                "--enable-prefill-context-parallel",
+                "--enable-prefill-cp",
+                "--cp-strategy",
+                "zigzag",
                 "--mem-fraction-static",
                 "0.8",
                 "--cuda-graph-max-bs-decode",