[Spec] Polish FutureMap after #25879: rename callback, async guard, cleanup (#25962)

python/sglang/srt/managers/overlap_utils.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Union
 
 import torch
 
@@ -67,8 +67,8 @@ class FutureMap:
             # Forward-only bufs are lazy (worker-dependent shape).
             self._forward_buf_initialized = False
 
-        # Fences the schedule-consumed buf fields.
+        # Fences schedule-consumed buf fields; lazy device.Event() (cuda/hip-agnostic).
-        self.publish_ready: Optional[torch.cuda.Event] = None
+        self.publish_ready = None
 
     def _lazy_init_forward_buf(self, draft_input: EagleDraftInput):
         self._forward_buf_initialized = True
@@ -115,6 +115,9 @@ class FutureMap:
             draft_input.new_seq_lens = self.new_seq_lens_buf[indices]
             # Resolve seq_lens placeholder (-indices) to the post-verify view.
             batch.seq_lens = draft_input.new_seq_lens
+            # Async guard: catches a (-indices) sentinel slipping through if
+            # publish_ready fencing or buf indexing is wrong.
+            torch._assert_async((batch.seq_lens > 0).all())
             if spec_need_hidden_states():
                 draft_input.hidden_states = self.hidden_states_buf[indices]
 
@@ -141,11 +144,16 @@ class FutureMap:
             self.publish_ready = torch.get_device_module(self.device).Event()
         self.publish_ready.record()
 
-    def stash(self, future_indices: FutureIndices, payload) -> None:
+    def stash(
+        self,
+        future_indices: FutureIndices,
+        payload: Union[torch.Tensor, EagleDraftInput],
+    ) -> None:
         """Store forward-only fields for the next forward batch to pick up."""
         indices = future_indices.indices
         if indices.shape[0] == 0:
-            return  # DP idle
+            # DP idle: payload is empty stub; lazy-init shape peek would IndexError.
+            return
         if self.spec_algo.is_none():
             # next_token_ids is int32; buf is int64. Advanced indexing requires
             # an explicit cast.

python/sglang/srt/managers/scheduler.py

@@ -2842,12 +2842,9 @@ class Scheduler(
         # Run forward
         if self.is_generation:
             if self.enable_overlap:
-                # Spec v2 pre-isolation CPU mirror prep: D2H new_seq_lens_buf
+                # Self-gates on batch.spec_info.future_indices; non-spec_v2
-                # into batch.seq_lens_cpu + set seq_lens_sum. For non-spec_v2,
+                # no-ops (ForwardBatch.init_new lazily computes the sum).
-                # ForwardBatch.init_new lazily computes the sum.
+                self.future_map.resolve_seq_lens_cpu(batch)
-                if batch.is_spec_v2:
-                    # FIXME: make this optional to different backends.
-                    self.future_map.resolve_seq_lens_cpu(batch)
 
                 with self._overlap_forward_isolation(batch):
                     future_indices = FutureIndices(indices=batch.req_pool_indices)
@@ -2856,11 +2853,7 @@ class Scheduler(
                     # draft_extend; publish moves the fence to verify-end so
                     # schedule prep can overlap with draft_extend.
                     fwd_kwargs = (
-                        {
+                        {"on_publish": partial(self.future_map.publish, future_indices)}
-                            "on_verify_complete": partial(
-                                self.future_map.publish, future_indices
-                            )
-                        }
                         if batch.is_spec_v2
                         else {}
                     )

python/sglang/srt/speculative/eagle_worker_v2.py

@@ -756,7 +756,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
         # allocator and kv cache pool are shared with target worker, which are cleared in scheduler
         pass
 
-    def forward_batch_generation(self, batch: ScheduleBatch, on_verify_complete=None):
+    def forward_batch_generation(self, batch: ScheduleBatch, on_publish=None):
         if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
             # Target prefill
             target_capture_mode = (
@@ -768,8 +768,8 @@ class EAGLEWorkerV2(BaseSpecWorker):
             batch_output = self.target_worker.forward_batch_generation(batch)
 
             # Publish before draft_extend so the fence is at target-end.
-            if on_verify_complete is not None:
+            if on_publish is not None:
-                on_verify_complete(batch.seq_lens)
+                on_publish(batch.seq_lens)
 
             # Draft prefill
             with (
@@ -814,8 +814,8 @@ class EAGLEWorkerV2(BaseSpecWorker):
             batch.spec_info = verify_input
             batch_output = self.verify(batch)
             # Publish before draft_extend so the fence is at verify-end.
-            if on_verify_complete is not None:
+            if on_publish is not None:
-                on_verify_complete(batch_output.next_draft_input.new_seq_lens)
+                on_publish(batch_output.next_draft_input.new_seq_lens)
             with (
                 self.draft_worker.draft_tp_context(
                     self.draft_worker.draft_runner.tp_group

python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py

@@ -669,7 +669,7 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
         # allocator and kv cache pool are shared with target worker, which are cleared in scheduler
         pass
 
-    def forward_batch_generation(self, batch: ScheduleBatch, on_verify_complete=None):
+    def forward_batch_generation(self, batch: ScheduleBatch, on_publish=None):
         if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
             # Target prefill
             target_capture_mode = (
@@ -681,8 +681,8 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
             batch_output = self.target_worker.forward_batch_generation(batch)
 
             # Publish before draft_extend so the fence is at target-end.
-            if on_verify_complete is not None:
+            if on_publish is not None:
-                on_verify_complete(batch.seq_lens)
+                on_publish(batch.seq_lens)
 
             # Chain-style MTP needs FULL to get all-token hidden states;
             # non-chain only needs LAST (the target model's hidden states).
@@ -711,8 +711,8 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
             batch.spec_info = verify_input
             batch_output = self.verify(batch)
             # Publish before draft_extend so the fence is at verify-end.
-            if on_verify_complete is not None:
+            if on_publish is not None:
-                on_verify_complete(batch_output.next_draft_input.new_seq_lens)
+                on_publish(batch_output.next_draft_input.new_seq_lens)
             self.draft_worker._draft_extend_for_decode(batch, batch_output)
             return batch_output