diff --git a/vllm/envs.py b/vllm/envs.py index 38ae7fd01b..e09defd278 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -511,6 +511,7 @@ if TYPE_CHECKING: VLLM_GPU_NIC_PCIE_MAPPING: str = "" VLLM_NIC_SELECTION_VARS: str = "" VLLM_ENABLE_HPC_OPS: bool = False + MERLIN_TOPK_EXACT_COUNTERS_LOG: int = 0 def get_default_cache_root(): @@ -2538,6 +2539,13 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_SKIP_VERSION_SUFFIX": lambda: bool( int(os.getenv("VLLM_SKIP_VERSION_SUFFIX", "0")) ), + # Seconds between `topk-exact-counters` lines from each GPU worker + # (0, the default, disables the reader entirely). The line reports the + # medium-path exact-restart counters; see + # vllm/models/qwen4_exp/nvidia/ops/topk_counters.py. + "MERLIN_TOPK_EXACT_COUNTERS_LOG": lambda: int( + os.getenv("MERLIN_TOPK_EXACT_COUNTERS_LOG", "0") + ), } @@ -2750,6 +2758,10 @@ def compile_factors() -> dict[str, object]: # configuration factor and is hashed conservatively. "VLLM_E2_PLE_FLAG_CHECK", "VLLM_E2_MAMBA_SLOT_OWNERSHIP", + # Logging cadence for an already-compiled kernel's counters: it gates a + # host-side timer and a snapshot op, and is read by no compiled code, so + # hashing it would invalidate every compile cache for a log interval. + "MERLIN_TOPK_EXACT_COUNTERS_LOG", # ------------------------------------------------------------------ "FORCE_COLOR", } diff --git a/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py b/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py new file mode 100644 index 0000000000..02c1720f11 --- /dev/null +++ b/vllm/models/qwen4_exp/nvidia/ops/topk_counters.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Python reader for the medium-path exact-restart top-k counters. + +The counters themselves are device-global uint32s owned by +`csrc/libtorch_stable/topk.cu`; `torch.ops._C.persistent_topk_exact_counters` +snapshots them into an int32 CUDA tensor. Without this module nothing in a +running engine ever calls that op, so the instrument records into a buffer no +one reads -- which measures nothing. + +SLOT LAYOUT is build-dependent, so nothing here hard-codes five. `slot_count()` +asks the build first (a `..._slots` op, should a later build expose one) and +otherwise reads `MERLIN_TOPK_COUNTER_SLOTS` (default 5); names come from a +table keyed by that length: + + 5 rows overflow fallbacks faults switch + 10 rows overflow fallbacks faults rows_large ties_cut ties_dropped + faults_large switch switch_large + +An unknown length degrades to `slot` names rather than mislabelling. + +HOW THE SMOKE READS IT. Set `MERLIN_TOPK_EXACT_COUNTERS_LOG=` on the +engine; each GPU worker then emits, at most that often and once more at +shutdown, one line into the ordinary engine log: + + topk-exact-counters rank=PP0_TP1 rows=... overflow=... fallbacks=... + faults=... switch=... + +Counts are cumulative per worker process, so the smoke takes the LATEST line +per rank -- `grep topk-exact-counters | ...` keyed on the `rank=` +field -- and reads `faults` as the pass/fail number (it must stay 0) with +`rows` as the denominator. Earlier lines are only useful as a rate. + +COST. `read()` passes `sync=1`, i.e. a full device synchronize, because a +stream-ordered snapshot can miss increments from work in flight on another +stream. That is acceptable ONLY because the cadence is coarse (>= 30 s in +practice): the time check happens before the op call, so a step that is not +due for a log touches nothing but `time.monotonic()`. The read is also refused +inside a cudagraph capture -- the op allocates and synchronizes, which would +abort the capture. +""" + +import os +import time + +import torch + +import vllm.envs as envs +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_OP_NAME = "persistent_topk_exact_counters" + +_MEDIUM_NAMES = ("rows", "overflow", "fallbacks", "faults", "switch") +_LARGEDET_NAMES = ( + "rows", + "overflow", + "fallbacks", + "faults", + "rows_large", + "ties_cut", + "ties_dropped", + "faults_large", + "switch", + "switch_large", +) +_SLOT_NAMES: dict[int, tuple[str, ...]] = {5: _MEDIUM_NAMES, 10: _LARGEDET_NAMES} + +# The medium-path line contract. Any further name is appended after these, in slot +# order, so a 10-slot build still starts with exactly the medium-path line. +_LINE_ORDER = _MEDIUM_NAMES + +_warned_absent = False +_warned_failed = False +_last_log_monotonic = 0.0 + + +def _get_op(name: str): + try: + return getattr(torch.ops._C, name, None) + except (AttributeError, RuntimeError): + return None + + +def slot_count() -> int: + """How many int32 slots this build's counter array carries.""" + query = _get_op(_OP_NAME + "_slots") + if query is not None: + return int(query()) + return int(os.getenv("MERLIN_TOPK_COUNTER_SLOTS", "5")) + + +def slot_names(count: int) -> tuple[str, ...]: + names = _SLOT_NAMES.get(count) + return names if names is not None else tuple(f"slot{i}" for i in range(count)) + + +def read(reset: bool = False, sync: bool = True) -> dict[str, int]: + """Snapshot the counters as {slot name: value}; {} when unreadable.""" + global _warned_absent + op = _get_op(_OP_NAME) + if op is None: + if not _warned_absent: + _warned_absent = True + logger.warning( + "torch.ops._C.%s is not in this build; the top-k exact-restart " + "counters cannot be read.", + _OP_NAME, + ) + return {} + if reset and not sync: + raise ValueError("reset requires sync: zeroing races in-flight atomics") + if not (torch.cuda.is_available() and torch.cuda.is_initialized()): + return {} + if torch.cuda.is_current_stream_capturing(): + # The op allocates and synchronizes; calling it here aborts the capture. + return {} + count = slot_count() + device = torch.device("cuda", torch.cuda.current_device()) + out = torch.zeros(count, dtype=torch.int32, device=device) + op(out, 1 if reset else 0, 1 if sync else 0) + return dict(zip(slot_names(count), out.tolist(), strict=True)) + + +def rank_tag() -> str: + from vllm.distributed.parallel_state import get_pp_group, get_tp_group + + try: + return f"PP{get_pp_group().rank_in_group}_TP{get_tp_group().rank_in_group}" + except Exception: + # Pre-init or single-process: say so rather than claim rank 0. + return "PP?_TP?" + + +def log_line(tag: str) -> str: + """The one-line form the smoke greps; "" when the counters are unreadable.""" + counters = read() + if not counters: + return "" + order = [name for name in _LINE_ORDER if name in counters] + order += [name for name in counters if name not in _LINE_ORDER] + fields = " ".join(f"{name}={counters[name]}" for name in order) + return f"topk-exact-counters rank={tag} {fields}" + + +def maybe_log(force: bool = False) -> str: + """Emit the line if due (or forced). Time check first: when not due this + costs one clock read and never touches the device.""" + global _last_log_monotonic, _warned_failed + period = envs.MERLIN_TOPK_EXACT_COUNTERS_LOG + if period <= 0: + return "" + now = time.monotonic() + if not force and now - _last_log_monotonic < period: + return "" + _last_log_monotonic = now + try: + line = log_line(rank_tag()) + except Exception as exc: # never let an instrument break a step or shutdown + if not _warned_failed: + _warned_failed = True + logger.warning("top-k exact-restart counter read failed: %s", exc) + return "" + if line: + logger.info("%s", line) + return line diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 187c6925b8..a9d87e4d2b 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -53,6 +53,7 @@ from vllm.model_executor.offloader import ( set_offloader, ) from vllm.model_executor.warmup.jit_warmup import JitWarmupRegistry +from vllm.models.qwen4_exp.nvidia.ops import topk_counters from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.encoder_budget import ( MultiModalBudget, @@ -1853,6 +1854,15 @@ class GPUModelRunner(LoRAModelRunnerMixin): # A Python bool; adds no device sync to the step. Late and # rank-local -- see guard_telemetry.check_pending. guard_telemetry.check_pending() + # Post-step read of the top-k exact-restart counters, i.e. + # everything the previous steps accumulated. Off unless + # MERLIN_TOPK_EXACT_COUNTERS_LOG > 0, and then a clock read on + # every step but the one per cadence that is actually due. It sits + # here rather than at the tail of sample_tokens because that method + # returns early on non-last PP ranks, so a tail hook would silently + # never fire on them. Outside any capture: this is the engine + # thread between steps (read() re-checks). + topk_counters.maybe_log() # Update the request states. self.update_pp_decode_requests() self.finish_requests(scheduler_output) @@ -2516,9 +2526,6 @@ class GPUModelRunner(LoRAModelRunnerMixin): return None input_batch = self.execute_model_state.input_batch - if self._ple_offload_connector is not None: - self._ple_offload_connector.close() - self._ple_offload_connector = None hidden_states = self.execute_model_state.hidden_states finished_req_ids = self.execute_model_state.finished_req_ids ec_connector_output = self.execute_model_state.ec_connector_output @@ -2551,6 +2558,9 @@ class GPUModelRunner(LoRAModelRunnerMixin): main_stream=self.main_stream, copy_stream=self.output_copy_stream, ) + # The last window of a run is the one a cadence log always + # misses; emit it before the device state goes away. + topk_counters.maybe_log(force=True) self.postprocess_num_computed_tokens(input_batch) return async_output @@ -2567,6 +2577,12 @@ class GPUModelRunner(LoRAModelRunnerMixin): """Release GPU tensors (model weights, KV caches, workspace) so that memory is reclaimable when running in the same process.""" torch.accelerator.synchronize() + # The last window of a run is the one a cadence log always + # misses; emit it before the device state goes away. + topk_counters.maybe_log(force=True) + if self._ple_offload_connector is not None: + self._ple_offload_connector.close() + self._ple_offload_connector = None if self.aux_output_connector is not None: self.aux_output_connector.close() set_offloader(None)