From f86e9b48e825ac5dba9e5e70e1d508f91314810b Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Wed, 17 Jun 2026 00:59:58 -0700 Subject: [PATCH] [Perf] Make spec-decode penalty H2D non-blocking and share decode cumulate path (#28500) --- python/sglang/srt/managers/schedule_batch.py | 37 +++++++++++-------- .../sglang/srt/speculative/eagle_info_v2.py | 16 +------- 2 files changed, 23 insertions(+), 30 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 5990e0cb1..02eef3460 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2572,6 +2572,27 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): pool.set_mamba_ping_pong_slot(req, other_idx, new_slot[0]) req.mamba_next_track_idx = other_idx + def cumulate_penalty_output_tokens(self): + # Under overlap batch.input_ids is just a placeholder here -- the + # real token is relayed via future_map and resolved at forward + # entry. So take the last output token from Req directly + # (origin_input_ids[-1] on the first decode, before any output). + last_tokens = [ + req.output_ids[-1] if len(req.output_ids) else req.origin_input_ids[-1] + for req in self.reqs + ] + # Non-blocking H2D so this per-step copy doesn't sync behind the forward. + # pin_memory (matching the prefill-path tensors) keeps the copy async; + # is_pin_memory_available falls back to pageable on unsupported devices. + latest_output_ids = torch.tensor( + last_tokens, + dtype=torch.int64, + pin_memory=is_pin_memory_available(self.device), + ).to(self.device, non_blocking=True) + self.sampling_info.penalizer_orchestrator.cumulate_output_tokens( + latest_output_ids + ) + def prepare_for_decode(self): self.forward_mode = ForwardMode.DECODE bs = len(self.reqs) @@ -2591,21 +2612,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): return if self.sampling_info.penalizer_orchestrator.is_required: - # Under overlap batch.input_ids is just a placeholder here -- the - # real token is relayed via future_map and resolved at forward - # entry. So take the last output token from Req directly - # (origin_input_ids[-1] on the first decode, before any output). - last_tokens = [ - req.output_ids[-1] if len(req.output_ids) else req.origin_input_ids[-1] - for req in self.reqs - ] - # Non-blocking H2D so this per-step copy doesn't sync behind the forward. - latest_output_ids = torch.tensor(last_tokens, dtype=torch.int64).to( - self.device, non_blocking=True - ) - self.sampling_info.penalizer_orchestrator.cumulate_output_tokens( - latest_output_ids - ) + self.cumulate_penalty_output_tokens() # input_ids is set at end of previous run_batch (placeholder for # overlap; next_token_ids cast for non-overlap). diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index c75589356..3723a5673 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -37,21 +37,7 @@ class EagleDraftInputV2Mixin: # Accumulate penalty # This is a relaxed version of penalties for speculative decoding. if batch.sampling_info.penalizer_orchestrator.is_required: - output_ids = torch.tensor( - [ - ( - req.output_ids[-1] - if len(req.output_ids) - else req.origin_input_ids[-1] - ) - for req in batch.reqs - ], - dtype=torch.int64, - device=batch.device, - ) - batch.sampling_info.penalizer_orchestrator.cumulate_output_tokens( - output_ids - ) + batch.cumulate_penalty_output_tokens() page_size = batch.token_to_kv_pool_allocator.page_size double_alloc = get_alloc_reserve_per_decode()