From d50c33aff024ad9ff0371656691771f88b39e2c7 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:50:37 +0300 Subject: [PATCH] fix: actually eliminate capture()'s per-block host sync, not just re-add .clone() Round 2 of adversarial review: my earlier "fix" still called bool(mask.any()) inside h3_repr_steering.capture() and h3_q_steering's _override on every one of the 50 blocks, every step -- despite the mask itself being cached per seq_len, the .any() reduction was recomputed and host-synced every call. The av_decouple mechanism already had the right pattern (cache .any() once per seq_len) and its own comment falsely claimed capture() had been fixed the same way; it hadn't. Both real callers (h3_repr_steering's hook, h3_q_steering's override) build their mask via video_mask_from_mod_segments, which by construction only ever returns non-None when it found at least one True row -- so both can pass has_rows=True and skip capture()'s internal check entirely, with zero syncs in the hot path. capture()'s own default (has_rows=None) still self-checks for any other/future/test caller that can't make that guarantee. Also dropped the .cpu() added in the first commit: a 50-descriptor-per-step [hidden_size] vector is a few hundred KB, not a real VRAM pressure source, and forcing that transfer synchronously was itself a sync in the exact loop this whole fix exists to keep sync-free. Reverted to leaving it on GPU (original behavior) -- torch.save/load already handles device-tagged tensors fine downstream. q_steering's unconditional q.clone() was flagged this round as likely pure waste under ComfyUI's single-owner attention-tensor-container contract, but left in place: verifying that contract myself with confidence is out of scope here, and the cost of being wrong (reintroducing the aliasing bug just fixed for h3_repr_steering's out) outweighs the small clone it would save. Co-Authored-By: Claude Sonnet 5 --- h3_repr_steering.py | 27 +++++++++++++++++---------- samplers.py | 15 ++++++++++++--- 2 files changed, 29 insertions(+), 13 deletions(-) diff --git a/h3_repr_steering.py b/h3_repr_steering.py index 32088367..cb8dd798 100644 --- a/h3_repr_steering.py +++ b/h3_repr_steering.py @@ -112,21 +112,28 @@ def audio_mask_from_mod_segments(mod_segments, seq_len, device): return _mask_from_mod_segments(mod_segments, seq_len, device, 2) -def capture(hidden_state, video_mask): +def capture(hidden_state, video_mask, has_rows=None): """[seq_len, hidden] + a video-row mask -> a [hidden] descriptor, or None. Mean over rows, nothing more -- hidden_size is architecturally fixed (H3 is one width throughout), so there is no varying dimension to pool away the way a latent's spatial - size needs adaptive_avg_pool1d for.""" - if video_mask is None or not bool(video_mask.any()): + size needs adaptive_avg_pool1d for. + + `has_rows`: pass True/False when the caller already knows (from a mask it cached and + already ran .any() on ONCE per seq_len, e.g. h3_repr_steering's/h3_q_steering's own + hooks) whether `video_mask` has at least one True row -- this skips re-deriving it here + with a SECOND host sync (bool() of a CUDA tensor) on every one of the 50 candidate + blocks, every step, which is exactly the kind of stall this hot path needs to avoid (see + project_reward_model_rework memory: implicated in a silent segfault under comfy_aimdo's + dynamic-VRAM block prefetch). Leave it None (the default, used by direct/test callers + with no cache of their own) to get the fully self-contained, host-synced check.""" + if video_mask is None: return None - # .cpu() immediately -- these accumulate across all 50 CANDIDATE_BLOCKS, every - # denoising step, and otherwise sit on the GPU for the rest of the run competing - # for VRAM headroom with the block-weight prefetch pipeline (comfy_aimdo). See - # project_reward_model_rework memory: that pressure was implicated in a silent - # segfault under dynamic VRAM offload. The descriptor only ever gets saved to - # disk (torch.save) or diffed against another CPU tensor -- never used on-GPU. - return hidden_state[video_mask].detach().float().mean(dim=0).cpu() + if has_rows is None: + has_rows = bool(video_mask.any()) + if not has_rows: + return None + return hidden_state[video_mask].detach().float().mean(dim=0) # --- persistence ------------------------------------------------------------------------- diff --git a/samplers.py b/samplers.py index 2ac54c25..16ec58b2 100644 --- a/samplers.py +++ b/samplers.py @@ -6926,7 +6926,11 @@ class FunPackLTXAVSceneChainSampler: # logged descriptor depend on that run's own strength setting instead # of reflecting the network's natural behavior. A block that never # injects (direction is None here) is unaffected by this ordering. - desc = _rs.capture(out, mask) + # has_rows=True: video_mask_from_mod_segments (above) only ever + # returns non-None when it found at least one True row -- passing + # that along skips capture()'s own .any() re-check, which would + # otherwise be a second host sync on every block, every step. + desc = _rs.capture(out, mask, has_rows=True) if desc is not None: capture_holder[0][block] = desc if direction is not None and _strength > 0.0: @@ -7342,13 +7346,18 @@ class FunPackLTXAVSceneChainSampler: vmask = _rs.video_mask_from_mod_segments( _seg_holder["mod_segments"], seq_len, q.device) _mask_cache[seq_len] = vmask - if vmask is None or not bool(vmask.any()): + if vmask is None: + # No .any() re-check: video_mask_from_mod_segments only ever returns + # non-None when it found at least one True row -- re-verifying with + # .any() would force a needless host sync on every steered block, every + # attention call, every step (see h3_repr_steering.capture()'s has_rows + # param and project_reward_model_rework memory). return _next(func, q, k, v, heads, mask=mask, skip_reshape=skip_reshape, **kwargs) # (B=1, H, S, D) -> (S, H*D) for the masked rows, so h3_repr_steering.capture() # sees the same shape it already expects from a hidden state. q_flat = q[0].permute(1, 0, 2).reshape(seq_len, -1) - desc = _rs.capture(q_flat, vmask) + desc = _rs.capture(q_flat, vmask, has_rows=True) if desc is not None: # Captured BEFORE injection, same reasoning as h3_repr_steering: a steered # block's descriptor must reflect the network's natural Q, not this run's