AttnRead

Recover attention weights without changing the attention backend.

AttnRead is a lightweight Python toolkit for extracting attention patterns from optimized transformer attention implementations, including PyTorch SDPA and FlashAttention 2.

Using identity-value passes and activation restoration, AttnRead recovers attention patterns while preserving the model's downstream computation. No fallback to eager attention required.

Version 0.1.1, distributed under the MIT License. The distribution and Python import name are both attnread.

Installation

Use Python 3.11 or newer in a venv. Install your preferred CPU/CUDA PyTorch build first, then install AttnRead:

python -m pip install attnread==0.1.1

Usage

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from attnread import extract_attention

checkpoint = "Qwen/Qwen3-4B"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval().requires_grad_(False)
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer("Describe the colors of an autumn leaf.",
                   return_tensors="pt").to(model.device)

patterns, checks = extract_attention(model, inputs, heads=[(0, 3), (12, 7)])
head = patterns[("text", 0, 3)]  # CPU tensor: [batch, query tokens, key tokens].
assert checks["probe_logit_max_abs"] == 0

Use the same call with AutoModelForImageTextToText and processor outputs:

# (layer, head) means text attention; qualify a vision head with its stream.
patterns, checks = extract_attention(
    vlm, processed_inputs, heads=[(12, 3), ("vision", 4, 2)],
)
text = patterns[("text", 12, 3)]
vision = patterns[("vision", 4, 2)]

Selection. heads=None extracts every discovered text head. A sequence of (layer, head) or (stream, layer, head) addresses extracts only those heads; returned keys always use the full triple, in request order. Layers and heads are zero-based within each stream, following discovered module order (reported in checks["modules"]). Use explicit paths for sparse/hybrid stacks or unusual module naming. Different streams can have different sequence lengths, so results are a dictionary. Probes are grouped by layer and save/transfer only selected heads. The underlying model still computes all heads; head selection reduces saved memory, not the number of identity probes. Start with one short prefix and one head. Unknown, duplicate or out-of-range addresses raise ValueError.

Model interface. Pass an ordinary HF model in evaluation mode and its preprocessed tensor input mapping. Discovery recognizes separate q_proj, k_proj, v_proj with o_proj/out_proj, GPT-2's fused c_attn/c_proj, and Qwen3-VL vision qkv/proj. Vision module paths identify the vision stream; others belong to text. For unusual or ambiguous naming, supply explicit paths:

patterns, checks = extract_attention(model, inputs, heads=[("text", 0, 1)],
    attention_modules={"text": ["model.layers.7.self_attn"]})
# With explicit lists, layer 0 denotes the first entry in that list.

Dependencies. AttnRead requires Python 3.11 or newer, NNsight 0.7.0, Transformers 5.19.x, and PyTorch 2.14.x. Install a PyTorch build suitable for your CPU or GPU. FlashAttention 2 is optional and installed separately.

Development

From the standalone source directory:

python -m venv .venv
source .venv/bin/activate
# Install your preferred CPU/CUDA PyTorch build first.
python -m pip install -e '.[dev]'
python -m pytest -q test_attnread.py
python -m build
python -m twine check --strict dist/*

The source distribution includes the package tests and release guide. The example notebook uses in-memory inputs and is excluded from release archives.

License

MIT License. Copyright (c) 2026 Etha Tianze Hua. The full license text is included in the source archive and wheel.