Metadata-Version: 2.4
Name: attnread
Version: 0.1.0
Summary: Recover attention weights without changing the attention backend.
License-Expression: MIT
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.11
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.11
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: nnsight==0.7.0
Requires-Dist: torch<2.15,>=2.14
Requires-Dist: transformers<5.20,>=5.19
Provides-Extra: dev
Requires-Dist: pytest<10,>=9; extra == "dev"
Requires-Dist: build<2,>=1.2; extra == "dev"
Requires-Dist: twine<7,>=6; extra == "dev"
Dynamic: license-file

# 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:

```bash
python -m pip install attnread==0.1.0
```

## Usage

```python
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:

```python
# (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:

```python
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:

```bash
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.
