[Runtime] Add decode CUDA graph hooks for eager logits processing (#40222)

Co-authored-by: mxz <mxz@fb.com>

python/sglang/srt/batch_overlap/two_batch_overlap.py

@@ -760,6 +760,7 @@ class TboForwardBatchPreparer:
             "is_prefill_only",
             "spec_algorithm",
             "capture_hidden_mode",
+            "defer_logits_to_eager",  # forward-level flag, inherited by both child batches
             "split_index",  # for split prefill
             "orig_seq_lens",  # only used by qwen-1m, thus not care
             "return_pooled_hidden_states",

python/sglang/srt/model_executor/forward_batch_info.py

@@ -576,6 +576,9 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
 
     # === Per-forward overrides passed explicitly to init_new ===
     capture_hidden_mode: CaptureHiddenMode = None
+    # CUDA-graph runners may capture only the transformer body and execute the
+    # logits processor eagerly after replay.
+    defer_logits_to_eager: bool = False
     # For hidden states before normal
     return_hidden_states_before_norm: bool = False
 

python/sglang/srt/model_executor/graph_shared_output.py

@@ -14,7 +14,11 @@ if TYPE_CHECKING:
 
 
 class GraphSharedOutput:
-    """``(max_rows, vocab)`` logits buffer, shared by every cuda-graph runner."""
+    """Persistent ``(max_rows, vocab)`` output shared by graph runners.
+
+    The producer need not be captured in a CUDA graph. A runner may capture only
+    the transformer body and reuse this buffer for an eager logits tail.
+    """
 
     _process_shared: Optional[GraphSharedOutput] = None
 
@@ -39,7 +43,7 @@ class GraphSharedOutput:
         max_rows = 0
         decode = cuda_graph_config.decode
         if decode.backend != Backend.DISABLED and decode.bs:
-            max_rows = max(max_rows, model_runner.max_decode_logits_rows())
+            max_rows = max(max_rows, model_runner.max_shared_logits_buffer_rows())
 
         if max_rows <= 0:
             return None

python/sglang/srt/model_executor/model_runner.py

@@ -892,6 +892,14 @@ class ModelRunner:
             device=self.device,
         )
 
+    def max_shared_logits_buffer_rows(self) -> int:
+        """Maximum rows in the persistent logits buffer used by graph runners.
+
+        This includes outputs produced inside a graph as well as eager logits
+        tails that reuse the runner-owned buffer after graph replay.
+        """
+        return self.max_decode_logits_rows()
+
     def alloc_memory_pool(self, memory_pool_config: Optional[MemoryPoolConfig] = None):
         """Allocate KV cache memory pools only (no backends or cuda graphs)."""
         if memory_pool_config is not None:

python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py

@@ -20,7 +20,7 @@ import gc
 import logging
 from abc import abstractmethod
 from contextlib import contextmanager
-from typing import TYPE_CHECKING, Any, List, Sequence, Tuple
+from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple
 
 from sglang.srt.model_executor.runner.base_runner import BaseRunner
 from sglang.srt.runtime_context import (
@@ -132,6 +132,27 @@ class BaseCudaGraphRunner(BaseRunner):
     buffers: ForwardInputBuffers
     backend: BaseCudaGraphBackend
 
+    def cuda_graph_output_rows(self, output: Any) -> Optional[int]:
+        """Rows of graph output that must be preserved for post-replay work.
+
+        The default graph key is a request count, which is also the output row
+        count for ordinary decode. A graph that returns per-token hidden states
+        for an eager tail can instead produce ``requests * tokens_per_request``
+        rows. Such a runner must return that actual row count here. ``None``
+        keeps the backend's default request-count behavior.
+        """
+        return None
+
+    def cuda_graph_output_capacity_rows(self, output: Any) -> Optional[int]:
+        """Capacity required by the output buffer shared across graph keys.
+
+        The breakable backend allocates this buffer once, while capturing its
+        first shape. A runner whose output uses token rows rather than request
+        rows must return the largest possible output here so later graph shapes
+        fit. ``None`` uses the current graph key as the capacity.
+        """
+        return None
+
     @staticmethod
     def _pad_to_bucket(raw_size: int, buckets: Sequence[int]) -> int:
         """Return the smallest buckets[i] >= raw_size.

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -26,6 +26,7 @@ Backend selection comes from cuda_graph_config.decode:
 from __future__ import annotations
 
 import contextlib
+import dataclasses
 import inspect
 import logging
 import os
@@ -384,13 +385,16 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             assert self.require_mlp_tp_gather or self.require_attn_tp_gather
 
         # --- buffers ---------------------------------------------------
+        logits_buffer_rows = self._next_token_logits_buffer_capacity_rows(
+            self.max_num_token
+        )
         self.buffers: DecodeInputBuffers = DecodeInputBuffers.create(
             device=self.device,
             max_bs=self.max_bs,
             max_num_token=self.max_num_token,
             hidden_size=self.model_runner.model_config.hidden_size,
             next_token_logits_buffer=self.model_runner.graph_shared_output.get_logits_buffer(
-                self.model_runner.model_config.vocab_size, rows=self.max_num_token
+                self.model_runner.model_config.vocab_size, rows=logits_buffer_rows
             ),
             dtype=self.model_runner.model_config.dtype,
             dp_size=self.dp_size,
@@ -468,6 +472,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
             )
 
+    def _next_token_logits_buffer_capacity_rows(self, max_num_tokens: int) -> int:
+        """Rows reserved for the largest shared logits output."""
+        return max_num_tokens
+
     def _record_in_graph_metadata_prep_done(self):
         # Purely a marker at this point in the graph; where the shared reads
         # actually end is the attn backend's call.
@@ -1457,6 +1465,8 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             if shared_read_ends is SharedReadEnds.POST_REPLAY:
                 self._publish_read_done(in_graph=False)
 
+            output = self._process_output_after_replay(output, forward_batch)
+
         if isinstance(output, LogitsProcessorOutput):
             if self.is_dllm:
                 next_token_logits = None
@@ -1473,7 +1483,9 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                     else None
                 )
 
-            return LogitsProcessorOutput(
+            # Preserve extension fields produced by the eager output processor.
+            return dataclasses.replace(
+                output,
                 next_token_logits=next_token_logits,
                 full_logits=full_logits,
                 hidden_states=(
@@ -1481,7 +1493,6 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                     if output.hidden_states is not None
                     else None
                 ),
-                customized_info=output.customized_info,
             )
         else:
             assert isinstance(output, PPProxyTensors)
@@ -1495,6 +1506,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 }
             )
 
+    def _process_output_after_replay(self, output, forward_batch: ForwardBatch):
+        """Optional eager tail executed after graph replay and inside its timer."""
+        return output
+
     def get_spec_info(self, num_tokens: int):
         spec_info = None
         if (

python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py

@@ -68,6 +68,7 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         enable_memory_saver: bool = False,
         debug_eager: bool = False,
     ) -> None:
+        self._cuda_graph_runner = cuda_graph_runner
         self._model_runner = cuda_graph_runner.model_runner
         self._graphs: Dict[Any, BreakableCUDAGraph] = {}
         self._outputs: Dict[Any, Any] = {}
@@ -130,7 +131,14 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         )
         size = shape_key.size
         if self._shared_output_buffer is None:
-            self._shared_output_buffer = self._alloc_full_buffer(warmup_out, size)
+            capacity_rows = self._cuda_graph_runner.cuda_graph_output_capacity_rows(
+                warmup_out
+            )
+            if capacity_rows is None:
+                capacity_rows = size
+            self._shared_output_buffer = self._alloc_full_buffer(
+                warmup_out, capacity_rows
+            )
         with (
             graph_pool_capture_scope(),
             BreakableCUDAGraphCapture(
@@ -157,6 +165,9 @@ class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
         A body that shards or prunes its output along dim 0 returns fewer than
         ``cap`` rows; everything else returns exactly ``cap``.
         """
+        runner_rows = self._cuda_graph_runner.cuda_graph_output_rows(output)
+        if runner_rows is not None:
+            return runner_rows
         if torch.is_tensor(output):
             return min(cap, output.shape[0])
         if isinstance(output, PPProxyTensors):

test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py

@@ -80,6 +80,17 @@ class TestTboFilterBatchMarker(CustomTestCase):
         self.assertFalse(child.forward_metadata_ready)
         self.assertFalse(child.forward_metadata_replan_equivalent)
 
+    def test_filter_batch_children_inherit_deferred_logits(self):
+        """A graph that defers logits to an eager tail defers them for both
+        TBO halves; a child that reset the flag would run the logits processor
+        inside the captured body."""
+        for deferred in (False, True):
+            with self.subTest(deferred=deferred):
+                parent = _make_target_verify_batch(8)
+                parent.defer_logits_to_eager = deferred
+                child = _filter(parent, lo=0, hi=4)
+                self.assertEqual(child.defer_logits_to_eager, deferred)
+
 
 def _make_valued_batch(bs: int) -> ForwardBatch:
     # Distinct per-position values so a filtered slice is unambiguous.