From c9a0e55eb597f3fe1a29fbeae546b2b294c16355 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 21 May 2026 13:56:22 -0700 Subject: [PATCH] [Spec] Polish FutureMap after #25879: rename callback, async guard, cleanup (#25962) --- python/sglang/srt/managers/overlap_utils.py | 18 +++++++++++++----- python/sglang/srt/managers/scheduler.py | 15 ++++----------- .../sglang/srt/speculative/eagle_worker_v2.py | 10 +++++----- .../speculative/multi_layer_eagle_worker_v2.py | 10 +++++----- 4 files changed, 27 insertions(+), 26 deletions(-) diff --git a/python/sglang/srt/managers/overlap_utils.py b/python/sglang/srt/managers/overlap_utils.py index 99b951fab..f75576687 100644 --- a/python/sglang/srt/managers/overlap_utils.py +++ b/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. - self.publish_ready: Optional[torch.cuda.Event] = None + # Fences schedule-consumed buf fields; lazy device.Event() (cuda/hip-agnostic). + 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. diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f69a698c2..3101b833d 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/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 - # into batch.seq_lens_cpu + set seq_lens_sum. For non-spec_v2, - # ForwardBatch.init_new lazily computes the sum. - if batch.is_spec_v2: - # FIXME: make this optional to different backends. - self.future_map.resolve_seq_lens_cpu(batch) + # Self-gates on batch.spec_info.future_indices; non-spec_v2 + # no-ops (ForwardBatch.init_new lazily computes the sum). + 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_verify_complete": partial( - self.future_map.publish, future_indices - ) - } + {"on_publish": partial(self.future_map.publish, future_indices)} if batch.is_spec_v2 else {} ) diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 6e81c5c25..4a9969113 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/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: - on_verify_complete(batch.seq_lens) + if on_publish is not None: + 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: - on_verify_complete(batch_output.next_draft_input.new_seq_lens) + if on_publish is not None: + on_publish(batch_output.next_draft_input.new_seq_lens) with ( self.draft_worker.draft_tp_context( self.draft_worker.draft_runner.tp_group diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 35f189c05..c8da2727e 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/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: - on_verify_complete(batch.seq_lens) + if on_publish is not None: + 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: - on_verify_complete(batch_output.next_draft_input.new_seq_lens) + if on_publish is not None: + on_publish(batch_output.next_draft_input.new_seq_lens) self.draft_worker._draft_extend_for_decode(batch, batch_output) return batch_output