[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.

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
2026-10-09 20:05:28 +07:00
co-authored by Hinata
parent ef55e095c3
commit 5df639bb4e
+30
View File
@@ -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