325 lines
14 KiB
Diff
325 lines
14 KiB
Diff
From 22b00ba2e2239b52923cb555ae6e24198bb59835 Mon Sep 17 00:00:00 2001
|
|
From: halt95 <halt95@users.noreply.github.com>
|
|
Date: Thu, 1 Oct 2026 05:08:33 -0400
|
|
Subject: [PATCH 069/103] qsa: start the clip-counter reader after capture, not
|
|
on the first forward
|
|
|
|
The FP8 KV clip counter's reader thread used to start on the first eager forward, before CUDA graph capture, so
|
|
its device-to-host copies could overlap a capture. _qsa_clip_count now only allocates the per-layer counter; the
|
|
reader is started by a hook that guard_telemetry runs after warmup and capture (or at once if capture is already
|
|
done). Counting is unchanged. Tests cover the order, idempotence and a failing control that restores the old start.
|
|
---
|
|
.../test_qsa_clip_reader_lifecycle.py | 170 ++++++++++++++++++
|
|
vllm/models/qwen4_exp/nvidia/qsa.py | 48 +++--
|
|
vllm/v1/worker/guard_telemetry.py | 30 ++++
|
|
3 files changed, 234 insertions(+), 14 deletions(-)
|
|
create mode 100644 tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py
|
|
|
|
diff --git a/tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py b/tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py
|
|
new file mode 100644
|
|
index 0000000..a81feb7
|
|
--- /dev/null
|
|
+++ b/tests/models/qwen4_exp/test_qsa_clip_reader_lifecycle.py
|
|
@@ -0,0 +1,170 @@
|
|
+# SPDX-License-Identifier: Apache-2.0
|
|
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
+"""The QSA clip-counter READER starts after warmup + cudagraph capture, not on the first forward (since v2.5).
|
|
+
|
|
+Before, `_qsa_clip_count` started the `qsa-clip-reader` thread on the first eager forward -- before capture -- so its
|
|
+device-to-host copies could run on a second thread while a capture was active (the call class
|
|
+cudaStreamCaptureModeGlobal rejects). That is why `COUNTERS` shipped off. Now the starter is registered with
|
|
+guard_telemetry.register_post_capture_hook and run by guard_telemetry.start_reader() (gpu_worker, after capture).
|
|
+
|
|
+qsa.py cannot be imported on a machine without zmq/CUDA, so the clip-counter block is cut out of the source
|
|
+text, written to a temporary module (with a small prelude) and imported next to the real guard_telemetry; the reader
|
|
+loop itself is replaced by a stub that records "I ran". The failing control re-inserts the OLD behaviour into the same
|
|
+block and must flip.
|
|
+"""
|
|
+
|
|
+import ast
|
|
+import importlib.util
|
|
+import itertools
|
|
+import os
|
|
+import threading
|
|
+import time
|
|
+
|
|
+import pytest
|
|
+import torch
|
|
+
|
|
+_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), *([".."] * 3)))
|
|
+_QSA = os.path.join(_ROOT, "vllm", "models", "qwen4_exp", "nvidia", "qsa.py")
|
|
+_WORKER = os.path.join(_ROOT, "vllm", "v1", "worker", "gpu_worker.py")
|
|
+_GT_PATH = os.path.join(_ROOT, "vllm", "v1", "worker", "guard_telemetry.py")
|
|
+_PRELUDE = (
|
|
+ "import logging, math, os, threading, time\n"
|
|
+ "import torch\n"
|
|
+ "logger = logging.getLogger('qsa-clip-test')\n"
|
|
+ "Qwen4ExpQSAAttention = object\n"
|
|
+)
|
|
+_seq = itertools.count()
|
|
+
|
|
+
|
|
+def _src(path):
|
|
+ with open(path, encoding="utf-8") as fh:
|
|
+ return fh.read()
|
|
+
|
|
+
|
|
+def _block(src):
|
|
+ a = src.index('_QSA_CLIP_ENV = "VLLM_QSA_KV_CLIP_COUNT"')
|
|
+ b = src.index("_QSA_FP8_CACHE_DTYPES = (")
|
|
+ return src[a:b]
|
|
+
|
|
+
|
|
+@pytest.fixture
|
|
+def gt(monkeypatch):
|
|
+ pytest.importorskip("vllm.v1.worker.guard_telemetry", reason="guard_telemetry needs torch + a vllm import")
|
|
+ from vllm.v1.worker import guard_telemetry as g
|
|
+
|
|
+ monkeypatch.setattr(g, "_post_capture_hooks", [])
|
|
+ monkeypatch.setattr(g, "_post_capture_done", False)
|
|
+ monkeypatch.setattr(g, "_MODE", "off") # no guard reader thread in this test
|
|
+ monkeypatch.setattr(g, "_reader_started", False)
|
|
+ return g
|
|
+
|
|
+
|
|
+class _Layer:
|
|
+ layer_name = "layers.0.attn"
|
|
+ _k_scale = torch.tensor(1.0 / 448.0) # ceiling = 1.0
|
|
+ _v_scale = torch.tensor(1.0 / 448.0)
|
|
+
|
|
+
|
|
+def _load_block(monkeypatch, tmp_path, env, transform=None):
|
|
+ if env is None:
|
|
+ monkeypatch.delenv("VLLM_QSA_KV_CLIP_COUNT", raising=False)
|
|
+ else:
|
|
+ monkeypatch.setenv("VLLM_QSA_KV_CLIP_COUNT", env)
|
|
+ text = _block(_src(_QSA))
|
|
+ if transform:
|
|
+ text = transform(text)
|
|
+ path = tmp_path / ("qsa_clip_block_%d.py" % next(_seq))
|
|
+ path.write_text(_PRELUDE + text, encoding="utf-8")
|
|
+ spec = importlib.util.spec_from_file_location(path.stem, str(path))
|
|
+ mod = importlib.util.module_from_spec(spec)
|
|
+ spec.loader.exec_module(mod)
|
|
+ ran = threading.Event()
|
|
+ mod._qsa_clip_reader = ran.set # the loop itself is not under test
|
|
+ return mod, ran
|
|
+
|
|
+
|
|
+def _forward(mod):
|
|
+ key = torch.tensor([[[2.0, 0.5]], [[0.1, 3.0]], [[5.0, 5.0]]]) # 3 tokens
|
|
+ value = torch.zeros_like(key)
|
|
+ slot = torch.tensor([0, 1, -1]) # the last row is padding: never counted
|
|
+ mod._qsa_clip_count(_Layer(), key, value, slot)
|
|
+
|
|
+
|
|
+def test_first_forward_does_not_start_the_reader(monkeypatch, tmp_path, gt):
|
|
+ mod, ran = _load_block(monkeypatch, tmp_path, "5")
|
|
+ _forward(mod)
|
|
+ time.sleep(0.05)
|
|
+ assert not ran.is_set()
|
|
+ assert mod._qsa_clip_reader_started is False
|
|
+ assert "_qsa_clip_start_reader" in [f.__name__ for f in gt._post_capture_hooks]
|
|
+ assert set(mod._qsa_clip_counters) == {"layers.0.attn"}
|
|
+
|
|
+
|
|
+def test_reader_starts_once_after_capture(monkeypatch, tmp_path, gt):
|
|
+ mod, ran = _load_block(monkeypatch, tmp_path, "5")
|
|
+ _forward(mod)
|
|
+ gt.start_reader() # what gpu_worker calls after capture
|
|
+ assert ran.wait(2.0)
|
|
+ assert mod._qsa_clip_reader_started is True
|
|
+ gt.start_reader() # idempotent
|
|
+ assert len([x for x in threading.enumerate() if x.name == "qsa-clip-reader"]) <= 1
|
|
+
|
|
+
|
|
+def test_counting_is_unchanged(monkeypatch, tmp_path, gt):
|
|
+ mod, _ = _load_block(monkeypatch, tmp_path, "5")
|
|
+ _forward(mod)
|
|
+ # ceiling 1.0: row0 has one element over (2.0), row1 has one (3.0), row2 is padding (slot -1), not counted
|
|
+ assert mod._qsa_clip_counters["layers.0.attn"].tolist() == [2, 0]
|
|
+
|
|
+
|
|
+def test_no_hook_when_the_env_is_unset(monkeypatch, tmp_path, gt):
|
|
+ mod, _ = _load_block(monkeypatch, tmp_path, None)
|
|
+ assert mod._qsa_clip_interval == 0.0
|
|
+ assert gt._post_capture_hooks == []
|
|
+
|
|
+
|
|
+def test_late_registration_starts_at_once(monkeypatch, tmp_path, gt):
|
|
+ gt.run_post_capture_hooks() # capture already finished
|
|
+ mod, _ = _load_block(monkeypatch, tmp_path, "5") # qsa imported late: the starter runs during the import
|
|
+ assert mod._qsa_clip_reader_started is True # (the thread was the real reader, asleep for its 5 s interval)
|
|
+
|
|
+
|
|
+def test_a_failing_starter_does_not_stop_the_others(gt):
|
|
+ seen = []
|
|
+
|
|
+ def boom():
|
|
+ raise RuntimeError("boom")
|
|
+
|
|
+ gt.register_post_capture_hook(boom)
|
|
+ gt.register_post_capture_hook(lambda: seen.append(1))
|
|
+ gt.start_reader()
|
|
+ assert seen == [1]
|
|
+
|
|
+
|
|
+def test_failing_control_the_old_path_started_the_reader_before_capture(monkeypatch, tmp_path, gt):
|
|
+ """re-insert the earlier behaviour (reader started inside the first forward): the observation must flip"""
|
|
+
|
|
+ def old_path(text):
|
|
+ marker = " _qsa_clip_counters[layer.layer_name] = ctr"
|
|
+ assert marker in text
|
|
+ return text.replace(marker, marker + "\n _qsa_clip_start_reader()", 1)
|
|
+
|
|
+ mod, ran = _load_block(monkeypatch, tmp_path, "5", transform=old_path)
|
|
+ _forward(mod) # first (pre-capture) forward, no start_reader() call
|
|
+ assert ran.wait(2.0), "control did not flip: the old path should start the reader on the first forward"
|
|
+ assert mod._qsa_clip_reader_started is True
|
|
+
|
|
+
|
|
+def test_source_pins():
|
|
+ qsa = _src(_QSA)
|
|
+ tree = ast.parse(qsa)
|
|
+ fn = next(n for n in ast.walk(tree) if isinstance(n, ast.FunctionDef) and n.name == "_qsa_clip_count")
|
|
+ assert "threading.Thread" not in ast.unparse(fn), "the clip reader must not be started from the forward"
|
|
+ assert "register_post_capture_hook(_qsa_clip_start_reader)" in qsa
|
|
+ gtree = ast.parse(_src(_GT_PATH))
|
|
+ sr = next(n for n in ast.walk(gtree) if isinstance(n, ast.FunctionDef) and n.name == "start_reader")
|
|
+ stmts = [s for s in sr.body if not (isinstance(s, ast.Expr) and isinstance(s.value, ast.Constant))]
|
|
+ assert isinstance(stmts[0], ast.Global)
|
|
+ assert ast.unparse(stmts[1]) == "run_post_capture_hooks()", "hooks must run before the guard-mode early return"
|
|
+ w = _src(_WORKER) # the worker calls start_reader after capture_model
|
|
+ assert w.index("self.model_runner.capture_model(") < w.index("guard_telemetry.start_reader()")
|
|
diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py
|
|
index 73dd07a..467f7ad 100644
|
|
--- a/vllm/models/qwen4_exp/nvidia/qsa.py
|
|
+++ b/vllm/models/qwen4_exp/nvidia/qsa.py
|
|
@@ -150,6 +150,13 @@ def _qsa_collect_absmax(
|
|
# graph replay keeps counting. A daemon reader thread .item()s the counters on
|
|
# a cadence and logs: WARNING on any new clips, INFO heartbeat every 10th pass
|
|
# so a quiet log still proves the instrument is alive.
|
|
+#
|
|
+# Since v2.5 the reader thread is started AFTER warmup and cudagraph capture
|
|
+# (guard_telemetry.register_post_capture_hook, run by start_reader() from
|
|
+# gpu_worker), never from the first eager forward. Before, it started inside the
|
|
+# first forward, i.e. before capture, and its device-to-host copies could run on
|
|
+# a second thread while a capture was active -- the call class
|
|
+# cudaStreamCaptureModeGlobal rejects. That is why COUNTERS shipped off.
|
|
_QSA_CLIP_ENV = "VLLM_QSA_KV_CLIP_COUNT"
|
|
_qsa_clip_raw = os.environ.get(_QSA_CLIP_ENV, "").strip()
|
|
try:
|
|
@@ -222,31 +229,44 @@ def _qsa_clip_reader() -> None:
|
|
logger.exception("QSA clip reader pass failed; thread continues")
|
|
|
|
|
|
+def _qsa_clip_start_reader() -> None:
|
|
+ """Start the reader thread. Runs after warmup + capture (post-capture hook); idempotent."""
|
|
+ global _qsa_clip_reader_started
|
|
+ with _qsa_clip_lock:
|
|
+ if _qsa_clip_reader_started:
|
|
+ return
|
|
+ _qsa_clip_reader_started = True
|
|
+ threading.Thread(
|
|
+ target=_qsa_clip_reader,
|
|
+ name="qsa-clip-reader",
|
|
+ daemon=True,
|
|
+ ).start()
|
|
+ logger.info(
|
|
+ "QSA clip counter reader started after capture (read interval %.0fs, %d layer(s) registered)",
|
|
+ _qsa_clip_interval,
|
|
+ len(_qsa_clip_counters),
|
|
+ )
|
|
+
|
|
+
|
|
+if _qsa_clip_interval > 0.0:
|
|
+ from vllm.v1.worker.guard_telemetry import register_post_capture_hook # noqa: E402
|
|
+
|
|
+ register_post_capture_hook(_qsa_clip_start_reader)
|
|
+
|
|
+
|
|
def _qsa_clip_count(
|
|
layer: Qwen4ExpQSAAttention,
|
|
key: torch.Tensor,
|
|
value: torch.Tensor,
|
|
slot_mapping: torch.Tensor | None = None,
|
|
) -> None:
|
|
- global _qsa_clip_reader_started
|
|
ctr = _qsa_clip_counters.get(layer.layer_name)
|
|
if ctr is None:
|
|
# First forward is eager (profiling/warmup precede capture), so the
|
|
- # allocation never happens inside cudagraph capture.
|
|
+ # allocation never happens inside cudagraph capture. The READER is not
|
|
+ # started here: _qsa_clip_start_reader runs after capture.
|
|
ctr = torch.zeros(2, dtype=torch.int64, device=key.device)
|
|
_qsa_clip_counters[layer.layer_name] = ctr # registered with the reader by the reader thread (never during capture)
|
|
- with _qsa_clip_lock:
|
|
- if not _qsa_clip_reader_started:
|
|
- _qsa_clip_reader_started = True
|
|
- threading.Thread(
|
|
- target=_qsa_clip_reader,
|
|
- name="qsa-clip-reader",
|
|
- daemon=True,
|
|
- ).start()
|
|
- logger.info(
|
|
- "QSA clip counter armed (read interval %.0fs)",
|
|
- _qsa_clip_interval,
|
|
- )
|
|
# Thresholds are 1-D so the comparison promotes to float32; a 0-dim float32
|
|
# tensor promotes like a Python scalar and the threshold would be rounded to
|
|
# the BF16 grid of K/V, undercounting values just above the calibrated ceiling.
|
|
diff --git a/vllm/v1/worker/guard_telemetry.py b/vllm/v1/worker/guard_telemetry.py
|
|
index b02a7c3..19cf966 100644
|
|
--- a/vllm/v1/worker/guard_telemetry.py
|
|
+++ b/vllm/v1/worker/guard_telemetry.py
|
|
@@ -486,9 +486,38 @@ def prealloc(device: torch.device) -> None:
|
|
_nb_reader.register(device, _buffers[device])
|
|
|
|
|
|
+# Readers that other modules own but that must obey the same rule: nothing reads a device tensor from a second
|
|
+# thread while a cudagraph capture is active. A module registers its starter here at import; start_reader() runs
|
|
+# them after capture, and one registered after that point (a late import) starts at once. Before v2.5 the QSA
|
|
+# clip-counter reader started on the first eager forward, i.e. BEFORE capture.
|
|
+_post_capture_hooks: list = []
|
|
+_post_capture_done = False
|
|
+
|
|
+
|
|
+def register_post_capture_hook(fn) -> None:
|
|
+ """Run fn() once after warmup and cudagraph capture (immediately if that already happened)."""
|
|
+ if _post_capture_done:
|
|
+ fn()
|
|
+ else:
|
|
+ _post_capture_hooks.append(fn)
|
|
+
|
|
+
|
|
+def run_post_capture_hooks() -> None:
|
|
+ global _post_capture_done
|
|
+ _post_capture_done = True
|
|
+ hooks, _post_capture_hooks[:] = list(_post_capture_hooks), []
|
|
+ for fn in hooks:
|
|
+ try:
|
|
+ fn()
|
|
+ except Exception: # noqa: BLE001 -- one broken starter must not stop the others or the engine
|
|
+ logger.exception("post-capture reader start failed: %r", fn)
|
|
+
|
|
+
|
|
def start_reader() -> None:
|
|
"""Start the cadence reader. Called AFTER warmup and cudagraph capture.
|
|
|
|
+ Also runs the post-capture hooks of the other readers (register_post_capture_hook), whatever the guard mode.
|
|
+
|
|
The reader must not exist during capture. It copies device counters to the
|
|
host, and a device-to-host copy from a second thread while a capture is
|
|
active is exactly the call class cudaStreamCaptureModeGlobal -- the default
|
|
@@ -500,6 +529,7 @@ def start_reader() -> None:
|
|
Between prealloc() and this call the kernels still count; nothing reads.
|
|
"""
|
|
global _reader_started
|
|
+ run_post_capture_hooks()
|
|
if _reader_started or _MODE == "off":
|
|
return
|
|
_reader_started = True
|
|
--
|
|
2.43.0
|
|
|