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

Support GPU execution time breakdown by forward mode metrics (#15396)

Browse Source
This commit is contained in:
fzyzcjy
2025-12-18 22:18:17 +08:00
committed by GitHub
parent 2c196f95a7
commit c5f4e20f2f
5 changed files with 98 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/managers/scheduler.py
+9 -7
View File
@@ -2105,10 +2105,11 @@ class Scheduler(
with self.forward_stream_ctx:
self.forward_stream.wait_stream(self.default_stream)
self.future_map.resolve_future(model_worker_batch)
batch_result = self.model_worker.forward_batch_generation(
model_worker_batch
# here pp is not compatible with overlap
)
with self.record_forward_metrics(batch):
batch_result = self.model_worker.forward_batch_generation(
model_worker_batch
# here pp is not compatible with overlap
)
# FIXME(lsyin): maybe move this to forward_batch_generation
batch_result.copy_done = self.device_module.Event()
if batch_result.delay_sample_func is None:
@@ -2144,9 +2145,10 @@ class Scheduler(
if self.spec_algorithm.is_none()
else {}
)
batch_result = self.model_worker.forward_batch_generation(
worker_batch_or_batch, **kwargs
)
with self.record_forward_metrics(batch):
batch_result = self.model_worker.forward_batch_generation(
worker_batch_or_batch, **kwargs
)
future_indices_or_next_token_ids = batch_result.next_token_ids
self.update_cache_from_scheduler(batch, batch_result)