From 01b8a37dffbe7c00c58fef36e86d2517cfbb4d7e Mon Sep 17 00:00:00 2001 From: Claude Date: Sun, 25 Oct 2026 22:10:19 +0000 Subject: [PATCH] Attention explorer: prompt weighting, prompt focus, per-prompt recipes, clip regions Regions rescale rows of q on a copy; checked against SLA's override contract. Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01RnrTBt9b8CdJwub5roffF3 --- CHANGELOG.md | 2 +- modules/sampling/attention_bandit/__init__.py | 295 ++++++++++++++---- .../tests/test_attention_bandit.py | 168 +++++++++- 3 files changed, 408 insertions(+), 57 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3338b711..bd538bb5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,7 +4,7 @@ ### Added (experimental) -- **Attention explorer** (needs a Taste key): each run tries a small attention change (sharper, as is or looser) per depth group of blocks, and your ratings decide which changes stay; how far a change goes learns from ratings too (or type your own). The Log names the exact blocks made sharper or looser. Zero extra model passes. Untested on a real H3. +- **Attention explorer** (needs a Taste key): each run tries a small attention change (sharper, as is or looser) per depth group of blocks, and your ratings decide which changes stay; how far a change goes learns from ratings too (or type your own). The Log names the exact blocks made sharper or looser. Also learns, per run: how strongly each group attends to your prompt ("Prompt focus"), and the clip's four time bands and its sound separately ("Clip regions"). A rating counts more for runs of a similar prompt, and a combination you liked on this prompt is reused (with some groups re-drawn). Zero extra model passes. Untested on a real H3. - **Noise merging** (Settings > Engine > Sampling): "Merge" can also run a whole second video beside yours from the second seed ("Second run": one extra pass at every step up to the last one named, blended at the steps you name), or mix the run's starting noise with a second seed's (no extra cost, variance kept). Or, at the steps you name, the step is answered a second time on the same picture with different noise (a second seed), and the two pictures are blended at your strength. Merge order picks which one is the base. The second seed is random each run unless you type one. Picture and sound are both merged (the sound can be left out with "Merge sound too"). "Second pass new noise" sets how much of the step's own noise the second pass swaps for the second seed's (1 all, lower keeps more of your own). Costs one extra model pass per merged step. Untested on a real H3. - **SLA rising sparsity** (diffusion model loader, `sla_sparsity_start`): skip fewer attention blocks on early steps and more on late ones, in a straight line up to `sla_sparsity`. Off at 0. One real H3 render (4 steps, 0.70 to 0.90): quality held, perhaps slightly better. Idea from GSRS (arXiv 2610.21129); only the schedule is taken, not its other parts. diff --git a/modules/sampling/attention_bandit/__init__.py b/modules/sampling/attention_bandit/__init__.py index 797533ae..6bac8c94 100644 --- a/modules/sampling/attention_bandit/__init__.py +++ b/modules/sampling/attention_bandit/__init__.py @@ -13,13 +13,27 @@ toward the size it used, a disliked one pushes it away (as Shot memory does for amount), with a little exploration. Or type your own size. Zero extra model calls; the speed cost is a division on the query. +Three more levers, all learned from the same ratings: +- Prompt focus: each group also draws a gain on the prompt tokens' keys, so attention to + the prompt can be sharpened or loosened separately from everything else. +- Prompt weighting: a rating counts more for runs of a similar prompt (pooled prompt + embedding, steep cosine) and less for other prompts, so what you tune on one prompt + shows up on that prompt first. +- Clip regions: the picture's four time bands and the sound each draw their own sharper or + looser setting, so a part you liked can stay while another changes. +- Recipes: when this prompt has a combination you liked, it is reused, with some groups + re-drawn so it keeps exploring. + Not proven: with a handful of ratings the per-group counts are noisy, and whether depth groups are the right unit is a guess. Experimental until it has been rated on a real model. """ +import math import random +import torch +import torch.nn.functional as F from comfy.patcher_extension import WrappersMP from ..._core import dit_hooks, log, registry @@ -54,6 +68,18 @@ SETTINGS = { "hint": "How much sharper or looser a change is. Keep it small.", "when": {"enabled": True, "mode": "manual"}, }, + "prompt_focus": { + "type": "bool", "default": True, + "label": "Prompt focus", + "hint": "Also learns how strongly each group attends to your prompt, apart from everything else.", + "when": {"enabled": True}, + }, + "regions": { + "type": "bool", "default": True, + "label": "Clip regions", + "hint": "Also learns the clip's four time bands and its sound separately, so a good part can stay while another changes.", + "when": {"enabled": True}, + }, "blocks": { "type": "text", "default": "2-10,29-46", "label": "Blocks", @@ -64,16 +90,26 @@ SETTINGS = { KIND = "attention_bandit" GROUPS = 4 -# Each arm is a temperature as 1 + RANGE * ARMS[i]: sharper, as is, looser. +# Each arm is a factor as 1 + size * ARMS[i]: lower, as is, higher. ARMS = (-1.0, 0.0, 1.0) STAY = 1 # index of "as is" STAY_HEAD_START = 2.0 # extra likes "as is" starts with, so a new key is stable -MAX_ROWS = 10 ** 9 # no cap: a rated run is under 1 KB +MAX_ROWS = 10 ** 9 # no cap: a rated run is a few KB at most START_SPREAD = 0.2 SPREAD_MIN, SPREAD_MAX = 0.05, 0.5 -LEARN_MAX = 0.35 # what ratings can grow the size to; 0.5x/1.5x on many blocks wrecked a picture on H3 +LEARN_MAX = 0.35 # what ratings can grow the size to; 0.5x/1.5x on many blocks wrecked a picture on H3 LR = 0.5 EXPLORE = 0.05 +# ponytail: pooled-prompt cosine sits high for most prompts, so the curve is steep and the +# constants are guesses. The log says how many rated runs count as this prompt. +SHARP = 100.0 # similarity = exp(SHARP * (cos - 1)); cos 0.99 -> 0.37 +FLOOR = 0.1 # what a different prompt's rating still counts for +UNKNOWN = 0.5 # a rating stored without its prompt +SAME = 0.5 # similarity from which a run counts as this prompt +RECIPE_BAR = 0.65 # a liked combination is reused when its draw beats this +REGIONS = ("start of the clip", "second quarter", "third quarter", "end of the clip", "sound") +BANDS = 4 # equal time bands of the picture; the sound is the last region +RECIPE_EXPLORE = 0.25 # per group: chance a reused recipe re-draws that group def _say(message, level=log.ALERT): @@ -92,17 +128,35 @@ def split(blocks, groups=GROUPS): return out -def counts(rows, signature, n_groups): - """[group][arm] -> [likes, dislikes] over the rated runs made with this same block layout.""" +def similarity(cond, stored): + """1 = the same prompt, toward 0 = a different one; None when either is unknown.""" + if cond is None or not torch.is_tensor(stored) or stored.numel() != cond.numel(): + return None + cos = float(F.cosine_similarity(cond.float().flatten(), stored.float().flatten(), dim=0)) + return math.exp(SHARP * (min(cos, 1.0) - 1.0)) + + +def weights_of(rows, cond): + """([how much each rated run counts now], [its similarity or None]).""" + sims = [similarity(cond, r["rows"].get("cond")) for r in rows] + return [UNKNOWN if s is None else FLOOR + (1.0 - FLOOR) * s for s in sims], sims + + +def _same_signature(p, signature): + return "sig" in p and [int(x) for x in p["sig"].tolist()] == signature + + +def counts(rows, signature, n_groups, field="arms", weights=None): + """[group][arm] -> [likes, dislikes] over the rated runs made with this block layout, + each counting its weight (1 when none given).""" table = [[[0.0, 0.0] for _ in ARMS] for _ in range(n_groups)] - for r in rows: + for i, r in enumerate(rows): p = r["rows"] - if "arms" not in p or "sig" not in p or r["reward"] == 0: + if field not in p or r["reward"] == 0 or not _same_signature(p, signature): continue - if [int(x) for x in p["sig"].tolist()] != signature: - continue - for g, arm in enumerate(int(a) for a in p["arms"].tolist()): - table[g][arm][0 if r["reward"] > 0 else 1] += 1.0 + w = 1.0 if weights is None else weights[i] + for g, arm in enumerate(int(a) for a in p[field].tolist()): + table[g][arm][0 if r["reward"] > 0 else 1] += w return table @@ -116,11 +170,9 @@ def learned_spread(rows, signature): m = START_SPREAD for r in rows: p = r["rows"] - if "spread" not in p or "sig" not in p or "arms" not in p or not r["reward"]: + if "spread" not in p or "arms" not in p or not r["reward"] or not _same_signature(p, signature): continue - if [int(x) for x in p["sig"].tolist()] != signature: - continue - if all(int(a) == STAY for a in p["arms"].tolist()): + if all(int(a) == STAY for f in ("arms", "tarms", "rarms") if f in p for a in p[f].tolist()): continue # nothing changed, so nothing to blame on the size used = float(p["spread"]) # A like pulls toward the size it used. A dislike can only say "too much": it pushes the @@ -137,17 +189,6 @@ def tally(rows): return f"{len(rows)} rated runs ({liked} liked, {disliked} disliked)" -def describe(groups, picks, spread): - """Which blocks went sharper or looser, exactly; groups left as is are not mentioned.""" - moved = {} - for g, a in zip(groups, picks): - if a != STAY: - moved.setdefault(a, []).extend(g) - parts = [f"{'sharper' if ARMS[a] < 0 else 'looser'} {1.0 + spread * ARMS[a]:g}x at blocks {ranges(sorted(bs))}" - for a, bs in sorted(moved.items())] - return "; ".join(parts) or "no change this run" - - def ranges(blocks): out, start = [], None for i, b in enumerate(blocks): @@ -159,6 +200,18 @@ def ranges(blocks): return ",".join(out) +def describe(groups, picks, spread, words=("sharper", "looser"), what="", blocks=True): + """Which blocks (or regions) went down or up, exactly; groups left as is are not mentioned.""" + moved = {} + for g, a in zip(groups, picks): + if a != STAY: + moved.setdefault(a, []).extend(g) + parts = [f"{what}{words[0] if ARMS[a] < 0 else words[1]} {1.0 + spread * ARMS[a]:g}x " + + (f"at blocks {ranges(sorted(bs))}" if blocks else f"in {', '.join(bs)}") + for a, bs in sorted(moved.items())] + return "; ".join(parts) + + def draw(table, rng=random): """One arm per group: the arm whose Beta draw (likes + 1, dislikes + 1) is highest.""" picks = [] @@ -173,8 +226,95 @@ def draw(table, rng=random): return picks +def pick_recipe(rows, sims, signature, focus, regions=False, rng=random): + """The arms of a combination liked on THIS prompt, or None. Each combination seen on this + prompt gets a Beta draw from its likes and dislikes; the best one is reused when it beats the bar.""" + combos = {} + for r, s in zip(rows, sims): + p = r["rows"] + if s is None or s < SAME or not r["reward"] or "arms" not in p or "spread" not in p \ + or not _same_signature(p, signature) or (focus and "tarms" not in p) \ + or (regions and "rarms" not in p): + continue + key = (tuple(int(a) for a in p["arms"].tolist()), + tuple(int(a) for a in p["tarms"].tolist()) if focus else None, + tuple(int(a) for a in p["rarms"].tolist()) if regions else None) + c = combos.setdefault(key, [0, 0, START_SPREAD]) + c[0 if r["reward"] > 0 else 1] += 1 + if r["reward"] > 0: + c[2] = float(p["spread"]) + best, best_theta = None, -1.0 + for key, (liked, disliked, spread) in combos.items(): + if not liked: + continue + theta = rng.betavariate(1.0 + liked, 1.0 + disliked) + if theta > best_theta: + best, best_theta = (key, spread, liked, disliked), theta + return best if best_theta > RECIPE_BAR else None + + +def make_plan(rows, cond, signature, n_groups, focus, manual_spread=None, regions=False, rng=random): + """What this run does: {arms, tarms | None, rarms | None, spread, recipe | None, same, unknown}.""" + weights, sims = weights_of(rows, cond) + arms = draw(counts(rows, signature, n_groups, "arms", weights), rng) + tarms = draw(counts(rows, signature, n_groups, "tarms", weights), rng) if focus else None + rarms = draw(counts(rows, signature, len(REGIONS), "rarms", weights), rng) if regions else None + spread = learned_spread(rows, signature) + recipe = pick_recipe(rows, sims, signature, focus, regions, rng) + if recipe is not None: + (r_arms, r_tarms, r_rarms), spread, liked, disliked = recipe + for mine, theirs in ((arms, r_arms), (tarms, r_tarms), (rarms, r_rarms)): + for i in range(len(mine or ())): + if rng.random() >= RECIPE_EXPLORE: + mine[i] = theirs[i] + recipe = (liked, disliked) + spread = manual_spread if manual_spread is not None else min(max( + spread + rng.gauss(0.0, EXPLORE), SPREAD_MIN), LEARN_MAX) + return {"arms": arms, "tarms": tarms, "rarms": rarms, "spread": spread, "recipe": recipe, + "same": sum(1 for s in sims if s is not None and s >= SAME), + "unknown": sum(1 for s in sims if s is None)} + + +def _prompt_of(guider): + try: + p = guider.conds["positive"][0]["cross_attn"] + except (AttributeError, KeyError, IndexError, TypeError): + return None + if not torch.is_tensor(p): + return None + p = p.detach().float() + while p.dim() > 1: + p = p.mean(dim=0) + return p.cpu() + + +def _segments(options): + """The packed sequence's {kind: (start, stop)}, or {} when this model does not say.""" + try: + return {kind: (int(a), int(b)) for a, b, kind in options["minimax_h3_layout"].segments} + except (KeyError, AttributeError, TypeError, ValueError): + return {} + + +def _text_rows(options): + """How many leading rows of the packed sequence are the prompt, or 0.""" + a, b = _segments(options).get("text", (0, 0)) + return b - a if a == 0 else 0 + + +def region_spans(options): + """[(start, stop)] per REGIONS entry: the target picture in equal time bands (its rows run + frame by frame), then the target sound. [] when the layout is unknown.""" + seg = _segments(options) + if "video" not in seg or "audio" not in seg: + return [] + a, b = seg["video"] + edges = [a + (b - a) * i // BANDS for i in range(BANDS + 1)] + return [(edges[i], edges[i + 1]) for i in range(BANDS)] + [seg["audio"]] + + # Carried on the per-block copy of transformer_options, as Randomizer does, so an -# attention call knows its block's temperature without a shared flag. +# attention call knows its block's factors without a shared flag. _SCOPE = "funui_attention_bandit" @@ -187,6 +327,8 @@ def install(patcher, values, key): return None manual = values.get("mode") == "manual" manual_spread = min(max(float(values.get("range", START_SPREAD)), SPREAD_MIN), SPREAD_MAX) + focus = bool(values.get("prompt_focus", True)) + regions = bool(values.get("regions", True)) blocks, problems = dit_hooks.parse_blocks(values.get("blocks"), dit_hooks.block_count(patcher)) for problem in problems: _say(f"ignored {problem}") @@ -194,30 +336,18 @@ def install(patcher, values, key): raise ValueError(f"no usable blocks in {values.get('blocks')!r}, so nothing would change") groups = split(blocks) signature = [len(groups), blocks[0], blocks[-1]] - live = {"temps": {}, "scaled": 0, "calls": 0} + live = {"factors": {}, "regions": [1.0] * len(REGIONS), "scaled": 0, "calls": 0, + "text": 0, "text_wanted": False, "region_hits": 0, "region_wanted": False} - def fresh(): - import torch - rows = taste.rows(KIND) - picks = draw(counts(rows, signature, len(groups))) - spread = manual_spread if manual else min(max( - learned_spread(rows, signature) + random.gauss(0.0, EXPLORE), SPREAD_MIN), LEARN_MAX) - live["scaled"] = live["calls"] = 0 - live["temps"] = {b: 1.0 + spread * ARMS[a] for g, a in zip(groups, picks) for b in g} - captured.update({"arms": torch.tensor(picks), "sig": torch.tensor(signature), - "spread": torch.tensor(float(spread))}) - log.once(f"{ID}:run", log.INFO, "FunUI Attention explorer", - f"key {taste.key!r}, {tally(rows)}: {describe(groups, picks, spread)}") - - captured = taste.collect(patcher, key, KIND, keep=MAX_ROWS, fresh=fresh) + captured = taste.collect(patcher, key, KIND, keep=MAX_ROWS) def scoped_for(index): def scoped(args, extra): - t = live["temps"].get(index) - if t is None or t == 1.0: + f = live["factors"].get(index) + if f is None or (f == (1.0, 1.0) and not live["region_wanted"]): return extra["original_block"](args) to = dict(args.get("transformer_options") or {}) - to[_SCOPE] = t + to[_SCOPE] = f return extra["original_block"]({**args, "transformer_options": to}) return scoped @@ -225,27 +355,82 @@ def install(patcher, values, key): dit_hooks.add_block_hook(patcher, key, index, scoped_for(index)) def override(func, q, k, v, *args, **kwargs): - t = (kwargs.get("transformer_options") or {}).get(_SCOPE) + options = kwargs.get("transformer_options") or {} + scope = options.get(_SCOPE) live["calls"] += 1 - live["scaled"] += t is not None - return func(q if t is None else q / t, k, v, *args, **kwargs) + if scope is None: + return func(q, k, v, *args, **kwargs) + temperature, prompt = scope + live["scaled"] += 1 + if temperature != 1.0: + q = q / temperature + if live["region_wanted"] and q.ndim == 4: + spans = region_spans(options) + if spans and spans[-1][1] <= q.shape[2]: + q = q.clone() + for (a, b), factor in zip(spans, live["regions"]): + if factor != 1.0: + q[:, :, a:b] /= factor # only those rows ask differently: the rest of the clip is untouched + live["region_hits"] += 1 + if prompt != 1.0 and kwargs.get("skip_reshape") and k.ndim == 4: + n = _text_rows(options) + if 0 < n < k.shape[2]: + k = k.clone() + k[:, :, :n] *= prompt # the prompt's keys only: its share of attention sharpens or loosens + live["text"] += 1 + return func(q, k, v, *args, **kwargs) dit_hooks.add_attention_override(patcher, key, override) - def report(executor, *args, **kwargs): + def run(executor, *args, **kwargs): + rows = taste.rows(KIND) + cond = _prompt_of(getattr(executor, "class_obj", None)) + plan = make_plan(rows, cond, signature, len(groups), focus, manual_spread if manual else None, regions) + arms, tarms, rarms, spread = plan["arms"], plan["tarms"], plan["rarms"], plan["spread"] + tfac = [1.0 + spread * ARMS[a] for a in tarms] if focus else [1.0] * len(groups) + rfac = [1.0 + spread * ARMS[a] for a in rarms] if regions else [1.0] * len(REGIONS) + live.update(scaled=0, calls=0, text=0, region_hits=0, regions=rfac, + region_wanted=any(f != 1.0 for f in rfac), text_wanted=any(t != 1.0 for t in tfac), + factors={b: (1.0 + spread * ARMS[a], tf) + for g, a, tf in zip(groups, arms, tfac) for b in g}) + captured.update({"arms": torch.tensor(arms), "sig": torch.tensor(signature), + "spread": torch.tensor(float(spread))}) + if focus: + captured["tarms"] = torch.tensor(tarms) + if regions: + captured["rarms"] = torch.tensor(rarms) + if cond is not None: + captured["cond"] = cond.half() + known = "prompt unknown" if cond is None else f"{plan['same']} of them this prompt" + how = f"recipe liked {plan['recipe'][0]}x, disliked {plan['recipe'][1]}x" if plan["recipe"] else "exploring" + changes = "; ".join(x for x in ( + describe(groups, arms, spread), + describe(groups, tarms, spread, ("weaker", "stronger"), "prompt ") if focus else "", + describe([[n] for n in REGIONS], rarms, spread, what="clip regions ", blocks=False) if regions else "") if x) + log.once(f"{ID}:run", log.INFO, "FunUI Attention explorer", + f"key {taste.key!r}, {tally(rows)}, {known}, {how}: {changes or 'no change this run'}") out = executor(*args, **kwargs) - changed = any(t != 1.0 for t in live["temps"].values()) - if changed and not live["scaled"]: + wanted = bool(changes) + if wanted and not live["scaled"]: _say("changed nothing this run: no attention call reached the changed blocks " f"({live['calls']} calls seen). Another feature may have replaced attention.", log.WARNING) - elif changed: + elif live["text_wanted"] and not live["text"]: + _say("prompt focus changed nothing this run: this model does not say where the prompt sits in " + "its sequence, so only the plain attention changes were applied.", log.WARNING) + elif live["region_wanted"] and not live["region_hits"]: + _say("clip regions changed nothing this run: this model does not say where the picture and " + "sound sit in its sequence.", log.WARNING) + elif wanted: log.once(f"{ID}:applied", log.INFO, "FunUI Attention explorer", - f"applied to {live['scaled']} of {live['calls']} attention calls") + f"applied to {live['scaled']} of {live['calls']} attention calls" + + (f", prompt keys scaled in {live['text']}" if live["text_wanted"] else "") + + (f", clip regions in {live['region_hits']}" if live["region_wanted"] else "")) return out - patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, f"{key}.report", report) + patcher.add_wrapper_with_key(WrappersMP.OUTER_SAMPLE, f"{key}.run", run) size = f"size {manual_spread:g}" if manual else "size learned from ratings" - return f"{len(groups)} depth groups over {len(blocks)} blocks, {size}" + return (f"{len(groups)} depth groups over {len(blocks)} blocks, {size}" + + (", prompt focus on" if focus else "") + (", clip regions on" if regions else "")) PROVIDES = {"modifier": install} diff --git a/modules/sampling/attention_bandit/tests/test_attention_bandit.py b/modules/sampling/attention_bandit/tests/test_attention_bandit.py index 016bdf02..89c02660 100644 --- a/modules/sampling/attention_bandit/tests/test_attention_bandit.py +++ b/modules/sampling/attention_bandit/tests/test_attention_bandit.py @@ -74,7 +74,7 @@ def test_dislikes_of_small_changes_never_grow_the_size(): def test_the_log_names_exact_blocks_and_skips_the_unchanged(): from modules.sampling.attention_bandit import describe groups = [[2, 3, 4], [9, 10, 29, 30], [34, 35], [41, 42]] - assert describe(groups, [1, 1, 1, 1], 0.2) == "no change this run" + assert describe(groups, [1, 1, 1, 1], 0.2) == "" assert describe(groups, [0, 1, 2, 0], 0.2) == "sharper 0.8x at blocks 2-4,41-42; looser 1.2x at blocks 34-35" @@ -176,3 +176,169 @@ def test_the_log_says_how_many_attention_calls_were_changed(tiny_h3, monkeypatch log.new_run() _run(_outer(patched), lambda: tiny_h3.run(patched)) assert any("applied to" in e["message"] and "attention calls" in e["message"] for e in log.history()) + + +# ── prompt weighting, recipes, prompt focus, clip regions ───────────────────── + +def _cond(seed, n=16): + g = torch.Generator().manual_seed(seed) + return torch.randn(n, generator=g) + + +def test_the_same_prompt_counts_fully_another_prompt_counts_for_the_floor_and_unknown_in_between(): + from modules.sampling.attention_bandit import FLOOR, UNKNOWN, similarity, weights_of + a, b = _cond(1), _cond(2) + assert similarity(a, a) == pytest.approx(1.0) + assert similarity(a, b) < 0.01 and similarity(a, None) is None and similarity(None, a) is None + rows = [{"reward": 1.0, "rows": {"cond": a.half()}}, {"reward": 1.0, "rows": {"cond": b.half()}}, + {"reward": 1.0, "rows": {}}] + w, sims = weights_of(rows, a) + assert w[0] == pytest.approx(1.0, abs=1e-3) and w[1] == pytest.approx(FLOOR, abs=0.02) and w[2] == UNKNOWN + + +def test_counts_use_the_weights(): + from modules.sampling.attention_bandit import counts + rows = [_row(1.0, [0], [1, 0, 0]), _row(1.0, [0], [1, 0, 0])] + assert counts(rows, [1, 0, 0], 1, weights=[1.0, 0.1])[0][0][0] == pytest.approx(1.1) + rows[0]["rows"]["tarms"] = torch.tensor([2]) + assert counts(rows, [1, 0, 0], 1, "tarms")[0][2] == [1.0, 0.0] # only the row that has the field + + +def _combo_row(reward, cond, arms=(0, 1), tarms=(2, 1), rarms=(1, 1, 1, 1, 1), spread=0.3, sig=(2, 0, 9)): + return {"reward": reward, "rows": {"arms": torch.tensor(list(arms)), "tarms": torch.tensor(list(tarms)), + "rarms": torch.tensor(list(rarms)), "sig": torch.tensor(list(sig)), + "spread": torch.tensor(spread), "cond": cond.half()}} + + +def test_a_combination_liked_on_this_prompt_is_reused_and_one_liked_on_another_is_not(): + from modules.sampling.attention_bandit import make_plan + here, other = _cond(1), _cond(2) + rows = [_combo_row(1.0, here) for _ in range(8)] + rng = random.Random(0) + plans = [make_plan(rows, here, [2, 0, 9], 2, True, regions=True, rng=rng) for _ in range(40)] + reused = [p for p in plans if p["recipe"]] + assert len(reused) > 30 and all(p["same"] == 8 for p in plans) + assert sum(p["arms"][0] == 0 for p in reused) > 0.6 * len(reused) # mostly the recipe, some groups re-drawn + away = [make_plan(rows, other, [2, 0, 9], 2, True, regions=True, rng=rng) for _ in range(40)] + assert not any(p["recipe"] for p in away) and all(p["same"] == 0 for p in away) + + +def test_disliked_combinations_are_not_a_recipe(): + from modules.sampling.attention_bandit import make_plan + here = _cond(1) + rows = [_combo_row(-1.0, here) for _ in range(8)] + assert not any(make_plan(rows, here, [2, 0, 9], 2, True, regions=True, rng=random.Random(i))["recipe"] for i in range(30)) + + +def _force(monkeypatch, arms=1, tarms=1, rarms=1): + from modules.sampling import attention_bandit as m + + def fixed(table, rng=random): + n = len(table) + return [{4: arms, 5: rarms}.get(n, arms)] * n if n == 5 else [arms] * n + seq = iter([]) + calls = [] + + def draw(table, rng=random): + calls.append(len(table)) + # the first draw is the plain groups, the second the prompt groups, the third the regions + pick = (arms, tarms, rarms)[min(len(calls) - 1, 2)] + return [pick] * len(table) + monkeypatch.setattr(m, "draw", draw) + + +def test_prompt_focus_alone_changes_the_output_and_says_where(tiny_h3, monkeypatch): + from core import log + base = tiny_h3.run() + patched, status = _load(tiny_h3, mode="manual", range=0.5, blocks="0-3", regions=False) + assert "prompt focus on" in status + _force(monkeypatch, arms=1, tarms=2) + log.new_run() + outs = _run(_outer(patched), lambda: tiny_h3.run(patched)) + assert not torch.equal(outs[0], base[0]) + messages = [e["message"] for e in log.history()] + assert any("prompt stronger" in m and "blocks 0-3" in m for m in messages) + assert any("prompt keys scaled in" in m for m in messages) + + +def test_with_prompt_focus_off_nothing_about_the_prompt_is_drawn_or_stored(tiny_h3, monkeypatch): + from modules.system.taste import store + patched, _ = _load(tiny_h3, blocks="0-3", prompt_focus=False, regions=False) + _force(monkeypatch, arms=1) + outs = _run(_outer(patched), lambda: tiny_h3.run(patched)) + assert torch.equal(outs[0], tiny_h3.run()[0]) + store.rate("run-1", "liked") + assert "tarms" not in store.load("fox", "attention_bandit")["rows"][0]["rows"] + + +def test_clip_regions_change_only_the_rows_they_name(tiny_h3, monkeypatch): + from modules.sampling.attention_bandit import REGIONS, region_spans + captured = {} + + def spy(func, q, k, v, *a, **kw): + options = kw.get("transformer_options") or {} + # the prompt refiner also runs with the layout in its options, on the prompt rows alone + if "minimax_h3_layout" in options and q.shape[2] > captured.get("rows", 0): + captured["spans"] = region_spans(options) + captured["rows"] = q.shape[2] + return func(q, k, v, *a, **kw) + + from core import dit_hooks + probe = tiny_h3.patcher.clone() + dit_hooks.add_attention_override(probe, "test.spy", spy) + tiny_h3.run(probe) + spans = captured["spans"] + assert len(spans) == len(REGIONS) + flat = sorted(spans) + assert all(a[1] <= b[0] for a, b in zip(flat, flat[1:])) # no row in two regions + assert all(0 <= a < b <= captured["rows"] for a, b in spans) + + +def test_clip_regions_change_the_output_and_stay_off_when_all_as_is(tiny_h3, monkeypatch): + from core import log + base = tiny_h3.run() + patched, _ = _load(tiny_h3, mode="manual", range=0.5, blocks="0-3", prompt_focus=False) + _force(monkeypatch, arms=1, rarms=0) + # prompt focus is off, so the draws are groups then regions + from modules.sampling import attention_bandit as m + picks = iter([1, 0]) + monkeypatch.setattr(m, "draw", lambda table, rng=random: [next(picks)] * len(table)) + log.new_run() + outs = _run(_outer(patched), lambda: tiny_h3.run(patched)) + assert not torch.equal(outs[0], base[0]) + messages = [e["message"] for e in log.history()] + assert any("clip regions sharper" in m and "sound" in m for m in messages) + assert any("clip regions in" in m for m in messages) + + +def test_the_real_sla_override_underneath_still_gets_a_normal_call(tiny_h3, monkeypatch): + """SLA hands anything off its fast path to dense and records a kernel failure if a call is odd. + Regions rescale rows of q in place on a copy, so shape, dtype and layout must be what SLA saw before.""" + from core import dit_hooks + from modules.loaders import sla_attention as sla + seen = [] + + def spy(func, q, k, v, *a, **kw): + # strides are left out on purpose: the copy of q is compact, which SLA reads as well as the fused-buffer view + seen.append((tuple(q.shape), q.dtype, tuple(k.shape), kw.get("skip_reshape"), kw.get("mask"))) + return func(q, k, v, *a, **kw) + + def run_with(on): + probe = tiny_h3.patcher.clone() + state = sla.new_state() + to = dit_hooks._transformer_options(probe) + cfg = {**sla.SLA_DEFAULTS, "engine": "triton", "dense_steps": frozenset(), "min_seq_len": 0} + to["optimized_attention_override"] = sla.make_override(state, cfg) + dit_hooks.add_attention_override(probe, "test.spy", spy) # sits on top: sees exactly what SLA would get + patched, _ = _load(type("T", (), {"patcher": probe})(), mode="manual", range=0.5, blocks="0-3") + picks = iter([1, 1, 0 if on else 1]) + from modules.sampling import attention_bandit as m + monkeypatch.setattr(m, "draw", lambda table, rng=random: [next(picks)] * len(table)) + seen.clear() + _run(_outer(patched), lambda: tiny_h3.run(patched)) + return list(seen), state + + plain, state_plain = run_with(False) + scaled, state_scaled = run_with(True) + assert plain and scaled == plain # the same calls, same shapes, strides and flags + assert state_scaled["failed"] is None and state_scaled["calls"] == state_plain["calls"]