attnread 0.1.0__tar.gz

attnread-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Etha
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

attnread-0.1.0/MANIFEST.in

@@ -0,0 +1,11 @@
+include LICENSE README.md PACKAGE.md RELEASE.md
+recursive-include src/attnread *.py py.typed
+include test_attnread.py
+exclude BLOG.md AGENTS.md compare.py test_compare.py test_value_patch.py value_patch.py
+prune figures
+prune outputs
+prune .venv
+prune .git
+prune __pycache__
+prune .pytest_cache
+global-exclude *.ipynb *.pdf *.png *.jpg *.jpeg *.gif *.webp

attnread-0.1.0/PACKAGE.md

@@ -0,0 +1,151 @@
+# 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`.
+
+```bash
+python -m venv .venv
+# Activate the venv and install your preferred CPU/CUDA PyTorch build first.
+python -m pip install .
+```
+
+```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.
+The expanded BF16 GPU validation covers 240 input/backend cases: 20 COCO images,
+20 standalone Qwen prompts, 20 GPT-2 prompts, and the same 20 prompts through the
+VLM text decoder, each with eager, ordinary SDPA and package FA2. It checks selected
+decoder heads at early, middle and final layers, plus packed vision layer/head 0
+for image inputs, against original and independent random V. Direct native eager
+weights provide a separate reference for text heads. These checks do not establish
+every head or every backend combination for every family. Generic adapters are
+also checked in small real HF models on CPU. The research repository
+contains the blog, numerical errors and provenance.
+
+See [the first-release guide](RELEASE.md).
+
+**Independent validation.** `verify_attention.py` is a research script, outside
+this package. It captures projected V and the native pre-output-projection
+activation with NNsight, then compares that activation with recovered A × V in
+float64. Three held-out Gaussian V replacements per layer additionally test the
+recovered matrix on values never used for extraction. Q/K and restored logits
+are checked for equality. Reversing A's key columns supplies a deliberately
+incorrect pattern as a sensitivity control. Small FP32 models also compare
+against directly returned eager weights. Original-V agreement alone cannot
+uniquely identify A when V has insufficient rank; finite random trials add
+evidence rather than a proof of equality for every possible V.
+
+From a separate research repository checkout, on a Slurm GPU allocation with
+that repository's venv and saved benchmark inputs:
+
+```bash
+.venv/bin/python verify_attention.py --output outputs/attention-validation
+```
+
+The JSON report records errors per stream/layer/head, seeds, input hashes,
+checkpoint revisions and package versions. Adjacent `.pt` files contain A,
+original V/output tensors, and the held-out V/output tensors for auditing.
+
+Each backend is additionally compared with the **identity-extracted eager
+pattern** on the same input: both A-versus-A errors and predictions using eager
+A with that backend's actual original/held-out V. This distinguishes within-
+backend recovery consistency from disagreement with eager. For later decoder
+layers, cross-backend matrix differences can include inherited upstream
+activation differences; this comparison does not isolate the local kernel.

attnread-0.1.0/PKG-INFO

@@ -0,0 +1,162 @@
+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.

attnread-0.1.0/README.md

@@ -0,0 +1,140 @@
+# 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.

attnread-0.1.0/RELEASE.md

@@ -0,0 +1,54 @@
+# First PyPI release
+
+Release version: **0.1.0**. License: **MIT**. No release/upload is automatic.
+
+1. Verify that the chosen distribution name, `attnread`, is
+   available on both PyPI and TestPyPI. Hyphens/underscores normalize to the
+   same distribution name; the Python import is `attnread`.
+2. Confirm that `LICENSE` is present in both release archives and that their
+   metadata contains `License-Expression: MIT` and `License-File: LICENSE`.
+   The PyPI README is self-contained. Add repository/documentation URLs only
+   once their destinations are publicly accessible. Review the files in the
+   wheel/sdist: these are public even if GitHub stays private.
+3. Keep supported models/backends and dependency bounds honest. The first adapter
+   recognizes specific separate/fused projection layouts; HF task classes alone
+   do not guarantee compatibility. Expand bounds after testing; keep
+   exact research environment pins separate from package dependencies.
+4. Run tests and build both artifacts in a venv. Test the installed wheel from a
+   directory outside the checkout; also build a wheel from the source archive.
+
+```bash
+python -m pip install -e '.[dev]'
+python -m pytest -q
+python -m build
+python -m twine check --strict dist/*
+```
+
+On this HPC system, installs/builds/tests belong in a CPU Slurm allocation.
+Run these commands from this directory. Verify the installed wheel outside the
+source checkout in a separate environment with the declared dependencies; avoid
+changing the research environment. `python -m build` builds the source archive
+and then builds the wheel from that archive.
+
+5. Create separate PyPI and TestPyPI accounts, verify email, enable 2FA, and save
+   recovery codes. Prefer GitHub Actions Trusted Publishing (OIDC) with a protected
+   release environment; it avoids storing a long-lived PyPI token. For a manual
+   first upload, use a token through Twine's prompt, never in source or commands.
+6. Upload version 0.1.0 to TestPyPI first, then install that exact artifact.
+   Preinstall dependencies from normal PyPI and use `--no-deps` for TestPyPI so
+   a test index does not become the source of your dependency graph.
+
+```bash
+python -m twine upload --repository testpypi dist/*
+python -m pip install --index-url https://test.pypi.org/simple/ --no-deps attnread==0.1.0
+```
+
+7. Once reviewed, tag the corresponding commit and upload the same verified
+   artifacts to PyPI. Every changed artifact needs a new version; uploaded
+   filenames cannot be reused, even after deletion. Version 0.1.0 is the first public release. Fix broken releases with a new
+   version (and consider yanking the broken one).
+
+Official guides: [packaging](https://packaging.python.org/en/latest/tutorials/packaging-projects/),
+[TestPyPI](https://packaging.python.org/en/latest/guides/using-testpypi/),
+[Trusted Publishing](https://docs.pypi.org/trusted-publishers/),
+[PyPI accounts and immutable uploads](https://pypi.org/help/).

attnread-0.1.0/pyproject.toml

@@ -0,0 +1,34 @@
+[build-system]
+requires = ["setuptools>=77.0.3", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "attnread"
+version = "0.1.0"
+description = "Recover attention weights without changing the attention backend."
+readme = "README.md"
+requires-python = ">=3.11"
+license = "MIT"
+license-files = ["LICENSE"]
+dependencies = [
+    "nnsight==0.7.0",
+    "torch>=2.14,<2.15",
+    "transformers>=5.19,<5.20",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=9,<10", "build>=1.2,<2", "twine>=6,<7"]
+
+[tool.setuptools.packages.find]
+where = ["src"]
+include = ["attnread*"]
+
+[tool.setuptools.package-data]
+attnread = ["py.typed"]

attnread-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

attnread-0.1.0/src/attnread/__init__.py

@@ -0,0 +1,4 @@
+"""Extract selected attention heads without switching their backend."""
+from ._core import extract_attention
+
+__all__ = ["extract_attention"]

attnread-0.1.0/src/attnread/_core.py

@@ -0,0 +1,215 @@
+"""Local attention recovery through NNsight partial-identity value probes."""
+import inspect
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from nnsight import NNsight
+
+Head = tuple[str, int, int]
+
+
+@dataclass(frozen=True)
+class _Spec:
+    path: str
+    value: str
+    output: str
+    heads: int
+    kv_heads: int
+    dim: int
+    fused: bool = False
+
+
+def _spec(path, module):
+    if getattr(module, 'is_cross_attention', False) or any(
+        part in path.split('.') for part in ('cross_attn', 'encoder_attn', 'crossattention')
+    ):
+        return None
+    dim = getattr(module, 'head_dim', None)
+    if not isinstance(dim, int) or dim <= 0:
+        return None
+    for output in ('o_proj', 'out_proj'):
+        v = getattr(module, 'v_proj', None)
+        o = getattr(module, output, None)
+        if all(hasattr(module, name) for name in ('q_proj', 'k_proj')) and v is not None and o is not None:
+            n, kv = getattr(o, 'in_features', 0), getattr(v, 'out_features', 0)
+            if n > 0 and kv > 0 and n % dim == kv % dim == 0:
+                return _Spec(path, 'v_proj', output, n // dim, kv // dim, dim)
+    # Fused QKV layouts are explicit: Q, then K, then V, with equal head counts.
+    if type(module).__name__ == 'GPT2Attention' and module.c_attn.nf == 3 * module.num_heads * dim:
+        return _Spec(path, 'c_attn', 'c_proj', module.num_heads, module.num_heads, dim, True)
+    if type(module).__name__ == 'Qwen3VLVisionAttention':
+        return _Spec(path, 'qkv', 'proj', module.num_heads, module.num_heads, dim, True)
+    return None
+
+
+def _discover(model, attention_modules=None):
+    groups = {}
+    if attention_modules is not None:
+        for stream, paths in attention_modules.items():
+            groups[stream] = []
+            for path in paths:
+                try:
+                    spec = _spec(path, model.get_submodule(path))
+                except AttributeError:
+                    spec = None
+                if spec is None:
+                    raise ValueError(f'No supported attention projection layout at {path!r}')
+                groups[stream].append(spec)
+    else:
+        for path, module in model.named_modules():
+            spec = _spec(path, module)
+            if spec is not None:
+                stream = 'vision' if any(part == 'visual' or part.startswith('vision')
+                                         for part in path.split('.')) else 'text'
+                groups.setdefault(stream, []).append(spec)
+    if not any(groups.values()):
+        raise ValueError('No supported attention projection layout found; see package compatibility')
+    return groups
+
+
+def _validate_inputs(model, inputs):
+    if model.training or any(module.training for module in model.modules()):
+        raise ValueError('Use evaluation mode: call model.eval()')
+    if inputs.get('past_key_values') is not None:
+        raise ValueError('Use an uncached full prefix; past_key_values is unsupported')
+    config = getattr(model, 'config', None)
+    configs = [config, getattr(config, 'text_config', None), getattr(config, 'vision_config', None)]
+    if inputs.get('output_attentions') or any(getattr(c, 'output_attentions', False) for c in configs):
+        raise ValueError('output_attentions=True can switch attention backends; disable it')
+    ids, embeds = inputs.get('input_ids'), inputs.get('inputs_embeds')
+    if not ((isinstance(ids, torch.Tensor) and ids.ndim == 2 and min(ids.shape) > 0) or
+            (isinstance(embeds, torch.Tensor) and embeds.ndim == 3 and min(embeds.shape) > 0)):
+        raise ValueError('Use nonempty input_ids [batch, tokens] or inputs_embeds [batch, tokens, dim]')
+    if hasattr(model, '_orig_mod'):
+        raise ValueError('Compiled models are outside the supported scope')
+
+
+def _envoy(wrapped, path):
+    node = wrapped
+    for part in path.split('.'):
+        node = node[int(part)] if part.isdigit() else getattr(node, part)
+    return node
+
+
+def extract_attention(
+    model: torch.nn.Module,
+    inputs: Mapping[str, Any],
+    heads: Sequence[tuple[int, int] | Head] | None = None,
+    *,
+    attention_modules: Mapping[str, Sequence[str]] | None = None,
+) -> tuple[dict[Head, torch.Tensor], dict[str, Any]]:
+    """Recover requested heads; return CPU [B,T,T] tensors and preservation checks.
+
+    (L,H) addresses text heads; (stream,L,H) also addresses vision/custom streams.
+    None selects every text head. Ordinary HF evaluation models only, uncached.
+    Partial-identity V probes preserve Q/K, masks, backend and downstream values.
+    Vision's packed tokens have a singleton batch axis. Reconstruction is finite
+    precision and quadratic in token count; head selection limits saved memory.
+    """
+    _validate_inputs(model, inputs)
+    specs = _discover(model, attention_modules)
+    if heads is None:
+        heads = [('text', layer, head) for layer, spec in enumerate(specs.get('text', []))
+                 for head in range(spec.heads)]
+    selected, grouped = [], {}
+    for address in heads:
+        if not isinstance(address, (tuple, list)) or len(address) not in (2, 3):
+            raise ValueError('heads must contain (layer, head) or (stream, layer, head) addresses')
+        stream, layer, head = ('text', *address) if len(address) == 2 else address
+        if (not isinstance(stream, str) or type(layer) is not int or type(head) is not int or
+            stream not in specs or not 0 <= layer < len(specs[stream]) or
+            not 0 <= head < specs[stream][layer].heads):
+            raise ValueError(f'Invalid heads address: {address!r}')
+        key = (stream, layer, head)
+        if key in selected:
+            raise ValueError(f'Duplicate heads address: {key!r}')
+        selected.append(key)
+        grouped.setdefault((stream, layer), []).append(head)
+    if not selected:
+        raise ValueError('heads selection is empty or no text attention was discovered')
+    paths = [specs[stream][layer].path for stream, layer in grouped]
+    if len(set(paths)) != len(paths):
+        raise ValueError('Do not address the same module through multiple stream aliases')
+    kwargs = {**inputs, 'use_cache': False, 'return_dict': True}
+    if 'logits_to_keep' in inspect.signature(type(model).forward).parameters:
+        kwargs['logits_to_keep'] = 1
+    checks = {'baseline_logit_max_abs': 0.0, 'probe_logit_max_abs': 0.0,
+              'unpatched_logit_max_abs': 0.0, 'heads': selected, 'probes': 0,
+              'modules': {stream: [{'path': s.path, 'heads': s.heads, 'head_dim': s.dim,
+                                   'packed': None} for s in group]
+                          for stream, group in specs.items()}}
+    patterns = {}
+    with torch.no_grad():
+        original = model(**kwargs).logits.detach().clone()
+        wrapped = NNsight(model)
+        for (stream, layer), chosen in grouped.items():
+            spec = specs[stream][layer]
+            attn = _envoy(wrapped, spec.path)
+            value, output = getattr(attn, spec.value), getattr(attn, spec.output)
+            # Save the genuine result first, for restoration before the output projection.
+            with wrapped.trace(**kwargs):
+                saved_z = output.input.detach().clone().save()
+                saved_logits = wrapped.output.logits.detach().clone().save()
+            baseline_z = saved_z
+            packed = baseline_z.ndim == 2
+            if baseline_z.ndim not in (2, 3) or baseline_z.shape[-1] != spec.heads * spec.dim:
+                raise ValueError(f'Unsupported merged-head output layout at {spec.path}')
+            batch, length = (1, baseline_z.shape[0]) if packed else baseline_z.shape[:2]
+            checks['modules'][stream][layer].update(packed=packed, tokens=length, batch=batch)
+            checks['baseline_logit_max_abs'] = max(checks['baseline_logit_max_abs'],
+                                                   (saved_logits-original).abs().max().item())
+            pieces = []
+            for offset in range(0, length, spec.dim):
+                width = min(spec.dim, length-offset)
+                basis = torch.zeros(length, spec.dim, device=baseline_z.device, dtype=baseline_z.dtype)
+                basis[offset:offset+width, :width] = torch.eye(width, device=basis.device, dtype=basis.dtype)
+                fake_v = basis[:, None, :].expand(length, spec.kv_heads, spec.dim).reshape(length, -1)
+                if not packed:
+                    fake_v = fake_v[None].expand(batch, -1, -1).contiguous()
+                with wrapped.trace(**kwargs):
+                    if spec.fused:
+                        # Preserve the original Q and K portions of a fused projection.
+                        value.output = torch.cat([value.output[..., :2*spec.heads*spec.dim], fake_v], dim=-1)
+                    else:
+                        value.output = fake_v
+                    # Slice heads before saving, so unrequested patterns are never transferred.
+                    block = output.input.reshape(batch, length, spec.heads, spec.dim).permute(
+                        0, 2, 1, 3)[:, chosen, :, :width].detach().clone().save()
+                    output.input = baseline_z
+                    probe_logits = wrapped.output.logits.detach().clone().save()
+                pieces.append(block.cpu())
+                checks['probe_logit_max_abs'] = max(checks['probe_logit_max_abs'],
+                                                   (probe_logits-original).abs().max().item())
+                checks['probes'] += 1
+            weights = torch.cat(pieces, dim=-1)
+            for index, head in enumerate(chosen):
+                patterns[(stream, layer, head)] = weights[:, index]
+        untouched = model(**kwargs).logits
+        checks['unpatched_logit_max_abs'] = (untouched-original).abs().max().item()
+    return {key: patterns[key] for key in selected}, checks
+
+
+def reconstruct_attention(model, inputs, layers=None, *, heads=None, attention_modules=None):
+    """Research compatibility adapter: return stacked text layers and checks."""
+    _validate_inputs(model, inputs)
+    if attention_modules is not None and layers is not None:
+        raise ValueError('Choose layers or attention_modules, not both')
+    explicit = {'text': list(attention_modules)} if attention_modules is not None else None
+    specs = _discover(model, explicit).get('text', [])
+    layers = list(range(len(specs))) if layers is None else list(layers)
+    if not layers or any(type(i) is not int or not 0 <= i < len(specs) for i in layers) or len(set(layers)) != len(layers):
+        raise ValueError('layers must contain distinct integer text indices in range')
+    if heads is None:
+        chosen = list(range(specs[layers[0]].heads))
+    else:
+        chosen = [heads] if type(heads) is int else list(heads)
+    if not chosen or any(type(h) is not int or not 0 <= h < specs[i].heads for i in layers for h in chosen) or len(set(chosen)) != len(chosen):
+        raise ValueError('heads must contain distinct integer head indices in range')
+    patterns, checks = extract_attention(model, inputs,
+        heads=[(i, h) for i in layers for h in chosen], attention_modules=explicit)
+    weights = torch.stack([torch.stack([patterns[('text', i, h)] for h in chosen], dim=1)
+                           for i in layers], dim=1)
+    checks.update(layers=layers, heads=chosen, modules=[specs[i].path for i in layers])
+    return weights, checks

attnread-0.1.0/src/attnread.egg-info/PKG-INFO

@@ -0,0 +1,162 @@
+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.

attnread-0.1.0/src/attnread.egg-info/SOURCES.txt

@@ -0,0 +1,15 @@
+LICENSE
+MANIFEST.in
+PACKAGE.md
+README.md
+RELEASE.md
+pyproject.toml
+test_attnread.py
+src/attnread/__init__.py
+src/attnread/_core.py
+src/attnread/py.typed
+src/attnread.egg-info/PKG-INFO
+src/attnread.egg-info/SOURCES.txt
+src/attnread.egg-info/dependency_links.txt
+src/attnread.egg-info/requires.txt
+src/attnread.egg-info/top_level.txt

attnread-0.1.0/src/attnread.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

attnread-0.1.0/src/attnread.egg-info/requires.txt

@@ -0,0 +1,8 @@
+nnsight==0.7.0
+torch<2.15,>=2.14
+transformers<5.20,>=5.19
+
+[dev]
+pytest<10,>=9
+build<2,>=1.2
+twine<7,>=6

attnread-0.1.0/src/attnread.egg-info/top_level.txt

@@ -0,0 +1 @@
+attnread

attnread-0.1.0/test_attnread.py

@@ -0,0 +1,196 @@
+import pytest
+import torch
+from transformers import Qwen3Config, Qwen3ForCausalLM
+from attnread._core import reconstruct_attention
+
+
+@pytest.mark.parametrize('operation', ['extract_attention'])
+@pytest.mark.parametrize('case, message', [
+    ('training', 'evaluation'),
+    ('cache', 'uncached'),
+    ('attentions', 'output_attentions'),
+    ('config_attentions', 'output_attentions'),
+    ('empty', 'nonempty'),
+])
+def test_unsafe_inputs_are_rejected_before_tracing(operation, case, message):
+    import attnread
+    model = Qwen3ForCausalLM(Qwen3Config(vocab_size=32, hidden_size=32,
+        intermediate_size=48, num_hidden_layers=1, num_attention_heads=2,
+        num_key_value_heads=1, head_dim=16)).eval()
+    model.set_attn_implementation('eager')
+    inputs = {'input_ids': torch.tensor([[1, 2, 3]])}
+    if case == 'training':
+        model.train()
+    elif case == 'cache':
+        inputs['past_key_values'] = object()
+    elif case == 'attentions':
+        inputs['output_attentions'] = True
+    elif case == 'config_attentions':
+        model.config.output_attentions = True
+    elif case == 'empty':
+        inputs['input_ids'] = torch.empty((1, 0), dtype=torch.long)
+    with pytest.raises(ValueError, match=message):
+        getattr(attnread, operation)(model, inputs)
+
+
+@pytest.mark.parametrize('layers', [[], [-1], [1], [0, 0], [0.5], [True]])
+def test_invalid_layer_selection_is_rejected(layers):
+    model = Qwen3ForCausalLM(Qwen3Config(vocab_size=32, hidden_size=32,
+        intermediate_size=48, num_hidden_layers=1, num_attention_heads=2,
+        num_key_value_heads=1, head_dim=16)).eval()
+    with pytest.raises(ValueError, match='layers'):
+        reconstruct_attention(model, {'input_ids': torch.tensor([[1, 2]])}, layers)
+
+
+
+def make_hf_model(family):
+    import transformers as hf
+    common = dict(vocab_size=40, hidden_size=48, intermediate_size=64,
+                  num_hidden_layers=2, num_attention_heads=4,
+                  num_key_value_heads=2, head_dim=16,
+                  bos_token_id=1, eos_token_id=2, pad_token_id=0)
+    if family == 'gpt2':
+        config = hf.GPT2Config(vocab_size=40, n_embd=32, n_head=2, n_layer=2,
+                              n_positions=32, bos_token_id=1, eos_token_id=2,
+                              resid_pdrop=0, embd_pdrop=0, attn_pdrop=0)
+    elif family == 'llava':
+        config = hf.LlavaConfig(
+            text_config=hf.LlamaConfig(**common),
+            vision_config=hf.CLIPVisionConfig(hidden_size=16, intermediate_size=32,
+                num_hidden_layers=2, num_attention_heads=2, image_size=8, patch_size=4),
+            image_token_id=39, image_seq_length=4, vision_feature_layer=-1)
+    else:
+        cls = {'llama': hf.LlamaConfig, 'mistral': hf.MistralConfig,
+               'gemma2': hf.Gemma2Config}[family]
+        extra = {'sliding_window': 4} if family in {'mistral', 'gemma2'} else {}
+        if family == 'gemma2':
+            extra['attn_logit_softcapping'] = 5.0
+        config = cls(**common, **extra)
+    config._attn_implementation = 'eager'
+    factory = hf.AutoModelForImageTextToText if family == 'llava' else hf.AutoModelForCausalLM
+    model = factory.from_config(config, attn_implementation='eager').eval()
+    inputs = {'input_ids': torch.tensor([[1, 3, 5, 6, 7, 8, 9, 2, 4, 5, 6]])}
+    if family == 'llava':
+        inputs = {'input_ids': torch.tensor([[1, 39, 39, 39, 39, 3, 5, 2]]),
+                  'pixel_values': torch.randn(1, 3, 8, 8)}
+    return model, inputs
+
+
+@pytest.mark.parametrize('family', ['llama', 'mistral', 'gemma2', 'gpt2', 'llava'])
+def test_hf_task_models_recover_native_attention(family):
+    torch.manual_seed(7)
+    model, inputs = make_hf_model(family)
+    with torch.no_grad():
+        expected = model(**inputs, use_cache=False, output_attentions=True)
+        weights, checks = reconstruct_attention(model, inputs, layers=[1, 0])
+    for output_index, layer in enumerate([1, 0]):
+        torch.testing.assert_close(weights[:, output_index], expected.attentions[layer],
+                                   atol=3e-7, rtol=3e-6)
+    assert checks['probe_logit_max_abs'] == 0
+    assert checks['baseline_logit_max_abs'] == 0
+    assert checks['unpatched_logit_max_abs'] == 0
+    assert len(checks['modules']) == 2
+    assert weights.device.type == 'cpu'
+
+
+@pytest.mark.parametrize('family', ['llama', 'gpt2', 'llava'])
+@pytest.mark.parametrize('heads', [1, [1, 0]])
+def test_selected_heads_are_recovered_in_requested_order(family, heads):
+    model, inputs = make_hf_model(family)
+    expected = model(**inputs, use_cache=False, output_attentions=True).attentions[0]
+    chosen = [heads] if isinstance(heads, int) else heads
+    weights, checks = reconstruct_attention(model, inputs, layers=[0], heads=heads)
+    torch.testing.assert_close(weights[:, 0], expected[:, chosen], atol=3e-7, rtol=3e-6)
+    assert weights.shape[2] == len(chosen)
+    assert checks['heads'] == chosen
+    assert checks['probe_logit_max_abs'] == 0
+
+
+@pytest.mark.parametrize('heads', [[], [-1], [4], [0, 0], [True]])
+def test_invalid_head_selection_is_rejected(heads):
+    model, inputs = make_hf_model('llama')
+    with pytest.raises(ValueError, match='heads'):
+        reconstruct_attention(model, inputs, layers=[0], heads=heads)
+
+
+def test_explicit_attention_module_paths_select_the_named_layer():
+    model, inputs = make_hf_model('llama')
+    expected = model(**inputs, use_cache=False, output_attentions=True).attentions[1]
+    path = 'model.layers.1.self_attn'
+    weights, checks = reconstruct_attention(model, inputs, attention_modules=[path], heads=0)
+    torch.testing.assert_close(weights[:, 0, 0], expected[:, 0], atol=3e-7, rtol=3e-6)
+    assert checks['modules'] == [path]
+
+
+def test_unsupported_projection_structure_fails_explicitly():
+    from transformers import BertConfig, BertForMaskedLM
+    model = BertForMaskedLM(BertConfig(vocab_size=40, hidden_size=16,
+        intermediate_size=32, num_hidden_layers=1, num_attention_heads=2)).eval()
+    with pytest.raises(ValueError, match='supported attention'):
+        reconstruct_attention(model, {'input_ids': torch.tensor([[1, 2]])})
+
+
+@pytest.mark.parametrize('family', ['llama', 'llava'])
+def test_public_head_addressing_returns_only_requested_patterns(family):
+    import attnread
+    model, inputs = make_hf_model(family)
+    text = model(**inputs, use_cache=False, output_attentions=True).attentions[1][:, 2]
+    selection = [(1, 2)]
+    expected = {('text', 1, 2): text}
+    if family == 'llava':
+        vision = model.model.vision_tower(inputs['pixel_values'], output_attentions=True).attentions[0][:, 1]
+        selection.append(('vision', 0, 1))
+        expected[('vision', 0, 1)] = vision
+    assert callable(getattr(attnread, 'extract_attention', None)), 'missing public extractor'
+    patterns, checks = attnread.extract_attention(model, inputs, heads=selection)
+    assert list(patterns) == list(expected)
+    for key in expected:
+        torch.testing.assert_close(patterns[key], expected[key], atol=3e-7, rtol=3e-6)
+    assert checks['probe_logit_max_abs'] == 0
+    assert checks['unpatched_logit_max_abs'] == 0
+
+
+def test_qwen_packed_vision_keeps_image_boundaries_and_preserves_logits():
+    import attnread
+    from transformers import Qwen3VLConfig, Qwen3VLForConditionalGeneration
+    config = Qwen3VLConfig(
+        text_config=dict(vocab_size=40, hidden_size=32, intermediate_size=48,
+            num_hidden_layers=2, num_attention_heads=2, num_key_value_heads=1,
+            head_dim=16, bos_token_id=1, eos_token_id=2, pad_token_id=0,
+            rope_parameters={'rope_type': 'default', 'rope_theta': 10000,
+                             'mrope_section': [2, 3, 3]}),
+        vision_config=dict(depth=2, hidden_size=16, intermediate_size=32,
+            num_heads=2, patch_size=2, temporal_patch_size=2, spatial_merge_size=2,
+            out_hidden_size=32, num_position_embeddings=16, deepstack_visual_indexes=[0, 1]),
+        image_token_id=39, video_token_id=38, vision_start_token_id=37, vision_end_token_id=36)
+    model = Qwen3VLForConditionalGeneration(config).eval()
+    model.set_attn_implementation('eager')
+    inputs = {'input_ids': torch.tensor([[1, 37, 39, 37, 39, 2]]),
+              'pixel_values': torch.randn(8, 24),
+              'mm_token_type_ids': torch.tensor([[0, 0, 1, 0, 1, 0]]),
+              'image_grid_thw': torch.tensor([[1, 2, 2], [1, 2, 2]])}
+    assert callable(getattr(attnread, 'extract_attention', None)), 'missing public extractor'
+    values, outputs = [], []
+    attn = model.model.visual.blocks[0].attn
+    h1 = attn.qkv.register_forward_hook(lambda m, a, o: values.append(o.detach()))
+    h2 = attn.proj.register_forward_pre_hook(lambda m, a: outputs.append(a[0].detach()))
+    try:
+        with torch.no_grad():
+            original = model(**inputs, use_cache=False).logits
+    finally:
+        h1.remove()
+        h2.remove()
+    patterns, checks = attnread.extract_attention(model, inputs,
+        heads=[('vision', 0, 1), (0, 0)])
+    matrix = patterns[('vision', 0, 1)]
+    assert matrix.shape == (1, 8, 8)  # Packed images share an axis, with disconnected blocks.
+    assert torch.count_nonzero(matrix[:, :4, 4:]) == 0
+    assert torch.count_nonzero(matrix[:, 4:, :4]) == 0
+    torch.testing.assert_close(matrix.sum(-1), torch.ones(1, 8), atol=3e-7, rtol=3e-6)
+    v = values[0].reshape(8, 3, 2, 8)[:, 2, 1]
+    z = outputs[0].reshape(8, 2, 8)[:, 1]
+    torch.testing.assert_close(matrix[0] @ v, z, atol=3e-7, rtol=3e-6)
+    assert checks['probe_logit_max_abs'] == 0
+    assert checks['unpatched_logit_max_abs'] == 0
+    assert checks['modules']['vision'][0]['packed'] is True
+    torch.testing.assert_close(model(**inputs, use_cache=False).logits, original, atol=0, rtol=0)