[Perf] Make spec-decode penalty H2D non-blocking and share decode cumulate path (#28500)

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
+            self.cumulate_penalty_output_tokens()
-            # 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
-            )
 
         # input_ids is set at end of previous run_batch (placeholder for
         # overlap; next_token_ids cast for non-overlap).

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(
+            batch.cumulate_penalty_output_tokens()
-                [
-                    (
-                        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
-            )
 
         page_size = batch.token_to_kv_pool_allocator.page_size
         double_alloc = get_alloc_reserve_per_decode()