From 77a27e728c6423453707ba02cb16b19937aa7976 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Fri, 30 Jan 2026 16:56:53 +0800 Subject: [PATCH] Add cuda graph status to prefill log (#17836) --- python/sglang/srt/managers/scheduler.py | 12 +++----- .../srt/managers/scheduler_metrics_mixin.py | 25 +++++++++++----- .../scheduler_output_processor_mixin.py | 20 +++++++++++++ .../sglang/srt/model_executor/model_runner.py | 30 ++++++++++++------- 4 files changed, 62 insertions(+), 25 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 3a06f92d9..e818deaa4 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2088,14 +2088,10 @@ class Scheduler( if self.dllm_staging_reqs.non_empty(): self.dllm_staging_reqs.update_chunked_status() - # Print stats - if self.current_scheduler_metrics_enabled: - self.log_prefill_stats( - adder, - can_run_list, - running_bs=len(self.running_batch.reqs), - running_bs_offline_batch=0, - ) + # Record for logging prefill stats after forward + self.adder = adder + self.can_run_list = can_run_list + self.running_bs = len(self.running_batch.reqs) # Record metrics for req in can_run_list: diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index d2f0df499..1a3ba69a1 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -150,6 +150,7 @@ class SchedulerMetricsMixin: can_run_list: List[Req], running_bs: int, running_bs_offline_batch: int, + can_run_cuda_graph: bool, ): gap_latency = time.perf_counter() - self.last_prefill_stats_tic self.last_prefill_stats_tic = time.perf_counter() @@ -204,7 +205,7 @@ class SchedulerMetricsMixin: self.stats.new_token_ratio = adder.new_token_ratio iter_msg = f" [{self.forward_ct + 1}]" if LOG_FORWARD_ITERS else "" - f = ( + msg = ( f"Prefill batch{iter_msg}, " f"#new-seq: {len(can_run_list)}, " f"#new-token: {adder.log_input_tokens}, " @@ -215,13 +216,23 @@ class SchedulerMetricsMixin: ) if self.disaggregation_mode == DisaggregationMode.PREFILL: - f += f"#prealloc-req: {len(self.disagg_prefill_bootstrap_queue.queue)}, " - f += f"#inflight-req: {len(self.disagg_prefill_inflight_queue)}, " - f += f"input throughput (token/s): {self.last_input_throughput:.2f}, " + msg += f"#prealloc-req: {len(self.disagg_prefill_bootstrap_queue.queue)}, " + msg += f"#inflight-req: {len(self.disagg_prefill_inflight_queue)}, " + msg += f"input throughput (token/s): {self.last_input_throughput:.2f}, " else: - f += f"input throughput (token/s): {self.last_input_throughput:.2f}, " + msg += f"input throughput (token/s): {self.last_input_throughput:.2f}, " - logger.info(f) + graph_backend = defaultdict( + lambda: "cuda graph", + { + "cpu": "cpu graph", + "npu": "npu graph", + }, + ) + + msg += f"{graph_backend[self.device]}: {can_run_cuda_graph}" + + logger.info(msg) if self.enable_metrics: # Basics @@ -395,7 +406,7 @@ class SchedulerMetricsMixin: msg += ( f"{graph_backend[self.device]}: {can_run_cuda_graph}, " f"gen throughput (token/s): {self.last_gen_throughput:.2f}, " - f"#queue-req: {len(self.waiting_queue)}, " + f"#queue-req: {len(self.waiting_queue)}" ) logger.info(msg) diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index 2c19dd3e8..d79f929d3 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -301,6 +301,16 @@ class SchedulerOutputProcessorMixin: self.stream_output(batch.reqs, batch.return_logprob, skip_stream_req) + if self.current_scheduler_metrics_enabled: + can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) + self.log_prefill_stats( + adder=self.adder, + can_run_list=self.can_run_list, + running_bs=self.running_bs, + running_bs_offline_batch=0, + can_run_cuda_graph=can_run_cuda_graph, + ) + def _resolve_spec_overlap_token_ids( self: Scheduler, result: GenerationBatchResult, batch: ScheduleBatch ) -> List[List[int]]: @@ -370,6 +380,16 @@ class SchedulerOutputProcessorMixin: self.stream_output(batch.reqs, batch.return_logprob) self.token_to_kv_pool_allocator.free_group_end() + if self.current_scheduler_metrics_enabled: + can_run_cuda_graph = getattr(result, "can_run_cuda_graph", False) + self.log_prefill_stats( + adder=self.adder, + can_run_list=self.can_run_list, + running_bs=self.running_bs, + running_bs_offline_batch=0, + can_run_cuda_graph=can_run_cuda_graph, + ) + def process_batch_result_decode( self: Scheduler, batch: ScheduleBatch, diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 4dcf981fa..11d60729d 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -2180,7 +2180,9 @@ class ModelRunner(ModelRunnerKVCacheMixin): forward_batch: ForwardBatch, skip_attn_backend_init: bool = False, pp_proxy_tensors=None, - ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]: + ) -> Tuple[ + Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput], bool + ]: kwargs = {} if self.support_pp: kwargs["pp_proxy_tensors"] = pp_proxy_tensors @@ -2189,20 +2191,28 @@ class ModelRunner(ModelRunnerKVCacheMixin): if not self.is_generation: kwargs["get_embedding"] = True - if ( + can_run_graph = ( self.piecewise_cuda_graph_runner is not None and self.piecewise_cuda_graph_runner.can_run(forward_batch) - ): - return self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs) + ) + + if can_run_graph: + return ( + self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs), + can_run_graph, + ) if not skip_attn_backend_init: self.attn_backend.init_forward_metadata(forward_batch) - return self.model.forward( - forward_batch.input_ids, - forward_batch.positions, - forward_batch, - **kwargs, + return ( + self.model.forward( + forward_batch.input_ids, + forward_batch.positions, + forward_batch, + **kwargs, + ), + can_run_graph, ) def forward_idle( @@ -2358,7 +2368,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): forward_count=split_forward_count, ) elif forward_batch.forward_mode.is_extend(include_draft_extend_v2=True): - ret = self.forward_extend( + ret, can_run_graph = self.forward_extend( forward_batch, skip_attn_backend_init=skip_attn_backend_init, pp_proxy_tensors=pp_proxy_tensors,