attnread 0.1.0__py3-none-any.whl

attnread/__init__.py

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

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.dist-info/METADATA

@@ -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.dist-info/RECORD

@@ -0,0 +1,8 @@
+attnread/__init__.py,sha256=q7PDqt5zmBTn2vhP0GXJ4TTVNmcYMROENKlFtASi-kI,142
+attnread/_core.py,sha256=4GIxsbqapW7_drM3bBDOK2a7sW8vkHyIQQ8f03j8CaY,11438
+attnread/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+attnread-0.1.0.dist-info/licenses/LICENSE,sha256=Xi4piPMZc-0ukHN4eM_sq8DGZcEfzPRr8Y2cIo1d8pM,1061
+attnread-0.1.0.dist-info/METADATA,sha256=D48LqpRQnnelzX0LPbDMYhVpb82xQyg0RgILX3sQmk4,7581
+attnread-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+attnread-0.1.0.dist-info/top_level.txt,sha256=0WrqukxxD6qUq7m_z28avBhVrgWOdeP7lQOFI7sq7yk,9
+attnread-0.1.0.dist-info/RECORD,,

attnread-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

attnread-0.1.0.dist-info/licenses/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.dist-info/top_level.txt

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