chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Add cuda graph status to prefill log (#17836)

Browse Source
This commit is contained in:
Ke Bao
2026-01-30 16:56:53 +08:00
committed by GitHub
parent c8dc543dc5
commit 77a27e728c
4 changed files with 62 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/sglang/srt/managers/scheduler.py
View File

@@ -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:

25
python/sglang/srt/managers/scheduler_metrics_mixin.py
View File

@@ -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)

20
python/sglang/srt/managers/scheduler_output_processor_mixin.py
View File

@@ -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,

30
python/sglang/srt/model_executor/model_runner.py
View File

@@ -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,