diff --git a/vllm/v1/worker/guard_telemetry.py b/vllm/v1/worker/guard_telemetry.py index b02a7c3d43..19cf966ac4 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