diff --git a/python/sglang/srt/batch_overlap/two_batch_overlap.py b/python/sglang/srt/batch_overlap/two_batch_overlap.py index 1cf69daa3..2a110ca08 100644 --- a/python/sglang/srt/batch_overlap/two_batch_overlap.py +++ b/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", diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 7cc421b14..94b7b597f 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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 diff --git a/python/sglang/srt/model_executor/graph_shared_output.py b/python/sglang/srt/model_executor/graph_shared_output.py index 1f153e2de..ea36c28c4 100644 --- a/python/sglang/srt/model_executor/graph_shared_output.py +++ b/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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 87dc0bfb0..b7856f1a3 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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: diff --git a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py index 880ea3136..8cbd23148 100644 --- a/python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py +++ b/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. diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 2bd702619..63f24d61f 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/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 ( diff --git a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py b/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py index 368db2ebc..32aa6a9b9 100644 --- a/python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py +++ b/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): diff --git a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py b/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py index b18d00b49..cba5a44ef 100644 --- a/test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py +++ b/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.