From 833be86c15044b312c09b5e3b0ceb73248049ce3 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Mon, 24 Aug 2026 20:15:34 -0700 Subject: [PATCH] [CP V1 Deprecation 1/5] Migrate tests to strategy-based prefill CP (#36222) --- test/manual/dsv4/test_b200_flash.py | 6 +- test/manual/dsv4/test_b200_pro.py | 6 +- test/manual/dsv4/test_b300_flash.py | 6 +- test/manual/dsv4/test_b300_pro.py | 6 +- test/manual/dsv4/test_gb300_flash.py | 6 +- test/manual/dsv4/test_gb300_pro.py | 6 +- test/manual/dsv4/test_h200_fp8_flash.py | 6 +- test/manual/test_kv_events.py | 4 +- test/manual/test_qwen3_235b.py | 3 +- test/registered/cp/test_cp_strategy_unit.py | 15 +- .../cp/test_deepseek_v3_cp_single_node.py | 6 +- .../cp/test_deepseek_v4_flash_fp4_b200_cp.py | 4 - test/registered/cp/test_dsa_prefill_cp.py | 1 - .../cp/test_dsa_prefill_cp_legacy.py | 157 ------------- test/registered/cp/test_gqa_prefill_cp.py | 2 - .../test_hicache_storage_mooncake_backend.py | 3 +- .../test_cp_prefix_len_fa3_parity.py | 143 ------------ .../ops/attention/test_mla_cp_fa3_parity.py | 210 ------------------ test/registered/pp/test_pp_parallel_compat.py | 1 - .../test_unified_radix_cache_kl_cp.py | 4 +- 20 files changed, 37 insertions(+), 558 deletions(-) delete mode 100644 test/registered/cp/test_dsa_prefill_cp_legacy.py delete mode 100644 test/registered/kernels/ops/attention/test_cp_prefix_len_fa3_parity.py delete mode 100644 test/registered/kernels/ops/attention/test_mla_cp_fa3_parity.py diff --git a/test/manual/dsv4/test_b200_flash.py b/test/manual/dsv4/test_b200_flash.py index bcf854aba..30df91e80 100644 --- a/test/manual/dsv4/test_b200_flash.py +++ b/test/manual/dsv4/test_b200_flash.py @@ -87,9 +87,9 @@ class TestB200FlashCP(DSV4FlashAime25TestBase): "4", "--moe-a2a-backend", "deepep", - "--enable-dsa-prefill-context-parallel", - "--dsa-prefill-cp-mode", - "round-robin-split", + "--enable-prefill-cp", + "--cp-strategy", + "interleave", "--chunked-prefill-size", "16384", "--mem-fraction-static", diff --git a/test/manual/dsv4/test_b200_pro.py b/test/manual/dsv4/test_b200_pro.py index b246413b3..f95f29020 100644 --- a/test/manual/dsv4/test_b200_pro.py +++ b/test/manual/dsv4/test_b200_pro.py @@ -101,9 +101,9 @@ class TestB200ProCP(DSV4ProAime25TestBase): "8", "--moe-a2a-backend", "deepep", - "--enable-dsa-prefill-context-parallel", - "--dsa-prefill-cp-mode", - "round-robin-split", + "--enable-prefill-cp", + "--cp-strategy", + "interleave", "--chunked-prefill-size", "16384", "--mem-fraction-static", diff --git a/test/manual/dsv4/test_b300_flash.py b/test/manual/dsv4/test_b300_flash.py index 72279b9f0..422299718 100644 --- a/test/manual/dsv4/test_b300_flash.py +++ b/test/manual/dsv4/test_b300_flash.py @@ -89,9 +89,9 @@ class TestB300FlashCP(DSV4FlashAime25TestBase): "4", "--moe-a2a-backend", "deepep", - "--enable-dsa-prefill-context-parallel", - "--dsa-prefill-cp-mode", - "round-robin-split", + "--enable-prefill-cp", + "--cp-strategy", + "interleave", "--chunked-prefill-size", "16384", "--mem-fraction-static", diff --git a/test/manual/dsv4/test_b300_pro.py b/test/manual/dsv4/test_b300_pro.py index 3a81d497f..0825e8510 100644 --- a/test/manual/dsv4/test_b300_pro.py +++ b/test/manual/dsv4/test_b300_pro.py @@ -103,9 +103,9 @@ class TestB300ProCP(DSV4ProAime25TestBase): "8", "--moe-a2a-backend", "deepep", - "--enable-dsa-prefill-context-parallel", - "--dsa-prefill-cp-mode", - "round-robin-split", + "--enable-prefill-cp", + "--cp-strategy", + "interleave", "--chunked-prefill-size", "16384", "--mem-fraction-static", diff --git a/test/manual/dsv4/test_gb300_flash.py b/test/manual/dsv4/test_gb300_flash.py index 5553b8c00..e19abae2f 100644 --- a/test/manual/dsv4/test_gb300_flash.py +++ b/test/manual/dsv4/test_gb300_flash.py @@ -87,9 +87,9 @@ class TestGB300FlashCP(DSV4FlashAime25TestBase): "4", "--moe-a2a-backend", "deepep", - "--enable-dsa-prefill-context-parallel", - "--dsa-prefill-cp-mode", - "round-robin-split", + "--enable-prefill-cp", + "--cp-strategy", + "interleave", "--chunked-prefill-size", "16384", "--mem-fraction-static", diff --git a/test/manual/dsv4/test_gb300_pro.py b/test/manual/dsv4/test_gb300_pro.py index f87b6e9b6..363296c47 100644 --- a/test/manual/dsv4/test_gb300_pro.py +++ b/test/manual/dsv4/test_gb300_pro.py @@ -103,9 +103,9 @@ class TestGB300ProCP(DSV4ProAime25TestBase): "4", "--moe-a2a-backend", "deepep", - "--enable-dsa-prefill-context-parallel", - "--dsa-prefill-cp-mode", - "round-robin-split", + "--enable-prefill-cp", + "--cp-strategy", + "interleave", "--chunked-prefill-size", "16384", "--mem-fraction-static", diff --git a/test/manual/dsv4/test_h200_fp8_flash.py b/test/manual/dsv4/test_h200_fp8_flash.py index 5e3b8809d..684272b9d 100644 --- a/test/manual/dsv4/test_h200_fp8_flash.py +++ b/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", - "--dsa-prefill-cp-mode", - "round-robin-split", + "--enable-prefill-cp", + "--cp-strategy", + "interleave", "--chunked-prefill-size", "16384", "--mem-fraction-static", diff --git a/test/manual/test_kv_events.py b/test/manual/test_kv_events.py index d4b9dbc8c..01607eb62 100644 --- a/test/manual/test_kv_events.py +++ b/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, diff --git a/test/manual/test_qwen3_235b.py b/test/manual/test_qwen3_235b.py index 19606a819..187c849ac 100644 --- a/test/manual/test_qwen3_235b.py +++ b/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 = [ diff --git a/test/registered/cp/test_cp_strategy_unit.py b/test/registered/cp/test_cp_strategy_unit.py index 1e8956892..4b8d6f8bd 100644 --- a/test/registered/cp/test_cp_strategy_unit.py +++ b/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 ): diff --git a/test/registered/cp/test_deepseek_v3_cp_single_node.py b/test/registered/cp/test_deepseek_v3_cp_single_node.py index 0c68500fd..2faa11005 100644 --- a/test/registered/cp/test_deepseek_v3_cp_single_node.py +++ b/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) diff --git a/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py b/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py index 56c1f0748..c618e0d31 100644 --- a/test/registered/cp/test_deepseek_v4_flash_fp4_b200_cp.py +++ b/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 diff --git a/test/registered/cp/test_dsa_prefill_cp.py b/test/registered/cp/test_dsa_prefill_cp.py index 3d82f33d7..357360d0e 100644 --- a/test/registered/cp/test_dsa_prefill_cp.py +++ b/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 diff --git a/test/registered/cp/test_dsa_prefill_cp_legacy.py b/test/registered/cp/test_dsa_prefill_cp_legacy.py deleted file mode 100644 index 34053f798..000000000 --- a/test/registered/cp/test_dsa_prefill_cp_legacy.py +++ /dev/null @@ -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() diff --git a/test/registered/cp/test_gqa_prefill_cp.py b/test/registered/cp/test_gqa_prefill_cp.py index f4a939d85..a1570fadb 100644 --- a/test/registered/cp/test_gqa_prefill_cp.py +++ b/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 diff --git a/test/registered/hicache/test_hicache_storage_mooncake_backend.py b/test/registered/hicache/test_hicache_storage_mooncake_backend.py index eec88bcaf..d3fa80654 100644 --- a/test/registered/hicache/test_hicache_storage_mooncake_backend.py +++ b/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, diff --git a/test/registered/kernels/ops/attention/test_cp_prefix_len_fa3_parity.py b/test/registered/kernels/ops/attention/test_cp_prefix_len_fa3_parity.py deleted file mode 100644 index d44f99552..000000000 --- a/test/registered/kernels/ops/attention/test_cp_prefix_len_fa3_parity.py +++ /dev/null @@ -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() diff --git a/test/registered/kernels/ops/attention/test_mla_cp_fa3_parity.py b/test/registered/kernels/ops/attention/test_mla_cp_fa3_parity.py deleted file mode 100644 index e4d841bb2..000000000 --- a/test/registered/kernels/ops/attention/test_mla_cp_fa3_parity.py +++ /dev/null @@ -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"])) diff --git a/test/registered/pp/test_pp_parallel_compat.py b/test/registered/pp/test_pp_parallel_compat.py index 39c77a03d..be848cc4d 100644 --- a/test/registered/pp/test_pp_parallel_compat.py +++ b/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): diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py index 40b8c733e..ad196a8bc 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_cp.py +++ b/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",