iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[TopK] Python reader for the exact-restart counters

Browse Source 
Reads persistent_topk_exact_counters into {slot: value}, logs one line on a
cadence, and degrades to a one-time warning when the op is absent.

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
MOVZX
2026-10-09 19:55:59 +07:00
co-authored by Hinata
parent 7f4385a422
commit 66c3d11016
3 changed files with 198 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/envs.py
+12
View File
@@ -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",
}
vllm/models/qwen4_exp/nvidia/ops/topk_counters.py
+167
View File
@@ -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<i>` names rather than mislabelling.
HOW THE SMOKE READS IT. Set `MERLIN_TOPK_EXACT_COUNTERS_LOG=<seconds>` 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 <engine log> | ...` 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
vllm/v1/worker/gpu/model_runner.py
+19 -3
View File
@@ -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)