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

Gate CP split validation to eligible extend forwards

Browse Source 
CP shared-KV marks all forwards with uses_cp_shared_kv, but TARGET_VERIFY/decode style forwards may legitimately carry no extend prefix page-plan. The CP split helper previously ran the hard page-plan validator before checking context-parallel extend eligibility, so non-extend spec paths could fail before the later no-split decision.

Hoist the eligibility check and only validate page-plan metadata for real context-parallel EXTEND or shared-KV draft extend forwards. Add a regression that TARGET_VERIFY with no prefix metadata returns no-split instead of failing.

Constraint: Absorb syh 7fea88278 only; MQA logits chunk and per-forward budget changes are intentionally excluded.
Rejected: Relax the validator globally | real prefill page-plan violations must remain fail-fast.
Confidence: high
Scope-risk: narrow
Directive: Do not run CP shared-KV page-plan validation for non-context-parallel forward modes without proving those modes own extend_prefix_lens_cpu.
Tested: Remote g0034 container py_compile for utils/test file.
Tested: Remote g0034 container pytest test_nsa_cp_utils.py -k can_cp_split: 8 passed, 92 deselected.
Not-tested: Full ETE speculative non-deepep MoE path.
This commit is contained in:
laoyao0822
2026-06-09 19:54:59 +08:00
parent 81eb138a26
commit df2a3696cd
2 changed files with 31 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
python/sglang/srt/layers/attention/nsa/utils.py
View File

@@ -1820,6 +1820,14 @@ def _validate_cp_shared_kv_cp_split_plan_inputs(
def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
_fail_if_cp_shared_kv_round_robin(forward_batch, op="can_cp_split")
# Only genuine context-parallel EXTEND forwards are eligible for in-seq CP
# splitting. CP shared-KV spec/decode batches can be flagged
# uses_cp_shared_kv=True while legitimately carrying no prefix page-plan,
# so the hard page-plan validator below must be gated on this eligibility.
is_context_parallel_extend = (
forward_batch.forward_mode.is_context_parallel_extend()
or _is_cp_shared_kv_draft_extend(forward_batch)
)
if is_nsa_prefill_cp_round_robin_split():
cur_cp_seq_len = seq_len // cp_size
assert seq_len % cp_size == 0, (
@@ -1834,9 +1842,8 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
if _is_cp_shared_kv_forward_batch(forward_batch):
cur_cp_seq_len = (
1
if _validate_cp_shared_kv_cp_split_plan_inputs(
forward_batch, cp_size
)
if is_context_parallel_extend
and _validate_cp_shared_kv_cp_split_plan_inputs(forward_batch, cp_size)
else 0
)
else:
@@ -1846,10 +1853,6 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
min_extend_token_count = 1
else:
min_extend_token_count = cp_size
is_context_parallel_extend = (
forward_batch.forward_mode.is_context_parallel_extend()
or _is_cp_shared_kv_draft_extend(forward_batch)
)
if (
cur_cp_seq_len != 0
and cp_size > 1

21
test/registered/unit/layers/test_nsa_cp_utils.py
View File

@@ -427,6 +427,27 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
):
can_cp_split(1089, 8, True, forward_batch)
def test_can_cp_split_skips_page_plan_validator_for_target_verify(self):
forward_batch = SimpleNamespace(
uses_cp_shared_kv=True,
extend_seq_lens_cpu=[64],
extend_prefix_lens_cpu=None,
token_to_kv_pool=SimpleNamespace(page_size=64),
forward_mode=ForwardMode.TARGET_VERIFY,
)
with (
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
return_value=False,
),
patch(
"sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
return_value=True,
),
):
self.assertFalse(can_cp_split(64, 8, True, forward_batch))
def test_can_cp_split_keeps_cp_for_radix_hit_suffix_with_one_page_per_rank(self):
class Mode:
def is_context_parallel_extend(self):