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.0, 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.0

Usage

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

checkpoint = "Qwen/Qwen3-4B"
revision = "1cfa9a7208912126459214e8b04321603b3df60c"
model = AutoModelForCausalLM.from_pretrained(
    checkpoint, revision=revision, dtype=torch.bfloat16,
    attn_implementation="sdpa",
).to("cuda").eval().requires_grad_(False)
tokenizer = AutoTokenizer.from_pretrained(checkpoint, revision=revision)
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.

This is structural compatibility, not a promise that every HF architecture works. The supported operation must be linear in V, with attention weights independent of V, and a conventional merged-head input to the output projection. Tested families include Qwen3, Llama, Mistral (sliding window), Gemma2 (score softcapping), GPT-2, LLaVA/CLIP vision with pixels, and Qwen3-VL packed vision/text. Unsupported projection layouts fail explicitly. Cross-attention, nonlinear/value-dependent attention, quantized/compiled models, distributed/sharded execution and remote NNsight are outside the verified scope. No weights or inputs are downloaded by the extraction call. Model weights stay unchanged; each probe restores the genuine pre-output-projection activation before the rest of the forward runs.

Inputs and results. Use nonempty input_ids or inputs_embeds, no KV cache, and output_attentions=False. Sequence length comes from the actual attention activation, accounting for expanded image tokens. Parameter gradients should be disabled for inference to reduce NNsight memory. Do not run concurrent calls on the same model. Qwen-VL vision packs patches from all images/frames on one axis; its result is [1, packed_tokens, packed_tokens], with separate images/frames forming disconnected blocks. Interpret coordinates with the processor's grids; these are patch tokens, not merged decoder image tokens.

Checks contain module paths/layouts, selected addresses, probe count, and maximum absolute logit deltas for traced baselines, probes and a final unpatched pass. Inspect the deltas: nondeterministic execution may produce nonzero values even without a probe. Recovered weights are finite-precision measurements and are not renormalized or guaranteed bit-exact native softmax probabilities. Masks and Q/K-only score modifications stay inside the original attention implementation.

Backend. sdpa uses normal PyTorch backend selection. Force math by wrapping extraction in torch.nn.attention.sdpa_kernel(SDPBackend.MATH). Package FA2 requires a separately installed compatible flash-attn and attn_implementation="flash_attention_2". The library does not install FA2 or override precision/backend.

Cost. For each selected layer with prefix T and head dimension d, reconstruction needs 1 + ceil(T/d) full forwards, plus two global unpatched forwards. Selected-head CPU storage scales as B×selected_heads×T² per layer; this remains quadratic. NNsight brings dependencies, but the package itself contains only the extractor.

Compatibility. Initially constrained to NNsight 0.7.0, Transformers 5.19.x and PyTorch 2.14.x; tested on Python 3.11. Experiment pins live in the research repository’s requirements.txt; the package has no CUDA URLs or plotting dependencies.

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/*

On Oscar, run installs, builds and tests in a CPU Slurm allocation; pretrained GPU examples require a GPU allocation. The source checkout includes the release guide and review notebook; the notebook uses in-memory example inputs and is excluded from release archives. No publishing step runs automatically.

License

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