diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index 7838f1023..cf76f34b2 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -152,6 +152,16 @@ class AttentionBackend(ABC): return SharedReadEnds.IN_REPLAY return SharedReadEnds.UNKNOWN + def prepare_prefill_shared_read_snapshot( + self, forward_batch: ForwardBatch, *, num_qo_tokens: int + ) -> None: + """Snapshot late prefill reads before a PRE_REPLAY event is published. + + Runners call this only after the actual eager/replay query geometry is + known. Backends that retain scheduler-shared reads into the model + forward keep the default no-op and must not declare PRE_REPLAY. + """ + # Chunked-prefix FullCG capture has a second model topology and stable # prefix buffers. Backends must opt in explicitly so the runner does not # assume that generic ForwardBatch metadata is sufficient for every diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index ac2345c7f..315496a16 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -408,10 +408,10 @@ class DSV4Metadata: c4_compress_metadata: Optional[FusedCompressMetadata] = None c128_compress_metadata: Optional[FusedCompressMetadata] = None - # Lazily populated on the first call to ``_forward_prefill_sparse`` and - # reused across every layer in the chunk. Reset to ``None`` when graph - # metadata is refreshed so replay rebuilds it from the live batch. + # Built at the runner's prefill WAR boundary when the fast path is on, + # otherwise lazily by ``_forward_prefill_sparse``. sparse_prefill_cache: Optional[SparsePrefillChunkCache] = None + prefill_shared_reads_snapshotted: bool = False @property def core_metadata(self) -> DSV4AttnMetadata: @@ -425,6 +425,7 @@ class DSV4Metadata: self.c128_compress_metadata, src=other.c128_compress_metadata ) self.sparse_prefill_cache = None + self.prefill_shared_reads_snapshotted = False def refresh_for_breakable_cuda_graph_replay_(self, static_metadata: DSV4Metadata): self.core_attn_metadata.refresh_for_breakable_cuda_graph_replay_( @@ -444,6 +445,7 @@ class DSV4Metadata: src=static_metadata.c128_compress_metadata, ) self.sparse_prefill_cache = None + self.prefill_shared_reads_snapshotted = False @dataclass @@ -519,6 +521,13 @@ class DeepseekV4AttnBackend( if self.model_runner.spec_algorithm.is_dspark(): return SharedReadEnds.IN_REPLAY return SharedReadEnds.POST_REPLAY + metadata = self.forward_metadata + if ( + fm == ForwardMode.EXTEND + and isinstance(metadata, DSV4Metadata) + and metadata.prefill_shared_reads_snapshotted + ): + return SharedReadEnds.PRE_REPLAY return super().shared_read_ends(fm) def __init__( @@ -1352,6 +1361,66 @@ class DeepseekV4AttnBackend( self.forward_metadata = self._build_forward_metadata(forward_batch) self.init_forward_metadata_in_graph(forward_batch) + def prepare_prefill_shared_read_snapshot( + self, forward_batch: ForwardBatch, *, num_qo_tokens: int + ) -> None: + # Sparse prefill otherwise reads req_to_token/full_to_swa lazily in its + # first layer. DFLASH/DSPARK have no later prefill draft-extend reader; + # CP-v2 shards the query layout that this global snapshot assumes. + metadata = self.forward_metadata + if isinstance(metadata, DSV4Metadata): + metadata.prefill_shared_reads_snapshotted = False + snapshot_shared_prefill_reads = ( + envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.get() + and forward_batch.forward_mode == ForwardMode.EXTEND + and self.model_runner.spec_algorithm.is_dflash_family() + and not is_cp_v2_active(forward_batch) + ) + if not snapshot_shared_prefill_reads: + return + + assert isinstance(metadata, DSV4Metadata) + use_sparse_prefill = not _is_sm120 and ( + num_qo_tokens > _LARGE_INDEXER_QUERY_THRESHOLD + or envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.get() + ) + if use_sparse_prefill: + metadata.sparse_prefill_cache = self._build_sparse_prefill_chunk_cache( + forward_batch, num_qo_tokens=num_qo_tokens + ) + # Marked for dense prefill too: that path reads only core_attn_metadata, + # which init_forward_metadata already snapshotted. + metadata.prefill_shared_reads_snapshotted = True + + def _build_sparse_prefill_chunk_cache( + self, forward_batch: ForwardBatch, *, num_qo_tokens: int + ) -> SparsePrefillChunkCache: + seq_lens_cpu = forward_batch.seq_lens_cpu + assert seq_lens_cpu is not None + extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu + assert extend_seq_lens_cpu is not None + seq_lens_cpu_list = seq_lens_cpu.tolist() + total_swa = sum( + min(int(seq_len), int(extend_len) + SWA_WINDOW - 1) + for seq_len, extend_len in zip( + seq_lens_cpu_list, extend_seq_lens_cpu, strict=True + ) + ) + # ``swa_window_size`` on the pool is its storage page size, not the + # model's SWA window, so pass both explicitly. + return SparsePrefillChunkCache.build( + seq_lens=forward_batch.seq_lens.to(torch.int32), + extend_seq_lens=forward_batch.extend_seq_lens.to(torch.int32), + req_pool_indices=forward_batch.req_pool_indices.to(torch.int32), + req_to_token=self.req_to_token, + full_to_swa=self.token_to_kv_pool.full_to_swa_index_mapping, + swa_window_size=SWA_WINDOW, + swa_page_size=self.token_to_kv_pool.swa_window_size, + num_qo_tokens=num_qo_tokens, + max_seq_len=max(seq_lens_cpu_list), + total_swa=total_swa, + ) + def _build_forward_metadata( self, forward_batch: ForwardBatch, @@ -1791,29 +1860,8 @@ class DeepseekV4AttnBackend( cache = self.forward_metadata.sparse_prefill_cache if cache is None: - seq_lens_cpu = forward_batch.seq_lens_cpu - assert seq_lens_cpu is not None - extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu - assert extend_seq_lens_cpu is not None - total_swa = sum( - min(int(seq_len), int(extend_len) + SWA_WINDOW - 1) - for seq_len, extend_len in zip( - seq_lens_cpu.tolist(), extend_seq_lens_cpu, strict=True - ) - ) - # ``swa_window_size`` on the pool is its storage page size, not - # the model's SWA window — pass both explicitly. - cache = SparsePrefillChunkCache.build( - seq_lens=forward_batch.seq_lens.to(torch.int32), - extend_seq_lens=forward_batch.extend_seq_lens.to(torch.int32), - req_pool_indices=forward_batch.req_pool_indices.to(torch.int32), - req_to_token=self.req_to_token, - full_to_swa=token_to_kv_pool.full_to_swa_index_mapping, - swa_window_size=SWA_WINDOW, - swa_page_size=token_to_kv_pool.swa_window_size, - num_qo_tokens=q_flat.shape[0], - max_seq_len=int(seq_lens_cpu.max().item()), - total_swa=total_swa, + cache = self._build_sparse_prefill_chunk_cache( + forward_batch, num_qo_tokens=q_flat.shape[0] ) self.forward_metadata.sparse_prefill_cache = cache @@ -1987,27 +2035,8 @@ class DeepseekV4AttnBackend( cache = self.forward_metadata.sparse_prefill_cache if cache is None: - seq_lens_cpu = forward_batch.seq_lens_cpu - assert seq_lens_cpu is not None - extend_seq_lens_cpu = forward_batch.extend_seq_lens_cpu - assert extend_seq_lens_cpu is not None - total_swa = sum( - min(int(seq_len), int(extend_len) + SWA_WINDOW - 1) - for seq_len, extend_len in zip( - seq_lens_cpu.tolist(), extend_seq_lens_cpu, strict=True - ) - ) - cache = SparsePrefillChunkCache.build( - seq_lens=forward_batch.seq_lens.to(torch.int32), - extend_seq_lens=forward_batch.extend_seq_lens.to(torch.int32), - req_pool_indices=forward_batch.req_pool_indices.to(torch.int32), - req_to_token=self.req_to_token, - full_to_swa=token_to_kv_pool.full_to_swa_index_mapping, - swa_window_size=SWA_WINDOW, - swa_page_size=token_to_kv_pool.swa_window_size, - num_qo_tokens=q_flat.shape[0], - max_seq_len=int(seq_lens_cpu.max().item()), - total_swa=total_swa, + cache = self._build_sparse_prefill_chunk_cache( + forward_batch, num_qo_tokens=q_flat.shape[0] ) self.forward_metadata.sparse_prefill_cache = cache diff --git a/python/sglang/srt/model_executor/runner/eager_runner.py b/python/sglang/srt/model_executor/runner/eager_runner.py index 83caa44d6..5f706128f 100644 --- a/python/sglang/srt/model_executor/runner/eager_runner.py +++ b/python/sglang/srt/model_executor/runner/eager_runner.py @@ -49,6 +49,9 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo enable_tc_piecewise_cuda_graph, set_tc_piecewise_forward_context, ) +from sglang.srt.model_executor.runner_utils import ( + maybe_publish_prefill_shared_read_done, +) from sglang.srt.runtime_context import ( get_parallel, get_spec, @@ -305,6 +308,15 @@ class EagerRunner(BaseRunner): # e.g. Moss-VL's prefill cross-attention custom mask. model_runner.model.prepare_forward_batch(forward_batch) model_runner.attn_backend.init_forward_metadata(forward_batch) + model_runner.attn_backend.prepare_prefill_shared_read_snapshot( + forward_batch, + num_qo_tokens=len(forward_batch.input_ids), + ) + maybe_publish_prefill_shared_read_done( + model_runner, + forward_batch, + torch.get_device_module(model_runner.device), + ) if not cp_v2_active: forward_batch.attn_cp_metadata = None diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 12e12c6eb..922b2d21a 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -1044,6 +1044,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): return if not self.use_captured_attn_metadata: attn_backend.init_forward_metadata(forward_batch) + attn_backend.prepare_prefill_shared_read_snapshot( + forward_batch, num_qo_tokens=num_tokens + ) return assert self.attn_metadata_buffers is not None metadata = self.attn_metadata_buffers[num_tokens] diff --git a/python/sglang/srt/model_executor/runner_utils/shared_read_event.py b/python/sglang/srt/model_executor/runner_utils/shared_read_event.py index f59e36361..4d9f36bb6 100644 --- a/python/sglang/srt/model_executor/runner_utils/shared_read_event.py +++ b/python/sglang/srt/model_executor/runner_utils/shared_read_event.py @@ -1,4 +1,4 @@ -"""Shared-read-done event utilities for CUDA graph runners.""" +"""Shared-read-done event utilities for graph and eager runners.""" import logging from typing import Optional @@ -31,9 +31,13 @@ def maybe_publish_prefill_shared_read_done( return if forward_batch.forward_mode != ForwardMode.EXTEND: return - # TODO(Jialin): Relax this gate for speculative decoding after its prefill - # WAR boundaries are validated. - if not model_runner.spec_algorithm.is_none(): + # TODO(Jialin): Relax for EAGLE/MTP after validating the later + # draft-extend reader's WAR boundary. + if ( + not model_runner.spec_algorithm.is_none() + and not model_runner.spec_algorithm.is_dflash_family() + ): + # Other speculative algorithms may have a later draft-extend reader. return # The record lands right after replay prep, so PRE_REPLAY only. declared = model_runner.attn_backend.shared_read_ends(forward_batch.forward_mode) diff --git a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py index 4e31d313b..979225dfb 100644 --- a/test/registered/attention/unittests/dsv4/test_deepseek_v4.py +++ b/test/registered/attention/unittests/dsv4/test_deepseek_v4.py @@ -391,6 +391,113 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase): DeepseekV4AttnBackend.use_captured_forward_metadata_for_breakable_cuda_graph ) + def test_prefill_snapshot_declares_pre_replay_boundary(self): + from sglang.srt.layers.attention.base_attn_backend import SharedReadEnds + from sglang.srt.layers.attention.deepseek_v4_backend import ( + DeepseekV4AttnBackend, + DSV4Metadata, + ) + + backend = object.__new__(DeepseekV4AttnBackend) + backend.forward_metadata = DSV4Metadata( + self._make_core_metadata(0), indexer_metadata=None + ) + self.assertIs( + backend.shared_read_ends(ForwardMode.EXTEND), + SharedReadEnds.UNKNOWN, + ) + + backend.forward_metadata.prefill_shared_reads_snapshotted = True + self.assertIs( + backend.shared_read_ends(ForwardMode.EXTEND), + SharedReadEnds.PRE_REPLAY, + ) + + def test_snapshot_builds_cache_only_for_sparse_prefill(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.deepseek_v4_backend import ( + _LARGE_INDEXER_QUERY_THRESHOLD, + DeepseekV4AttnBackend, + DSV4Metadata, + ) + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + batch = SimpleNamespace(forward_mode=ForwardMode.EXTEND) + cache = object() + for num_qo_tokens, builds in ( + (_LARGE_INDEXER_QUERY_THRESHOLD, False), + (_LARGE_INDEXER_QUERY_THRESHOLD + 1, True), + ): + with self.subTest(num_qo_tokens=num_qo_tokens): + backend = object.__new__(DeepseekV4AttnBackend) + backend.model_runner = SimpleNamespace( + spec_algorithm=SpeculativeAlgorithm.DFLASH + ) + backend.forward_metadata = DSV4Metadata( + self._make_core_metadata(0), indexer_metadata=None + ) + backend._build_sparse_prefill_chunk_cache = mock.Mock( + return_value=cache + ) + + with ( + envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True), + envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.override(False), + mock.patch( + "sglang.srt.layers.attention.deepseek_v4_backend._is_sm120", + False, + ), + ): + backend.prepare_prefill_shared_read_snapshot( + batch, num_qo_tokens=num_qo_tokens + ) + + metadata = backend.forward_metadata + if builds: + backend._build_sparse_prefill_chunk_cache.assert_called_once_with( + batch, num_qo_tokens=num_qo_tokens + ) + self.assertIs(metadata.sparse_prefill_cache, cache) + else: + backend._build_sparse_prefill_chunk_cache.assert_not_called() + self.assertIsNone(metadata.sparse_prefill_cache) + # Dense declares the boundary too; it reads only the metadata + # that init_forward_metadata already snapshotted. + self.assertTrue(metadata.prefill_shared_reads_snapshotted) + + def test_sparse_prefill_snapshot_marks_success_only_after_build(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.deepseek_v4_backend import ( + DeepseekV4AttnBackend, + DSV4Metadata, + ) + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + backend = object.__new__(DeepseekV4AttnBackend) + backend.model_runner = SimpleNamespace( + spec_algorithm=SpeculativeAlgorithm.DFLASH + ) + backend.forward_metadata = DSV4Metadata( + self._make_core_metadata(0), indexer_metadata=None + ) + backend.forward_metadata.prefill_shared_reads_snapshotted = True + backend._build_sparse_prefill_chunk_cache = mock.Mock( + side_effect=RuntimeError("snapshot failed") + ) + batch = SimpleNamespace(forward_mode=ForwardMode.EXTEND) + + with ( + envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True), + envs.SGLANG_OPT_FLASHMLA_SPARSE_PREFILL.override(True), + mock.patch( + "sglang.srt.layers.attention.deepseek_v4_backend._is_sm120", False + ), + self.assertRaisesRegex(RuntimeError, "snapshot failed"), + ): + backend.prepare_prefill_shared_read_snapshot(batch, num_qo_tokens=12288) + + self.assertFalse(backend.forward_metadata.prefill_shared_reads_snapshotted) + def test_refresh_replay_metadata_preserves_captured_tensor_storage(self): capture_metadata = self._make_core_metadata(0) replay_metadata = self._make_core_metadata(1000) @@ -459,6 +566,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase): self._make_core_metadata(0), indexer_metadata=None ) capture_metadata.sparse_prefill_cache = object() + capture_metadata.prefill_shared_reads_snapshotted = True replay_metadata = DSV4Metadata( self._make_core_metadata(1000), indexer_metadata=None ) @@ -487,6 +595,7 @@ class TestDSV4BreakableCudaGraphMetadataContract(CustomTestCase): self.assertTrue(calls[0][2]) self.assertIs(backend.forward_metadata, capture_metadata) self.assertIsNone(capture_metadata.sparse_prefill_cache) + self.assertFalse(capture_metadata.prefill_shared_reads_snapshotted) self.assertTrue( torch.equal( capture_metadata.core_attn_metadata.seq_lens_casual, diff --git a/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py b/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py index 36183c9b2..8a04996df 100644 --- a/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py +++ b/test/registered/unit/model_executor/runner/test_prefill_cuda_graph_padding.py @@ -1,5 +1,6 @@ import unittest from types import SimpleNamespace +from unittest import mock from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.model_executor.forward_batch_info import ( @@ -52,6 +53,25 @@ class TestPrefillCudaGraphPadding(CustomTestCase): self.assertTrue(runner.can_run_graph(self._make_forward_batch(8))) + def test_replay_snapshot_uses_padded_token_count(self): + runner = self._make_runner() + runner.use_captured_attn_metadata = False + attn_backend = mock.Mock() + runner.model_runner = SimpleNamespace(attn_backend=attn_backend) + forward_batch = self._make_forward_batch(8) + static_forward_batch = self._make_forward_batch(16) + + runner._prepare_forward_metadata_for_replay( + forward_batch, + static_forward_batch, + num_tokens=16, + ) + + attn_backend.init_forward_metadata.assert_called_once_with(forward_batch) + attn_backend.prepare_prefill_shared_read_snapshot.assert_called_once_with( + forward_batch, num_qo_tokens=16 + ) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py b/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py index c36e2dd40..6842bff49 100644 --- a/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py +++ b/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py @@ -60,6 +60,17 @@ def test_disabled_when_flag_is_false(): assert runner.shared_read_done_event is None +@pytest.mark.parametrize( + "algorithm", (SpeculativeAlgorithm.DFLASH, SpeculativeAlgorithm.DSPARK) +) +def test_dflash_family_target_prefill_publishes(algorithm): + runner = _model_runner(spec_algorithm=algorithm) + with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True): + maybe_publish_prefill_shared_read_done(runner, _batch(), _DEVICE_MODULE) + published = runner.shared_read_done_event + assert isinstance(published, _Event) and published.recorded + + def test_gates_exclude_non_prefill_unsupported_algorithm_and_noncompliant_backend(): with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True): for runner, batch in (