From f5729a96df8e2dd200091e457e12a4f3377791c1 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 19 Sep 2026 18:34:50 +0300 Subject: [PATCH] fix: stop REINS/q-steering/av-decouple forcing host syncs in the H3 block loop comfy_aimdo's dynamic-VRAM block prefetch pipelines weight cast/unpin against CUDA stream timing; a CPU sync (bool()/item() of a CUDA tensor) or a surprise allocation mid-loop can race it into reclaiming a weight buffer a still-queued kernel needs -- a silent segfault, no Python exception possible. All three H3 rating-driven dit hooks did this on every block, every step: - h3_repr_steering's injection guard and h3_q_steering's injection guard both did `if torch.isfinite(x) and x > 0:` (forces a host sync) before writing the steered rows, and h3_repr_steering cloned the whole block output first. Replaced with a GPU-only torch.where guard; dropped the now-unnecessary clone in h3_repr_steering (out is this hook's own private tensor already). - h3_repr_steering.capture() now .cpu()s its descriptor immediately instead of leaving it on GPU competing for VRAM headroom for the rest of the run. - h3_av_decouple's attention override re-checked `.any()` on its video/audio/ other masks every single attention call; now decided once per unique seq_len alongside the mask itself (typically once per generation). Unvalidated against a real GPU (none available here) -- worth confirming on the rental that REINS no longer crashes silently under dynamic VRAM offload. Co-Authored-By: Claude Sonnet 5 --- h3_repr_steering.py | 8 ++++++- samplers.py | 51 ++++++++++++++++++++++++++++++++------------- 2 files changed, 43 insertions(+), 16 deletions(-) diff --git a/h3_repr_steering.py b/h3_repr_steering.py index df21e9db..32088367 100644 --- a/h3_repr_steering.py +++ b/h3_repr_steering.py @@ -120,7 +120,13 @@ def capture(hidden_state, video_mask): size needs adaptive_avg_pool1d for.""" if video_mask is None or not bool(video_mask.any()): return None - return hidden_state[video_mask].detach().float().mean(dim=0) + # .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() # --- persistence ------------------------------------------------------------------------- diff --git a/samplers.py b/samplers.py index dcdc4987..c63f27d3 100644 --- a/samplers.py +++ b/samplers.py @@ -6932,10 +6932,18 @@ class FunPackLTXAVSceneChainSampler: if direction is not None and _strength > 0.0: rows = out[mask] row_norm = rows.detach().float().norm(dim=-1).mean() - if torch.isfinite(row_norm) and row_norm > 0: - out = out.clone() - out[mask] = rows + (direction.to(out.dtype).to(out.device) - * _strength * row_norm).to(rows.dtype) + # GPU-only guard (no .item()/bool() host sync): a host sync here + # forces the CPU thread to wait mid-block-loop, which throws off + # the async weight-prefetch pipeline's own stream-wait timing + # (comfy_aimdo/model_prefetch.py) enough to segfault on a + # VRAM-constrained dynamic-offload run -- see project_reward_model + # -rework memory. torch.where no-ops the update when row_norm is + # 0/NaN/inf instead of skipping it, at the cost of the multiply + # always running -- cheaper than a stall that can crash the run. + bump = (direction.to(out.dtype).to(out.device) + * _strength * row_norm.to(out.dtype)) + safe = torch.isfinite(row_norm) & (row_norm > 0) + out[mask] = torch.where(safe, rows + bump.to(rows.dtype), rows) return {"img": out} return _tag_dit_hook(_hook, inner) @@ -7062,10 +7070,20 @@ class FunPackLTXAVSceneChainSampler: seq_len = q.shape[-2] cached = _mask_cache.get(seq_len) if cached is None: - cached = (_rs.video_mask_from_mod_segments(seg, seq_len, q.device), - _rs.audio_mask_from_mod_segments(seg, seq_len, q.device)) + _v = _rs.video_mask_from_mod_segments(seg, seq_len, q.device) + _a = _rs.audio_mask_from_mod_segments(seg, seq_len, q.device) + # .any() decided ONCE per unique seq_len (one host sync, typically + # once for the whole generation) instead of on every attention call, + # every block, every step -- see h3_repr_steering's capture() fix and + # project_reward_model_rework memory for why a host sync in this loop + # risks a silent segfault under comfy_aimdo's dynamic-VRAM prefetch. + _other = ~(_v | _a) if (_v is not None and _a is not None) else None + cached = (_v, _a, _other, + bool(_v.any()) if _v is not None else False, + bool(_a.any()) if _a is not None else False, + bool(_other.any()) if _other is not None else False) _mask_cache[seq_len] = cached - vmask, amask = cached + vmask, amask, other, has_v, has_a, has_other = cached if vmask is None or amask is None or k.shape[-2] != seq_len: # Captured mod_segments but couldn't turn it into rows to bias -- the # penalty is silently NOT applying this generation. Loud once, not once @@ -7079,7 +7097,6 @@ class FunPackLTXAVSceneChainSampler: f"generation is running the ORDINARY joint attention. The " f"penalty is not being applied; this is not saturation.") return _next(func, q, k, v, heads, mask=mask, skip_reshape=skip_reshape, **kwargs) - other = ~(vmask | amask) # `strength` IS the raw softmax logit penalty, not a normalized 0-1 fraction -- # there's no local way to know H3's actual trained attention-logit scale (the # checkpoint isn't on this machine, see the method docstring), so a 0-1 dial @@ -7094,19 +7111,19 @@ class FunPackLTXAVSceneChainSampler: f"seq_len={seq_len}. (Printed once per generation, not per block.)") out = torch.empty(q.shape[0], seq_len, heads * q.shape[-1], dtype=q.dtype, device=q.device) - if amask.any(): + if has_a: bias = torch.zeros(seq_len, dtype=q.dtype, device=q.device) bias[vmask] = bias_val out[:, amask, :] = _next(func, q[:, :, amask, :], k, v, heads, mask=bias.view(1, 1, 1, seq_len), skip_reshape=skip_reshape, **kwargs) - if vmask.any(): + if has_v: bias = torch.zeros(seq_len, dtype=q.dtype, device=q.device) bias[amask] = bias_val out[:, vmask, :] = _next(func, q[:, :, vmask, :], k, v, heads, mask=bias.view(1, 1, 1, seq_len), skip_reshape=skip_reshape, **kwargs) - if other.any(): + if has_other: out[:, other, :] = _next(func, q[:, :, other, :], k, v, heads, mask=None, skip_reshape=skip_reshape, **kwargs) return out @@ -7359,10 +7376,14 @@ class FunPackLTXAVSceneChainSampler: **kwargs) rows = q[:, :, vmask, :] row_norm = rows.detach().float().norm(dim=-1).mean() - if torch.isfinite(row_norm) and row_norm > 0: - dir_per_head = direction.to(q.dtype).to(q.device).view(heads, 1, head_dim) - q = q.clone() - q[:, :, vmask, :] = rows + dir_per_head * _strength * row_norm + # GPU-only guard (no .item()/bool() host sync) -- same fix as + # h3_repr_steering's hook, see that one for why a host sync here can + # segfault under comfy_aimdo's dynamic-VRAM block prefetch. + dir_per_head = direction.to(q.dtype).to(q.device).view(heads, 1, head_dim) + bump = dir_per_head * _strength * row_norm.to(q.dtype) + safe = torch.isfinite(row_norm) & (row_norm > 0) + q = q.clone() + q[:, :, vmask, :] = torch.where(safe, rows + bump, rows) return _next(func, q, k, v, heads, mask=mask, skip_reshape=skip_reshape, **kwargs)