fix: H3 attention decoupling was scaling K magnitude, which barely moves softmax

Live 0-1 strength sweep came back flat -- confirmed the mechanism, not just the
value, was wrong. H3's Q/K get RMSNorm'd before attention, so scaling a key
toward zero pushes its logit toward the *mean* of an already near-zero-centered
distribution, not toward excluded. Switched to what actually suppresses a
softmax entry: an additive negative bias on the targeted logits, still passed
through the same per-query-group split (still no S^2 mask, still same total
FLOPs -- the bias is a per-key vector broadcast over the group's queries).

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
digital-garbage
2026-09-05 02:23:30 +03:00
co-authored by Claude Sonnet 5
parent 9bbdfc8b4b
commit 1e4553e658
3 changed files with 62 additions and 33 deletions
+1 -1
View File
@@ -336,7 +336,7 @@
detail: "Blocks 40-49 carry the most residual movement (48+49 alone about a third of it), so a steering push has the most to act on there. Blocks 0-1 are structurally distinct rather than just 'quiet early blocks' — 0 is the raw input projection, 1 does the first cross-patch mixing — expect steering there to behave differently in kind, not just 'weaker', than the middle of the stack." },
{ name: "h3_av_decouple", label: "Video/audio attention decoupling (H3, experimental)", kind: "float", default: 0.0, min: 0.0, max: 1.0, step: 0.01,
hint: "EXPERIMENTAL, unvalidated, H3 only. Weakens how much video and audio attend to each other, for the failure pattern where one comes out clean and the other doesn't. 0 = off. Not rating-driven — a manual dial, same as Repeat block(s).",
detail: "H3 has no separate cross-attention to mask — video, text and audio share one joint self-attention. Instead of building a mask over that whole sequence (unaffordable at H3's token count), this runs the same attention as separate video-query and audio-query passes against a version of the keys where the other modality is scaled down by this amount. Same total cost as the ordinary single pass. 1.0 cuts cross-modal attention entirely, which likely also kills legitimate audio/video sync — search down from a small value, not up to it." },
detail: "H3 has no separate cross-attention to mask — video, text and audio share one joint self-attention. Instead of building a mask over that whole sequence (unaffordable at H3's token count), this runs the same attention as separate video-query and audio-query passes, each carrying a penalty on the other modality's attention weight. Same total cost as the ordinary single pass. 1.0 cuts cross-modal attention entirely, which likely also kills legitimate audio/video sync — search down from a small value, not up to it. (v2: the first version tried this by shrinking the other modality's keys instead of penalizing attention directly, and live testing showed almost no effect across the whole 0-1 range — fixed to actually penalize.)" },
{ name: "identity_transfer_enabled", label: "Best-FaceID compatibility", kind: "bool", default: false,
hint: "Feeds the identity pin image the way Best-FaceID identity LoRAs expect it. Needs an Identity pin set.",
detail: "Replaces Continuity's identity-pin guide with separate, non-rendered reference tokens plus an optional ArcFace projector below. Load the LoRA itself in Models. No effect without an identity pin." },
+33 -14
View File
@@ -2808,7 +2808,7 @@ class FunPackLTXAVSceneChainSampler:
}),
"h3_av_decouple": ("FLOAT", {
"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01,
"tooltip": "EXPERIMENTAL, unvalidated (H3 only). Damps how much video and audio attend to EACH OTHER inside H3's joint self-attention, without touching any hidden state -- for the failure pattern where one comes out clean and the other doesn't. 0 = off (installs nothing, zero cost). Higher = weaker cross-modal influence; 1.0 fully cuts it, which likely also kills legitimate AV sync, so this is a knob to search down from, not up to.",
"tooltip": "EXPERIMENTAL, unvalidated (H3 only). Damps how much video and audio attend to EACH OTHER inside H3's joint self-attention, without touching any hidden state -- for the failure pattern where one comes out clean and the other doesn't. 0 = off (installs nothing, zero cost). Higher = weaker cross-modal influence via a negative bias on the other modality's attention logits (v2: the first version scaled key magnitude instead and measured flat 0-1 live -- fixed). 1.0 fully cuts it, which likely also kills legitimate AV sync, so this is a knob to search down from, not up to.",
}),
"alg_guide_blur_strength": ("FLOAT", {
"default": 2.0, "min": 1.0, "max": 4.0, "step": 0.1,
@@ -6831,12 +6831,20 @@ class FunPackLTXAVSceneChainSampler:
sequence). This does not build that S×S mask at all. It uses comfy's own
`optimized_attention_override` extension point (comfy/ldm/modules/attention.py's
`wrap_attn` -- already wired into every backend, nothing to monkeypatch) to run the
SAME joint attention as separate video-query and audio-query passes against a K
where the OTHER modality's rows are scaled down by `strength`; text/ref/cond rows
SAME joint attention as separate video-query and audio-query passes, each carrying
an additive negative bias on the OTHER modality's logits (a per-key vector, not a
[seq,seq] matrix -- broadcasts over every query in the group); text/ref/cond rows
stay on one ordinary pass, untouched. Splitting the query dimension like this costs
nothing extra -- same total Q*K*V work as the single joint call, just partitioned --
and the only new memory is two full-length clones of K (audio and video zeroed out
by contrast, at most `strength`), not a dense attention matrix.
nothing extra -- same total Q*K*V work as the single joint call, just partitioned.
First cut of this scaled K's magnitude toward zero instead of biasing the logit --
measured flat live across the whole 0-1 strength range. After H3's own per-head
RMSNorm on Q/K, attention logits already sit close to 0, so shrinking a key's
magnitude pushes its logit toward "average", not toward "excluded" -- softmax barely
notices. An additive penalty actually pushes the targeted logit negative, which is
what exclusion requires. Still unvalidated whether the *hypothesis* (cross-modal
attention causes the clean-video/corrupted-audio split) is even correct -- this only
fixes the mechanism not doing what the strength dial implied.
`strength` in [0, 1]; 0 or blank installs nothing (exact pass-through, not just a
no-op override) so an unset knob costs literally zero.
@@ -6876,6 +6884,14 @@ class FunPackLTXAVSceneChainSampler:
to["patches_replace"] = patches_replace
_mask_cache = {} # seq_len -> (video_mask, audio_mask); one live resolution per gen
# Logit bias scale, not a key-magnitude scale (see the method docstring for why
# the first cut of this used K-scaling and measured flat across the whole 0-1
# range live). After H3's own per-head RMSNorm on Q/K, attention logits already
# sit close to 0 -- an additive penalty pushes a targeted key toward genuinely
# excluded, where a multiplicative one only pushes it toward "average".
# ponytail: fixed heuristic scale, unvalidated against H3's real logit spread --
# revisit if the live strength sweep is still flat, or saturates immediately.
_BIAS_SCALE = 12.0
def _override(func, q, k, v, heads, mask=None, skip_reshape=True,
transformer_options=None, **kwargs):
@@ -6892,18 +6908,21 @@ class FunPackLTXAVSceneChainSampler:
if vmask is None or amask is None or k.shape[-2] != seq_len:
return func(q, k, v, heads, mask=mask, skip_reshape=skip_reshape, **kwargs)
other = ~(vmask | amask)
bias_val = -_strength * _BIAS_SCALE
out = torch.empty(q.shape[0], seq_len, heads * q.shape[-1],
dtype=q.dtype, device=q.device)
if amask.any():
k_damped = k.clone()
k_damped[:, :, vmask, :] *= (1.0 - _strength)
out[:, amask, :] = func(q[:, :, amask, :], k_damped, v, heads,
mask=None, skip_reshape=skip_reshape, **kwargs)
bias = torch.zeros(seq_len, dtype=q.dtype, device=q.device)
bias[vmask] = bias_val
out[:, amask, :] = func(q[:, :, amask, :], k, v, heads,
mask=bias.view(1, 1, 1, seq_len),
skip_reshape=skip_reshape, **kwargs)
if vmask.any():
k_damped = k.clone()
k_damped[:, :, amask, :] *= (1.0 - _strength)
out[:, vmask, :] = func(q[:, :, vmask, :], k_damped, v, heads,
mask=None, skip_reshape=skip_reshape, **kwargs)
bias = torch.zeros(seq_len, dtype=q.dtype, device=q.device)
bias[amask] = bias_val
out[:, vmask, :] = func(q[:, :, vmask, :], k, v, heads,
mask=bias.view(1, 1, 1, seq_len),
skip_reshape=skip_reshape, **kwargs)
if other.any():
out[:, other, :] = func(q[:, :, other, :], k, v, heads,
mask=None, skip_reshape=skip_reshape, **kwargs)
+28 -18
View File
@@ -51,8 +51,11 @@ def _qkv(seq_len=6, heads=1, head_dim=2):
def _fake_func_factory(calls):
def fake_func(q, k, v, heads, mask=None, skip_reshape=True, **kw):
calls.append({"q_len": q.shape[-2], "k_sum": float(k.sum()), "mask": mask})
return torch.full((q.shape[0], q.shape[-2], heads * q.shape[-1]), float(k.sum()))
# value encodes the bias so assertions can tell which group produced which output
# without needing a real attention computation.
marker = float(mask.sum()) if mask is not None else 0.0
calls.append({"q_len": q.shape[-2], "mask": mask, "marker": marker})
return torch.full((q.shape[0], q.shape[-2], heads * q.shape[-1]), marker)
return fake_func
@@ -80,8 +83,8 @@ def test_before_mod_segments_captured_override_is_a_passthrough():
q, k, v = _qkv()
out = override(_fake_func_factory(calls), q, k, v, heads=1, mask=None, skip_reshape=True)
assert len(calls) == 1, "no mod_segments yet -> single unmodified call, not split"
assert calls[0]["k_sum"] == float(k.sum())
assert torch.allclose(out, torch.full_like(out, float(k.sum())))
assert calls[0]["mask"] is None
assert torch.allclose(out, torch.zeros_like(out))
def test_an_explicit_mask_is_never_overridden():
@@ -96,7 +99,7 @@ def test_an_explicit_mask_is_never_overridden():
assert len(calls) == 1 and calls[0]["mask"] is real_mask
def test_splits_into_three_groups_with_correctly_damped_keys():
def test_splits_into_three_groups_with_a_bias_penalizing_the_other_modality():
node = S()
patched = node._install_h3_av_decouple(_FakeModel(), 0.5)
_fire_capture_hook(patched)
@@ -109,20 +112,27 @@ def test_splits_into_three_groups_with_correctly_damped_keys():
by_len = {c["q_len"]: c for c in calls}
assert set(by_len) == {len(_VIDEO_IDX), len(_AUDIO_IDX), len(_OTHER_IDX)}
full_k_sum = float(k.sum()) # 6 rows * 2 dims * 2.0 = 24
# audio-query pass: video rows (3 of them) halved in K -> k_sum drops by half of their share
per_row = 2 * 2.0 # head_dim * key value
audio_call = by_len[len(_AUDIO_IDX)]
assert audio_call["k_sum"] == full_k_sum - len(_VIDEO_IDX) * per_row * 0.5
video_call = by_len[len(_VIDEO_IDX)]
assert video_call["k_sum"] == full_k_sum - len(_AUDIO_IDX) * per_row * 0.5
other_call = by_len[len(_OTHER_IDX)]
assert other_call["k_sum"] == full_k_sum, "text/other rows ride the ordinary, undamped K"
bias_val = -0.5 * 12.0 # strength * _BIAS_SCALE
# scatter landed each group's output back at its own rows
assert torch.allclose(out[0, _AUDIO_IDX], torch.full((len(_AUDIO_IDX), 2), audio_call["k_sum"]))
assert torch.allclose(out[0, _VIDEO_IDX], torch.full((len(_VIDEO_IDX), 2), video_call["k_sum"]))
assert torch.allclose(out[0, _OTHER_IDX], torch.full((len(_OTHER_IDX), 2), other_call["k_sum"]))
audio_call = by_len[len(_AUDIO_IDX)]
audio_bias = audio_call["mask"].view(-1)
assert torch.allclose(audio_bias[_VIDEO_IDX], torch.full((len(_VIDEO_IDX),), bias_val))
assert torch.allclose(audio_bias[_AUDIO_IDX + _OTHER_IDX],
torch.zeros(len(_AUDIO_IDX) + len(_OTHER_IDX)))
video_call = by_len[len(_VIDEO_IDX)]
video_bias = video_call["mask"].view(-1)
assert torch.allclose(video_bias[_AUDIO_IDX], torch.full((len(_AUDIO_IDX),), bias_val))
assert torch.allclose(video_bias[_VIDEO_IDX + _OTHER_IDX],
torch.zeros(len(_VIDEO_IDX) + len(_OTHER_IDX)))
other_call = by_len[len(_OTHER_IDX)]
assert other_call["mask"] is None, "text/other rows ride the ordinary, unbiased pass"
# scatter landed each group's output back at its own rows (marker = mask.sum())
assert torch.allclose(out[0, _AUDIO_IDX], torch.full((len(_AUDIO_IDX), 2), audio_call["marker"]))
assert torch.allclose(out[0, _VIDEO_IDX], torch.full((len(_VIDEO_IDX), 2), video_call["marker"]))
assert torch.allclose(out[0, _OTHER_IDX], torch.zeros(len(_OTHER_IDX), 2))
def test_chains_through_an_already_installed_block0_hook():