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, 13:15 UTC
[Misc] Add more scoping for improved trace (#28329)
Signed-off-by: Wei Wei <wwei6@meta.com>
Wei Wei committed
Nov 10, 2025 at 13:03 UTC
bf6a3d0ff5a69e0a30567f2ad417530c002eaa4e
4 files changed
+192
-148
vllm/v1/core/sched/scheduler.py
+61
-55
@@ -38,6 +38,7 @@ from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
38
from vllm.v1.request import Request, RequestStatus
39
from vllm.v1.spec_decode.metrics import SpecDecodingStats
40
from vllm.v1.structured_output import StructuredOutputManager
41
+from vllm.v1.utils import record_function_or_nullcontext
42
43
logger = init_logger(__name__)
44
@@ -259,49 +260,52 @@ class Scheduler(SchedulerInterface):
260
continue
261
262
# Schedule newly needed KV blocks for the request.
262
- while True:
263
- new_blocks = self.kv_cache_manager.allocate_slots(
264
- request,
265
- num_new_tokens,
266
- num_lookahead_tokens=self.num_lookahead_tokens,
267
- )
268
-
269
- if new_blocks is not None:
270
- # The request can be scheduled.
271
- break
272
-
273
- # The request cannot be scheduled.
274
- # Preempt the lowest-priority request.
275
- if self.policy == SchedulingPolicy.PRIORITY:
276
- preempted_req = max(
277
- self.running,
278
- key=lambda r: (r.priority, r.arrival_time),
263
+ with record_function_or_nullcontext("schedule: allocate_slots"):
264
+ while True:
265
+ new_blocks = self.kv_cache_manager.allocate_slots(
266
+ request,
267
+ num_new_tokens,
268
+ num_lookahead_tokens=self.num_lookahead_tokens,
269
)
280
- self.running.remove(preempted_req)
281
- if preempted_req in scheduled_running_reqs:
282
- scheduled_running_reqs.remove(preempted_req)
283
- token_budget += num_scheduled_tokens[preempted_req.request_id]
284
- req_to_new_blocks.pop(preempted_req.request_id)
285
- num_scheduled_tokens.pop(preempted_req.request_id)
286
- req_index -= 1
287
- else:
288
- preempted_req = self.running.pop()
270
290
- self.kv_cache_manager.free(preempted_req)
291
- self.encoder_cache_manager.free(preempted_req)
292
- preempted_req.status = RequestStatus.PREEMPTED
293
- preempted_req.num_computed_tokens = 0
294
- preempted_req.num_preemptions += 1
295
- if self.log_stats:
296
- preempted_req.record_event(
297
- EngineCoreEventType.PREEMPTED, scheduled_timestamp
298
- )
271
+ if new_blocks is not None:
272
+ # The request can be scheduled.
273
+ break
274
300
- self.waiting.prepend_request(preempted_req)
301
- preempted_reqs.append(preempted_req)
302
- if preempted_req == request:
303
- # No more request to preempt. Cannot schedule this request.
304
- break
275
+ # The request cannot be scheduled.
276
+ # Preempt the lowest-priority request.
277
+ if self.policy == SchedulingPolicy.PRIORITY:
278
+ preempted_req = max(
279
+ self.running,
280
+ key=lambda r: (r.priority, r.arrival_time),
281
+ )
282
+ self.running.remove(preempted_req)
283
+ if preempted_req in scheduled_running_reqs:
284
+ scheduled_running_reqs.remove(preempted_req)
285
+ token_budget += num_scheduled_tokens[
286
+ preempted_req.request_id
287
+ ]
288
+ req_to_new_blocks.pop(preempted_req.request_id)
289
+ num_scheduled_tokens.pop(preempted_req.request_id)
290
+ req_index -= 1
291
+ else:
292
+ preempted_req = self.running.pop()
293
+
294
+ self.kv_cache_manager.free(preempted_req)
295
+ self.encoder_cache_manager.free(preempted_req)
296
+ preempted_req.status = RequestStatus.PREEMPTED
297
+ preempted_req.num_computed_tokens = 0
298
+ preempted_req.num_preemptions += 1
299
+ if self.log_stats:
300
+ preempted_req.record_event(
301
+ EngineCoreEventType.PREEMPTED, scheduled_timestamp
302
+ )
303
+
304
+ self.waiting.prepend_request(preempted_req)
305
+ preempted_reqs.append(preempted_req)
306
+ if preempted_req == request:
307
+ # No more request to preempt. Cannot schedule this request.
308
+ break
309
310
if new_blocks is None:
311
# Cannot schedule this request.
@@ -599,13 +603,14 @@ class Scheduler(SchedulerInterface):
603
# Get the longest common prefix among all requests in the running queue.
604
# This can be potentially used for cascade attention.
605
num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
602
- if self.running:
603
- any_request = self.running[0]
604
- num_common_prefix_blocks = (
605
- self.kv_cache_manager.get_num_common_prefix_blocks(
606
- any_request.request_id
606
+ with record_function_or_nullcontext("schedule: get_num_common_prefix_blocks"):
607
+ if self.running:
608
+ any_request = self.running[0]
609
+ num_common_prefix_blocks = (
610
+ self.kv_cache_manager.get_num_common_prefix_blocks(
611
+ any_request.request_id
612
+ )
613
)
608
- )
614
615
# Construct the scheduler output.
616
new_reqs_data = [
@@ -614,13 +619,14 @@ class Scheduler(SchedulerInterface):
619
)
620
for req in scheduled_new_reqs
621
]
617
- cached_reqs_data = self._make_cached_request_data(
618
- scheduled_running_reqs,
619
- scheduled_resumed_reqs,
620
- num_scheduled_tokens,
621
- scheduled_spec_decode_tokens,
622
- req_to_new_blocks,
623
- )
622
+ with record_function_or_nullcontext("schedule: make_cached_request_data"):
623
+ cached_reqs_data = self._make_cached_request_data(
624
+ scheduled_running_reqs,
625
+ scheduled_resumed_reqs,
626
+ num_scheduled_tokens,
627
+ scheduled_spec_decode_tokens,
628
+ req_to_new_blocks,
629
+ )
630
631
# Record the request ids that were scheduled in this step.
632
self.prev_step_scheduled_req_ids.clear()
@@ -649,8 +655,8 @@ class Scheduler(SchedulerInterface):
655
if self.connector is not None:
656
meta = self.connector.build_connector_meta(scheduler_output)
657
scheduler_output.kv_connector_metadata = meta
652
-
653
- self._update_after_schedule(scheduler_output)
658
+ with record_function_or_nullcontext("schedule: update_after_schedule"):
659
+ self._update_after_schedule(scheduler_output)
660
return scheduler_output
661
662
def _update_after_schedule(
vllm/v1/engine/core.py
+73
-44
@@ -61,6 +61,7 @@ from vllm.v1.outputs import ModelRunnerOutput
61
from vllm.v1.request import Request, RequestStatus
62
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
63
from vllm.v1.structured_output import StructuredOutputManager
64
+from vllm.v1.utils import record_function_or_nullcontext
65
from vllm.version import __version__ as VLLM_VERSION
66
67
logger = init_logger(__name__)
@@ -315,17 +316,21 @@ class EngineCore:
316
# or finished and not yet removed from the batch.
317
if not self.scheduler.has_requests():
318
return {}, False
318
- scheduler_output = self.scheduler.schedule()
319
- future = self.model_executor.execute_model(scheduler_output, non_block=True)
320
- grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output)
321
- with self.log_error_detail(scheduler_output):
322
- model_output = future.result()
323
- if model_output is None:
324
- model_output = self.model_executor.sample_tokens(grammar_output)
325
-
326
- engine_core_outputs = self.scheduler.update_from_output(
327
- scheduler_output, model_output
328
- )
319
+ with record_function_or_nullcontext("core step: schedule"):
320
+ scheduler_output = self.scheduler.schedule()
321
+
322
+ with record_function_or_nullcontext("core step: execute_model"):
323
+ future = self.model_executor.execute_model(scheduler_output, non_block=True)
324
+ grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output)
325
+ with self.log_error_detail(scheduler_output):
326
+ model_output = future.result()
327
+ if model_output is None:
328
+ model_output = self.model_executor.sample_tokens(grammar_output)
329
+
330
+ with record_function_or_nullcontext("core step: update_from_output"):
331
+ engine_core_outputs = self.scheduler.update_from_output(
332
+ scheduler_output, model_output
333
+ )
334
335
return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0
336
@@ -363,32 +368,49 @@ class EngineCore:
368
model_executed = False
369
deferred_scheduler_output = None
370
if self.scheduler.has_requests():
366
- scheduler_output = self.scheduler.schedule()
367
- exec_future = self.model_executor.execute_model(
368
- scheduler_output, non_block=True
369
- )
371
+ with record_function_or_nullcontext("core step_with_batch_queue: schedule"):
372
+ scheduler_output = self.scheduler.schedule()
373
+ with record_function_or_nullcontext(
374
+ "core step_with_batch_queue: execute_model"
375
+ ):
376
+ exec_future = self.model_executor.execute_model(
377
+ scheduler_output, non_block=True
378
+ )
379
model_executed = scheduler_output.total_num_scheduled_tokens > 0
380
381
if scheduler_output.pending_structured_output_tokens:
373
- # We need to defer sampling until we have processed the model output
374
- # from the prior step.
375
- deferred_scheduler_output = scheduler_output
376
- # Block-wait for execute to return (continues running async on the GPU).
377
- with self.log_error_detail(scheduler_output):
378
- exec_result = exec_future.result()
379
- assert exec_result is None
382
+ with record_function_or_nullcontext(
383
+ "core step_with_batch_queue: pending_structured_output_tokens"
384
+ ):
385
+ # We need to defer sampling until we have processed the model output
386
+ # from the prior step.
387
+ deferred_scheduler_output = scheduler_output
388
+ # Block-wait for execute to return
389
+ # (continues running async on the GPU).
390
+ with self.log_error_detail(scheduler_output):
391
+ exec_result = exec_future.result()
392
+ assert exec_result is None
393
else:
381
- # We aren't waiting for any tokens, get any grammar output immediately.
382
- grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output)
394
+ with record_function_or_nullcontext(
395
+ "core step_with_batch_queue: get_grammar_bitmask"
396
+ ):
397
+ # We aren't waiting for any tokens, get any grammar
398
+ # output immediately.
399
+ grammar_output = self.scheduler.get_grammar_bitmask(
400
+ scheduler_output
401
+ )
402
# Block-wait for execute to return (continues running async on the GPU).
403
with self.log_error_detail(scheduler_output):
404
exec_result = exec_future.result()
405
406
if exec_result is None:
388
- # Call sample tokens.
389
- future = self.model_executor.sample_tokens(
390
- grammar_output, non_block=True
391
- )
407
+ with record_function_or_nullcontext(
408
+ "core step_with_batch_queue: sample_tokens"
409
+ ):
410
+ # Call sample tokens.
411
+ future = self.model_executor.sample_tokens(
412
+ grammar_output, non_block=True
413
+ )
414
else:
415
# No sampling required (e.g. all requests finished).
416
future = cast(Future[ModelRunnerOutput], exec_future)
@@ -408,27 +430,34 @@ class EngineCore:
430
# only be called when the scheduler contains requests or the queue
431
# is non-empty.
432
return None, False
411
-
412
- # Block until the next result is available.
413
- future, scheduler_output = batch_queue.pop()
414
- with self.log_error_detail(scheduler_output):
415
- model_output = future.result()
416
-
417
- engine_core_outputs = self.scheduler.update_from_output(
418
- scheduler_output, model_output
419
- )
433
+ with record_function_or_nullcontext("core step_with_batch_queue: model_output"):
434
+ # Block until the next result is available.
435
+ future, scheduler_output = batch_queue.pop()
436
+ with self.log_error_detail(scheduler_output):
437
+ model_output = future.result()
438
+ with record_function_or_nullcontext(
439
+ "core step_with_batch_queue: update_from_output"
440
+ ):
441
+ engine_core_outputs = self.scheduler.update_from_output(
442
+ scheduler_output, model_output
443
+ )
444
445
# NOTE(nick): We can either handle the deferred tasks here or save
446
# in a field and do it immediately once step_with_batch_queue is
447
# re-called. The latter slightly favors TTFT over TPOT/throughput.
448
if deferred_scheduler_output:
425
- # We now have the tokens needed to compute the bitmask for the
426
- # deferred request. Get the bitmask and call sample tokens.
427
- grammar_output = self.scheduler.get_grammar_bitmask(
428
- deferred_scheduler_output
429
- )
430
- future = self.model_executor.sample_tokens(grammar_output, non_block=True)
431
- batch_queue.appendleft((future, deferred_scheduler_output))
449
+ with record_function_or_nullcontext(
450
+ "core step_with_batch_queue: deferred_scheduler_output"
451
+ ):
452
+ # We now have the tokens needed to compute the bitmask for the
453
+ # deferred request. Get the bitmask and call sample tokens.
454
+ grammar_output = self.scheduler.get_grammar_bitmask(
455
+ deferred_scheduler_output
456
+ )
457
+ future = self.model_executor.sample_tokens(
458
+ grammar_output, non_block=True
459
+ )
460
+ batch_queue.appendleft((future, deferred_scheduler_output))
461
462
return engine_core_outputs, model_executed
463
vllm/v1/engine/llm_engine.py
+21
-16
@@ -36,6 +36,7 @@ from vllm.v1.executor import Executor
36
from vllm.v1.metrics.loggers import StatLoggerFactory, StatLoggerManager
37
from vllm.v1.metrics.reader import Metric, get_metrics_snapshot
38
from vllm.v1.metrics.stats import IterationStats
39
+from vllm.v1.utils import record_function_or_nullcontext
40
from vllm.v1.worker.worker_base import WorkerBase
41
42
logger = init_logger(__name__)
@@ -280,28 +281,32 @@ class LLMEngine:
281
return []
282
283
# 1) Get EngineCoreOutput from the EngineCore.
283
- outputs = self.engine_core.get_output()
284
+ with record_function_or_nullcontext("llm_genine step: get_output"):
285
+ outputs = self.engine_core.get_output()
286
287
# 2) Process EngineCoreOutputs.
286
- iteration_stats = IterationStats() if self.log_stats else None
287
- processed_outputs = self.output_processor.process_outputs(
288
- outputs.outputs,
289
- engine_core_timestamp=outputs.timestamp,
290
- iteration_stats=iteration_stats,
291
- )
292
- self.output_processor.update_scheduler_stats(outputs.scheduler_stats)
288
+ with record_function_or_nullcontext("llm_genine step: process_outputs"):
289
+ iteration_stats = IterationStats() if self.log_stats else None
290
+ processed_outputs = self.output_processor.process_outputs(
291
+ outputs.outputs,
292
+ engine_core_timestamp=outputs.timestamp,
293
+ iteration_stats=iteration_stats,
294
+ )
295
+ self.output_processor.update_scheduler_stats(outputs.scheduler_stats)
296
297
# 3) Abort any reqs that finished due to stop strings.
295
- self.engine_core.abort_requests(processed_outputs.reqs_to_abort)
298
+ with record_function_or_nullcontext("llm_genine step: abort_requests"):
299
+ self.engine_core.abort_requests(processed_outputs.reqs_to_abort)
300
301
# 4) Record stats
298
- if self.logger_manager is not None and outputs.scheduler_stats is not None:
299
- self.logger_manager.record(
300
- scheduler_stats=outputs.scheduler_stats,
301
- iteration_stats=iteration_stats,
302
- mm_cache_stats=self.processor.stat_mm_cache(),
303
- )
304
- self.do_log_stats_with_interval()
302
+ with record_function_or_nullcontext("llm_genine step: record_stats"):
303
+ if self.logger_manager is not None and outputs.scheduler_stats is not None:
304
+ self.logger_manager.record(
305
+ scheduler_stats=outputs.scheduler_stats,
306
+ iteration_stats=iteration_stats,
307
+ mm_cache_stats=self.processor.stat_mm_cache(),
308
+ )
309
+ self.do_log_stats_with_interval()
310
311
return processed_outputs.request_outputs
312
vllm/v1/worker/gpu_model_runner.py
+37
-33
@@ -2525,7 +2525,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2525
"after execute_model() returns None."
2526
)
2527
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
2528
- with record_function_or_nullcontext("Preprocess"):
2528
+ with record_function_or_nullcontext("gpu_model_runner: preprocess"):
2529
with self.synchronize_input_prep():
2530
# Update persistent batch states.
2531
self._update_states(scheduler_output)
@@ -2648,7 +2648,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2648
batch_descriptor=batch_descriptor,
2649
ubatch_slices=ubatch_slices,
2650
),
2651
- record_function_or_nullcontext("Forward"),
2651
+ record_function_or_nullcontext("gpu_model_runner: forward"),
2652
self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
2653
):
2654
model_output = self._model_forward(
@@ -2659,7 +2659,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2659
**model_kwargs,
2660
)
2661
2662
- with record_function_or_nullcontext("Postprocess"):
2662
+ with record_function_or_nullcontext("gpu_model_runner: postprocess"):
2663
if self.use_aux_hidden_state_outputs:
2664
# True when EAGLE 3 is used.
2665
hidden_states, aux_hidden_states = model_output
@@ -2756,12 +2756,12 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2756
scheduler_output, grammar_output, self.input_batch, logits
2757
)
2758
2759
- with record_function_or_nullcontext("Sample"):
2759
+ with record_function_or_nullcontext("gpu_model_runner: sample"):
2760
sampler_output = self._sample(logits, spec_decode_metadata)
2761
2762
def propose_draft_token_ids(sampled_token_ids):
2763
assert spec_decode_common_attn_metadata is not None
2764
- with record_function_or_nullcontext("Draft"):
2764
+ with record_function_or_nullcontext("gpu_model_runner: draft"):
2765
self._draft_token_ids = self.propose_draft_token_ids(
2766
scheduler_output,
2767
sampled_token_ids,
@@ -2799,7 +2799,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2799
# as inputs, and does not need to wait for bookkeeping to finish.
2800
propose_draft_token_ids(sampler_output.sampled_token_ids)
2801
2802
- with record_function_or_nullcontext("Bookkeep"):
2802
+ with record_function_or_nullcontext("gpu_model_runner: bookkeep"):
2803
(
2804
num_nans_in_logits,
2805
logprobs_lists,
@@ -2826,37 +2826,41 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2826
# tokens on the CPU, so they are run after bookkeeping.
2827
propose_draft_token_ids(valid_sampled_token_ids)
2828
2829
- with record_function_or_nullcontext("EPLB"):
2829
+ with record_function_or_nullcontext("gpu_model_runner: eplb"):
2830
self.eplb_step()
2831
-
2832
- output = ModelRunnerOutput(
2833
- req_ids=req_ids_output_copy,
2834
- req_id_to_index=req_id_to_index_output_copy,
2835
- sampled_token_ids=valid_sampled_token_ids,
2836
- logprobs=logprobs_lists,
2837
- prompt_logprobs_dict=prompt_logprobs_dict,
2838
- pooler_output=[],
2839
- kv_connector_output=kv_connector_output,
2840
- num_nans_in_logits=num_nans_in_logits,
2841
- )
2831
+ with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"):
2832
+ output = ModelRunnerOutput(
2833
+ req_ids=req_ids_output_copy,
2834
+ req_id_to_index=req_id_to_index_output_copy,
2835
+ sampled_token_ids=valid_sampled_token_ids,
2836
+ logprobs=logprobs_lists,
2837
+ prompt_logprobs_dict=prompt_logprobs_dict,
2838
+ pooler_output=[],
2839
+ kv_connector_output=kv_connector_output,
2840
+ num_nans_in_logits=num_nans_in_logits,
2841
+ )
2842
2843
if not self.use_async_scheduling:
2844
return output
2845
-
2846
- async_output = AsyncGPUModelRunnerOutput(
2847
- model_runner_output=output,
2848
- sampled_token_ids=sampler_output.sampled_token_ids,
2849
- logprobs_tensors=sampler_output.logprobs_tensors,
2850
- invalid_req_indices=invalid_req_indices,
2851
- async_output_copy_stream=self.async_output_copy_stream,
2852
- )
2853
-
2854
- # Save ref of sampled_token_ids CPU tensor if the batch contains
2855
- # any requests with sampling params that that require output ids.
2856
- self.input_batch.set_async_sampled_token_ids(
2857
- async_output.sampled_token_ids_cpu,
2858
- async_output.async_copy_ready_event,
2859
- )
2845
+ with record_function_or_nullcontext(
2846
+ "gpu_model_runner: AsyncGPUModelRunnerOutput"
2847
+ ):
2848
+ async_output = AsyncGPUModelRunnerOutput(
2849
+ model_runner_output=output,
2850
+ sampled_token_ids=sampler_output.sampled_token_ids,
2851
+ logprobs_tensors=sampler_output.logprobs_tensors,
2852
+ invalid_req_indices=invalid_req_indices,
2853
+ async_output_copy_stream=self.async_output_copy_stream,
2854
+ )
2855
+ with record_function_or_nullcontext(
2856
+ "gpu_model_runner: set_async_sampled_token_ids"
2857
+ ):
2858
+ # Save ref of sampled_token_ids CPU tensor if the batch contains
2859
+ # any requests with sampling params that that require output ids.
2860
+ self.input_batch.set_async_sampled_token_ids(
2861
+ async_output.sampled_token_ids_cpu,
2862
+ async_output.async_copy_ready_event,
2863
+ )
2864
2865
return async_output
2866