feat: H3 query steering -- REINS-style direction learned and applied on Q, not hidden state

Same liked-minus-disliked mean-difference as h3_repr_steering, captured off Q inside
optimized_attention_override instead of a block's hidden state, own kind="q_steer" sidecar
so the two vector spaces never collide. Only Q is biased: a uniform additive bias on K is a
provable softmax no-op, and on V it's mathematically identical to h3_repr_steering's own
mechanism, so neither is a genuinely different lever. Off by default (blank block spec).

h3_repr_steering.py's persistence functions take an optional kind= param (default unchanged)
so this reuses capture()/direction()/commit() verbatim instead of a parallel implementation.

conditioning.py

@@ -10216,6 +10216,14 @@ class FunPackVideoRefinerV2(FunPackVideoRefiner):
                     }.get(_rs_outcome, _rs_outcome)
                     print(f"[FunPackRefiner] H3 representation steering: weight "
                           f"{_rs_reward:+.3f} for rating '{rating_label}' — {_rs_note}")
+                    # Query steering rides the same rating, its own kind="q_steer" store —
+                    # see h3_repr_steering.py / _install_h3_q_steering. Unconditional and
+                    # silent-on-no-pending for the same reason as REINS above: a run with the
+                    # sampler's h3_q_steer_block blank never captured anything to commit.
+                    _rs_q_outcome = _rs.commit(refinement_key, _rs_reward, kind="q_steer")
+                    if _rs_q_outcome == "recorded":
+                        print(f"[FunPackRefiner] H3 query steering: weight "
+                              f"{_rs_reward:+.3f} for rating '{rating_label}' — recorded")
                     # Block-influence probe rides the same rating. Measurement only -- it
                     # steers nothing, so it is committed unconditionally alongside REINS
                     # rather than behind its own admissibility gate. DELIBERATELY SILENT:

h3_repr_steering.py

@@ -70,12 +70,12 @@ def _log():
     return fl
 
 
-def state_path(refinement_key):
+def state_path(refinement_key, kind="repr_steer"):
     try:
         from .conditioning import refinement_state_path
     except ImportError:
         from conditioning import refinement_state_path
-    return refinement_state_path(refinement_key, "repr_steer", prefix="refine_v2", extension="pt")
+    return refinement_state_path(refinement_key, kind, prefix="refine_v2", extension="pt")
 
 
 def _mask_from_mod_segments(mod_segments, seq_len, device, tag):
@@ -133,8 +133,8 @@ def capture(hidden_state, video_mask):
 # one that steers, so block_sweep() can be answered from data already being collected instead
 # of needing its own separate collection pass.
 
-def _load(refinement_key):
-    path = state_path(refinement_key)
+def _load(refinement_key, kind="repr_steer"):
+    path = state_path(refinement_key, kind)
     if not os.path.exists(path):
         return {"rows": [], "pending": None, "slots": {}}
     try:
@@ -149,33 +149,38 @@ def _load(refinement_key):
     return data
 
 
-def _save(refinement_key, data):
-    path = state_path(refinement_key)
+def _save(refinement_key, data, kind="repr_steer"):
+    path = state_path(refinement_key, kind)
     os.makedirs(os.path.dirname(path), exist_ok=True)
     tmp = path + ".tmp"
     torch.save(data, tmp)
     os.replace(tmp, path)
 
 
-def save_pending(refinement_key, descriptors):
+def save_pending(refinement_key, descriptors, kind="repr_steer"):
     """Written on EVERY captured run, win or lose -- same fix trajectory_probe needed: a
     pending capture must not survive to be paired with the NEXT run's rating instead of its
     own, and the only way to guarantee that is to overwrite it every time, not just when a
     rating happens to follow. `descriptors` is {block: tensor}, whatever candidate blocks the
     sampler actually captured this run -- not required to cover every entry in
-    CANDIDATE_BLOCKS (a shape the hook has never seen just means that block is missing)."""
+    CANDIDATE_BLOCKS (a shape the hook has never seen just means that block is missing).
+
+    `kind` picks which sidecar file this is read/written under (see state_path) -- a second
+    mechanism capturing a DIFFERENT vector space at the same block index (h3_q_steering
+    captures Q, not hidden state) uses its own kind so the two never share rows under the
+    same block key."""
     if not refinement_key or not descriptors:
         return
-    data = _load(refinement_key)
+    data = _load(refinement_key, kind)
     data["pending"] = {int(b): d.detach().float().cpu() for b, d in descriptors.items()}
     try:
-        _save(refinement_key, data)
+        _save(refinement_key, data, kind)
     except OSError as e:
         _log().failed("H3 representation steering", "save pending", e,
                        "this run's capture is not kept")
 
 
-def commit(refinement_key, reward):
+def commit(refinement_key, reward, kind="repr_steer"):
     """Pairs the pending capture with a rating's WEIGHT and appends it to the log. `reward`
     is the caller's already-resolved, already-admissibility-filtered reward -- this function
     trusts it rather than re-deriving a verdict, same as the value functions do with the same
@@ -187,12 +192,12 @@ def commit(refinement_key, reward):
     reports the reward it is ABOUT to commit cannot tell that apart from a real commit."""
     if not refinement_key:
         return "no_key"
-    data = _load(refinement_key)
+    data = _load(refinement_key, kind)
     pending = data.get("pending")
     data["pending"] = None
     if pending is None:
         try:
-            _save(refinement_key, data)
+            _save(refinement_key, data, kind)
         except OSError:
             pass
         return "no_pending"
@@ -200,12 +205,12 @@ def commit(refinement_key, reward):
         data["rows"].append({"desc": pending, "weight": float(reward)})
     except (TypeError, ValueError):
         try:
-            _save(refinement_key, data)
+            _save(refinement_key, data, kind)
         except OSError:
             pass
         return "no_pending"
     try:
-        _save(refinement_key, data)
+        _save(refinement_key, data, kind)
     except OSError as e:
         _log().failed("H3 representation steering", "commit", e, "rating not recorded")
         return "no_pending"
@@ -274,16 +279,16 @@ def discard_slot(refinement_key, slot_id):
             pass
 
 
-def clear_all(refinement_key):
+def clear_all(refinement_key, kind="repr_steer"):
     if not refinement_key:
         return
     try:
-        os.remove(state_path(refinement_key))
+        os.remove(state_path(refinement_key, kind))
     except FileNotFoundError:
         pass
 
 
-def direction(refinement_key, block=None):
+def direction(refinement_key, block=None, kind="repr_steer"):
     """-> (unit_vector, n_positive, n_negative), or (None, n_positive, n_negative) if either
     side is under MIN_PER_GROUP. One Awful and one Perfect is not a direction, it is two
     points -- doesn't matter how many barely-positive rows sit between them.
@@ -303,7 +308,7 @@ def direction(refinement_key, block=None):
     that block was in CANDIDATE_BLOCKS (or missing it for any other reason) are skipped, not
     treated as zero."""
     block = DEFAULT_BLOCK if block is None else int(block)
-    data = _load(refinement_key)
+    data = _load(refinement_key, kind)
     rows = [r for r in data["rows"]
             if isinstance(r, dict) and "weight" in r and isinstance(r.get("desc"), dict)
             and block in r["desc"]]

samplers.py

@@ -3006,6 +3006,17 @@ class FunPackLTXAVSceneChainSampler:
                     "default": 0, "min": -50, "max": 50,
                     "tooltip": "Only repeat during part of the schedule. 0 = every step (the original behaviour). Positive N = the FINAL N steps -- structure is already settled, so this only refines. Negative N = the FIRST |N| steps instead -- apply while structure is still forming, then get out of the way and let the untouched rest of the schedule resolve it, the same way the model ordinarily turns a rough early layout into a finished image. Counted in STEPS, not sigma, so it lands on the same point of the schedule regardless of step count or sampler.",
                 }),
+                # Appended here, not next to h3_repr_steering_* above — same widgets_values
+                # positional-mapping trap noted near h3_repr_steering_block. New widgets go at
+                # the end, always.
+                "h3_q_steer_strength": ("FLOAT", {
+                    "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01,
+                    "tooltip": "EXPERIMENTAL, unvalidated (H3 only). Same liked-minus-disliked learned direction as h3_repr_steering, but captured and applied on the QUERY vectors going into attention at h3_q_steer_block, not on the block's hidden state. Steers WHICH keys the video rows attend to -- it does not add anything to the residual stream the way h3_repr_steering does, so it is a genuinely different lever, not the same mechanism twice. 0 = off, model not cloned. Needs 2+ liked and 2+ disliked ratings at the SAME block before it steers -- captures read-only until then. Its learned direction lives in its own store, separate from h3_repr_steering's.",
+                }),
+                "h3_q_steer_block": ("STRING", {
+                    "default": "",
+                    "tooltip": "Which block(s) query-steering is captured AND applied at. Blank = off (nowhere to apply is not a lesser version of this, it costs nothing and captures nothing). Same syntax as h3_block_repeat -- a block, a range, or a comma list; try 0,1, the early blocks where h3_repr_steering's own stylistic effect turned out to live. Each named block learns and steers with its OWN direction, independent of every other block and of h3_repr_steering_block (different store, different vector space).",
+                }),
                 # A connection socket, never a widget — safe at the end, and it must stay after
                 # every widget above (see the widgets_values note at the top of this block).
                 "second_pass_sigmas": ("SIGMAS", {
@@ -3720,6 +3731,7 @@ class FunPackLTXAVSceneChainSampler:
                       h3_block_repeat="", h3_block_repeat_times=1,
                       h3_block_repeat_video_only=False, h3_block_repeat_span_loop=False,
                       h3_block_repeat_last_steps=0,
+                      h3_q_steer_strength=0.0, h3_q_steer_block="",
                       refinement_key="",):
         if sampler is None:
             raise ValueError("sampler input is required.")
@@ -3809,6 +3821,11 @@ class FunPackLTXAVSceneChainSampler:
             model = self._install_h3_repr_steering(
                 model, refinement_key, h3_repr_steering_strength, _repr_capture,
                 steer_block=h3_repr_steering_block)
+        _q_steer_capture = [{}]
+        if refinement_key and self._parse_block_spec(h3_q_steer_block):
+            model = self._install_h3_q_steering(
+                model, refinement_key, h3_q_steer_strength, _q_steer_capture,
+                steer_block=h3_q_steer_block)
         model = self._install_h3_av_decouple(model, h3_av_decouple)
         model = self._install_h3_attn_temperature(
             model, h3_explore_temperature, h3_explore_temperature_block)
@@ -3828,6 +3845,12 @@ class FunPackLTXAVSceneChainSampler:
                     _rs.save_capture_slot(refinement_key, h3_repr_capture_slot, _repr_capture[0])
                 else:
                     _rs.save_pending(refinement_key, _repr_capture[0])
+            if _q_steer_capture[0]:
+                try:
+                    from . import h3_repr_steering as _rs
+                except ImportError:
+                    import h3_repr_steering as _rs
+                _rs.save_pending(refinement_key, _q_steer_capture[0], kind="q_steer")
             if _influence_capture[0]:
                 try:
                     from . import block_influence as _bi
@@ -7064,6 +7087,156 @@ class FunPackLTXAVSceneChainSampler:
                         "not applied this run")
             return model
 
+    def _install_h3_q_steering(self, model, refinement_key, strength, capture_holder,
+                                steer_block=None):
+        """EXPERIMENTAL, unvalidated (H3 only). Same liked-minus-disliked mean-difference
+        mechanism as h3_repr_steering (see that module's docstring for the shared math), but
+        captured and applied on QUERY vectors going into attention instead of on the block's
+        hidden state. Reshapes Q from (heads, head_dim) back into one flat hidden-sized
+        vector per video row so it can reuse h3_repr_steering.capture()/direction() verbatim.
+        The learned direction lives under its own kind="q_steer" sidecar file -- entirely
+        separate from h3_repr_steering's hidden-state directions, same refinement_key -- a
+        block index means a different vector space in each store, so nothing here reads or
+        writes the other's rows.
+
+        Only Q is biased, never K or V. Adding a fixed vector to every KEY uniformly is a
+        provable no-op through softmax: for a fixed query the shift q.c is the same constant
+        added to every key's logit in that row, and softmax is invariant to a constant added
+        across the whole row -- biasing K would look installed and cost real compute while
+        changing nothing. Biasing Q does not have that problem (q'.k_j = q.k_j + c.k_j, and
+        c.k_j varies per key), so Q is the only side of this that can actually move the
+        output. V was not picked either: a uniform additive V bias is mathematically
+        identical to adding straight to the block's output (attention weights sum to 1), which
+        is just h3_repr_steering's own mechanism wearing an attention-shaped costume, not a
+        different lever.
+
+        Only VIDEO-row queries are touched, matching h3_repr_steering's masking rationale
+        (H3 has no separate audio-attention path to spare -- see that module's docstring).
+        `steer_block` -- same range/list syntax as h3_block_repeat -- defaults to
+        h3_repr_steering.DEFAULT_BLOCK if blank; each named block captures and steers with
+        its own independently-learned direction. Assumes batch size 1 (this sampler's only
+        supported shape); no-ops with a one-time notice otherwise rather than guessing at a
+        batched reshape.
+        """
+        try:
+            _strength = float(strength or 0.0)
+        except (TypeError, ValueError):
+            _strength = 0.0
+        try:
+            try:
+                from . import h3_repr_steering as _rs
+            except ImportError:
+                import h3_repr_steering as _rs
+            steer_blocks = (self._parse_block_spec(steer_block)
+                            if steer_block else {_rs.DEFAULT_BLOCK})
+            if not steer_blocks:
+                return model
+            directions = {}
+            _counts = {}
+            for _sb in sorted(steer_blocks):
+                direction, n_liked, n_disliked = _rs.direction(refinement_key, block=_sb,
+                                                                kind="q_steer")
+                directions[_sb] = direction
+                _counts[_sb] = (n_liked, n_disliked)
+                if direction is None:
+                    print(f"[FunPackSceneChain] H3 query steering: not enough data yet "
+                          f"({n_liked} liked / {n_disliked} disliked, need "
+                          f"{_rs.MIN_PER_GROUP} of each at block {_sb}) — capturing this run, "
+                          "not steering it.")
+
+            patched = model.clone()
+            to = patched.model_options.get("transformer_options", {}).copy()
+            patches_replace = dict(to.get("patches_replace", {}))
+            dit_patches = dict(patches_replace.get("dit", {}))
+
+            _seg_holder = {"mod_segments": None}
+            _active_block = {"cur": None}
+            _warned_batch = {"on": False}
+
+            def _make_hook(block):
+                inner = dit_patches.get(("double_block", block))
+
+                def _hook(args, extra):
+                    _seg_holder["mod_segments"] = args.get("mod_segments")
+                    _active_block["cur"] = block
+                    try:
+                        return (extra["original_block"](args) if inner is None
+                                else inner(args, extra))
+                    finally:
+                        _active_block["cur"] = None
+                return _hook
+
+            for _b in steer_blocks:
+                dit_patches[("double_block", _b)] = _make_hook(_b)
+            patches_replace["dit"] = dit_patches
+            to["patches_replace"] = patches_replace
+
+            _mask_cache = {}
+            _inner_attn_override = to.get("optimized_attention_override")
+
+            def _next(func, q_, k_, v_, heads_, **kw):
+                return (_inner_attn_override(func, q_, k_, v_, heads_, **kw)
+                        if _inner_attn_override is not None
+                        else func(q_, k_, v_, heads_, **kw))
+
+            def _override(func, q, k, v, heads, mask=None, skip_reshape=True,
+                          transformer_options=None, **kwargs):
+                block = _active_block["cur"]
+                if block is None or mask is not None or not skip_reshape:
+                    return _next(func, q, k, v, heads, mask=mask, skip_reshape=skip_reshape,
+                                **kwargs)
+                if q.shape[0] != 1:
+                    if not _warned_batch["on"]:
+                        _warned_batch["on"] = True
+                        print("[FunPackSceneChain] H3 query steering: batch size > 1 is not "
+                              "supported by this mechanism — running unmodified this "
+                              "generation.")
+                    return _next(func, q, k, v, heads, mask=mask, skip_reshape=skip_reshape,
+                                **kwargs)
+                seq_len = q.shape[-2]
+                vmask = _mask_cache.get(seq_len, "MISS")
+                if vmask == "MISS":
+                    vmask = _rs.video_mask_from_mod_segments(
+                        _seg_holder["mod_segments"], seq_len, q.device)
+                    _mask_cache[seq_len] = vmask
+                if vmask is None or not bool(vmask.any()):
+                    return _next(func, q, k, v, heads, mask=mask, skip_reshape=skip_reshape,
+                                **kwargs)
+                # (B=1, H, S, D) -> (S, H*D) for the masked rows, so h3_repr_steering.capture()
+                # sees the same shape it already expects from a hidden state.
+                q_flat = q[0].permute(1, 0, 2).reshape(seq_len, -1)
+                desc = _rs.capture(q_flat, vmask)
+                if desc is not None:
+                    # Captured BEFORE injection, same reasoning as h3_repr_steering: a steered
+                    # block's descriptor must reflect the network's natural Q, not this run's
+                    # own strength setting.
+                    capture_holder[0][block] = desc
+                direction = directions.get(block)
+                if direction is not None and _strength > 0.0:
+                    head_dim = q.shape[-1]
+                    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
+                return _next(func, q, k, v, heads, mask=mask, skip_reshape=skip_reshape,
+                            **kwargs)
+
+            to["optimized_attention_override"] = _override
+            patched.model_options["transformer_options"] = to
+            _steered = sorted(b for b, d in directions.items() if d is not None)
+            if _steered:
+                _per_block = ", ".join(f"{b} ({_counts[b][0]}/{_counts[b][1]})"
+                                       for b in _steered)
+                print(f"[FunPackSceneChain] H3 query steering: applying learned direction(s) "
+                      f"at block (liked/disliked): {_per_block}, strength {_strength:g}.")
+            return patched
+        except Exception as _e:  # noqa: BLE001
+            _log.failed("FunPackSceneChain", "H3 query steering", _e,
+                        "not applied and not capturing this run")
+            return model
+
     @staticmethod
     def _parse_block_spec(spec, n_blocks=50):
         """"40" | "38-42" | "10,40,44" | "" -> a set of block indices. Anything unparseable
@@ -7714,6 +7887,7 @@ class FunPackLTXAVSceneChainSampler:
                h3_block_repeat="", h3_block_repeat_times=1,
                h3_block_repeat_video_only=False, h3_block_repeat_span_loop=False,
                h3_block_repeat_last_steps=0,
+               h3_q_steer_strength=0.0, h3_q_steer_block="0,1",
                alg_guide_blur_strength=2.0, alg_guide_blur_sigma_threshold=0.975,
                alg_anchor=False, alg_anchor_strength=2.0, alg_anchor_sigma_threshold=0.975,
                identity_transfer_enabled=False, identity_projector="None", source_id=2.0,
@@ -8617,6 +8791,8 @@ class FunPackLTXAVSceneChainSampler:
                     h3_block_repeat_video_only=h3_block_repeat_video_only,
                     h3_block_repeat_span_loop=h3_block_repeat_span_loop,
                     h3_block_repeat_last_steps=h3_block_repeat_last_steps,
+                    h3_q_steer_strength=h3_q_steer_strength,
+                    h3_q_steer_block=h3_q_steer_block,
                     refinement_key=refinement_key_input,
                 )
                 # cut_opening_frames: let the real, untouched i2v anchor condition the scene

tests/test_h3_q_steering.py

@@ -0,0 +1,141 @@
+"""H3 query steering (_install_h3_q_steering): REINS-style liked-minus-disliked direction
+captured and applied on Q instead of the block's hidden state. Exercised at the install/
+override level with fake tensors -- no real attention backend or GPU needed.
+
+Covers the two places a silent bug would hide: the (heads, head_dim) <-> flat-vector
+reshape used to reuse h3_repr_steering.capture()/direction() on Q, and that capture happens
+on the PRE-injection Q (same ordering rule h3_repr_steering itself follows)."""
+import sys
+from pathlib import Path
+
+import torch
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+sys.path.insert(0, ".")
+import _comfy_stubs  # noqa: E402, F401
+import h3_repr_steering as _rs  # noqa: E402
+from samplers import FunPackLTXAVSceneChainSampler as S  # noqa: E402
+
+
+class _FakeModel:
+    def __init__(self):
+        self.model_options = {}
+
+    def clone(self):
+        m = _FakeModel()
+        m.model_options = dict(self.model_options)
+        return m
+
+
+# 4 rows: 0 and 2 are video, 1 is text, 3 is audio -- non-contiguous video, same as
+# test_h3_av_decouple's layout.
+_MOD_SEGMENTS = [(0, 1, 0), (1, 2, 1), (2, 3, 0), (3, 4, 2)]
+
+
+def _run_attention_inside_block(patched, block, attn_fn, mod_segments=_MOD_SEGMENTS, seq_len=4):
+    """The active-block flag _install_h3_q_steering uses is only true WHILE the block's own
+    forward is running (attention happens inside it) -- so `attn_fn` (whatever calls the
+    optimized_attention_override) must run from inside `original_block`, not after the hook
+    returns, or the override sees no active block at all (a real trap this test caught once)."""
+    dit = patched.model_options["transformer_options"]["patches_replace"]["dit"]
+    hook = dit[("double_block", block)]
+    result = {}
+
+    def _original_block(a):
+        result["out"] = attn_fn()
+        return {"img": a["img"]}
+
+    args = {"img": torch.ones(seq_len, 3), "mod_segments": mod_segments}
+    hook(args, {"original_block": _original_block})
+    return result["out"]
+
+
+def _install(monkeypatch, direction_vec, strength=1.0, block=0):
+    """direction_vec: flat (heads*head_dim,) tensor, or None for 'not enough data yet'."""
+    def _fake_direction(refinement_key, block=None, kind="repr_steer"):
+        assert kind == "q_steer"
+        if direction_vec is None:
+            return None, 0, 0
+        return direction_vec, 5, 5
+    monkeypatch.setattr(_rs, "direction", _fake_direction)
+    capture_holder = [{}]
+    patched = S()._install_h3_q_steering(
+        _FakeModel(), "fake_key", strength, capture_holder, steer_block=str(block))
+    return patched, capture_holder
+
+
+def _override_fn(patched):
+    return patched.model_options["transformer_options"]["optimized_attention_override"]
+
+
+def test_capture_and_inject_at_named_block(monkeypatch):
+    # heads=2, head_dim=1, values chosen so row_norm and the injected shift are round numbers.
+    q = torch.zeros(1, 2, 4, 1)
+    q[0, 0, 0, 0] = 3.0
+    q[0, 0, 2, 0] = 3.0
+    q[0, 1, 0, 0] = 1.0
+    q[0, 1, 2, 0] = 1.0
+    k = torch.zeros(1, 2, 4, 1)
+    v = torch.zeros(1, 2, 4, 1)
+
+    patched, capture_holder = _install(monkeypatch, torch.tensor([1.0, -1.0]), strength=1.0)
+
+    calls = []
+
+    def fake_func(q_, k_, v_, heads_, mask=None, skip_reshape=True, **kw):
+        calls.append(q_.clone())
+        return torch.zeros(q_.shape[0], q_.shape[-2], heads_ * q_.shape[-1])
+
+    _run_attention_inside_block(
+        patched, block=0,
+        attn_fn=lambda: _override_fn(patched)(fake_func, q, k, v, 2, skip_reshape=True))
+
+    # Captured BEFORE injection: q_flat[video_rows].mean(dim=0) == [3.0, 1.0] (head0, head1).
+    assert torch.allclose(capture_holder[0][0], torch.tensor([3.0, 1.0]))
+
+    # row_norm = mean(|3|, |3|, |1|, |1|) = 2.0; shift = direction * strength * row_norm =
+    # [1, -1] * 1.0 * 2.0 = [+2, -2], applied identically to both masked rows, K/V untouched.
+    q_after = calls[0]
+    assert torch.allclose(q_after[0, 0, 0, 0], torch.tensor(5.0))   # 3 + 2
+    assert torch.allclose(q_after[0, 0, 2, 0], torch.tensor(5.0))   # 3 + 2
+    assert torch.allclose(q_after[0, 1, 0, 0], torch.tensor(-1.0))  # 1 - 2
+    assert torch.allclose(q_after[0, 1, 2, 0], torch.tensor(-1.0))  # 1 - 2
+    # Unmasked (text/audio) rows are untouched.
+    assert torch.allclose(q_after[0, :, 1, :], q[0, :, 1, :])
+    assert torch.allclose(q_after[0, :, 3, :], q[0, :, 3, :])
+
+
+def test_not_enough_data_captures_without_steering(monkeypatch):
+    q = torch.ones(1, 2, 4, 1) * 7.0
+    patched, capture_holder = _install(monkeypatch, None, strength=1.0)
+
+    def fake_func(q_, k_, v_, heads_, mask=None, skip_reshape=True, **kw):
+        return q_.clone()
+
+    out = _run_attention_inside_block(
+        patched, block=0,
+        attn_fn=lambda: _override_fn(patched)(
+            fake_func, q, torch.zeros_like(q), torch.zeros_like(q), 2, skip_reshape=True))
+    # No direction -> Q passed through unmodified, but still captured for next time.
+    assert torch.allclose(out, q)
+    assert 0 in capture_holder[0]
+
+
+def test_batch_greater_than_one_is_a_noop(monkeypatch, capsys):
+    q = torch.ones(2, 2, 4, 1)
+    patched, capture_holder = _install(monkeypatch, torch.tensor([1.0, -1.0]), strength=1.0)
+
+    def fake_func(q_, k_, v_, heads_, mask=None, skip_reshape=True, **kw):
+        return q_.clone()
+
+    out = _run_attention_inside_block(
+        patched, block=0,
+        attn_fn=lambda: _override_fn(patched)(
+            fake_func, q, torch.zeros_like(q), torch.zeros_like(q), 2, skip_reshape=True))
+    assert torch.allclose(out, q)
+    assert capture_holder[0] == {}
+    assert "batch size > 1" in capsys.readouterr().out
+
+
+if __name__ == "__main__":
+    print("run via pytest -- these tests need the monkeypatch/capsys fixtures")

tests/test_h3_repr_steering.py

@@ -9,7 +9,8 @@ import h3_repr_steering as rs  # noqa: E402
 
 @pytest.fixture(autouse=True)
 def _isolate_state(monkeypatch, tmp_path):
-    monkeypatch.setattr(rs, "state_path", lambda key: str(tmp_path / f"{key}.pt"))
+    monkeypatch.setattr(rs, "state_path",
+                        lambda key, kind="repr_steer": str(tmp_path / f"{key}.{kind}.pt"))
 
 
 def _pend(v, block=None):