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