attention_routing.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
# This file is part of the Aquin Engine. Unauthorized copying, modification,
# distribution, or use of this file, via any medium, is strictly prohibited.
# Proprietary and confidential. See LICENSE for terms.
 
"""Per-head attention routing: sink + induction scores across all layers."""
 
from __future__ import annotations
 
from typing import Any
 
import torch
 
 
def _count_induction_pairs(token_ids: list[int]) -> int:
    """Positions (q, j) where token[j] repeats at q and j+1 is valid."""
    n = 0
    for q in range(len(token_ids)):
        for j in range(q):
            if token_ids[j] == token_ids[q] and j + 1 < len(token_ids):
                n += 1
    return n
 
 
def _score_head(
    pattern: torch.Tensor,
    token_ids: list[int],
) -> tuple[float, float]:
    """pattern: (seq, seq) causal attn weights for one head."""
    seq_len = pattern.shape[0]
 
    if seq_len > 1:
        sink = float(pattern[1:, 0].mean().item())
    else:
        sink = float(pattern[0, 0].item())
 
    ind_sum = 0.0
    ind_count = 0
    for q in range(seq_len):
        tq = token_ids[q]
        for j in range(q):
            if token_ids[j] == tq and j + 1 < seq_len:
                ind_sum += float(pattern[q, j + 1].item())
                ind_count += 1
    induction = ind_sum / max(ind_count, 1)
    return sink, induction
 
 
def _heads_from_pattern_stack(
    patterns: list[tuple[int, torch.Tensor]],
    token_ids: list[int],
) -> list[dict]:
    """patterns: list of (layer_idx, tensor [n_heads, seq, seq])."""
    heads: list[dict] = []
    for layer, attn in patterns:
        if attn is None:
            continue
        t = attn.detach().float()
        if t.ndim == 4:
            t = t[0]
        if t.ndim != 3:
            continue
        n_heads = t.shape[0]
        for h in range(n_heads):
            sink, induction = _score_head(t[h], token_ids)
            heads.append(
                {
                    "layer": int(layer),
                    "head": int(h),
                    "sink_score": round(sink, 4),
                    "induction_score": round(induction, 4),
                }
            )
    return heads
 
 
def _patterns_from_tl_cache(model: Any, tokens: torch.Tensor, n_layers: int) -> list[tuple[int, torch.Tensor]]:
    with torch.no_grad():
        _, cache = model.run_with_cache(
            tokens,
            names_filter=lambda name: name.endswith("attn.hook_pattern"),
        )
    out: list[tuple[int, torch.Tensor]] = []
    for layer in range(n_layers):
        key = f"blocks.{layer}.attn.hook_pattern"
        if key not in cache:
            continue
        attn = cache[key]
        if isinstance(attn, torch.Tensor):
            out.append((layer, attn[0] if attn.ndim == 4 else attn))
    return out
 
 
def _patterns_from_hf_output_attentions(
    model: Any, tokens: torch.Tensor
) -> tuple[list[tuple[int, torch.Tensor]], str | None]:
    """
    HF-native / HfLlmShim path.
    TransformerLens patterns are unavailable on hf_only models (LFM, …).
    """
    hf = getattr(model, "hf_model", None)
    if hf is None:
        return [], None
    try:
        with torch.no_grad():
            out = hf(
                input_ids=tokens,
                output_attentions=True,
                use_cache=False,
                return_dict=True,
            )
    except TypeError:
        try:
            with torch.no_grad():
                out = hf(input_ids=tokens, output_attentions=True, return_dict=True)
        except Exception as exc:  # noqa: BLE001
            return [], f"output_attentions failed: {exc}"
    except Exception as exc:  # noqa: BLE001
        return [], f"output_attentions failed: {exc}"
 
    attns = getattr(out, "attentions", None)
    if not attns:
        return [], (
            "Model returned no attentions (output_attentions unused or architecture "
            "does not expose multi-head attention weights)."
        )
 
    patterns: list[tuple[int, torch.Tensor]] = []
    for layer, layer_attn in enumerate(attns):
        if layer_attn is None:
            continue  # e.g. non-attention / conv hybrid layers
        patterns.append((layer, layer_attn[0] if layer_attn.ndim == 4 else layer_attn))
    if not patterns:
        return [], "All layer attentions were empty (hybrid/conv-only blocks?)."
    return patterns, None
 
 
def run_attention_routing(
    model: Any,
    prompt: str,
    *,
    top_k: int = 5,
    model_id: str = "",
) -> dict:
    tokens = model.to_tokens(prompt)
    if tokens.ndim == 1:
        tokens = tokens.unsqueeze(0)
    token_ids = tokens[0].tolist()
    n_layers = int(getattr(getattr(model, "cfg", None), "n_layers", 0) or 0)
    n_heads = int(getattr(getattr(model, "cfg", None), "n_heads", 0) or 0)
 
    source = "tl_cache"
    warning: str | None = None
    patterns = _patterns_from_tl_cache(model, tokens, n_layers) if n_layers else []
 
    if not patterns:
        patterns, warning = _patterns_from_hf_output_attentions(model, tokens)
        source = "hf_output_attentions"
 
    heads = _heads_from_pattern_stack(patterns, token_ids)
 
    if heads and (n_layers <= 0 or n_heads <= 0):
        n_layers = max(h["layer"] for h in heads) + 1
        n_heads = max(h["head"] for h in heads) + 1
 
    by_sink = sorted(heads, key=lambda x: x["sink_score"], reverse=True)
    by_ind = sorted(heads, key=lambda x: x["induction_score"], reverse=True)
    n_ind_pairs = _count_induction_pairs(token_ids)
 
    result: dict[str, Any] = {
        "prompt": prompt,
        "model_id": model_id,
        "n_layers": n_layers,
        "n_heads": n_heads,
        "n_induction_pairs": n_ind_pairs,
        "heads": heads,
        "top_sink_heads": by_sink[:top_k],
        "top_induction_heads": by_ind[:top_k] if n_ind_pairs else [],
        "source": source,
    }
    if not heads:
        result["warning"] = warning or (
            "No per-head attention patterns available for this model. "
            "TransformerLens hook_pattern cache was empty and HF output_attentions "
            "did not return usable weights."
        )
    elif warning:
        result["warning"] = warning
    return result