From 6ae95b6fbd3def4f08e65519ea81e2497c92f427 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sun, 23 Aug 2026 08:25:05 +0300 Subject: [PATCH] feat(h3): token weighting as an attention-logit bias MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- h3_token_weights.py | 162 +++++++++++++++++++++++++ tests/test_h3_token_weights.py | 212 +++++++++++++++++++++++++++++++++ 2 files changed, 374 insertions(+) create mode 100644 h3_token_weights.py create mode 100644 tests/test_h3_token_weights.py diff --git a/h3_token_weights.py b/h3_token_weights.py new file mode 100644 index 00000000..f8c41246 --- /dev/null +++ b/h3_token_weights.py @@ -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"(? (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 diff --git a/tests/test_h3_token_weights.py b/tests/test_h3_token_weights.py new file mode 100644 index 00000000..9c59eb73 --- /dev/null +++ b/tests/test_h3_token_weights.py @@ -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)]) == []