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 <noreply@anthropic.com>

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

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