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

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

Browse Source
This commit is contained in:
Baizhou Zhang
2026-06-10 14:11:38 -07:00
committed by GitHub
parent 73d0989d9b
commit 3c1b0fb226
5 changed files with 318 additions and 98 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/manual/test_dsa_alias_cli_registry_env.py
+6 -45
View File
@@ -2,11 +2,10 @@
Manual test for step 01: NSA → DSA user-facing alias layer.
Tests:
1. CLI: --dsa-* canonical flags write to dsa_* attrs
2. CLI: --nsa-* deprecated flags write to dsa_* attrs + log deprecation warning
3. Registry: "dsa" key creates the backend; "nsa" key triggers DeprecationWarning
4. Env: SGLANG_DSA_* canonical vars work
5. Env: SGLANG_NSA_* deprecated vars fall back to SGLANG_DSA_* with DeprecationWarning
1. CLI: --dsa-* non-CP canonical flags write to dsa_* attrs
2. Registry: "dsa" key creates the backend; "nsa" key triggers DeprecationWarning
3. Env: SGLANG_DSA_* canonical vars work
4. Env: SGLANG_NSA_* deprecated vars fall back to SGLANG_DSA_* with DeprecationWarning
Run:
python test/manual/test_dsa_alias_cli_registry_env.py
@@ -27,17 +26,13 @@ class TestDSAChoicesAndFields(unittest.TestCase):
def setUp(self):
from sglang.srt.server_args import (
DSA_CHOICES,
DSA_PREFILL_CP_SPLIT_CHOICES,
NSA_CHOICES,
NSA_PREFILL_CP_SPLIT_CHOICES,
ServerArgs,
)
self.ServerArgs = ServerArgs
self.DSA_CHOICES = DSA_CHOICES
self.NSA_CHOICES = NSA_CHOICES
self.DSA_PREFILL_CP_SPLIT_CHOICES = DSA_PREFILL_CP_SPLIT_CHOICES
self.NSA_PREFILL_CP_SPLIT_CHOICES = NSA_PREFILL_CP_SPLIT_CHOICES
def test_dsa_choices_is_canonical(self):
self.assertIn("fa3", self.DSA_CHOICES)
@@ -50,18 +45,10 @@ class TestDSAChoicesAndFields(unittest.TestCase):
"NSA_CHOICES must be the same object as DSA_CHOICES",
)
def test_nsa_cp_split_choices_is_alias(self):
self.assertIs(
self.NSA_PREFILL_CP_SPLIT_CHOICES,
self.DSA_PREFILL_CP_SPLIT_CHOICES,
)
def test_serverargs_has_dsa_fields(self):
sa = self.ServerArgs
self.assertTrue(hasattr(sa, "dsa_prefill_backend"))
self.assertTrue(hasattr(sa, "dsa_decode_backend"))
self.assertTrue(hasattr(sa, "enable_dsa_prefill_context_parallel"))
self.assertTrue(hasattr(sa, "dsa_prefill_cp_mode"))
def test_serverargs_no_nsa_fields(self):
"""The nsa_* attributes should no longer exist on ServerArgs."""
@@ -74,12 +61,10 @@ class TestDSAChoicesAndFields(unittest.TestCase):
hasattr(sa, "nsa_decode_backend"),
"nsa_decode_backend should have been renamed",
)
self.assertFalse(hasattr(sa, "enable_nsa_prefill_context_parallel"))
self.assertFalse(hasattr(sa, "nsa_prefill_cp_mode"))
class TestCLICanonicalFlags(unittest.TestCase):
"""--dsa-* canonical flags write to dsa_* attributes with no warning."""
"""Canonical flags write to canonical attributes with no warning."""
def setUp(self):
from sglang.srt.server_args import ServerArgs
@@ -98,20 +83,10 @@ class TestCLICanonicalFlags(unittest.TestCase):
args = self._parse(["--dsa-decode-backend", "tilelang"])
self.assertEqual(args.dsa_decode_backend, "tilelang")
def test_enable_dsa_prefill_cp_canonical(self):
args = self._parse(["--enable-dsa-prefill-context-parallel"])
self.assertTrue(args.enable_dsa_prefill_context_parallel)
def test_dsa_prefill_cp_mode_canonical(self):
args = self._parse(["--dsa-prefill-cp-mode", "in-seq-split"])
self.assertEqual(args.dsa_prefill_cp_mode, "in-seq-split")
def test_defaults_are_none_or_false(self):
args = self._parse([])
self.assertIsNone(args.dsa_prefill_backend)
self.assertIsNone(args.dsa_decode_backend)
self.assertFalse(args.enable_dsa_prefill_context_parallel)
self.assertEqual(args.dsa_prefill_cp_mode, "round-robin-split")
def test_attention_backend_dsa_key_in_choices(self):
args = self._parse(["--attention-backend", "dsa"])
@@ -119,7 +94,7 @@ class TestCLICanonicalFlags(unittest.TestCase):
class TestCLIDeprecatedFlags(unittest.TestCase):
"""--nsa-* deprecated flags write to dsa_* attributes and emit logger warning."""
"""Deprecated flags write to canonical attributes and emit logger warning."""
def setUp(self):
import logging
@@ -174,20 +149,6 @@ class TestCLIDeprecatedFlags(unittest.TestCase):
self.assertEqual(args.dsa_decode_backend, "tilelang")
self.assertIn("deprecated", log_output.lower())
def test_enable_nsa_prefill_cp_deprecated(self):
args, log_output = self._parse_capture_warnings(
["--enable-nsa-prefill-context-parallel"]
)
self.assertTrue(args.enable_dsa_prefill_context_parallel)
self.assertIn("deprecated", log_output.lower())
def test_nsa_prefill_cp_mode_deprecated(self):
args, log_output = self._parse_capture_warnings(
["--nsa-prefill-cp-mode", "in-seq-split"]
)
self.assertEqual(args.dsa_prefill_cp_mode, "in-seq-split")
self.assertIn("deprecated", log_output.lower())
def test_attention_backend_nsa_still_accepted(self):
"""attention_backend='nsa' still parses without error (registry handles the deprecation)."""
args = self._parse(["--attention-backend", "nsa"])
test/registered/unit/server_args/test_server_args.py
+169 -1
View File
@@ -97,6 +97,174 @@ class TestLoadBalanceMethod(unittest.TestCase):
self.assertIn("'fake'", str(context.exception))
class TestContextParallelServerArgs(CustomTestCase):
def setUp(self):
self.parser = server_args_module.argparse.ArgumentParser()
ServerArgs.add_cli_args(self.parser)
def _new_cp_args(self, **overrides):
server_args = object.__new__(ServerArgs)
defaults = dict(
enable_prefill_context_parallel=False,
enable_dsa_prefill_context_parallel=False,
enable_prefill_cp=False,
cp_strategy=None,
dsa_prefill_cp_mode="round-robin-split",
prefill_cp_mode="in-seq-split",
attn_cp_size=1,
tp_size=1,
dp_size=1,
moe_dp_size=1,
ep_size=1,
pp_size=1,
enable_aiter_allreduce_fusion=False,
)
defaults.update(overrides)
for key, value in defaults.items():
setattr(server_args, key, value)
return server_args
def test_canonical_prefill_cp_cli_sets_unified_fields(self):
args = self.parser.parse_args(
["--model", "dummy", "--enable-prefill-cp", "--cp-strategy", "interleave"]
)
self.assertTrue(args.enable_prefill_cp)
self.assertEqual(args.cp_strategy, "interleave")
def test_canonical_prefill_cp_requires_strategy(self):
args = self.parser.parse_args(["--model", "dummy", "--enable-prefill-cp"])
self.assertTrue(args.enable_prefill_cp)
self.assertIsNone(args.cp_strategy)
server_args = self._new_cp_args(
enable_prefill_cp=args.enable_prefill_cp,
cp_strategy=args.cp_strategy,
)
with self.assertRaisesRegex(ValueError, "--cp-strategy"):
server_args._handle_context_parallelism()
def test_deprecated_dsa_cp_mode_maps_to_unified_strategy(self):
args = self.parser.parse_args(
[
"--model",
"dummy",
"--enable-dsa-prefill-context-parallel",
"--dsa-prefill-cp-mode",
"round-robin-split",
]
)
server_args = self._new_cp_args(
enable_dsa_prefill_context_parallel=(
args.enable_dsa_prefill_context_parallel
),
dsa_prefill_cp_mode=args.dsa_prefill_cp_mode,
)
server_args._handle_legacy_cp_arguments()
self.assertTrue(server_args.enable_prefill_cp)
self.assertEqual(server_args.cp_strategy, "interleave")
self.assertEqual(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",
)
server_args._handle_legacy_cp_arguments()
server_args._handle_context_parallelism()
self.assertTrue(server_args.enable_dsa_prefill_context_parallel)
self.assertFalse(server_args.enable_prefill_context_parallel)
self.assertEqual(server_args.dsa_prefill_cp_mode, "round-robin-split")
self.assertEqual(server_args.prefill_cp_mode, "round-robin-split")
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)
server_args._handle_legacy_cp_arguments()
server_args._handle_context_parallelism()
self.assertTrue(server_args.enable_prefill_cp)
self.assertEqual(server_args.cp_strategy, strategy)
self.assertEqual(server_args.dsa_prefill_cp_mode, mode)
self.assertEqual(server_args.prefill_cp_mode, mode)
self.assertEqual(
server_args.enable_dsa_prefill_context_parallel, expect_dsa
)
self.assertEqual(
server_args.enable_prefill_context_parallel, expect_generic
)
class TestPortArgs(unittest.TestCase):
@patch("sglang.srt.server_args.get_free_port")
@patch("sglang.srt.server_args.tempfile.NamedTemporaryFile")
@@ -647,7 +815,7 @@ class TestPrefillOnlyDisableKvCache(unittest.TestCase):
ServerArgs(**self._base_kwargs(attn_cp_size=2, tp_size=2))
def test_rejects_prefill_context_parallel(self):
with self.assertRaisesRegex(ValueError, "--enable-prefill-context-parallel"):
with self.assertRaisesRegex(ValueError, "--enable-prefill-cp"):
ServerArgs(**self._base_kwargs(enable_prefill_context_parallel=True))
def test_rejects_hisparse(self):