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 10, 11:30 UTC
[Feature] Add iteration level logging and enhance nvtx marker (#31193)
Signed-off-by: Max Hu <maxhu@nvidia.com> Signed-off-by: Max Hu <hyoung2991@gmail.com> Co-authored-by: Max Hu <maxhu@nvidia.com>
Max Hu committed
Jan 8, 2026 at 19:13 UTC
6ebe34d6fab6ab8e6db6e37dcfd3117da1125d11
6 files changed
+138
-10
vllm/config/observability.py
+6
@@ -75,6 +75,12 @@ class ObservabilityConfig:
75
enable_mfu_metrics: bool = False
76
"""Enable Model FLOPs Utilization (MFU) metrics."""
77
78
+ enable_logging_iteration_details: bool = False
79
+ """Enable detailed logging of iteration details.
80
+ If set, vllm EngineCore will log iteration details
81
+ This includes number of context/generation requests and tokens
82
+ and the elapsed cpu time for the iteration."""
83
+
84
@cached_property
85
def collect_model_forward_time(self) -> bool:
86
"""Whether to collect model forward time for the request."""
vllm/engine/arg_utils.py
+8
@@ -525,6 +525,9 @@ class EngineArgs:
525
ObservabilityConfig.enable_layerwise_nvtx_tracing
526
)
527
enable_mfu_metrics: bool = ObservabilityConfig.enable_mfu_metrics
528
+ enable_logging_iteration_details: bool = (
529
+ ObservabilityConfig.enable_logging_iteration_details
530
+ )
531
enable_mm_processor_stats: bool = ObservabilityConfig.enable_mm_processor_stats
532
scheduling_policy: SchedulerPolicy = SchedulerConfig.policy
533
scheduler_cls: str | type[object] | None = SchedulerConfig.scheduler_cls
@@ -1059,6 +1062,10 @@ class EngineArgs:
1062
"--enable-mfu-metrics",
1063
**observability_kwargs["enable_mfu_metrics"],
1064
)
1065
+ observability_group.add_argument(
1066
+ "--enable-logging-iteration-details",
1067
+ **observability_kwargs["enable_logging_iteration_details"],
1068
+ )
1069
1070
# Scheduler arguments
1071
scheduler_kwargs = get_kwargs(SchedulerConfig)
@@ -1713,6 +1720,7 @@ class EngineArgs:
1720
enable_layerwise_nvtx_tracing=self.enable_layerwise_nvtx_tracing,
1721
enable_mfu_metrics=self.enable_mfu_metrics,
1722
enable_mm_processor_stats=self.enable_mm_processor_stats,
1723
+ enable_logging_iteration_details=self.enable_logging_iteration_details,
1724
)
1725
1726
# Compilation config overrides
vllm/v1/core/sched/output.py
+15
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from dataclasses import dataclass
5
+from functools import cached_property
6
from typing import TYPE_CHECKING
7
8
from vllm._bc_linter import bc_linter_include
@@ -151,6 +152,20 @@ class CachedRequestData:
152
def num_reqs(self) -> int:
153
return len(self.req_ids)
154
155
+ @cached_property
156
+ def _req_id_to_num_output_tokens(self) -> dict[str, int]:
157
+ """Cache mapping of req_id to num_output_tokens for O(1) lookup.
158
+
159
+ This cached property is safe because CachedRequestData instances
160
+ are created fresh each scheduling iteration and not mutated during
161
+ computation of iteration details.
162
+ """
163
+ return dict(zip(self.req_ids, self.num_output_tokens))
164
+
165
+ def is_context_phase(self, req_id: str) -> bool:
166
+ num_output_tokens = self._req_id_to_num_output_tokens.get(req_id)
167
+ return num_output_tokens is not None and num_output_tokens == 0
168
+
169
@classmethod
170
def make_empty(cls) -> "CachedRequestData":
171
return cls(
vllm/v1/engine/core.py
+39
-3
@@ -65,6 +65,7 @@ from vllm.v1.outputs import ModelRunnerOutput
65
from vllm.v1.request import Request, RequestStatus
66
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
67
from vllm.v1.structured_output import StructuredOutputManager
68
+from vllm.v1.utils import compute_iteration_details
69
from vllm.version import __version__ as VLLM_VERSION
70
71
logger = init_logger(__name__)
@@ -208,7 +209,6 @@ class EngineCore:
209
self.async_scheduling = vllm_config.scheduler_config.async_scheduling
210
211
self.aborts_queue = queue.Queue[list[str]]()
211
-
212
# Mark the startup heap as static so that it's ignored by GC.
213
# Reduces pause times of oldest generation collections.
214
freeze_gc_heap()
@@ -337,6 +337,36 @@ class EngineCore:
337
)
338
raise err
339
340
+ @contextmanager
341
+ def log_iteration_details(self, scheduler_output: SchedulerOutput):
342
+ if not self.vllm_config.observability_config.enable_logging_iteration_details:
343
+ yield
344
+ return
345
+ self._iteration_index = getattr(self, "_iteration_index", 0)
346
+ iteration_details = compute_iteration_details(scheduler_output)
347
+ before = time.monotonic()
348
+ yield
349
+ logger.info(
350
+ "".join(
351
+ [
352
+ "Iteration(",
353
+ str(self._iteration_index),
354
+ "): ",
355
+ str(iteration_details.num_ctx_requests),
356
+ " context requests, ",
357
+ str(iteration_details.num_ctx_tokens),
358
+ " context tokens, ",
359
+ str(iteration_details.num_generation_requests),
360
+ " generation requests, ",
361
+ str(iteration_details.num_generation_tokens),
362
+ " generation tokens, iteration elapsed time: ",
363
+ format((time.monotonic() - before) * 1000, ".2f"),
364
+ " ms",
365
+ ]
366
+ )
367
+ )
368
+ self._iteration_index += 1
369
+
370
def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]:
371
"""Schedule, execute, and make output.
372
@@ -351,7 +381,10 @@ class EngineCore:
381
scheduler_output = self.scheduler.schedule()
382
future = self.model_executor.execute_model(scheduler_output, non_block=True)
383
grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output)
354
- with self.log_error_detail(scheduler_output):
384
+ with (
385
+ self.log_error_detail(scheduler_output),
386
+ self.log_iteration_details(scheduler_output),
387
+ ):
388
model_output = future.result()
389
if model_output is None:
390
model_output = self.model_executor.sample_tokens(grammar_output)
@@ -447,7 +480,10 @@ class EngineCore:
480
481
# Block until the next result is available.
482
future, scheduler_output, exec_model_fut = batch_queue.pop()
450
- with self.log_error_detail(scheduler_output):
483
+ with (
484
+ self.log_error_detail(scheduler_output),
485
+ self.log_iteration_details(scheduler_output),
486
+ ):
487
model_output = future.result()
488
if model_output is None:
489
# None from sample_tokens() implies that the original execute_model()
vllm/v1/utils.py
+52
@@ -7,6 +7,7 @@ import time
7
import weakref
8
from collections.abc import Callable, Sequence
9
from contextlib import AbstractContextManager
10
+from dataclasses import dataclass
11
from multiprocessing import connection
12
from multiprocessing.process import BaseProcess
13
from typing import (
@@ -27,6 +28,7 @@ from vllm.logger import init_logger
28
from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message
29
from vllm.utils.network_utils import get_open_port, get_open_zmq_ipc_path, get_tcp_uri
30
from vllm.utils.system_utils import kill_process_tree
31
+from vllm.v1.core.sched.output import SchedulerOutput
32
33
if TYPE_CHECKING:
34
import numpy as np
@@ -412,3 +414,53 @@ def tensor_data(tensor: torch.Tensor) -> memoryview:
414
A memoryview of the tensor data as uint8.
415
"""
416
return tensor.flatten().contiguous().view(torch.uint8).numpy().data
417
+
418
+
419
+@dataclass
420
+class IterationDetails:
421
+ num_ctx_requests: int
422
+ num_ctx_tokens: int
423
+ num_generation_requests: int
424
+ num_generation_tokens: int
425
+
426
+ def __repr__(self) -> str:
427
+ return f"IterationDetails(num_ctx_requests={self.num_ctx_requests},\
428
+ num_ctx_tokens={self.num_ctx_tokens}, \
429
+ num_generation_requests={self.num_generation_requests}, \
430
+ num_generation_tokens={self.num_generation_tokens})"
431
+
432
+
433
+def compute_iteration_details(scheduler_output: SchedulerOutput) -> IterationDetails:
434
+ """
435
+ Compute the number of context/generation requests and tokens
436
+ for the current iteration's scheduler output. A requests is regarded
437
+ as a context request if its output tokens are still 0, an extended chunk
438
+ of chunked prefill falls into this category.
439
+
440
+ Args:
441
+ scheduler_output: The scheduler output for the current iteration.
442
+
443
+ Returns:
444
+ An IterationDetails object containing the number of
445
+ context/generation requests and tokens.
446
+ """
447
+ num_context_requests = 0
448
+ num_context_tokens = 0
449
+ num_generation_requests = 0
450
+ num_generation_tokens = 0
451
+ new_req_ids = {new_req.req_id for new_req in scheduler_output.scheduled_new_reqs}
452
+ for req_id, num_tokens in scheduler_output.num_scheduled_tokens.items():
453
+ if scheduler_output.scheduled_cached_reqs.is_context_phase(req_id) or (
454
+ req_id in new_req_ids
455
+ ):
456
+ num_context_requests += 1
457
+ num_context_tokens += num_tokens
458
+ else:
459
+ num_generation_requests += 1
460
+ num_generation_tokens += num_tokens
461
+ return IterationDetails(
462
+ num_context_requests,
463
+ num_context_tokens,
464
+ num_generation_requests,
465
+ num_generation_tokens,
466
+ )
vllm/v1/worker/gpu_worker.py
+18
-7
@@ -50,7 +50,7 @@ from vllm.v1.outputs import (
50
DraftTokenIds,
51
ModelRunnerOutput,
52
)
53
-from vllm.v1.utils import report_usage_stats
53
+from vllm.v1.utils import compute_iteration_details, report_usage_stats
54
from vllm.v1.worker.utils import is_residual_scattered_for_sp
55
from vllm.v1.worker.worker_base import WorkerBase
56
from vllm.v1.worker.workspace import init_workspace_manager
@@ -547,18 +547,29 @@ class Worker(WorkerBase):
547
548
def annotate_profile(self, scheduler_output):
549
# add trace annotation so that we can easily distinguish
550
- # new/cached request numbers in each iteration
550
+ # context/generation request numbers in each iteration.
551
+ # A context request is a request that has not yet generated any tokens
552
if not self.profiler:
553
return nullcontext()
554
555
self.profiler.step()
556
556
- num_new = len(scheduler_output.scheduled_new_reqs)
557
- num_cached = len(scheduler_output.scheduled_cached_reqs.req_ids)
558
-
559
- return self.profiler.annotate_context_manager(
560
- f"execute_new_{num_new}_cached_{num_cached}"
557
+ iteration_details = compute_iteration_details(scheduler_output)
558
+
559
+ annotation = "".join(
560
+ [
561
+ "execute_context_",
562
+ str(iteration_details.num_ctx_requests),
563
+ "(",
564
+ str(iteration_details.num_ctx_tokens),
565
+ ")_generation_",
566
+ str(iteration_details.num_generation_requests),
567
+ "(",
568
+ str(iteration_details.num_generation_tokens),
569
+ ")",
570
+ ]
571
)
572
+ return self.profiler.annotate_context_manager(annotation)
573
574
@torch.inference_mode()
575
def sample_tokens(