From bfc2b18e97ae9ae4cf1b149b25347b5c4c00a37d Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 12 Oct 2026 16:02:56 -0700 Subject: [PATCH 1/2] [DSV4][CP] Write low-ratio page indices when CP prefill runs paged attention #38603 excluded active CP batches from sparse prefill, but the low-ratio indexer's page-index gate still mirrored the old sparse gate. Eager CP prefill then ran paged attention without its c1/c2 page indices. Derive the gate from _use_sparse_prefill so the two stay in sync. --- .../layers/attention/deepseek_v4_backend.py | 7 ++- .../layers/cp/test_dsv4_prefill_cp_bcg.py | 44 +++++++++++++++++++ 2 files changed, 47 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 832c7678b089..8acf85216c4b 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -3323,15 +3323,14 @@ def _make_low_ratio_prefill_indexer_inputs( def _low_ratio_prefill_reads_page_indices(self, forward_batch: ForwardBatch): # The negation of the sparse-prefill gate in forward(), which reads only - # raw_indices; the query-count clause is dropped, so this errs to True. + # raw_indices; num_qo_tokens=0 drops the query-count clause, so this errs + # to True. Paged attention under active CP needs the page indices. return ( self.trtllm_attn - or not forward_batch.forward_mode.is_extend_without_speculative() - or get_platform().is_sm120 or self.tail_forward_metadata is not None or self.forward_metadata.late_layer_tail is not None or self.token_to_kv_pool.request_window is not None - or not envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get() + or not self._use_sparse_prefill(forward_batch, num_qo_tokens=0) ) def _publish_candidate_metadata(self, published: Optional[CandidateMetadata]): diff --git a/test/registered/unit/layers/cp/test_dsv4_prefill_cp_bcg.py b/test/registered/unit/layers/cp/test_dsv4_prefill_cp_bcg.py index ab32ef791b38..176b0ebe92fa 100644 --- a/test/registered/unit/layers/cp/test_dsv4_prefill_cp_bcg.py +++ b/test/registered/unit/layers/cp/test_dsv4_prefill_cp_bcg.py @@ -57,6 +57,50 @@ def test_sparse_prefill_never_uses_cp_local_queries(self): selected, not cp_active and (forced or num_tokens > 11673) ) + def test_low_ratio_indexer_writes_page_indices_for_paged_cp(self): + from sglang.srt.layers.attention.deepseek_v4_backend import ( + DeepseekV4AttnBackend, + ) + + backend = SimpleNamespace( + trtllm_attn=False, + tail_forward_metadata=None, + forward_metadata=SimpleNamespace(late_layer_tail=None), + token_to_kv_pool=SimpleNamespace(request_window=None), + ) + backend._use_sparse_prefill = lambda fb, *, num_qo_tokens: ( + DeepseekV4AttnBackend._use_sparse_prefill( + None, fb, num_qo_tokens=num_qo_tokens + ) + ) + batch = SimpleNamespace(forward_mode=ForwardMode.EXTEND) + for cp_active in (False, True): + for forced in (False, True): + with ( + patch( + "sglang.srt.layers.attention.deepseek_v4_backend.is_cp_active", + return_value=cp_active, + ), + patch( + "sglang.srt.layers.attention.deepseek_v4_backend.get_platform", + return_value=SimpleNamespace(is_sm120=False), + ), + patch( + "sglang.srt.layers.attention.deepseek_v4_backend.envs" + ) as envs, + ): + envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get.return_value = forced + reads = DeepseekV4AttnBackend._low_ratio_prefill_reads_page_indices( + backend, batch + ) + for num_tokens in (256, 11674): + # Any batch that runs paged attention must get page indices. + if not backend._use_sparse_prefill( + batch, num_qo_tokens=num_tokens + ): + self.assertTrue(reads) + self.assertEqual(reads, cp_active or not forced) + def test_sparse_metadata_rejects_short_query_buffer_before_launch(self): from sglang.srt.layers.attention.deepseek_v4_backend import ( DeepseekV4AttnBackend, From 62954207139e6691d2041a6907823dfae002d483 Mon Sep 17 00:00:00 2001 From: Yangmin Li Date: Mon, 12 Oct 2026 16:37:34 -0700 Subject: [PATCH 2/2] Drop the gate-agreement unit test --- .../layers/cp/test_dsv4_prefill_cp_bcg.py | 44 ------------------- 1 file changed, 44 deletions(-) diff --git a/test/registered/unit/layers/cp/test_dsv4_prefill_cp_bcg.py b/test/registered/unit/layers/cp/test_dsv4_prefill_cp_bcg.py index 176b0ebe92fa..ab32ef791b38 100644 --- a/test/registered/unit/layers/cp/test_dsv4_prefill_cp_bcg.py +++ b/test/registered/unit/layers/cp/test_dsv4_prefill_cp_bcg.py @@ -57,50 +57,6 @@ def test_sparse_prefill_never_uses_cp_local_queries(self): selected, not cp_active and (forced or num_tokens > 11673) ) - def test_low_ratio_indexer_writes_page_indices_for_paged_cp(self): - from sglang.srt.layers.attention.deepseek_v4_backend import ( - DeepseekV4AttnBackend, - ) - - backend = SimpleNamespace( - trtllm_attn=False, - tail_forward_metadata=None, - forward_metadata=SimpleNamespace(late_layer_tail=None), - token_to_kv_pool=SimpleNamespace(request_window=None), - ) - backend._use_sparse_prefill = lambda fb, *, num_qo_tokens: ( - DeepseekV4AttnBackend._use_sparse_prefill( - None, fb, num_qo_tokens=num_qo_tokens - ) - ) - batch = SimpleNamespace(forward_mode=ForwardMode.EXTEND) - for cp_active in (False, True): - for forced in (False, True): - with ( - patch( - "sglang.srt.layers.attention.deepseek_v4_backend.is_cp_active", - return_value=cp_active, - ), - patch( - "sglang.srt.layers.attention.deepseek_v4_backend.get_platform", - return_value=SimpleNamespace(is_sm120=False), - ), - patch( - "sglang.srt.layers.attention.deepseek_v4_backend.envs" - ) as envs, - ): - envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get.return_value = forced - reads = DeepseekV4AttnBackend._low_ratio_prefill_reads_page_indices( - backend, batch - ) - for num_tokens in (256, 11674): - # Any batch that runs paged attention must get page indices. - if not backend._use_sparse_prefill( - batch, num_qo_tokens=num_tokens - ): - self.assertTrue(reads) - self.assertEqual(reads, cp_active or not forced) - def test_sparse_metadata_rejects_short_query_buffer_before_launch(self): from sglang.srt.layers.attention.deepseek_v4_backend import ( DeepseekV4AttnBackend,