From f6f1cacbce0e39992d04843f652c1adda373ae43 Mon Sep 17 00:00:00 2001
From: Gahow Wang <gahow.wang@gmail.com>
Date: Sat, 11 Jul 2026 17:29:02 +0800
Subject: [PATCH 1/5] Add lightweight per-step OpProf telemetry

Assisted-by: OpenAI Codex
---
 vllm/envs.py                       |   4 +
 vllm/v1/core/sched/scheduler.py    |  28 +++
 vllm/v1/opprof.py                  | 337 +++++++++++++++++++++++++++++
 vllm/v1/worker/gpu_model_runner.py |   6 +-
 4 files changed, 374 insertions(+), 1 deletion(-)
 create mode 100644 vllm/v1/opprof.py

diff --git a/vllm/envs.py b/vllm/envs.py
index 27a85bb..b3093e9 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -45,6 +45,7 @@ if TYPE_CHECKING:
     VLLM_LOGGING_COLOR: str = "auto"
     NO_COLOR: bool = False
     VLLM_LOG_STATS_INTERVAL: float = 10.0
+    VLLM_OPPROF_DIR: str = ""
     VLLM_TRACE_FUNCTION: int = 0
     VLLM_USE_FLASHINFER_SAMPLER: bool = True
     VLLM_PP_LAYER_PARTITION: str | None = None
@@ -786,6 +787,9 @@ environment_variables: dict[str, Callable[[], Any]] = {
         if (val := float(os.getenv("VLLM_LOG_STATS_INTERVAL", "10."))) > 0.0
         else 10.0
     ),
+    # Directory for per-step OpProf JSONL telemetry.
+    # Empty disables OpProf.
+    "VLLM_OPPROF_DIR": lambda: os.getenv("VLLM_OPPROF_DIR", ""),
     # Trace function calls
     # If set to 1, vllm will trace function calls
     # Useful for debugging
diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py
index 90d93a1..303c562 100644
--- a/vllm/v1/core/sched/scheduler.py
+++ b/vllm/v1/core/sched/scheduler.py
@@ -7,6 +7,7 @@ from collections.abc import Iterable
 from dataclasses import replace
 from typing import Any
 
+import vllm.envs as envs
 from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.config import VllmConfig
 from vllm.distributed.ec_transfer.ec_connector.base import (
@@ -55,6 +56,7 @@ from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutp
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.metrics.perf import ModelMetrics, PerfStats
 from vllm.v1.metrics.stats import PrefixCacheStats, SchedulerStats
+from vllm.v1.opprof import OpProfRecorder
 from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput
 from vllm.v1.request import Request, RequestStatus, StreamingUpdate
 from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
@@ -271,6 +273,12 @@ class Scheduler(SchedulerInterface):
         if self.connector is not None:
             self.connector.bind_gpu_block_pool(self.kv_cache_manager.block_pool)
 
+        self.opprof = OpProfRecorder.create(
+            output_dir=envs.VLLM_OPPROF_DIR,
+            dp_rank=self.parallel_config.data_parallel_index,
+            log_stats=self.log_stats,
+        )
+
         self.use_pp = self.parallel_config.pipeline_parallel_size > 1
         self.use_v2_model_runner = vllm_config.use_v2_model_runner
         # Scheduler iteration counter. Drives the V2+PP+async decode-throttle
@@ -386,6 +394,9 @@ class Scheduler(SchedulerInterface):
         return num_new_tokens
 
     def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput:
+        opprof_start = (
+            self.opprof.capture_start(self) if self.opprof is not None else None
+        )
         self.current_step += 1
         # NOTE(woosuk) on the scheduling algorithm:
         # There's no "decoding phase" nor "prefill phase" in the scheduler.
@@ -1090,6 +1101,14 @@ class Scheduler(SchedulerInterface):
             )
             scheduler_output.ec_connector_metadata = ec_meta
 
+        if self.opprof is not None:
+            assert opprof_start is not None
+            self.opprof.begin(
+                scheduler=self,
+                output=scheduler_output,
+                start=opprof_start,
+            )
+
         # Advance the fence only for non-empty steps (those that actually
         # write KV and have their output processed later in update_from_output).
         if self.defer_block_free and total_num_scheduled_tokens > 0:
@@ -1800,6 +1819,12 @@ class Scheduler(SchedulerInterface):
                 engine_core_outputs[0] = eco = EngineCoreOutputs()
             eco.scheduler_stats = stats
 
+        if self.opprof is not None:
+            self.opprof.finalize(
+                output=scheduler_output,
+                cudagraph_stat=cudagraph_stats,
+            )
+
         return engine_core_outputs
 
     @staticmethod
@@ -2292,6 +2317,9 @@ class Scheduler(SchedulerInterface):
         if self.ec_connector is not None:
             self.ec_connector.shutdown()
 
+        if self.opprof is not None:
+            self.opprof.close()
+
         logger.debug_once("[shutdown] Scheduler: complete")
 
     ########################################################################
diff --git a/vllm/v1/opprof.py b/vllm/v1/opprof.py
new file mode 100644
index 0000000..f0330d0
--- /dev/null
+++ b/vllm/v1/opprof.py
@@ -0,0 +1,337 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import atexit
+import logging
+import os
+import queue
+import threading
+import time
+from bisect import bisect_left
+from pathlib import Path
+from typing import Any
+
+import msgspec
+
+logger = logging.getLogger(__name__)
+
+SCHEMA_VERSION = 1
+CONTEXT_LENGTH_EDGES = tuple(1 << exponent for exponent in range(7, 18))
+CHUNK_SIZE_EDGES = tuple(1 << exponent for exponent in range(4, 12))
+DEFAULT_QUEUE_CAPACITY = 8192
+_CLOSE_TIMEOUT_SECONDS = 1.0
+_STOP = object()
+_PREFIX_FIELDS = (  # noqa: SIM905
+    "requests queries hits preempted_requests preempted_queries preempted_hits"
+).split()
+
+
+ScheduleStart = tuple[int, int, dict[str, dict[str, int] | None]]
+
+
+def classify_chunk(was_chunk: bool, end: int, target: int) -> str:
+    assert 0 <= end <= target
+    if was_chunk:
+        return "middle" if end < target else "final"
+    return "first" if end < target else "unsplit"
+
+
+def _prefix_values(stats: Any | None) -> dict[str, int] | None:
+    if stats is None:
+        return None
+    return {name: int(getattr(stats, name, 0)) for name in _PREFIX_FIELDS}
+
+
+def _prefix_snapshot(scheduler: Any) -> dict[str, dict[str, int] | None]:
+    return {
+        "local": _prefix_values(scheduler.kv_cache_manager.prefix_cache_stats),
+        "external": _prefix_values(scheduler.connector_prefix_cache_stats),
+    }
+
+
+def _prefix_delta(
+    before: dict[str, int] | None, after: dict[str, int] | None
+) -> dict[str, int] | None:
+    if before is None and after is None:
+        return None
+    before = before or dict.fromkeys(_PREFIX_FIELDS, 0)
+    after = after or dict.fromkeys(_PREFIX_FIELDS, 0)
+    delta = {name: after[name] - before[name] for name in _PREFIX_FIELDS}
+    assert all(value >= 0 for value in delta.values())
+    return delta
+
+
+class JSONLWriter:
+    def __init__(
+        self,
+        path: Path,
+        capacity: int = DEFAULT_QUEUE_CAPACITY,
+        start: bool = True,
+    ) -> None:
+        self._queue: queue.Queue[Any] = queue.Queue(capacity)
+        self._encoder = msgspec.json.Encoder()
+        self._file = path.open("xb", buffering=1 << 20)
+        self._thread = threading.Thread(target=self._run, daemon=True)
+        self._failure_lock = threading.Lock()
+        self._started = self._closed = False
+        self.failed = False
+        self.failure: Exception | None = None
+        self.encoded_records = self.written_records = 0
+        self.dropped_records = self._unreported_drops = 0
+        if start:
+            self.start()
+
+    def start(self) -> None:
+        if not self._started:
+            self._started = True
+            self._thread.start()
+
+    def _record_failure(self, error: Exception) -> None:
+        with self._failure_lock:
+            if self.failed:
+                return
+            self.failed = True
+            self.failure = error
+            logger.error("OpProf writer failed: %s", error)
+
+    def _writer_unavailable(self) -> bool:
+        if self.failed:
+            return True
+        if self._started and not self._thread.is_alive():
+            self._record_failure(RuntimeError("writer thread stopped unexpectedly"))
+            return True
+        return False
+
+    def _drop(self, pending: int) -> bool:
+        self.dropped_records += 1
+        self._unreported_drops = pending + 1
+        return False
+
+    def submit(self, record: dict[str, Any]) -> bool:
+        if self._closed:
+            raise RuntimeError("OpProf writer is closed")
+        pending = self._unreported_drops
+        record["dropped_records_before"] = pending
+        if self._writer_unavailable():
+            return self._drop(pending)
+        payload = self._encoder.encode(record) + b"\n"
+        self.encoded_records += 1
+        if self._writer_unavailable():
+            return self._drop(pending)
+        try:
+            self._queue.put_nowait(payload)
+        except queue.Full:
+            return self._drop(pending)
+        self._unreported_drops = 0
+        return True
+
+    def _run(self) -> None:
+        buffered = 0
+        last_flush = time.monotonic()
+        try:
+            while True:
+                try:
+                    item = self._queue.get(timeout=1.0)
+                except queue.Empty:
+                    self._file.flush()
+                    buffered = 0
+                    last_flush = time.monotonic()
+                    continue
+                try:
+                    if item is _STOP:
+                        break
+                    self._file.write(item)
+                    buffered += len(item)
+                    self.written_records += 1
+                    now = time.monotonic()
+                    if buffered >= 1 << 20 or now - last_flush >= 1.0:
+                        self._file.flush()
+                        buffered = 0
+                        last_flush = now
+                finally:
+                    self._queue.task_done()
+        except Exception as error:
+            self._record_failure(error)
+        finally:
+            footer = dict(
+                schema=SCHEMA_VERSION,
+                record_type="footer",
+                encoded_records=self.encoded_records,
+                written_records=self.written_records,
+                dropped_records=self.dropped_records,
+            )
+            try:
+                self._file.write(self._encoder.encode(footer) + b"\n")
+                self._file.flush()
+            except Exception as error:
+                self._record_failure(error)
+            finally:
+                try:
+                    self._file.close()
+                except Exception as error:
+                    self._record_failure(error)
+
+    def close(self) -> None:
+        if self._closed:
+            return
+        self._closed = True
+        self.start()
+        if not self._thread.is_alive():
+            return
+        try:
+            self._queue.put(_STOP, timeout=_CLOSE_TIMEOUT_SECONDS)
+        except queue.Full:
+            if not self._thread.is_alive():
+                return
+            try:
+                self._queue.put_nowait(_STOP)
+            except queue.Full:
+                self._record_failure(
+                    TimeoutError("timed out enqueueing writer stop sentinel")
+                )
+                return
+        self._thread.join(timeout=_CLOSE_TIMEOUT_SECONDS)
+        if self._thread.is_alive():
+            self._record_failure(TimeoutError("timed out joining writer thread"))
+
+
+class OpProfRecorder:
+    def __init__(self, engine_id: str, writer: JSONLWriter) -> None:
+        self.engine_id = engine_id
+        self.writer = writer
+        self._next_step = 0
+        self._pending: dict[int, dict[str, Any]] = {}
+        atexit.register(self.close)
+
+    @classmethod
+    def create(
+        cls, output_dir: str, dp_rank: int, log_stats: bool
+    ) -> "OpProfRecorder | None":
+        if not output_dir:
+            return None
+        if not log_stats:
+            raise ValueError("VLLM_OPPROF_DIR requires log stats to be enabled")
+        directory = Path(output_dir).expanduser()
+        if not directory.is_absolute():
+            raise ValueError("VLLM_OPPROF_DIR must be an absolute path")
+        directory.mkdir(parents=True, exist_ok=True)
+        engine_id = f"dp{dp_rank}-pid{os.getpid()}"
+        name = f"opprof-v{SCHEMA_VERSION}-{engine_id}-{time.time_ns()}.jsonl"
+        return cls(engine_id, JSONLWriter(directory / name))
+
+    @staticmethod
+    def capture_start(scheduler: Any) -> ScheduleStart:
+        return time.time_ns(), time.monotonic_ns(), _prefix_snapshot(scheduler)
+
+    def begin(self, scheduler: Any, output: Any, start: ScheduleStart) -> None:
+        key = id(output)
+        assert key not in self._pending, "duplicate OpProf begin"
+        new_ids = {request.req_id for request in output.scheduled_new_reqs}
+        prefill_requests = prefill_tokens = decode_requests = decode_tokens = 0
+        context_length_hist = [0] * (len(CONTEXT_LENGTH_EDGES) + 1)
+        chunk_size_hist = [0] * (len(CHUNK_SIZE_EDGES) + 1)
+        chunks: dict[str, Any] = dict.fromkeys(
+            ("first", "middle", "final", "unsplit"), 0
+        )
+        for req_id, num_tokens in output.num_scheduled_tokens.items():
+            request = scheduler.requests[req_id]
+            end = request.num_computed_tokens + num_tokens
+            assert end >= 0
+            context_length_hist[bisect_left(CONTEXT_LENGTH_EDGES, end)] += 1
+            is_prefill = (
+                req_id in new_ids
+                or output.scheduled_cached_reqs.is_context_phase(req_id)
+            )
+            if is_prefill:
+                prefill_requests += 1
+                prefill_tokens += num_tokens
+                assert num_tokens >= 0
+                chunk_size_hist[bisect_left(CHUNK_SIZE_EDGES, num_tokens)] += 1
+                target = request.num_tokens + request.num_output_placeholders
+                chunks[classify_chunk(request.is_prefill_chunk, end, target)] += 1
+            else:
+                decode_requests += 1
+                decode_tokens += num_tokens
+        prefix_after = _prefix_snapshot(scheduler)
+        block_pool = scheduler.kv_cache_manager.block_pool
+        total_blocks = block_pool.num_gpu_blocks - 1
+        free_blocks = block_pool.get_num_free_blocks()
+        assert 0 <= free_blocks <= total_blocks
+        chunks["tokens"] = prefill_tokens
+        chunks["chunk_size_hist"] = chunk_size_hist
+        values = dict(
+            schema=SCHEMA_VERSION,
+            engine_id=self.engine_id,
+            step_index=self._next_step,
+            submit_wall_ns=start[0],
+            submit_mono_ns=start[1],
+            model_executed=output.total_num_scheduled_tokens > 0,
+            scheduled_requests=len(output.num_scheduled_tokens),
+            decode_batch_size=decode_requests,
+            prefill_requests=prefill_requests,
+            prefill_tokens=prefill_tokens,
+            decode_tokens=decode_tokens,
+            chunked_prefill=chunks,
+            context_length_hist=context_length_hist,
+            preemptions=len(output.preempted_req_ids or ()),
+            queues=dict(
+                running=len(scheduler.running),
+                waiting=len(scheduler.waiting),
+                deferred=len(scheduler.skipped_waiting),
+            ),
+            kv=dict(
+                total_blocks=total_blocks,
+                free_blocks=free_blocks,
+                used_blocks=total_blocks - free_blocks,
+                usage=scheduler.kv_cache_manager.usage,
+            ),
+            prefix=dict(
+                local=_prefix_delta(start[2]["local"], prefix_after["local"]),
+                external=_prefix_delta(start[2]["external"], prefix_after["external"]),
+            ),
+        )
+        assert prefill_tokens + decode_tokens == output.total_num_scheduled_tokens
+        self._pending[key] = values
+        self._next_step += 1
+
+    def finalize(self, output: Any, cudagraph_stat: Any | None) -> bool:
+        try:
+            values = self._pending.pop(id(output))
+        except KeyError:
+            raise AssertionError("missing or already finalized OpProf step") from None
+        if cudagraph_stat is None:
+            assert not values["model_executed"]
+            cudagraph = dict(
+                hit=False,
+                runtime_mode="NONE",
+                unpadded_tokens=0,
+                bucket_tokens=0,
+                padding_tokens=0,
+            )
+        else:
+            mode = str(cudagraph_stat.runtime_mode).rsplit(".", 1)[-1]
+            cudagraph = dict(
+                hit=mode != "NONE",
+                runtime_mode=mode,
+                unpadded_tokens=cudagraph_stat.num_unpadded_tokens,
+                bucket_tokens=cudagraph_stat.num_padded_tokens,
+                padding_tokens=cudagraph_stat.num_paddings,
+            )
+        record = dict(
+            values,
+            complete_mono_ns=time.monotonic_ns(),
+            cudagraph=cudagraph,
+            moe_expert_load=None,
+            dropped_records_before=0,
+        )
+        return self.writer.submit(record)
+
+    def close(self) -> None:
+        self.writer.close()
+
+    @property
+    def failed(self) -> bool:
+        return self.writer.failed
+
+    @property
+    def failure(self) -> Exception | None:
+        return self.writer.failure
diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py
index 74938a8..c11d773 100644
--- a/vllm/v1/worker/gpu_model_runner.py
+++ b/vllm/v1/worker/gpu_model_runner.py
@@ -437,6 +437,7 @@ class GPUModelRunner(
         self.scheduler_config = vllm_config.scheduler_config
         self.speculative_config = vllm_config.speculative_config
         self.observability_config = vllm_config.observability_config
+        self.opprof_enabled = bool(envs.VLLM_OPPROF_DIR)
 
         model_config = self.model_config
         cache_config = self.cache_config
@@ -3917,7 +3918,10 @@ class GPUModelRunner(
                 assert batch_descriptor.num_tokens == num_tokens_padded
 
         cudagraph_stats = None
-        if self.vllm_config.observability_config.cudagraph_metrics:
+        if (
+            self.vllm_config.observability_config.cudagraph_metrics
+            or self.opprof_enabled
+        ):
             cudagraph_stats = CUDAGraphStat(
                 num_unpadded_tokens=num_tokens,
                 num_padded_tokens=batch_descriptor.num_tokens,
-- 
2.43.0

