batch the eval: guided_rollout_batch + 4× speedup

Sequential eval was the bottleneck (~12s/vignette × 131 = 27 min/pass; with
bidirectional ±C × 14 methods, that projected to ~17 h). Three model calls
per row (phase1 generate, scoring forward, cosmetic continuation) became one
phase1 + one scoring per *batch*; continuation generate dropped (callers
only use p_true + pmass_format).

Parity smoke (scripts/smoke_batch_parity.py): float32 is bit-exact (max
Δp_true=0.0000 over 16 prompts, 4.3× speedup at limit=4). bf16 drifts on
individual rows — greedy argmax flips at near-tie tokens then phase1
diverges — but pmass agrees within 0.03 (scoring forward correct) and
aggregates over 131 vignettes will average out the per-row noise.

Other changes shipped in this commit:
- core.py: analyse() now returns raw_pmass dict alongside raw p_true (callers
  needed per-(vid,cond,frame) pmass for diagnostic warnings).
- guided.py guided_rollout: warn + log top-5 when pmass<0.9 (catches OOD
  steering / format-broken vignettes without a separate audit pass).
- eval.py: pre-tokenize a sample to log expected prompt+cache budget so OOM
  is predictable from the SHOULD line; group items by frame so each batch
  shares schema_hint + prefill.

scripts/smoke_batch_parity.py

@@ -0,0 +1,122 @@
+"""Parity smoke: guided_rollout vs guided_rollout_batch on a small vignette subset.
+
+Asserts p_true and pmass_format match within fp tolerance. Same chat template,
+same prompts, same model, same generation kwargs -- only batching differs.
+
+usage:
+    uv run python scripts/smoke_batch_parity.py --model Qwen/Qwen3-0.6B --limit 4
+"""
+from __future__ import annotations
+import argparse
+import time
+
+import torch
+from loguru import logger
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from tinymfv.core import CONDITIONS, FRAMES
+from tinymfv.data import load_vignettes
+from tinymfv.guided import guided_rollout, guided_rollout_batch, choice_token_ids_tf
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--model", default="Qwen/Qwen3-0.6B")
+    ap.add_argument("--limit", type=int, default=4)
+    ap.add_argument("--max-think-tokens", type=int, default=32)
+    ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    ap.add_argument("--dtype", default="bfloat16")
+    args = ap.parse_args()
+
+    rows = load_vignettes("")[: args.limit]
+    logger.info(f"{len(rows)} vignettes; testing parity")
+
+    dtype = getattr(torch, args.dtype)
+    tok = AutoTokenizer.from_pretrained(args.model)
+    if tok.pad_token is None:
+        tok.pad_token = tok.eos_token
+    tok.padding_side = "left"
+    model = AutoModelForCausalLM.from_pretrained(args.model, dtype=dtype).to(args.device).eval()
+
+    choice_ids = choice_token_ids_tf(tok)
+
+    # --- Sequential ---
+    t0 = time.time()
+    seq_results = []  # list of (vid, cond, frame, p_true, pmass)
+    for r in rows:
+        for cond in CONDITIONS:
+            for frame, fr in FRAMES.items():
+                res = guided_rollout(
+                    model, tok,
+                    user_prompt=r[cond],
+                    choice_token_ids=choice_ids,
+                    max_think_tokens=args.max_think_tokens,
+                    schema_hint=fr["q"],
+                    prefill=fr["prefill"],
+                )
+                seq_results.append((r["id"], cond, frame, res.p_true, res.pmass_format))
+    seq_elapsed = time.time() - t0
+    logger.info(f"sequential: {seq_elapsed:.1f}s ({len(seq_results)} prompts)")
+
+    # --- Batched ---
+    t0 = time.time()
+    batch_results = []
+    for frame, fr in FRAMES.items():
+        for cond in CONDITIONS:
+            user_prompts = [r[cond] for r in rows]
+            outs = guided_rollout_batch(
+                model, tok,
+                user_prompts=user_prompts,
+                choice_token_ids=choice_ids,
+                max_think_tokens=args.max_think_tokens,
+                schema_hint=fr["q"],
+                prefill=fr["prefill"],
+            )
+            for r, o in zip(rows, outs):
+                batch_results.append((r["id"], cond, frame, o.p_true, o.pmass_format))
+    batch_elapsed = time.time() - t0
+    logger.info(f"batched: {batch_elapsed:.1f}s (speedup={seq_elapsed/batch_elapsed:.1f}x)")
+
+    # --- Compare ---
+    seq_d = {(vid, c, f): (pt, pm) for vid, c, f, pt, pm in seq_results}
+    batch_d = {(vid, c, f): (pt, pm) for vid, c, f, pt, pm in batch_results}
+    assert set(seq_d) == set(batch_d), "key mismatch"
+
+    n = 0
+    max_pt_diff, max_pm_diff = 0.0, 0.0
+    rows_out = []
+    for k in seq_d:
+        spt, spm = seq_d[k]
+        bpt, bpm = batch_d[k]
+        d_pt = abs(spt - bpt)
+        d_pm = abs(spm - bpm)
+        max_pt_diff = max(max_pt_diff, d_pt)
+        max_pm_diff = max(max_pm_diff, d_pm)
+        rows_out.append((k, spt, bpt, d_pt, spm, bpm, d_pm))
+        n += 1
+
+    from tabulate import tabulate
+    print()
+    print(tabulate(
+        [(f"{k[0][:8]}|{k[1]}|{k[2]}", spt, bpt, d_pt, spm, bpm, d_pm)
+         for (k, spt, bpt, d_pt, spm, bpm, d_pm) in rows_out],
+        headers=["key", "p_true_seq", "p_true_bat", "Δp_true", "pm_seq", "pm_bat", "Δpm"],
+        floatfmt="+.4f", tablefmt="tsv",
+    ))
+
+    # bf16 batched greedy decoding can pick different argmax than per-row greedy
+    # when two tokens tie within bf16 precision. The phase1 think rollout then
+    # diverges and per-row p_true drifts. float32 is bit-exact (use --dtype float32
+    # to verify the batching logic itself). At aggregate eval (131 vignettes
+    # averaged) the bf16 drift averages out; we accept it.
+    TOL = 0.20 if args.dtype != "float32" else 0.001
+    cue = "🟢" if (max_pt_diff < TOL and max_pm_diff < TOL) else "🔴"
+    print(f"\n{cue} max Δp_true={max_pt_diff:.4f}  max Δpmass={max_pm_diff:.4f}  (tol={TOL})")
+    print(f"speedup: {seq_elapsed/batch_elapsed:.1f}x ({len(seq_results)} prompts)")
+
+    if max_pt_diff >= TOL or max_pm_diff >= TOL:
+        raise SystemExit(f"PARITY FAILED: Δp_true={max_pt_diff:.4f} Δpmass={max_pm_diff:.4f}")
+
+
+if __name__ == "__main__":
+    main()

src/tinymfv/core.py

@@ -214,10 +214,15 @@ def analyse(
     if bool_mass is not None:
         info["bool_mass_mean"] = float(sum(map(float, bool_mass)) / len(bool_mass))
 
+    raw_pmass = (
+        {f"{vid}|{cond}|{frame}": float(b) for (vid, _, cond, frame, _), b in zip(meta, bool_mass)}
+        if bool_mass is not None else {}
+    )
     return {
         "wrongness": float(df["s_other_violate"].mean()),
         "gap": float(df["gap"].mean()),
         "table": df,
         "raw": {f"{vid}|{cond}|{frame}": p for (vid, _, cond, frame, _), p in zip(meta, p_true)},
+        "raw_pmass": raw_pmass,
         "info": info,
     }

src/tinymfv/eval.py

@@ -9,7 +9,7 @@ from tqdm.auto import tqdm
 
 from .core import format_prompts, next_token_logits, score_prompts, analyse, CONDITIONS, FRAMES
 from .data import load_vignettes
-from .guided import guided_rollout, choice_token_ids_tf
+from .guided import guided_rollout_batch, choice_token_ids_tf
 
 
 def evaluate(
@@ -39,36 +39,63 @@ def evaluate(
     t0 = time.time()
     
     if max_think_tokens > 0:
-        logger.info(f"Using guided_rollout with {max_think_tokens} max_think_tokens (sequential)")
-        p_true_list = []
-        meta = []
-        bool_mass_list = []
+        logger.info(f"Using guided_rollout_batch with {max_think_tokens} max_think_tokens, batch_size={batch_size}")
         choice_ids = choice_token_ids_tf(tokenizer)
-        
-        for r in tqdm(vignettes, desc="Evaluating"):
+
+        # Build all (vid, cond, frame) items, grouped by frame so each batch
+        # shares schema_hint + prefill (collapses the per-row branching).
+        items_per_frame: dict[str, list[tuple]] = {f: [] for f in FRAMES}
+        for r in vignettes:
             for cond in CONDITIONS:
-                for frame, fr in FRAMES.items():
-                    user_prompt = f"{r[cond]}"
-                    schema_hint = fr["q"]
-                    prefill = fr["prefill"]
-                    
-                    res = guided_rollout(
+                for frame in FRAMES:
+                    items_per_frame[frame].append(
+                        (r["id"], r["foundation_coarse"], cond, frame, r.get("wrong"), r[cond])
+                    )
+
+        # Pretokenize a sample to log expected prompt length / cache budget.
+        sample_user = items_per_frame[next(iter(FRAMES))][0][5]
+        sample_q = FRAMES[next(iter(FRAMES))]["q"]
+        sample_full = f"{sample_user}\n\n{sample_q}"
+        sample_msgs = [{"role": "user", "content": sample_full}]
+        try:
+            sample_p = tokenizer.apply_chat_template(sample_msgs, tokenize=False, add_generation_prompt=True)
+        except TypeError:
+            sample_p = tokenizer.apply_chat_template(sample_msgs, tokenize=False)
+        sample_p = sample_p + "<think>\n"
+        sample_len = len(tokenizer(sample_p).input_ids)
+        logger.info(
+            f"SHOULD: prompt_len≈{sample_len} tok; max cache ≈ {sample_len + max_think_tokens} per row × "
+            f"batch_size={batch_size}. If OOM, lower batch_size."
+        )
+
+        p_true_list, meta, bool_mass_list = [], [], []
+        total = sum(len(v) for v in items_per_frame.values())
+        with tqdm(total=total, desc="Evaluating") as pbar:
+            for frame, items in items_per_frame.items():
+                fr = FRAMES[frame]
+                schema_hint = fr["q"]
+                prefill = fr["prefill"]
+                for i in range(0, len(items), batch_size):
+                    chunk = items[i:i + batch_size]
+                    user_prompts = [it[5] for it in chunk]
+                    results = guided_rollout_batch(
                         model, tokenizer,
-                        user_prompt=user_prompt,
+                        user_prompts=user_prompts,
                         choice_token_ids=choice_ids,
                         max_think_tokens=max_think_tokens,
                         schema_hint=schema_hint,
                         prefill=prefill,
-                        verbose=False
                     )
-                    
-                    p_true_list.append(res.p_true)
-                    meta.append((r["id"], r["foundation_coarse"], cond, frame, r.get("wrong")))
-                    bool_mass_list.append(res.pmass_format)
-        
+                    for it, res in zip(chunk, results):
+                        vid, found, cond, fr_name, wrong, _ = it
+                        p_true_list.append(res.p_true)
+                        meta.append((vid, found, cond, fr_name, wrong))
+                        bool_mass_list.append(res.pmass_format)
+                    pbar.update(len(chunk))
+
         elapsed = time.time() - t0
         logger.info(f"guided eval: {elapsed:.1f}s ({len(p_true_list)/elapsed:.1f} prompts/s)")
-        
+
         report = analyse(p_true_list, meta, bool_mass=bool_mass_list)
         
     else:

src/tinymfv/guided.py

@@ -108,7 +108,17 @@ def guided_rollout(
         
     all_ids = torch.tensor(a_ids + b_ids, device=device, dtype=torch.long)
     pmass_format = float(logp[all_ids].exp().sum().item())
-    
+
+    # SHOULD: pmass≈1 (model picks one of the JSON-bool tokens). pmass<0.9
+    # means the model is leaking probability to other tokens -> the schema
+    # is being ignored or the steering vector has pushed the model OOD.
+    if pmass_format < 0.9:
+        topk = torch.topk(logp.exp(), k=5)
+        toks = [tok.decode([i]) for i in topk.indices.tolist()]
+        probs = topk.values.tolist()
+        top5 = ", ".join(f"{repr(t)}={p:.3f}" for t, p in zip(toks, probs))
+        logger.warning(f"pmass={pmass_format:.3f}<0.9 — top-5: {top5}")
+
     if a_ids and b_ids:
         a_t = torch.tensor(a_ids, device=device, dtype=torch.long)
         b_t = torch.tensor(b_ids, device=device, dtype=torch.long)
@@ -143,6 +153,128 @@ def guided_rollout(
         p_true=p_true,
     )
 
+@torch.no_grad()
+def guided_rollout_batch(
+    model, tok,
+    user_prompts: list[str],
+    choice_token_ids: list,
+    max_think_tokens: int = 128,
+    schema_hint: str = _DEFAULT_SCHEMA_HINT,
+    prefill: str = '\n{"choice": ',
+) -> list[GuidedResult]:
+    """Batched guided rollout. Same logic as guided_rollout but over a list of
+    user_prompts that share schema_hint + prefill (so prefill cases collapse).
+
+    Skips the cosmetic answer-continuation generate (caller only needs p_true,
+    pmass_format, think_text). Two model calls per batch instead of 3 per row:
+    one phase1 generate (think) + one scoring forward.
+
+    Tokenizer must already have padding_side='left' and pad_token set."""
+    if tok.padding_side != "left":
+        raise ValueError("tok.padding_side must be 'left' for batched rollout")
+    device = next(model.parameters()).device
+
+    prompts = []
+    for up in user_prompts:
+        full_user = f"{up}\n\n{schema_hint}" if schema_hint else up
+        msgs = [{"role": "user", "content": full_user}]
+        try:
+            p = tok.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+        except TypeError:
+            p = tok.apply_chat_template(msgs, tokenize=False)
+        prompts.append(p + "<think>\n")
+
+    think_end_id = tok.convert_tokens_to_ids("</think>")
+    if think_end_id in (None, getattr(tok, "unk_token_id", None)):
+        think_end_id = tok.eos_token_id
+    pad_id = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id
+
+    enc = tok(prompts, return_tensors="pt", padding=True).to(device)
+    prompt_len = enc.input_ids.shape[1]
+
+    phase1 = model.generate(
+        **enc,
+        max_new_tokens=max_think_tokens,
+        do_sample=False,
+        eos_token_id=think_end_id,
+        pad_token_id=pad_id,
+    )
+
+    scoring_texts = []
+    per_row = []  # (think_text, emitted_close, emitted_prefill, n_think_tokens)
+    for i, p in enumerate(prompts):
+        gen_ids = phase1[i, prompt_len:]
+        keep = gen_ids != pad_id
+        gen_ids = gen_ids[keep] if keep.any() else gen_ids[:0]
+        gen_text = tok.decode(gen_ids, skip_special_tokens=True)
+        n_think = int(gen_ids.shape[0])
+
+        emitted_close = _CLOSE_MARKER in gen_text
+        if emitted_close:
+            think_text, after = gen_text.split(_CLOSE_MARKER, 1)
+            if prefill.lstrip() in after:
+                emitted_prefill = True
+                before_value = after.split(prefill.lstrip(), 1)[0]
+                scoring_text = p + think_text + _CLOSE_MARKER + before_value + prefill.lstrip()
+            else:
+                emitted_prefill = False
+                scoring_text = p + think_text + _CLOSE_MARKER + prefill
+        else:
+            think_text = gen_text
+            emitted_prefill = False
+            force_suffix = "\nI should answer now." + _CLOSE_MARKER + prefill
+            scoring_text = p + gen_text + force_suffix
+
+        scoring_texts.append(scoring_text)
+        per_row.append((think_text, emitted_close, emitted_prefill, n_think))
+
+    score_enc = tok(scoring_texts, return_tensors="pt", padding=True,
+                    add_special_tokens=False).to(device)
+    score_logits = model(**score_enc).logits[:, -1].float()
+    score_logp = F.log_softmax(score_logits, dim=-1)
+
+    if (len(choice_token_ids) == 2 and all(isinstance(x, (list, tuple)) for x in choice_token_ids)):
+        a_ids, b_ids = list(choice_token_ids[0]), list(choice_token_ids[1])
+    else:
+        a_ids, b_ids = list(choice_token_ids), []
+    all_ids = torch.tensor(a_ids + b_ids, device=device, dtype=torch.long)
+    a_t = torch.tensor(a_ids, device=device, dtype=torch.long) if a_ids else None
+    b_t = torch.tensor(b_ids, device=device, dtype=torch.long) if b_ids else None
+
+    results = []
+    for i, (up, (think_text, emitted_close, emitted_prefill, n_think)) in enumerate(zip(user_prompts, per_row)):
+        logp = score_logp[i]
+        pmass_format = float(logp[all_ids].exp().sum().item())
+        if pmass_format < 0.9:
+            topk = torch.topk(logp.exp(), k=5)
+            toks = [tok.decode([j]) for j in topk.indices.tolist()]
+            probs = topk.values.tolist()
+            top5 = ", ".join(f"{repr(t)}={pp:.3f}" for t, pp in zip(toks, probs))
+            logger.warning(f"pmass={pmass_format:.3f}<0.9 — top-5: {top5}")
+        if a_t is not None and b_t is not None:
+            la = torch.logsumexp(logp[a_t], dim=0)
+            lb = torch.logsumexp(logp[b_t], dim=0)
+            logratio = float((la - lb).item())
+            p_true = float(torch.softmax(torch.stack([la, lb]), dim=0)[0].item())
+        else:
+            logratio = float("nan")
+            p_true = float("nan")
+        results.append(GuidedResult(
+            user_prompt=up,
+            think_text=think_text,
+            answer_text="",
+            raw_full_text="",
+            pmass_format=pmass_format,
+            logratio_ab=logratio,
+            rep_ratio_think=_ngram_rep_ratio(think_text, n=4),
+            think_tokens=n_think,
+            emitted_close=emitted_close,
+            emitted_prefill=emitted_prefill,
+            p_true=p_true,
+        ))
+    return results
+
+
 def choice_token_ids_tf(tok) -> list[list[int]]:
     def _variants(words):
         seen = []