From 8f98f46dacc084cdc93580fe467253d2fb2f54ca Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:55:41 +0300 Subject: [PATCH] fix: cache a real .any(), not an unenforced construction guarantee Round 3 of adversarial review, executed and confirmed live: the "mask is not None implies non-empty" reasoning behind has_rows=True was false. _mask_from_mod_segments' scalar branch sets found=True without checking the segment has nonzero length, so a zero-length scalar-tagged segment produces an all-False mask that is not None -- has_rows=True would then make capture() silently return an all-NaN descriptor instead of None, poisoning that block's learned direction() on disk with no error anywhere. No reachable path into a zero-length segment was found in this repo's own code (the real video segment is never zero-length; keyframe/ref segments come from real encoded latents here), so this doesn't fire today, but the invariant it depended on isn't enforced or asserted anywhere, so a future upstream ComfyUI change could trip it silently. Fixed the right way, which was available at the same cost av_decouple already pays (h3_repr_steering.capture()'s has_rows param exists for exactly this): cache a REAL bool(mask.any()) once per unique seq_len, alongside the mask itself, in both h3_repr_steering's hook and h3_q_steering's override -- same single-sync-per-generation cost class as av_decouple's has_v/has_a/has_other, zero per-block/per-step syncs, and correct regardless of what video_mask_from_mod_segments actually guarantees. Co-Authored-By: Claude Sonnet 5 --- samplers.py | 43 ++++++++++++++++++++++++++----------------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/samplers.py b/samplers.py index 16ec58b2..8aa5b14d 100644 --- a/samplers.py +++ b/samplers.py @@ -6915,22 +6915,27 @@ class FunPackLTXAVSceneChainSampler: out = extra["original_block"](args)["img"] if inner is None else \ inner(args, extra)["img"] seq_len = int(out.shape[0]) - mask = mask_cache.get(seq_len, "MISS") - if mask == "MISS": + cached = mask_cache.get(seq_len, "MISS") + if cached == "MISS": mask = _rs.video_mask_from_mod_segments( args.get("mod_segments"), seq_len, out.device) - mask_cache[seq_len] = mask + # Real .any(), not "mask is not None implies non-empty" -- + # _mask_from_mod_segments' scalar branch sets found=True without + # checking the segment has nonzero length, so that implication isn't + # actually enforced. Decided ONCE per unique seq_len (same cost class + # av_decouple already pays for its own has_v/has_a/has_other), not + # per block/step -- see project_reward_model_rework memory. + has_rows = bool(mask.any()) if mask is not None else False + cached = (mask, has_rows) + mask_cache[seq_len] = cached + mask, has_rows = cached if mask is not None: # Captured BEFORE injection, deliberately -- a steered block is the # only kind where capturing the post-injection state would make the # 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. - # 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) + desc = _rs.capture(out, mask, has_rows=has_rows) if desc is not None: capture_holder[0][block] = desc if direction is not None and _strength > 0.0: @@ -7341,23 +7346,27 @@ class FunPackLTXAVSceneChainSampler: return _next(func, q, k, v, heads, mask=mask, skip_reshape=skip_reshape, **kwargs) seq_len = q.shape[-2] - vmask = _mask_cache.get(seq_len, "MISS") - if vmask == "MISS": + cached = _mask_cache.get(seq_len, "MISS") + if cached == "MISS": vmask = _rs.video_mask_from_mod_segments( _seg_holder["mod_segments"], seq_len, q.device) - _mask_cache[seq_len] = vmask + # Real .any(), decided ONCE per unique seq_len, not per attention call -- + # "vmask is not None implies non-empty" is NOT actually enforced by + # video_mask_from_mod_segments' scalar branch (a zero-length segment can + # set found=True with nothing marked), so this can't be skipped. Same + # cost class av_decouple already pays for has_v/has_a/has_other -- see + # h3_repr_steering's mask_cache and project_reward_model_rework memory. + has_rows = bool(vmask.any()) if vmask is not None else False + cached = (vmask, has_rows) + _mask_cache[seq_len] = cached + vmask, has_rows = cached 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, has_rows=True) + desc = _rs.capture(q_flat, vmask, has_rows=has_rows) 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