nodes/olivv-cs-ComfyUI-FunPack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat(h3): token weighting as an attention-logit bias

Browse Source 
Qwen never applies (word:1.2): ComfyUI's H3 tokenizer hardcodes every
weight to 1.0 and runs the raw tokenizer over the literal string, so the
syntax reaches the encoder as punctuation inside the sentence and
sd1_clip's weighting arm never runs.

Scaling embeddings would not have worked either — Qwen's rows are
contextual hidden states, not per-word vectors. But H3 is one packed
self-attention stream with the prompt's tokens in it, so log(w) added to
the logits at that phrase's key positions multiplies its share of the
attention mass by w, which is what a weight means.

Module and tests only; nothing is wired to a node yet.
This commit is contained in:
digital-garbage
2026-08-23 08:25:05 +03:00
parent 6308578bd7
commit 6ae95b6fbd
2 changed files with 374 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
h3_token_weights.py
+162
View File
@@ -0,0 +1,162 @@
"""Real token weighting for MiniMax H3.
Qwen ignores `(word:1.2)`. ComfyUI's `MiniMaxH3Tokenizer.tokenize_with_weights` hardcodes
every weight to 1.0 and runs the raw HF tokenizer over the literal string, so the syntax
does not weight anything — it reaches the encoder as punctuation inside your sentence, and
`sd1_clip`'s weighting arm never runs because no weight is ever != 1.0.
Scaling the embeddings would not have worked anyway. Qwen's conditioning rows are
CONTEXTUAL hidden states — row i is "having read the prompt this far", not "the vector for
this word" — so there is no per-word magnitude to turn up.
What H3 does have is a single packed self-attention stream with the prompt's tokens IN it.
Every video token attends over those positions, so the weighting CLIP only ever
approximated can be done exactly: add `log(w)` to the attention logits at that phrase's key
positions. softmax(logits + log w) is softmax(logits) with that key's share of the
attention mass multiplied by w, which is what a weight is supposed to mean.
Costs one broadcast add per attention call. The catch is stated where it is installed: a
biased call cannot run on a block-sparse kernel, so SLA goes dense for the run.
"""
from __future__ import annotations
import math
import re
from typing import Optional
# (phrase:1.3) / (phrase:-0.4). Escaped \( is left alone so a literal bracket still works.
_WEIGHTED = re.compile(r"(?<!\\)\(([^():]*?):\s*(-?\d+(?:\.\d+)?)\s*\)")
# Below this a weight is treated as "remove entirely" rather than log(0) = -inf, which
# produces NaN the moment a query attends to nothing else.
MIN_WEIGHT = 1e-3
MASKED_BIAS = -30.0
# A guard, not a preference: bias is added to raw logits, and past this the phrase wins
# every query regardless of content, which reads as the prompt collapsing onto one word.
MAX_ABS_BIAS = 6.0
def parse(text: str):
"""Strip weight syntax. -> (clean_text, [(start_char, end_char, weight), ...]).
Spans index into `clean_text`, so they survive the removal of the brackets. Returns the
text unchanged and an empty list when there is nothing to parse, which is the common
case and must cost nothing.
"""
if not text or "(" not in text:
return text, []
out = []
spans = []
pos = 0
for m in _WEIGHTED.finditer(text):
out.append(text[pos:m.start()])
phrase = m.group(1)
start = sum(len(p) for p in out)
out.append(phrase)
spans.append((start, start + len(phrase), float(m.group(2))))
pos = m.end()
out.append(text[pos:])
return "".join(out), spans
def token_spans(tokenizer, clean_text: str, char_spans):
"""Map character spans to token index ranges. -> [(start_tok, end_tok, weight), ...].
Uses the tokenizer's own offset mapping rather than re-tokenizing each phrase on its
own: "cat" tokenized alone and "cat" inside a sentence are not always the same ids, and
a span that is off by one token weights the wrong word.
"""
if not char_spans:
return []
try:
enc = tokenizer(clean_text, add_special_tokens=False, return_offsets_mapping=True)
offsets = enc["offset_mapping"]
except Exception: # noqa: BLE001 — a slow tokenizer has no offsets; weighting is off
return []
out = []
for c_start, c_end, weight in char_spans:
toks = [i for i, (a, b) in enumerate(offsets) if a < c_end and b > c_start]
if toks:
out.append((toks[0], toks[-1] + 1, weight))
return out
def bias_value(weight: float) -> float:
"""The additive logit bias a multiplicative weight corresponds to."""
if weight <= MIN_WEIGHT:
return MASKED_BIAS
return max(-MAX_ABS_BIAS, min(MAX_ABS_BIAS, math.log(weight)))
def build_bias(spans, prompt_tokens: int, cond_len: int, seq_len: int, device, dtype):
"""An additive [1, 1, 1, seq_len] attention bias, or None if nothing applies.
The prompt is the LAST thing the H3 tokenizer appends — reference pictures, audio and
video labels all go in front of it — so the prompt's tokens are the final
`prompt_tokens` rows of the conditioning block, and the conditioning block leads the
packed sequence. That places prompt token j at packed position
`cond_len - prompt_tokens + j`, without needing to know what the references added.
"""
import torch
if not spans or prompt_tokens <= 0 or cond_len <= 0 or seq_len < cond_len:
return None
base = cond_len - prompt_tokens
if base < 0:
return None
bias = None
for start, end, weight in spans:
b = bias_value(weight)
if b == 0.0:
continue
lo, hi = base + max(0, start), base + min(prompt_tokens, end)
if hi <= lo or hi > cond_len:
continue
if bias is None:
bias = torch.zeros(1, 1, 1, seq_len, device=device, dtype=dtype)
bias[..., lo:hi] += b
return bias
def make_override(spans, prompt_tokens: int, cond_len: int, inner=None, on_apply=None):
"""An `optimized_attention_override` that adds the bias, then delegates.
`inner` is the override this one displaces (SLA's, or the mask-safe one) so installing
weighting never silently discards the backend the user chose. The bias is cached per
(sequence length, device, dtype) — it is the same tensor on every call and every block.
"""
cache: dict = {}
def override(func, q, k, v, heads, mask=None, attn_precision=None,
skip_reshape=False, skip_output_reshape=False, **kwargs):
def run(m):
if inner is not None:
return inner(func, q, k, v, heads, mask=m, attn_precision=attn_precision,
skip_reshape=skip_reshape,
skip_output_reshape=skip_output_reshape, **kwargs)
return func(q, k, v, heads, mask=m, attn_precision=attn_precision,
skip_reshape=skip_reshape,
skip_output_reshape=skip_output_reshape, **kwargs)
try:
# Only the packed self-attention: k and q are the same sequence there. The
# 2-block token refiner and any cross-shaped call are left alone.
if not skip_reshape or k.ndim != 4 or q.shape[2] != k.shape[2]:
return run(mask)
seq_len = int(k.shape[2])
if seq_len < cond_len:
return run(mask)
key = (seq_len, k.device, k.dtype)
if key not in cache:
cache[key] = build_bias(spans, prompt_tokens, cond_len, seq_len,
k.device, k.dtype)
if on_apply is not None and cache[key] is not None:
on_apply(seq_len)
bias = cache[key]
if bias is None:
return run(mask)
return run(bias if mask is None else mask + bias)
except Exception: # noqa: BLE001 — weighting must never cost the step
return run(mask)
return override
tests/test_h3_token_weights.py
+212
View File
@@ -0,0 +1,212 @@
"""Attention-logit token weighting for MiniMax H3.
Qwen never applies (word:1.2) — ComfyUI's H3 tokenizer hardcodes every weight to 1.0 — so
the syntax has to be stripped before encoding and re-applied where the model can act on it:
the prompt's key positions in the packed self-attention stream.
"""
import math
import sys
import pytest
import torch
sys.path.insert(0, ".")
import h3_token_weights as tw # noqa: E402
# --- parsing ---------------------------------------------------------------
def test_plain_text_is_returned_untouched():
assert tw.parse("a cat on a mat") == ("a cat on a mat", [])
def test_the_syntax_is_removed_from_the_text():
"""Left in, Qwen reads the brackets and digits as part of the sentence — worse than no
weighting at all."""
clean, spans = tw.parse("a (fluffy:1.4) cat")
assert clean == "a fluffy cat"
assert [(s, e) for s, e, _ in spans] == [(2, 8)]
assert clean[2:8] == "fluffy"
def test_spans_index_the_cleaned_text_not_the_original():
clean, spans = tw.parse("(red:1.2) car and a (blue:0.5) boat")
assert clean == "red car and a blue boat"
assert [clean[s:e] for s, e, _ in spans] == ["red", "blue"]
assert [w for _, _, w in spans] == [1.2, 0.5]
def test_a_negative_weight_is_parsed():
_, spans = tw.parse("(blurry:-0.4) shot")
assert spans[0][2] == -0.4
def test_an_escaped_bracket_is_left_alone():
clean, spans = tw.parse(r"a \(not a weight:1.2\) thing")
assert spans == []
assert clean == r"a \(not a weight:1.2\) thing"
# --- weight -> bias --------------------------------------------------------
def test_a_weight_of_one_is_no_bias_at_all():
assert tw.bias_value(1.0) == 0.0
def test_the_bias_is_the_log_of_the_weight():
"""softmax(logits + log w) multiplies that key's attention share by w. That identity is
the whole reason this is the right place to weight."""
assert tw.bias_value(2.0) == pytest.approx(math.log(2.0))
assert tw.bias_value(0.5) == pytest.approx(math.log(0.5))
def test_zero_masks_instead_of_producing_negative_infinity():
"""log(0) is -inf, and a query whose every key is -inf comes out NaN."""
assert tw.bias_value(0.0) == tw.MASKED_BIAS
assert math.isfinite(tw.bias_value(0.0))
def test_an_extreme_weight_is_clamped():
assert tw.bias_value(10_000.0) == tw.MAX_ABS_BIAS
# --- placing the bias in the packed sequence -------------------------------
def test_the_bias_lands_on_the_prompt_tail_of_the_conditioning():
"""References are prepended by the tokenizer, so the prompt occupies the END of the
conditioning block, and the conditioning block leads the packed sequence."""
bias = tw.build_bias([(0, 2, 2.0)], prompt_tokens=5, cond_len=9, seq_len=100,
device="cpu", dtype=torch.float32)
assert bias.shape == (1, 1, 1, 100)
assert torch.allclose(bias[0, 0, 0, 4:6], torch.full((2,), math.log(2.0)))
assert bias[0, 0, 0, :4].abs().sum() == 0 # the reference tokens are untouched
assert bias[0, 0, 0, 6:].abs().sum() == 0 # so is audio and video
def test_no_spans_means_no_tensor():
assert tw.build_bias([], 5, 9, 100, "cpu", torch.float32) is None
def test_a_span_reaching_past_the_prompt_is_dropped_not_wrapped():
assert tw.build_bias([(4, 99, 2.0)], prompt_tokens=5, cond_len=9, seq_len=100,
device="cpu", dtype=torch.float32) is not None
assert tw.build_bias([(50, 60, 2.0)], prompt_tokens=5, cond_len=9, seq_len=100,
device="cpu", dtype=torch.float32) is None
def test_a_sequence_shorter_than_the_conditioning_is_refused():
assert tw.build_bias([(0, 2, 2.0)], 5, 9, 4, "cpu", torch.float32) is None
# --- the override ----------------------------------------------------------
def _q(seq_len, heads=2, dim=8):
return torch.zeros(1, heads, seq_len, dim)
def test_the_override_biases_the_packed_self_attention():
seen = {}
def func(q, k, v, heads, mask=None, **kw):
seen["mask"] = mask
return q
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9)
ov(func, _q(100), _q(100), _q(100), 2, skip_reshape=True)
assert seen["mask"] is not None
assert seen["mask"][0, 0, 0, 4] == pytest.approx(math.log(2.0))
def test_an_existing_mask_is_added_to_not_replaced():
"""Stomping a real padding mask would corrupt attention rather than weight it."""
seen = {}
def func(q, k, v, heads, mask=None, **kw):
seen["mask"] = mask
return q
existing = torch.full((1, 1, 1, 100), -1.0)
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9)
ov(func, _q(100), _q(100), _q(100), 2, mask=existing, skip_reshape=True)
assert seen["mask"][0, 0, 0, 4] == pytest.approx(-1.0 + math.log(2.0))
assert seen["mask"][0, 0, 0, 0] == pytest.approx(-1.0)
def test_the_token_refiner_and_cross_shaped_calls_are_left_alone():
"""q and k are the same sequence only in the packed self-attention."""
seen = {}
def func(q, k, v, heads, mask=None, **kw):
seen["mask"] = mask
return q
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9)
ov(func, _q(100), _q(40), _q(40), 2, skip_reshape=True)
assert seen["mask"] is None
ov(func, _q(100), _q(100), _q(100), 2, skip_reshape=False)
assert seen["mask"] is None
def test_the_displaced_override_still_runs():
"""There is one override slot. Taking it without chaining would silently discard the
attention backend the user selected."""
calls = []
def inner(func, q, k, v, heads, mask=None, **kw):
calls.append(mask)
return q
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9, inner=inner)
ov(lambda *a, **kw: None, _q(100), _q(100), _q(100), 2, skip_reshape=True)
assert len(calls) == 1 and calls[0] is not None
def test_a_failure_inside_the_override_keeps_the_step():
"""Weighting is a refinement. It must never be the reason a render dies."""
def func(q, k, v, heads, mask=None, **kw):
return "ran"
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9)
broken = torch.zeros(1) # no .shape[2]; build_bias will raise
assert ov(func, broken, broken, broken, 2, skip_reshape=True) == "ran"
def test_the_bias_is_built_once_per_sequence_shape():
built = []
def func(q, k, v, heads, mask=None, **kw):
return q
ov = tw.make_override([(0, 2, 2.0)], prompt_tokens=5, cond_len=9,
on_apply=built.append)
for _ in range(4):
ov(func, _q(100), _q(100), _q(100), 2, skip_reshape=True)
assert built == [100]
# --- token spans -----------------------------------------------------------
class _FakeTokenizer:
"""Character-per-token, which makes the offset arithmetic checkable by hand."""
def __call__(self, text, add_special_tokens=False, return_offsets_mapping=False):
return {"offset_mapping": [(i, i + 1) for i in range(len(text))]}
def test_char_spans_become_token_ranges():
clean, spans = tw.parse("a (bc:1.5) d")
assert clean == "a bc d"
assert tw.token_spans(_FakeTokenizer(), clean, spans) == [(2, 4, 1.5)]
def test_a_tokenizer_without_offsets_turns_weighting_off_rather_than_guessing():
class _NoOffsets:
def __call__(self, *a, **kw):
raise ValueError("slow tokenizer")
assert tw.token_spans(_NoOffsets(), "a bc d", [(2, 4, 1.5)]) == []