A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 13, 17:00 UTC
[Debugging] Add annotation for easier trace analysis (#22496)
Dayeol Lee committed
Nov 5, 2025 at 16:52 UTC
17676585599934fe3f6e1bff61185a313c5da2cd
1 file changed
+19
-3
vllm/v1/worker/gpu_worker.py
+19
-3
@@ -509,6 +509,19 @@ class Worker(WorkerBase):
509
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
510
return self.model_runner.get_supported_tasks()
511
512
+ def annotate_profile(self, scheduler_output):
513
+ # add trace annotation so that we can easily distinguish
514
+ # new/cached request numbers in each iteration
515
+ if not self.profiler:
516
+ return nullcontext()
517
+
518
+ num_new = len(scheduler_output.scheduled_new_reqs)
519
+ num_cached = len(scheduler_output.scheduled_cached_reqs.req_ids)
520
+
521
+ return torch.profiler.record_function(
522
+ f"execute_new_{num_new}_cached_{num_cached}"
523
+ )
524
+
525
@torch.inference_mode()
526
def sample_tokens(
527
self, grammar_output: "GrammarOutput | None"
@@ -536,9 +549,12 @@ class Worker(WorkerBase):
549
)
550
)
551
539
- output = self.model_runner.execute_model(scheduler_output, intermediate_tensors)
540
- if isinstance(output, (ModelRunnerOutput, NoneType)):
541
- return output
552
+ with self.annotate_profile(scheduler_output):
553
+ output = self.model_runner.execute_model(
554
+ scheduler_output, intermediate_tensors
555
+ )
556
+ if isinstance(output, (ModelRunnerOutput, NoneType)):
557
+ return output
558
559
assert isinstance(output, IntermediateTensors)
560
parallel_config = self.vllm_config.parallel_config