iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
vLLM-4x-RTX-3090/patches/0040-topk-Python-reader-for-the-exact-restart-counters.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

347 lines
14 KiB
Diff
Raw Permalink Blame History

From 8be2e3ab6246beafa10aa7ed2de9591c1b4cfc37 Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
Date: Wed, 23 Sep 2026 08:31:06 -0400
Subject: [PATCH 040/103] topk: Python reader for the exact-restart counters
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.
---
tests/kernels/test_top_k_per_row.py | 60 +++++++
vllm/envs.py | 12 ++
.../qwen4_exp/nvidia/ops/topk_counters.py | 167 ++++++++++++++++++
vllm/v1/worker/gpu/model_runner.py | 13 ++
4 files changed, 252 insertions(+)
create mode 100644 vllm/models/qwen4_exp/nvidia/ops/topk_counters.py
diff --git a/tests/kernels/test_top_k_per_row.py b/tests/kernels/test_top_k_per_row.py
index 70332b8..b0564e8 100644
--- a/tests/kernels/test_top_k_per_row.py
+++ b/tests/kernels/test_top_k_per_row.py
@@ -1,11 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import re
import numpy as np
import pytest
import torch
+import vllm.envs as envs
from vllm.platforms import current_platform
from vllm.utils.torch_utils import set_random_seed
@@ -2090,3 +2092,61 @@ def test_persistent_topk_exact_counters_rejects_unsafe_output() -> None:
with pytest.raises(Exception):
op(ok, 1, 0) # reset without sync is not a complete handoff
+
+
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="This test requires CUDA")
+@torch.inference_mode()
+def test_topk_exact_counters_reader_log_line(monkeypatch) -> None:
+ """The Python reader must produce the exact line the smoke greps.
+
+ Runs the op path directly -- no engine, no model -- so a failure here is
+ the reader's, not a serving regression. Asserting the format with a regex
+ is the point: the smoke keys on these field names, and a reader that
+ prints something else is as good as no reader at all.
+ """
+ from vllm.models.qwen4_exp.nvidia.ops import topk_counters
+
+ torch.set_default_device("cuda:0")
+ direct = _read_exact_counters()
+ if direct is None:
+ pytest.skip(
+ "QUALIFICATION FAILURE if seen on the branch build: no "
+ "persistent_topk_exact_counters op. A skip is not a pass."
+ )
+
+ monkeypatch.setenv("MERLIN_TOPK_EXACT_COUNTERS_LOG", "1")
+ monkeypatch.setenv("MERLIN_TOPK_COUNTER_SLOTS", str(_EXACT_COUNTER_SLOTS))
+ assert envs.MERLIN_TOPK_EXACT_COUNTERS_LOG == 1
+
+ counters = topk_counters.read()
+ assert list(counters) == [
+ "rows",
+ "overflow",
+ "fallbacks",
+ "faults",
+ "switch",
+ ], f"unexpected slot names: {list(counters)}"
+ assert counters == dict(zip(counters, direct, strict=True)), (
+ "the reader disagrees with a direct op call"
+ )
+
+ line = topk_counters.log_line("PP0_TP0")
+ assert re.fullmatch(
+ r"topk-exact-counters rank=PP0_TP0 rows=\d+ overflow=\d+ "
+ r"fallbacks=\d+ faults=\d+ switch=\d+",
+ line,
+ ), f"log line does not match the documented format: {line!r}"
+
+ # The cadence gate must also be able to emit -- a reader that never fires
+ # from maybe_log() is the same defect one level up.
+ topk_counters._last_log_monotonic = 0.0
+ expected = line.replace("PP0_TP0", topk_counters.rank_tag())
+ assert topk_counters.maybe_log() == expected
+ # ... and must then stay silent until the period elapses.
+ assert topk_counters.maybe_log() == ""
+
+ # Unknown slot counts must degrade to positional names, never mislabel.
+ assert topk_counters.slot_names(7) == tuple(f"slot{i}" for i in range(7))
+ if getattr(torch.ops._C, "persistent_topk_exact_counters_slots", None) is None:
+ monkeypatch.setenv("MERLIN_TOPK_COUNTER_SLOTS", "7")
+ assert topk_counters.slot_count() == 7
diff --git a/vllm/envs.py b/vllm/envs.py
index 35047dd..5923dc5 100755
--- a/vllm/envs.py
+++ b/vllm/envs.py
@@ -506,6 +506,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():
@@ -2459,6 +2460,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
"VLLM_E2_MAMBA_SLOT_OWNERSHIP": lambda: _e2_diagnostic_mode(
"VLLM_E2_MAMBA_SLOT_OWNERSHIP"
),
+ # 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")
+ ),
}
@@ -2683,6 +2691,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",
# ------------------------------------------------------------------
}
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 0000000..02c1720
--- /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<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
diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py
index 67ba37f..b5b8a70 100644
--- a/vllm/v1/worker/gpu/model_runner.py
+++ b/vllm/v1/worker/gpu/model_runner.py
@@ -55,6 +55,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,
@@ -1784,6 +1785,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)
@@ -2479,6 +2489,9 @@ 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
--
2.43.0
Reference in New Issue View Git Blame Copy Permalink