minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CP V1 Deprecation 2/5] Make strategy prefill CP canonical (#36223)

Browse Source
This commit is contained in:
Baizhou Zhang
2026-09-03 20:24:19 -07:00
committed by GitHub
parent 59799a3687
commit ff1285cc28
9 changed files with 240 additions and 279 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/cp/test_cp_strategy_unit.py
+48 -41
View File
@@ -5,6 +5,7 @@ from unittest.mock import patch
import torch
from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.cp.base import (
ContextParallelStrategyKind,
get_cp_strategy,
@@ -94,12 +95,43 @@ class TestCPStrategyUnit(CustomTestCase):
cp_strategy="interleave",
)
with patch(
"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())
def test_hip_dsa_cp_uses_protected_legacy_runtime_flag(self):
parallel = SimpleNamespace(
enable_dsa_prefill_context_parallel=False,
attn_cp_size=2,
)
model_config = SimpleNamespace(hf_config=SimpleNamespace())
with (
patch(
"sglang.srt.layers.attention.dsa.utils.get_parallel",
return_value=parallel,
),
patch(
"sglang.srt.layers.attention.dsa.utils.process_model_config",
return_value=model_config,
),
patch("sglang.srt.layers.attention.dsa.utils.is_hip", return_value=True),
patch(
"sglang.srt.configs.model_config.is_deepseek_dsa",
return_value=True,
),
):
self.assertIsNotNone(get_cp_strategy())
self.assertTrue(is_cp_enabled())
self.assertTrue(is_interleave())
self.assertFalse(is_dsa_enable_prefill_cp())
@patch("sglang.srt.utils.is_npu", return_value=False)
@patch("sglang.srt.utils.is_hip", return_value=True)
def test_hip_keeps_strategy_cp_disabled(self, _mock_is_hip, _mock_is_npu):
self.assertFalse(enable_cp_v2())
@patch("sglang.srt.utils.is_npu", return_value=True)
@patch("sglang.srt.utils.is_hip", return_value=False)
def test_npu_keeps_strategy_cp_disabled(self, _mock_is_hip, _mock_is_npu):
self.assertFalse(enable_cp_v2())
class TestPrefillCPBCGReplay(CustomTestCase):
@@ -155,10 +187,6 @@ class TestPrefillCPBCGReplay(CustomTestCase):
self._enable_zigzag()
with (
patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get",
return_value=True,
),
patch(
"sglang.srt.layers.cp.bcg.get_cp_padding_align_size",
return_value=8,
@@ -209,10 +237,6 @@ class TestPrefillCPBCGReplay(CustomTestCase):
with (
get_parallel().override(attn_cp_rank=0, attn_cp_size=4),
patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get",
return_value=True,
),
patch(
"sglang.srt.layers.cp.bcg.get_cp_padding_align_size",
return_value=8,
@@ -243,10 +267,6 @@ class TestPrefillCPBCGReplay(CustomTestCase):
with (
get_parallel().override(attn_cp_rank=0, attn_cp_size=4),
patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get",
return_value=True,
),
patch(
"sglang.srt.layers.cp.bcg.get_cp_padding_align_size",
return_value=8,
@@ -298,9 +318,7 @@ class TestCPZigzagStrategy(CustomTestCase):
extend_seq_lens_cpu=[7],
)
with patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True
):
with patch.dict("os.environ", {"SGLANG_ENABLE_CP_V2": "0"}):
self.assertTrue(enable_cp_v2())
self.assertTrue(is_cp_v2_active(active_batch))
self.assertFalse(is_cp_v2_active(inactive_batch))
@@ -494,14 +512,11 @@ class TestCPZigzagStrategy(CustomTestCase):
local_x = strategy.shard_hidden_states(x, fb)
local_positions = strategy.shard_position_ids(positions, fb)
with patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get", return_value=True
):
helper_x, helper_positions = cp_split_before_forward(
x,
positions,
fb,
)
helper_x, helper_positions = cp_split_before_forward(
x,
positions,
fb,
)
self.assertTrue(torch.equal(local_x, expected_x))
self.assertTrue(torch.equal(local_positions, expected_positions))
@@ -916,10 +931,6 @@ class TestCPInterleaveStrategy(CustomTestCase):
with (
get_parallel().override(attn_cp_rank=2, attn_cp_size=4),
patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get",
return_value=True,
),
patch(
"sglang.srt.layers.cp.padding.get_cp_padding_align_size",
return_value=4,
@@ -965,15 +976,11 @@ class TestCPInterleaveStrategy(CustomTestCase):
local_x = strategy.shard_hidden_states(x, fb)
local_positions = strategy.shard_position_ids(positions, fb)
with patch(
"sglang.srt.environ.envs.SGLANG_ENABLE_CP_V2.get",
return_value=True,
):
helper_x, helper_positions = cp_split_before_forward(
x,
positions,
fb,
)
helper_x, helper_positions = cp_split_before_forward(
x,
positions,
fb,
)
self.assertTrue(torch.equal(local_x, expected_x))
self.assertTrue(torch.equal(local_positions, expected_positions))
test/registered/unit/server_args/test_server_args.py
+107 -131
View File
@@ -47,7 +47,8 @@ from sglang.srt.arg_groups.overrides import (
from sglang.srt.arg_groups.parallel_hook import (
handle_context_parallelism,
handle_data_parallelism,
handle_legacy_cp_arguments,
handle_legacy_cp_runtime_compatibility,
handle_platform_cp_compatibility,
)
from sglang.srt.arg_groups.pd_disaggregation_hook import handle_pd_disaggregation
from sglang.srt.arg_groups.serving_hook import (
@@ -1056,52 +1057,116 @@ class TestContextParallelServerArgs(CustomTestCase):
with self.assertRaisesRegex(ValueError, "--cp-strategy"):
handle_context_parallelism(server_args)
def test_deprecated_dsa_cp_mode_maps_to_unified_strategy(self):
@override_platform(is_hip=False, is_npu=False)
def test_deepseek_v32_prefill_cp_rejects_zigzag(self):
server_args = self._new_cp_args(
model_path="deepseek-ai/DeepSeek-V3.2",
enable_prefill_cp=True,
cp_strategy="zigzag",
)
server_args._model_config = SimpleNamespace(
hf_config=SimpleNamespace(architectures=["DeepseekV32ForCausalLM"]),
is_multimodal=False,
)
with self.assertRaisesRegex(ValueError, "DeepSeek V3.2.*interleave"):
handle_context_parallelism(server_args)
@override_platform(is_hip=False, is_npu=False)
def test_generic_canonical_cp_mirrors_to_transitional_runtime_fields(self):
cases = (
(
"zigzag_mla_or_gqa",
"zigzag",
"fa3",
True,
False,
"in-seq-split",
),
(
"interleave_dsa",
"interleave",
"dsa",
False,
True,
"round-robin-split",
),
)
for name, strategy, backend, expect_generic, expect_dsa, mode in cases:
with self.subTest(name=name):
server_args = self._new_cp_args(
enable_prefill_cp=True,
cp_strategy=strategy,
attention_backend=backend,
)
handle_platform_cp_compatibility(server_args)
self.assertFalse(
resolution_result(server_args, "enable_prefill_context_parallel")
)
self.assertFalse(
resolution_result(
server_args, "enable_dsa_prefill_context_parallel"
)
)
handle_legacy_cp_runtime_compatibility(server_args)
self.assertEqual(
resolution_result(server_args, "enable_prefill_context_parallel"),
expect_generic,
)
self.assertEqual(
resolution_result(
server_args, "enable_dsa_prefill_context_parallel"
),
expect_dsa,
)
self.assertEqual(
resolution_result(server_args, "dsa_prefill_cp_mode"), mode
)
self.assertEqual(
resolution_result(server_args, "prefill_cp_mode"), mode
)
@override_platform(is_hip=False, is_npu=False)
def test_non_platform_legacy_prefill_cp_is_rejected(self):
server_args = ServerArgs(
model_path="instance://127.0.0.1:8000/dummy",
enable_prefill_context_parallel=True,
)
with self.assertRaisesRegex(ValueError, "HIP or Ascend NPU"):
handle_platform_cp_compatibility(server_args)
def test_generic_v1_cp_options_are_not_public_cli(self):
removed_options = (
("--enable-dsa-prefill-context-parallel", []),
("--dsa-prefill-cp-mode", ["round-robin-split"]),
("--prefill-cp-mode", ["in-seq-split"]),
)
for option, values in removed_options:
with self.subTest(option=option), self.assertRaises(SystemExit):
self.parser.parse_args(["--model", "dummy", option, *values])
def test_npu_cp_compatibility_options_remain_public_cli(self):
args = self.parser.parse_args(
[
"--model",
"dummy",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"--enable-prefill-context-parallel",
"--enable-nsa-prefill-context-parallel",
"--nsa-prefill-cp-mode",
"round-robin-split",
]
)
server_args = self._new_cp_args(
enable_dsa_prefill_context_parallel=(
resolution_result(args, "enable_dsa_prefill_context_parallel")
),
dsa_prefill_cp_mode=resolution_result(args, "dsa_prefill_cp_mode"),
)
handle_legacy_cp_arguments(server_args)
self.assertTrue(resolution_result(server_args, "enable_prefill_cp"))
self.assertEqual(resolution_result(server_args, "cp_strategy"), "interleave")
self.assertTrue(resolution_result(args, "enable_prefill_context_parallel"))
self.assertTrue(resolution_result(args, "enable_dsa_prefill_context_parallel"))
self.assertEqual(
resolution_result(server_args, "dsa_prefill_cp_mode"), "round-robin-split"
)
def test_canonical_interleave_cp_mirrors_to_dsa_runtime_aliases(self):
server_args = self._new_cp_args(
enable_prefill_cp=True,
cp_strategy="interleave",
attention_backend="dsa",
)
handle_legacy_cp_arguments(server_args)
handle_context_parallelism(server_args)
self.assertTrue(
resolution_result(server_args, "enable_dsa_prefill_context_parallel")
)
self.assertFalse(
resolution_result(server_args, "enable_prefill_context_parallel")
)
self.assertEqual(
resolution_result(server_args, "dsa_prefill_cp_mode"), "round-robin-split"
)
self.assertEqual(
resolution_result(server_args, "prefill_cp_mode"), "round-robin-split"
resolution_result(args, "dsa_prefill_cp_mode"), "round-robin-split"
)
def test_context_parallel_handler_initializes_cp_strategy(self):
@@ -1117,97 +1182,6 @@ class TestContextParallelServerArgs(CustomTestCase):
self.assertTrue(is_cp_enabled())
self.assertTrue(is_interleave())
def test_registered_cp_legacy_args_map_to_unified_strategy(self):
cases = [
(
"deepseek_v3_mla_cp",
dict(enable_prefill_context_parallel=True),
"zigzag",
"in-seq-split",
False,
True,
),
(
"qwen3_gqa_cp",
dict(
enable_prefill_context_parallel=True,
tp_size=4,
attn_cp_size=2,
),
"zigzag",
"in-seq-split",
False,
True,
),
(
"deepseek_v32_dsa_in_seq_split",
dict(
enable_dsa_prefill_context_parallel=True,
dsa_prefill_cp_mode="in-seq-split",
tp_size=8,
dp_size=2,
attn_cp_size=4,
),
"zigzag",
"in-seq-split",
True,
False,
),
(
"deepseek_v32_dsa_round_robin_split",
dict(
enable_dsa_prefill_context_parallel=True,
tp_size=8,
attn_cp_size=8,
),
"interleave",
"round-robin-split",
True,
False,
),
(
"deepseek_v4_flash_fp4_b200_dsa_round_robin_split",
dict(
enable_dsa_prefill_context_parallel=True,
dsa_prefill_cp_mode="round-robin-split",
tp_size=4,
attn_cp_size=4,
),
"interleave",
"round-robin-split",
True,
False,
),
]
for name, overrides, strategy, mode, expect_dsa, expect_generic in cases:
with self.subTest(name=name):
server_args = self._new_cp_args(**overrides)
handle_legacy_cp_arguments(server_args)
handle_context_parallelism(server_args)
self.assertTrue(resolution_result(server_args, "enable_prefill_cp"))
self.assertEqual(
resolution_result(server_args, "cp_strategy"), strategy
)
self.assertEqual(
resolution_result(server_args, "dsa_prefill_cp_mode"), mode
)
self.assertEqual(
resolution_result(server_args, "prefill_cp_mode"), mode
)
self.assertEqual(
resolution_result(
server_args, "enable_dsa_prefill_context_parallel"
),
expect_dsa,
)
self.assertEqual(
resolution_result(server_args, "enable_prefill_context_parallel"),
expect_generic,
)
class TestPortArgs(unittest.TestCase):
@patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
@@ -1804,7 +1778,6 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase):
def _validate_prefill_only_args(self, **overrides):
sa = ServerArgs(**self._base_kwargs(**overrides))
handle_legacy_cp_arguments(sa)
validate_prefill_only_disable_kv_cache_args(sa)
return sa
@@ -1830,7 +1803,10 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase):
def test_rejects_prefill_context_parallel(self):
with self.assertRaisesRegex(ValueError, "--enable-prefill-cp"):
self._validate_prefill_only_args(enable_prefill_context_parallel=True)
self._validate_prefill_only_args(
enable_prefill_cp=True,
cp_strategy="zigzag",
)
def test_rejects_hisparse(self):
with self.assertRaisesRegex(ValueError, "--enable-hisparse"):