""" MRI Live — streaming rank-capture for mid-inference monitoring. Two projection methods, side by side: - lm : residual / per-head vectors -> final_norm -> lm_head (the standard logit-lens readout; expensive, so by default I only project the 'Layer' residual column). - interp : the same vectors linearly interpolated up to vocab length and softmaxed — a projection-free readout that needs NO lm_head. Cheap enough that by default I project EVERY head + Layer. For each generated token it yields a TraceStep carrying, per active method, the rank of each tracked token at each (layer, col) cell, plus the projection latency for that method (so the cost difference is visible). Interventions: pass an InterventionConfig. The trigger is screened against each active method; if any fires I suppress/stop the (single, real) generation and record which method(s) tripped. """ from __future__ import annotations import os import sys import time from dataclasses import dataclass, field from typing import Callable, Iterator, Optional sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) import numpy as np import torch import torch.nn.functional as F from projection_map import load_model, build_col_specs # --------------------------------------------------------------------------- # Configs # --------------------------------------------------------------------------- @dataclass class TrackedWord: """A word to track. Tokenizes to 1..N tokens; anchor picks which one.""" text: str token_ids: list[int] anchor_idx: int = 0 @property def anchor_id(self) -> int: return self.token_ids[self.anchor_idx] @property def label(self) -> str: if len(self.token_ids) == 1: return self.text return f"{self.text} [tok{self.anchor_idx}]" @dataclass class InterventionConfig: """When to intervene, and how. trigger_fn: given a MethodTrace, returns True if intervention fires. action: 'log' | 'suppress' | 'stop' """ trigger_fn: Callable[["MethodTrace"], bool] action: str = "log" @dataclass class CellTrace: """Per-cell trace data for a single (layer, col).""" layer: int col: str ranks: list[int] # rank per tracked word, 1-indexed probs: list[float] # prob per tracked word top30_ids: list[int] = field(default_factory=list) top30_probs: list[float] = field(default_factory=list) @dataclass class MethodTrace: """One projection method's view of a single generated token.""" method: str # 'lm' | 'interp' cells: list[CellTrace] layers_used: list[int] cols_used: list[str] proj_ms: float = 0.0 def rank_matrix(self, tracked_idx: int) -> np.ndarray: n_layers, n_cols = len(self.layers_used), len(self.cols_used) m = np.full((n_layers, n_cols), 1, dtype=np.int32) lr = {l: i for i, l in enumerate(self.layers_used)} cc = {c: j for j, c in enumerate(self.cols_used)} for cell in self.cells: if cell.layer in lr and cell.col in cc: m[lr[cell.layer], cc[cell.col]] = cell.ranks[tracked_idx] return m def prob_matrix(self, tracked_idx: int) -> np.ndarray: n_layers, n_cols = len(self.layers_used), len(self.cols_used) m = np.zeros((n_layers, n_cols), dtype=np.float32) lr = {l: i for i, l in enumerate(self.layers_used)} cc = {c: j for j, c in enumerate(self.cols_used)} for cell in self.cells: if cell.layer in lr and cell.col in cc: m[lr[cell.layer], cc[cell.col]] = cell.probs[tracked_idx] return m @dataclass class TraceStep: """One generated token's full trace across the active method(s).""" step: int token_id: int token_str: str token_prob: float fwd_ms: float methods: dict # 'lm'/'interp' -> MethodTrace intervened: bool = False intervention_action: str = "" intervention_methods: list = field(default_factory=list) # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- def tokenize_word(tokenizer, word: str, with_leading_space: bool = True) -> list[int]: text = (" " + word) if with_leading_space else word return tokenizer.encode(text, add_special_tokens=False) def word_preview(tokenizer, word: str) -> dict: with_space = tokenize_word(tokenizer, word, with_leading_space=True) no_space = tokenize_word(tokenizer, word, with_leading_space=False) return { "with_leading_space": [(tid, tokenizer.decode([tid])) for tid in with_space], "no_leading_space": [(tid, tokenizer.decode([tid])) for tid in no_space], } def _is_cuda(device) -> bool: if isinstance(device, str): return device.startswith("cuda") return getattr(device, "type", "") == "cuda" # --------------------------------------------------------------------------- # Streaming trace # --------------------------------------------------------------------------- def stream_trace( model, tokenizer, device, num_layers: int, num_heads: int, prompt: str, tracked_words: list[TrackedWord], max_new_tokens: int = 30, layer_mask: Optional[list[int]] = None, lm_cols: Optional[list[str]] = None, interp_cols: Optional[list[str]] = None, methods: tuple = ("lm",), scale_mode: str = "full", intervention: Optional[InterventionConfig] = None, apply_chat_template: bool = True, precise_timing: bool = False, ) -> Iterator[TraceStep]: """Generate tokens one-by-one, yielding a TraceStep after each. methods: any of ('lm',), ('interp',), ('lm','interp'). lm_cols / interp_cols: column labels each method projects. Defaults: lm -> ['Layer'] (residual only; lm_head is expensive) interp -> all cols (every head + Layer; interp is cheap) """ col_specs, head_indices, _, _ = build_col_specs(num_heads) default_cols = [c for c in col_specs if c[0] in ("single", "layer")] all_labels = [c[2] for c in default_cols] if lm_cols is None: lm_cols = ["Layer"] if interp_cols is None: interp_cols = all_labels # everything def _resolve(wanted): return [c for c in default_cols if c[2] in set(wanted)] method_cols = {} if "lm" in methods: method_cols["lm"] = _resolve(lm_cols) if "interp" in methods: method_cols["interp"] = _resolve(interp_cols) if layer_mask is None: layer_mask = list(range(num_layers)) # Do I need per-head decomposition at all? (only if some active col is a head) need_heads = any( ct != "layer" for cols in method_cols.values() for (ct, _, _) in cols ) final_norm = model.model.norm lm_head = model.lm_head vocab_size = lm_head.weight.shape[0] tracked_ids = [tw.anchor_id for tw in tracked_words] tids_t = torch.tensor(tracked_ids, device=device) # Hook o_proj inputs (per-head pre-projection) — only used if need_heads. captured: dict[int, torch.Tensor] = {} def make_hook(layer_idx): def hook_fn(module, args, output): captured[layer_idx] = args[0].detach() return hook_fn handles = [] if need_heads: for l in range(num_layers): handles.append( model.model.layers[l].self_attn.o_proj.register_forward_hook(make_hook(l)) ) # ---- project a list of (layer, col_label, vec[hidden]) through a method ---- def project_method(entries, method): if not entries: return [], -1.0 if precise_timing and _is_cuda(device): torch.cuda.synchronize() t0 = time.time() stack = torch.stack([v for (_, _, v) in entries], dim=0) # [N, hidden], model dtype with torch.no_grad(): if method == "lm": dist = lm_head(final_norm(stack)).float() # [N, vocab] logits else: # interp — linear upsample to vocab; align_corners=True == np.interp(linspace) dist = F.interpolate( stack.float().unsqueeze(1), size=vocab_size, mode="linear", align_corners=True ).squeeze(1) # [N, vocab] sel = dist.index_select(1, tids_t) # [N, T] ranks = torch.empty((dist.shape[0], len(tracked_ids)), dtype=torch.long, device=dist.device) for ti in range(len(tracked_ids)): ranks[:, ti] = (dist > sel[:, ti:ti + 1]).sum(dim=1) + 1 # prob of the tracked tokens without materializing the full [N,vocab] softmax lse = torch.logsumexp(dist, dim=1, keepdim=True) # [N, 1] probs = torch.exp(sel - lse) # [N, T] if precise_timing and _is_cuda(device): torch.cuda.synchronize() proj_ms = (time.time() - t0) * 1000.0 if precise_timing else -1.0 ranks = ranks.cpu().numpy() probs = probs.float().cpu().numpy() cells = [] for i, (layer, col_label, _) in enumerate(entries): cells.append(CellTrace( layer=layer, col=col_label, ranks=[int(x) for x in ranks[i]], probs=[float(x) for x in probs[i]], )) return cells, proj_ms # ---- build input ---- if prompt and apply_chat_template: try: full_prompt = tokenizer.apply_chat_template( [{"role": "user", "content": prompt}], tokenize=False, add_generation_prompt=True ) except Exception: full_prompt = prompt else: full_prompt = prompt or "" if full_prompt: step_input = tokenizer(full_prompt, return_tensors="pt").input_ids.to(device) else: step_input = torch.tensor([[tokenizer.bos_token_id]], device=device) # KV cache: first pass processes the whole prompt; every later pass feeds # ONLY the new token with past_key_values, so generation is O(n) not O(n^2). # The projection code only ever reads the last position, so it is unchanged. past = None try: for step_idx in range(max_new_tokens): captured.clear() if precise_timing and _is_cuda(device): torch.cuda.synchronize() t_fwd = time.time() with torch.no_grad(): outputs = model(step_input, past_key_values=past, use_cache=True, output_hidden_states=True) past = outputs.past_key_values if precise_timing and _is_cuda(device): torch.cuda.synchronize() fwd_ms = (time.time() - t_fwd) * 1000.0 if precise_timing else -1.0 logits = outputs.logits[0, -1, :] probs_final = F.softmax(logits.float(), dim=-1) # Gather per-layer residual + (if needed) per-head vectors at last pos. residual_by_layer = {} head_by_layer = {} for l_idx in layer_mask: residual_by_layer[l_idx] = outputs.hidden_states[l_idx + 1][:, -1, :][0] if need_heads and l_idx in captured: inp = captured[l_idx] ow = model.model.layers[l_idx].self_attn.o_proj.weight hidden_size, attn_out_dim = ow.shape per_head_dim = attn_out_dim // num_heads # all heads at the last position in one shot, then slice: # proj[h, o] = sum_d o_proj_weight[o, h, d] * head_out[h, d] mh_last = inp[0, -1].view(num_heads, per_head_dim) # [H, d] wv = ow.view(hidden_size, num_heads, per_head_dim) # [O, H, d] proj = torch.einsum("ohd,hd->ho", wv, mh_last) # [H, O] head_by_layer[l_idx] = {h: proj[h] for h in head_indices} # Build entries + project, per method. step_methods = {} for m, cols in method_cols.items(): entries = [] for l_idx in layer_mask: hv = head_by_layer.get(l_idx, {}) for (col_type, col_heads, col_label) in cols: if col_type == "layer": entries.append((l_idx, col_label, residual_by_layer[l_idx])) else: if not hv: continue combined = sum(hv[h] for h in col_heads) if scale_mode == "full": combined = combined * (num_heads / len(col_heads)) elif scale_mode == "mean": combined = combined / len(col_heads) entries.append((l_idx, col_label, combined)) cells, proj_ms = project_method(entries, m) cols_used = [lbl for (_, _, lbl) in ([(0, 0, c[2]) for c in cols])] step_methods[m] = MethodTrace( method=m, cells=cells, layers_used=list(layer_mask), cols_used=[c[2] for c in cols], proj_ms=proj_ms, ) # Pick next token (greedy, from the real logits). top_prob, top_id = torch.max(probs_final, dim=-1) chosen_id = int(top_id.item()) chosen_str = tokenizer.decode(chosen_id) chosen_prob = float(top_prob.item()) step = TraceStep( step=step_idx, token_id=chosen_id, token_str=chosen_str, token_prob=chosen_prob, fwd_ms=fwd_ms, methods=step_methods, ) # Intervention — screen each active method. if intervention is not None: fired = [m for m, mt in step_methods.items() if intervention.trigger_fn(mt)] if fired: step.intervened = True step.intervention_action = intervention.action step.intervention_methods = fired if intervention.action == "suppress": masked = logits.clone() masked[chosen_id] = float("-inf") np2 = F.softmax(masked.float(), dim=-1) p2, i2 = torch.max(np2, dim=-1) chosen_id = int(i2.item()) chosen_str = tokenizer.decode(chosen_id) chosen_prob = float(p2.item()) step.token_id, step.token_str, step.token_prob = chosen_id, chosen_str, chosen_prob yield step if step.intervened and intervention.action == "stop": break # Next pass: feed ONLY the new token (the cache holds the rest). step_input = torch.tensor([[chosen_id]], device=device) if chosen_id == tokenizer.eos_token_id: break finally: for h in handles: h.remove() # --------------------------------------------------------------------------- # Threshold-based trigger (screened per method) # --------------------------------------------------------------------------- def make_rank_threshold_trigger( tracked_idx: int, layers: list[int], cols: list[str], rank_below: Optional[int] = None, rank_above: Optional[int] = None, ) -> Callable[["MethodTrace"], bool]: """Fire if the tracked token's rank at any matching cell crosses a threshold. `rank_below=10` -> fires when rank <= 10 (token surfaced).""" layers_s = set(layers) cols_s = set(cols) def trigger(mt: "MethodTrace") -> bool: for cell in mt.cells: if cell.layer not in layers_s or cell.col not in cols_s: continue r = cell.ranks[tracked_idx] if rank_below is not None and r <= rank_below: return True if rank_above is not None and r >= rank_above: return True return False return trigger