Publish gated DSV4 DFLASH-family target-prefill read completion (#35947)

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>

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

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 (