[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:
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user