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

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

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

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]

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)