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

Co-authored-by: weireweire <20922698+weireweire@users.noreply.github.com>
2026-08-26 21:33:08 -07:00
parent 7324021e6c
commit 4d5d506486
8 changed files with 249 additions and 51 deletions
@@ -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