Gate CP split validation to eligible extend forwards

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

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):