interp-engine 1.2.7__tar.gz → 1.2.8__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {interp_engine-1.2.7 → interp_engine-1.2.8}/PKG-INFO +23 -34
- interp_engine-1.2.8/README.md +117 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/README.md +24 -6
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/bench_spec.py +6 -4
- {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/PERFORMANCE.md +30 -9
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/_hooks.py +25 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/freeze.py +22 -8
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/requests.py +1 -17
- {interp_engine-1.2.7 → interp_engine-1.2.8}/pyproject.toml +1 -1
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_hyper_connections.py +276 -0
- interp_engine-1.2.7/README.md +0 -128
- {interp_engine-1.2.7 → interp_engine-1.2.8}/.gitignore +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/LICENSE +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/__init__.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/cells.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/probe.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/publish.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/report_bench.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__eager.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__vllm-cudagraph.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__vllm-dspark-cudagraph.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__vllm-dspark.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__vllm-freeze.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__vllm.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/gemma-2-2b__eager.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/gemma-2-2b__vllm-cudagraph.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/gemma-2-2b__vllm-freeze.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/gemma-2-2b__vllm.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/llama-3.1-8b__eager.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/llama-3.1-8b__vllm-cudagraph.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/llama-3.1-8b__vllm-freeze.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/llama-3.1-8b__vllm.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3-4b__eager.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3-4b__vllm-cudagraph.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3-4b__vllm-freeze.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3-4b__vllm.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3.8-27b__eager.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3.8-27b__vllm-cudagraph.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3.8-27b__vllm-freeze.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3.8-27b__vllm.json +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results-latest.md +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/run_all.sh +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/run_bench.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/workloads.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/AGENT_INTEGRATION.md +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/ARCHITECTURE_QUIRKS.md +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/COMPATIBILITY.md +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/ENGINE_HOOK_MAPPINGS.md +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/GRADIENTS.md +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/INTERNALS.md +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/PORTING.md +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/README.md +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/SUPPORTED_POINTS.md +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/USAGE.md +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/__init__.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/_loop.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/address.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/arch.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/attn_config.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/attn_scores.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/autograd_support.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/capture.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/chat_compose.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/chat_conventions.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/chat_formatters.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/cuda_preflight.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/dispatch.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/facts.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/hooks.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/lens.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/load.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/mappers.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/model.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/moe_routing.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/points.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/protocol.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/residual_basis.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/select.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/steer.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/steer_specs.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/sync.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/tokenize.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_backend.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/__init__.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/_demux.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/_payload.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/_tree.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/attn.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/capture.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/graphs.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/lens/__init__.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/lens/intervene.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/lens/readout.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/lens/unembed.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/mhc.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/native.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/steering.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_plugin.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/conftest.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/harness.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/model_expectations.yaml +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/synthetic_families.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_address.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_attn_config_tripwire.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_attn_probs_indexing.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_attn_scores.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_attn_z_gqa.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_autograd_support.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_bench_workloads.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_capability_refusals.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_capture_addressing.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_chat_compose.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_chat_formatters.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_chat_templates.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_core.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_cuda_preflight.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_doc_code_fences.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_eager_autograd.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_facts.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_family_points.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_freeze_dsv4_gpu.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_freeze_parity_gpu.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_freeze_set.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_freeze_warmup.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_gated_attn_out.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_head_contributions.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_hook_call_conventions.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_layer_kinds.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_load.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_logit_transform.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_mappers.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_mlp_internals.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_model_expectations.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_moe.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_multimodal_arch.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_new_models_gpu.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_no_chat_template.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_normalized_hook.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_packaging.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_parity_gpt2.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_per_layer_attn_dims.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_points_registry.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_protocol.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_published_benchmarks.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_qk_norm.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_qkv_layout.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_reasoning_spans.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_release.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_resid_mid.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_residual_basis.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_sandwich_norms.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_select.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_sliding_window_attn.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_small_models_gpu.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_steer_context.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_steer_math_parity.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_sync_loop.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_sync_parity.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_unified_free_functions.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_unresolved_families.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_capture_gpu.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_capture_scales.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_graph_path.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_graphs_on_gpu.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_hook_availability.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_kv_isolation.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_new_points.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_only_families.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_plugin.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_wire_grammar.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vocabulary_boundary.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_worker_lens_capture_readout.py +0 -0
- {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_worker_lens_readout.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.5
|
|
2
2
|
Name: interp-engine
|
|
3
|
-
Version: 1.2.
|
|
3
|
+
Version: 1.2.8
|
|
4
4
|
Summary: A fast, standardized interpretability engine that supports most modern models and architectures. Powers Neuronpedia.
|
|
5
5
|
Project-URL: Homepage, https://github.com/decoderesearch/interp-engine
|
|
6
6
|
Project-URL: Repository, https://github.com/decoderesearch/interp-engine
|
|
@@ -31,33 +31,22 @@ Description-Content-Type: text/markdown
|
|
|
31
31
|
|
|
32
32
|
# interp-engine
|
|
33
33
|
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
🔗 <a href="https://interp-engine.org"><strong>interp-engine.org</strong></a>
|
|
39
|
-
</p>
|
|
40
|
-
<p align="center">
|
|
41
|
-
<a href="https://github.com/decoderesearch/interp-engine/actions/workflows/engine-tests.yml"><img src="https://github.com/decoderesearch/interp-engine/actions/workflows/engine-tests.yml/badge.svg?branch=main" alt="CI status"></a>
|
|
42
|
-
<a href="https://pypi.org/project/interp-engine/"><img src="https://img.shields.io/pypi/v/interp-engine.svg" alt="PyPI version"></a>
|
|
43
|
-
<a href="LICENSE"><img src="https://img.shields.io/pypi/l/interp-engine.svg" alt="Apache-2.0 license"></a>
|
|
44
|
-
<a href="https://join.slack.com/t/opensourcemechanistic/shared_invite/zt-3z9o0hxjl-MDX9pbATO2qESOazNDLpdQ"><img src="https://img.shields.io/badge/Slack-Open%20Source%20Mechanistic%20Interpretability-4A154B?logo=slack&logoColor=white" alt="Join the Slack"></a>
|
|
45
|
-
</p>
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
🔗 **[interp-engine.org](https://interp-engine.org)**
|
|
37
|
+
|
|
46
38
|
|
|
47
39
|
|
|
48
40
|
`interp-engine` is an interpretability engine that is fast (>40x tok/s vs HF eager), standardized (34 'points'/addresses across architectures), and easy to use and debug. It powers all of [Neuronpedia](https://neuronpedia.org)'s inference and is checked for accuracy against HF Transformers and other engines.
|
|
49
41
|
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
<p align="center">
|
|
54
|
-
<img src="https://neuronpedia.s3.amazonaws.com/site-assets/interp-engine-demo.gif" alt="interp-engine demo gif" width="100%">
|
|
55
|
-
</p>
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
|
|
56
45
|
|
|
57
46
|
This repo contains:
|
|
58
47
|
|
|
59
|
-
1. [
|
|
60
|
-
2. [
|
|
48
|
+
1. `[validator/](validator/)`, which compares/validates it against TransformerLens, and nnsight/nnterp on real architectures.
|
|
49
|
+
2. `[visualizer-web/](visualizer-web/)`, a "cheat sheet" hosted at [interp-engine.org](https://interp-engine.org) of each 'point' (eg `resid_post.16`), standardized across architectures.
|
|
61
50
|
|
|
62
51
|
## Installation
|
|
63
52
|
|
|
@@ -71,7 +60,15 @@ pip install interp-engine # eager backend only
|
|
|
71
60
|
```python
|
|
72
61
|
from interp_engine import Address, load_model, run_with_cache
|
|
73
62
|
|
|
74
|
-
|
|
63
|
+
# VLLM MODE (default): low VRAM, medium speed
|
|
64
|
+
model = load_model("Qwen/Qwen3-8B")
|
|
65
|
+
|
|
66
|
+
# VLLM-FREEZE MODE: high VRAM, high speed, only frozen points (default resid_post)
|
|
67
|
+
# model = load_model("Qwen/Qwen3-8B", freeze_points="auto")
|
|
68
|
+
|
|
69
|
+
# EAGER MODE: low VRAM, low speed
|
|
70
|
+
# model = load_model("Qwen/Qwen3-8B", backend="eager")
|
|
71
|
+
|
|
75
72
|
point = Address("resid_post", 10) # or string: "resid_post.10"
|
|
76
73
|
cache = run_with_cache(model, model.to_tokens("Hello, world"), [point])
|
|
77
74
|
cache[point] # [batch, pos, ...]
|
|
@@ -87,7 +84,7 @@ Add "use interp-engine" to your prompt and let your agent figure it out - everyt
|
|
|
87
84
|
|
|
88
85
|
## Performance / Speed
|
|
89
86
|
|
|
90
|
-
vLLM gives `interp-engine` high throughput via concurrency, and **graph freeze** adds CUDA-graph replay on top of that
|
|
87
|
+
vLLM gives `interp-engine` high throughput via concurrency, and **graph freeze** adds CUDA-graph replay on top of that *without giving up capture or steering*. Every column below is capture-capable.
|
|
91
88
|
|
|
92
89
|
<!-- THROUGHPUT:START -->
|
|
93
90
|
|
|
@@ -117,22 +114,14 @@ One stream (tok/s):
|
|
|
117
114
|
|
|
118
115
|
<!-- THROUGHPUT:END -->
|
|
119
116
|
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
**Graph freeze** is opt-in, via `freeze_points`. Ordinary hooks cannot survive CUDA graphs — replay never calls the Python forward a hook is attached to — so the tap becomes a preallocated buffer plus a `copy_` that graph capture records and replay re-executes. Steering rides the same wrap, so additive, orthogonal, projection-cap and j-lens steer/ablate/swap all keep working under graphs, position masks included.
|
|
123
|
-
|
|
124
|
-
```python
|
|
125
|
-
model = load_model("Qwen/Qwen3-8B", freeze_points="auto") # resid_post at every layer, under graphs
|
|
126
|
-
```
|
|
127
|
-
|
|
128
|
-
The trade is that a frozen engine serves _the set it froze_ rather than any point on request: `"auto"` covers `resid_post` at every layer on a conventional trunk, an explicit list covers whatever you name, and anything outside the set is refused rather than quietly returned empty. Omit `freeze_points` for today's hooked vLLM, which still serves every point. `qwen3.8-27b` is a hybrid trunk, so its freeze row runs prefill eagerly and keeps the decode graphs — freeze pins that mode, because breakable graphs turn `torch.compile` off and vLLM's mixed prefill-decode capture then miscomputes prefill on a gated-delta trunk. On a hyper-connection trunk such as `deepseek-v4-flash-0731`, whose block carries four parallel residual streams, `"auto"` freezes `resid_streams` instead — the whole stack per layer, which is four times the width and so four times the buffer, and is what that row's freeze cell prices. See [PERFORMANCE.md](docs/PERFORMANCE.md) for the full trade-off, and [`benchmarks/results-latest.md`](benchmarks/results-latest.md) for the sweep, including capture, steering and lens latencies.
|
|
117
|
+
**Graph freeze** is opt-in via `freeze_points`, and a frozen engine serves only the set it froze. [PERFORMANCE.md](docs/PERFORMANCE.md) has how it works and what it trades; `[benchmarks/results-latest.md](benchmarks/results-latest.md)` has every figure at full precision, including capture, steering and lens latencies; `[benchmarks/README.md](benchmarks/README.md)` has how the tables above are rounded.
|
|
129
118
|
|
|
130
119
|
## Correctness
|
|
131
120
|
|
|
132
121
|
We verify correctness in two main ways:
|
|
133
122
|
|
|
134
123
|
1. A test suite that checks results across several models - what each check is designed to catch is in [INTERNALS.md](docs/INTERNALS.md#correctness).
|
|
135
|
-
2. A full `validator` comparison engine that checks most hook points across 50+ models, at early, middle and late layers - fully reproducible, with detailed results saved in the git repo at [
|
|
124
|
+
2. A full `validator` comparison engine that checks most hook points across 50+ models, at early, middle and late layers - fully reproducible, with detailed results saved in the git repo at `[validator/](validator/)`.
|
|
136
125
|
|
|
137
126
|
## Why use an Interpretability Engine, instead of just having my AI code whatever it needs on the fly?
|
|
138
127
|
|
|
@@ -156,4 +145,4 @@ Bugs and feature requests belong in [issues](https://github.com/decoderesearch/i
|
|
|
156
145
|
|
|
157
146
|
## License
|
|
158
147
|
|
|
159
|
-
Apache 2.0
|
|
148
|
+
Apache 2.0
|
|
@@ -0,0 +1,117 @@
|
|
|
1
|
+
# interp-engine
|
|
2
|
+
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
🔗 **[interp-engine.org](https://interp-engine.org)**
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
`interp-engine` is an interpretability engine that is fast (>40x tok/s vs HF eager), standardized (34 'points'/addresses across architectures), and easy to use and debug. It powers all of [Neuronpedia](https://neuronpedia.org)'s inference and is checked for accuracy against HF Transformers and other engines.
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
This repo contains:
|
|
16
|
+
|
|
17
|
+
1. `[validator/](validator/)`, which compares/validates it against TransformerLens, and nnsight/nnterp on real architectures.
|
|
18
|
+
2. `[visualizer-web/](visualizer-web/)`, a "cheat sheet" hosted at [interp-engine.org](https://interp-engine.org) of each 'point' (eg `resid_post.16`), standardized across architectures.
|
|
19
|
+
|
|
20
|
+
## Installation
|
|
21
|
+
|
|
22
|
+
```bash
|
|
23
|
+
pip install 'interp-engine[vllm]' # preferred install: includes vLLM support (CUDA required)
|
|
24
|
+
pip install interp-engine # eager backend only
|
|
25
|
+
```
|
|
26
|
+
|
|
27
|
+
## Simple Usage
|
|
28
|
+
|
|
29
|
+
```python
|
|
30
|
+
from interp_engine import Address, load_model, run_with_cache
|
|
31
|
+
|
|
32
|
+
# VLLM MODE (default): low VRAM, medium speed
|
|
33
|
+
model = load_model("Qwen/Qwen3-8B")
|
|
34
|
+
|
|
35
|
+
# VLLM-FREEZE MODE: high VRAM, high speed, only frozen points (default resid_post)
|
|
36
|
+
# model = load_model("Qwen/Qwen3-8B", freeze_points="auto")
|
|
37
|
+
|
|
38
|
+
# EAGER MODE: low VRAM, low speed
|
|
39
|
+
# model = load_model("Qwen/Qwen3-8B", backend="eager")
|
|
40
|
+
|
|
41
|
+
point = Address("resid_post", 10) # or string: "resid_post.10"
|
|
42
|
+
cache = run_with_cache(model, model.to_tokens("Hello, world"), [point])
|
|
43
|
+
cache[point] # [batch, pos, ...]
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
### AI Agents
|
|
47
|
+
|
|
48
|
+
Add "use interp-engine" to your prompt and let your agent figure it out - everything is fully documented in this repo and open source.
|
|
49
|
+
|
|
50
|
+
## Supported Points ("Addresses")
|
|
51
|
+
|
|
52
|
+
`interp-engine` supports 34 standardized points ("Addresses") across architectures: every one of them on the eager backend, 28 of them on vLLM. Check [interp-engine.org](https://interp-engine.org) for the "cheat sheet", or [SUPPORTED_POINTS.md](docs/SUPPORTED_POINTS.md) for a markdown version with the per-backend detail.
|
|
53
|
+
|
|
54
|
+
## Performance / Speed
|
|
55
|
+
|
|
56
|
+
vLLM gives `interp-engine` high throughput via concurrency, and **graph freeze** adds CUDA-graph replay on top of that *without giving up capture or steering*. Every column below is capture-capable.
|
|
57
|
+
|
|
58
|
+
<!-- THROUGHPUT:START -->
|
|
59
|
+
|
|
60
|
+
<!-- Generated by `python -m benchmarks.report_bench`. Do not edit: rerun the sweep. -->
|
|
61
|
+
|
|
62
|
+
Measured on NVIDIA B200, bf16, 512-token prompt, 128 new tokens.
|
|
63
|
+
|
|
64
|
+
One stream (tok/s):
|
|
65
|
+
|
|
66
|
+
| model | eager | vLLM | vLLM + graph freeze |
|
|
67
|
+
| ------------------------ | ----- | ---------- | ------------------- |
|
|
68
|
+
| `gemma-2-2b` | 31 | 31 (1.0x) | **214 (6.9x)** |
|
|
69
|
+
| `qwen3-4b` | 24 | 47 (2.0x) | **296 (12.3x)** |
|
|
70
|
+
| `llama-3.1-8b` | 33 | 57 (1.7x) | **256 (7.9x)** |
|
|
71
|
+
| `qwen3.8-27b` | 9.9 | 12 (1.2x) | **63 (6.4x)** |
|
|
72
|
+
| `deepseek-v4-flash-0731` | 3.3 | 2.9 (0.9x) | **119 (36x)** |
|
|
73
|
+
|
|
74
|
+
8 concurrent requests (aggregate tok/s):
|
|
75
|
+
|
|
76
|
+
| model | eager | vLLM | vLLM + graph freeze |
|
|
77
|
+
| ------------------------ | ----- | ----------- | ------------------- |
|
|
78
|
+
| `gemma-2-2b` | 30 | 226 (7.5x) | **1,238 (41x)** |
|
|
79
|
+
| `qwen3-4b` | 24 | 333 (14.0x) | **1,018 (43x)** |
|
|
80
|
+
| `llama-3.1-8b` | 32 | 419 (13.0x) | **1,536 (48x)** |
|
|
81
|
+
| `qwen3.8-27b` | 9.5 | 87 (9.2x) | **386 (41x)** |
|
|
82
|
+
| `deepseek-v4-flash-0731` | 3.2 | 23 (7.2x) | **402 (127x)** |
|
|
83
|
+
|
|
84
|
+
<!-- THROUGHPUT:END -->
|
|
85
|
+
|
|
86
|
+
**Graph freeze** is opt-in via `freeze_points`, and a frozen engine serves only the set it froze. [PERFORMANCE.md](docs/PERFORMANCE.md) has how it works and what it trades; `[benchmarks/results-latest.md](benchmarks/results-latest.md)` has every figure at full precision, including capture, steering and lens latencies; `[benchmarks/README.md](benchmarks/README.md)` has how the tables above are rounded.
|
|
87
|
+
|
|
88
|
+
## Correctness
|
|
89
|
+
|
|
90
|
+
We verify correctness in two main ways:
|
|
91
|
+
|
|
92
|
+
1. A test suite that checks results across several models - what each check is designed to catch is in [INTERNALS.md](docs/INTERNALS.md#correctness).
|
|
93
|
+
2. A full `validator` comparison engine that checks most hook points across 50+ models, at early, middle and late layers - fully reproducible, with detailed results saved in the git repo at `[validator/](validator/)`.
|
|
94
|
+
|
|
95
|
+
## Why use an Interpretability Engine, instead of just having my AI code whatever it needs on the fly?
|
|
96
|
+
|
|
97
|
+
Software is easy to make in 2026, here's why an engine makes sense:
|
|
98
|
+
|
|
99
|
+
1. **Speed**: Get performance without sacrificing correctness.
|
|
100
|
+
2. **Standardization + Verification**: Eliminate ambiguity when referring to points, plus a full test suite included.
|
|
101
|
+
3. **Faster Dev / Fewer Tokens Used**: You could spend ten million tokens and have your AI write, test, and make production-ready an interpretability engine. Or you could just `pip install interp-engine[vllm]`.
|
|
102
|
+
|
|
103
|
+
## Development / Contributing
|
|
104
|
+
|
|
105
|
+
Activate the shared git hooks once per clone — they format staged Python, rebuild the generated files, and run CI's static checks before a push. Details in [CONTRIBUTING.md](CONTRIBUTING.md).
|
|
106
|
+
|
|
107
|
+
```bash
|
|
108
|
+
make hooks # or: git config core.hooksPath .githooks
|
|
109
|
+
```
|
|
110
|
+
|
|
111
|
+
## Contact
|
|
112
|
+
|
|
113
|
+
Bugs and feature requests belong in [issues](https://github.com/decoderesearch/interp-engine/issues). For anything else: [johnny@neuronpedia.org](mailto:johnny@neuronpedia.org).
|
|
114
|
+
|
|
115
|
+
## License
|
|
116
|
+
|
|
117
|
+
Apache 2.0
|
|
@@ -104,6 +104,21 @@ python -m benchmarks.publish # rewrite both
|
|
|
104
104
|
python -m benchmarks.publish --check # exit 1 if either has drifted, and name it
|
|
105
105
|
```
|
|
106
106
|
|
|
107
|
+
Both print the same cells in a display form of their own, applied by `publish.py` and pinned by
|
|
108
|
+
`tests/test_published_benchmarks.py`:
|
|
109
|
+
|
|
110
|
+
- **tok/s is whole at 10 and above, one decimal below.** A tenth beside a four-digit figure in the
|
|
111
|
+
next column claims a resolution the reader cannot use. At 3 tok/s that same tenth is worth several
|
|
112
|
+
percent, so the small rows keep it.
|
|
113
|
+
- **every comparison is a multiplier, never a percent**, with one decimal below 20x, where it is
|
|
114
|
+
still checkable against the two printed figures. `+20%` beside `27x` makes the reader convert one
|
|
115
|
+
of them.
|
|
116
|
+
- **multipliers are ratios of the unrounded metrics**, each against eager on the same workload, so
|
|
117
|
+
dividing two printed figures by hand can differ in the last place. Ratios of the rounded figures
|
|
118
|
+
would make the published win depend on the rounding, which is worse.
|
|
119
|
+
|
|
120
|
+
`results-latest.md` keeps every figure at full precision, so nothing is lost to those rules.
|
|
121
|
+
|
|
107
122
|
`tests/test_published_benchmarks.py` runs that check over the committed cells, so a stale copy is a
|
|
108
123
|
red suite rather than a claim nobody re-read. The visualizer's chatbot answers out of a bundle holding
|
|
109
124
|
the README verbatim, so a publish that changed the README also wants `make viz-knowledge` -- the
|
|
@@ -225,12 +240,15 @@ already refers to, so a replay serves capture and steering without a Python forw
|
|
|
225
240
|
`freeze_points="auto"` takes in production, and the row exists to price it against the
|
|
226
241
|
`enforce_eager=True` column capture would otherwise have to use.
|
|
227
242
|
|
|
228
|
-
`"auto"`
|
|
229
|
-
half the feature and reported the other half as `n/a` -- with a message that blamed graph
|
|
230
|
-
it, which is the thing freeze exists to work around.
|
|
231
|
-
|
|
232
|
-
|
|
233
|
-
|
|
243
|
+
`"auto"` once installed **reads** only, and a steering op needs a write tap to land in, so this row
|
|
244
|
+
priced half the feature and reported the other half as `n/a` -- with a message that blamed graph
|
|
245
|
+
replay for it, which is the thing freeze exists to work around. Auto now covers both halves, so the
|
|
246
|
+
cell would be a number either way, and this row still passes `freeze_writes` on purpose: an explicit
|
|
247
|
+
list *narrows* what auto would install, to the one mid-stack site the `steer` workload actually
|
|
248
|
+
writes. A row that priced a write buffer at every layer would not be comparable with the ones beside
|
|
249
|
+
it, which is the whole job of the column. Its value is the sentinel `run_bench.STEER_WRITES` rather
|
|
250
|
+
than a site, because that layer differs per model and a freeze write is a `load_model` argument, so
|
|
251
|
+
it has to be resolved from the config before a model exists to ask.
|
|
234
252
|
|
|
235
253
|
`VariantSpec.models` restricts the row to the checkpoints freeze has been shown correct on, so a model
|
|
236
254
|
missing from it renders `--` rather than a number nobody checked.
|
|
@@ -264,12 +264,14 @@ VARIANTS: tuple[VariantSpec, ...] = (
|
|
|
264
264
|
# Its write tap is `mlp_out`, not `resid_post`: the steer workload addresses the model's
|
|
265
265
|
# `capture_point`, and a hyper-connection trunk refuses the default name (`run_bench._steer_site`).
|
|
266
266
|
#
|
|
267
|
-
# `freeze_writes`
|
|
267
|
+
# `freeze_writes` was once what made the `steer` cell a number instead of `n/a`: `"auto"` installed
|
|
268
268
|
# reads, and a steering op needs a write tap to land in, so this row priced capture under replay
|
|
269
269
|
# and left the other half of the feature unmeasured -- with a message that blamed graph replay for
|
|
270
|
-
# it.
|
|
271
|
-
#
|
|
272
|
-
#
|
|
270
|
+
# it. Auto covers writes now, so the cell stands either way, and naming them here has become a
|
|
271
|
+
# *narrowing*: one write buffer at the site the workload steers rather than one per layer, which
|
|
272
|
+
# is what keeps this row's memory comparable with the columns beside it. The value is the sentinel
|
|
273
|
+
# `run_bench.STEER_WRITES`, resolved there to the mid-stack `resid_post` the workload steers,
|
|
274
|
+
# because the layer differs per model and a freeze write has to be named before the model exists.
|
|
273
275
|
VariantSpec(
|
|
274
276
|
"vllm-freeze",
|
|
275
277
|
"vllm",
|
|
@@ -124,15 +124,24 @@ There *is* a middle ground, which was an open question in this file and is now a
|
|
|
124
124
|
preallocated buffer plus a `copy_` that graph capture records, so replay re-executes the tap instead
|
|
125
125
|
of skipping it. That is `freeze_points` on `VLLMModel`, and it is opt-in:
|
|
126
126
|
|
|
127
|
+
```python
|
|
128
|
+
model = load_model("Qwen/Qwen3-8B", freeze_points="auto") # resid_post at every layer, read and write, under graphs
|
|
129
|
+
```
|
|
130
|
+
|
|
127
131
|
- **`freeze_points=None`** (omit it) — hooked vLLM, `enforce_eager=True`. The default, and the only
|
|
128
132
|
configuration that serves every point.
|
|
129
|
-
- **`freeze_points="auto"`** — `resid_post` at every layer,
|
|
130
|
-
replay and turns Dynamo off (`VLLM_USE_BREAKABLE_CUDAGRAPH=1`), so the
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
133
|
+
- **`freeze_points="auto"`** — `resid_post` at every layer, to read *and* to write, via static
|
|
134
|
+
`copy_` taps. Keeps graph replay and turns Dynamo off (`VLLM_USE_BREAKABLE_CUDAGRAPH=1`), so the
|
|
135
|
+
win is replay without compile. Because breakable `add_eager` keeps the wrap as ordinary PyTorch on
|
|
136
|
+
the live tensor, the write path is not limited to additive — `orthogonal`, `projection_cap` and the
|
|
137
|
+
j-lens `steer`/`ablate`/`swap` ops all ride the same wrap, per-request and with `position_mask`
|
|
138
|
+
honoured (`register_freeze_write`). An op outside that set is refused rather than silently skipped.
|
|
139
|
+
Auto covers the write because the two halves are one decision: a read tap alone serves the lens
|
|
140
|
+
read-out and refuses every steer, ablation and swap derived from it, at an address already tapped.
|
|
141
|
+
Pass `freeze_writes=[]` for the reads without the write buffers, and an explicit `freeze_points`
|
|
142
|
+
list to name both halves yourself — neither is filled in for a caller who said something.
|
|
143
|
+
- **`freeze_writes=[...]`** — write sites, and on their own (no `freeze_points`) an engine that
|
|
144
|
+
writes and does not read.
|
|
136
145
|
- **`freeze_points=[]`** — no taps, graphs and inductor both on. Generation-only, the same thing as
|
|
137
146
|
`enforce_eager=False` but reached through the flag that knows what it costs. Prefer it to setting
|
|
138
147
|
`enforce_eager` yourself.
|
|
@@ -150,11 +159,23 @@ What freeze does not do, which is why `enforce_eager=True` is still the default:
|
|
|
150
159
|
served short.
|
|
151
160
|
- **A narrower batch.** Freeze buffers have to fit alongside the graphs, so `max_num_batched_tokens`
|
|
152
161
|
is stepped down (16384 → 1024) to make room, and refuses rather than OOM-ing in graph capture.
|
|
162
|
+
`freeze_writes=[]` buys some of that width back when the write buffers are what stands between you
|
|
163
|
+
and a wide enough batch.
|
|
164
|
+
- **Four times the buffer on a hyper-connection trunk.** `deepseek-v4-flash-0731`'s block carries
|
|
165
|
+
four parallel residual streams, so `"auto"` freezes `resid_streams` — the whole stack per layer,
|
|
166
|
+
four times the width and so four times the buffer. That is what its freeze row in the sweep prices.
|
|
167
|
+
- **Graphs for decode only on a hybrid trunk.** `qwen3.8-27b`, and any other linear-attention or
|
|
168
|
+
unclassified trunk, runs prefill eagerly and keeps the decode graphs
|
|
169
|
+
(`freeze.decode_only_graphs_reason`). Freeze pins that mode rather than offering it: breakable
|
|
170
|
+
graphs turn `torch.compile` off, and vLLM's mixed prefill-decode graph capture then miscomputes
|
|
171
|
+
prefill on a gated-delta trunk — a whole wrong forward, not a bad tap, and reproducible on plain
|
|
172
|
+
`vllm.LLM`. Eager prefill costs little here, since the wraps are ordinary PyTorch and replay's win
|
|
173
|
+
is in decode.
|
|
153
174
|
- **One engine per process.** `VLLM_USE_BREAKABLE_CUDAGRAPH` is process-global, so a freeze engine
|
|
154
175
|
and a compiled engine cannot share a process.
|
|
155
176
|
- **Measured for throughput, not yet for capture latency.** The sweep runs `generate` and
|
|
156
|
-
`generate_x8` on the freeze variant and only
|
|
157
|
-
steering and lens cells are unmeasured. Capture *correctness* under freeze is covered separately —
|
|
177
|
+
`generate_x8` on the freeze variant, and the tables above compare only the three conventional
|
|
178
|
+
trunks; its capture, steering and lens cells are unmeasured. Capture *correctness* under freeze is covered separately —
|
|
158
179
|
`tests/test_freeze_parity_gpu.py` holds the harvest to cosine ≥ 0.999 and greedy token-id parity
|
|
159
180
|
against hooked capture, across concurrent requests, chunked prefill and decode rows — and the
|
|
160
181
|
backend self-tests each tap at startup, refusing to serve if a `copy_` produced a non-finite or
|
|
@@ -8,6 +8,8 @@ hooks (steering then capture) in :mod:`~interp_engine.vllm_capture.requests` ins
|
|
|
8
8
|
|
|
9
9
|
from __future__ import annotations
|
|
10
10
|
|
|
11
|
+
from collections.abc import Iterable
|
|
12
|
+
|
|
11
13
|
import torch
|
|
12
14
|
|
|
13
15
|
from interp_engine.hooks import hidden_arg_index
|
|
@@ -15,6 +17,29 @@ from interp_engine.hooks import hidden_from_call as _hidden_from_call
|
|
|
15
17
|
from interp_engine.vllm_capture._payload import select_stream
|
|
16
18
|
from interp_engine.vllm_capture._tree import LAYER_RETURN_INDEX
|
|
17
19
|
|
|
20
|
+
|
|
21
|
+
def position_mask(positions: Iterable[int], num_tokens: int, like: torch.Tensor) -> torch.Tensor:
|
|
22
|
+
"""A boolean row mask over ``positions``, shaped to broadcast against ``like``.
|
|
23
|
+
|
|
24
|
+
One trailing singleton axis per axis ``like`` has after the token one, rather than exactly one.
|
|
25
|
+
On a hyper-connection trunk a point is ``[tokens, streams, width]``, and a ``[tokens, 1]`` mask
|
|
26
|
+
would line its token axis up against the *stream* axis: a shape error for most prompts, and --
|
|
27
|
+
worse -- a silent masking of the wrong thing for a prompt whose length happens to equal the
|
|
28
|
+
stream count.
|
|
29
|
+
|
|
30
|
+
Lives here, in the leaf both write paths already import, because there is one right answer and
|
|
31
|
+
two callers: the hooked path in :mod:`~interp_engine.vllm_capture.requests` and the frozen one
|
|
32
|
+
in :mod:`~interp_engine.vllm_capture.freeze`. It was fixed in the first and not the second while
|
|
33
|
+
they were separate copies, which is how a jlens intervention on a stream stack came to crash on
|
|
34
|
+
a shape under CUDA graphs and work under hooks.
|
|
35
|
+
"""
|
|
36
|
+
mask = torch.zeros(num_tokens, *([1] * (like.dim() - 1)), dtype=torch.bool, device=like.device)
|
|
37
|
+
for position in positions:
|
|
38
|
+
if 0 <= position < num_tokens:
|
|
39
|
+
mask[position] = True
|
|
40
|
+
return mask
|
|
41
|
+
|
|
42
|
+
|
|
18
43
|
# vLLM's decoder-layer convention is `forward(positions, hidden, residual) -> (hidden, residual)`
|
|
19
44
|
# where the returned `hidden` is the part NOT yet added to the residual stream: the fused add+norm at
|
|
20
45
|
# the top of the *next* block does that add, so the stream is the sum of the pair.
|
|
@@ -26,7 +26,7 @@ from interp_engine.facts import is_linear_attention_layer, unclassified_layer_ki
|
|
|
26
26
|
from interp_engine.hooks import hidden_arg_index, hidden_from_call
|
|
27
27
|
from interp_engine.points import steer_refusal_reason
|
|
28
28
|
from interp_engine.vllm_capture._demux import _ensure_patched, _get_demux, _resolve_rid
|
|
29
|
-
from interp_engine.vllm_capture._hooks import layer_return_tensor, returns_full_residual
|
|
29
|
+
from interp_engine.vllm_capture._hooks import layer_return_tensor, position_mask, returns_full_residual
|
|
30
30
|
from interp_engine.vllm_capture._payload import attn_payload_key, decode_capture_payload, encode_tensor_payload
|
|
31
31
|
from interp_engine.vllm_capture._tree import (
|
|
32
32
|
_INPUT_POINTS,
|
|
@@ -1394,7 +1394,15 @@ def _is_fused_resid(module: torch.nn.Module, name: str, residual: torch.Tensor |
|
|
|
1394
1394
|
|
|
1395
1395
|
|
|
1396
1396
|
def _apply_lens_scope(delta: torch.Tensor, n: int, scope: dict[str, Any] | None) -> torch.Tensor | None:
|
|
1397
|
-
"""Prefill-vs-decode skip used by jlens. None means leave the live tensor alone.
|
|
1397
|
+
"""Prefill-vs-decode skip used by jlens. None means leave the live tensor alone.
|
|
1398
|
+
|
|
1399
|
+
The mask comes from :func:`~interp_engine.vllm_capture._hooks.position_mask`, the one the hooked
|
|
1400
|
+
path uses, rather than being built here: a ``[tokens, 1]`` mask is right for every point with one
|
|
1401
|
+
width axis and wrong for a hyper-connection trunk, where the delta is ``[tokens, streams,
|
|
1402
|
+
width]`` and broadcasting -- which pads on the LEFT -- lines the token axis up against the stream
|
|
1403
|
+
axis. There is nothing about the answer that differs between the two paths, so there is no
|
|
1404
|
+
longer a second construction of it here.
|
|
1405
|
+
"""
|
|
1398
1406
|
if not scope:
|
|
1399
1407
|
return delta
|
|
1400
1408
|
is_prefill = n > 1
|
|
@@ -1403,10 +1411,7 @@ def _apply_lens_scope(delta: torch.Tensor, n: int, scope: dict[str, Any] | None)
|
|
|
1403
1411
|
skip = scope.get("skip_positions") or []
|
|
1404
1412
|
prompt_len = int(scope.get("prompt_len") or 0)
|
|
1405
1413
|
if is_prefill and skip and n == prompt_len:
|
|
1406
|
-
mask =
|
|
1407
|
-
for i in skip:
|
|
1408
|
-
if 0 <= int(i) < n:
|
|
1409
|
-
mask[int(i)] = True
|
|
1414
|
+
mask = position_mask((int(i) for i in skip), n, delta)
|
|
1410
1415
|
delta = torch.where(mask, torch.zeros_like(delta), delta)
|
|
1411
1416
|
return delta
|
|
1412
1417
|
|
|
@@ -1607,6 +1612,13 @@ def worker_set_freeze_delta(
|
|
|
1607
1612
|
Additive ``op="add"`` without a lens scope fills the static ``delta`` buffer. Orthogonal,
|
|
1608
1613
|
projection_cap, and lens ops attach a live ``modify`` that reads the residual each
|
|
1609
1614
|
forward (breakable ``add_eager``). ``lens_scope`` is the jlens prefill/decode skip.
|
|
1615
|
+
|
|
1616
|
+
A ``stream`` disqualifies the static buffer too. The buffer is added whole to a ``[tokens,
|
|
1617
|
+
streams, width]`` activation, so it has no way to say "this stream and not the others"; taking
|
|
1618
|
+
the fast path anyway would steer all four streams of a hyper-connection trunk on a request that
|
|
1619
|
+
asked for one, and report success. :func:`~interp_engine.vllm_capture.steering._one_stream`
|
|
1620
|
+
knows how, so a stream is served by the modifier path like every other op that is not a plain
|
|
1621
|
+
constant.
|
|
1610
1622
|
"""
|
|
1611
1623
|
freeze = _state(worker)
|
|
1612
1624
|
if freeze is None:
|
|
@@ -1621,7 +1633,7 @@ def worker_set_freeze_delta(
|
|
|
1621
1633
|
raise ValueError(f"graph freeze cannot serve op={op!r}; supported ops are {sorted(_FREEZE_WRITE_OPS)}")
|
|
1622
1634
|
site = _write_site(freeze, Address(str(spec["point"]), int(spec["layer"])))
|
|
1623
1635
|
assert site.delta is not None
|
|
1624
|
-
if op == "add" and lens_scope is None:
|
|
1636
|
+
if op == "add" and lens_scope is None and spec.get("stream") is None:
|
|
1625
1637
|
vec = torch.tensor(spec["vector"], dtype=torch.float32, device=site.delta.device)
|
|
1626
1638
|
vec = (vec * float(spec["coeff"])).to(dtype=site.delta.dtype)
|
|
1627
1639
|
site.delta.copy_(vec.reshape(1, -1).expand_as(site.delta))
|
|
@@ -1647,7 +1659,9 @@ def _compile_write_req(
|
|
|
1647
1659
|
raise ValueError(f"graph freeze cannot serve op={op!r}; supported ops are {sorted(_FREEZE_WRITE_OPS)}")
|
|
1648
1660
|
assert site.delta is not None
|
|
1649
1661
|
device, dtype = site.delta.device, site.delta.dtype
|
|
1650
|
-
|
|
1662
|
+
# A constant `[1, width]` vector broadcasts over a stream axis and so cannot exclude one; see
|
|
1663
|
+
# `worker_set_freeze_delta` for why a `stream` therefore has to go the modifier way.
|
|
1664
|
+
if op == "add" and spec.get("stream") is None:
|
|
1651
1665
|
vec = torch.tensor(spec["vector"], dtype=torch.float32, device=device)
|
|
1652
1666
|
vec = (vec * float(spec["coeff"])).to(dtype=dtype).reshape(1, -1)
|
|
1653
1667
|
return _WriteReq(
|
|
@@ -14,7 +14,6 @@ cycle around :mod:`~interp_engine.vllm_capture._demux`.
|
|
|
14
14
|
from __future__ import annotations
|
|
15
15
|
|
|
16
16
|
import inspect
|
|
17
|
-
from collections.abc import Iterable
|
|
18
17
|
from typing import cast
|
|
19
18
|
|
|
20
19
|
import torch
|
|
@@ -37,6 +36,7 @@ from interp_engine.vllm_capture._hooks import (
|
|
|
37
36
|
layer_return_tensor,
|
|
38
37
|
returns_full_residual,
|
|
39
38
|
)
|
|
39
|
+
from interp_engine.vllm_capture._hooks import position_mask as _position_mask
|
|
40
40
|
from interp_engine.vllm_capture._payload import (
|
|
41
41
|
attn_payload_key,
|
|
42
42
|
encode_tensor_payload,
|
|
@@ -64,22 +64,6 @@ from interp_engine.vllm_capture.steering import _make_steer_modifier
|
|
|
64
64
|
# --- the per-request combined hook body ---------------------------------------
|
|
65
65
|
|
|
66
66
|
|
|
67
|
-
def _position_mask(positions: Iterable[int], num_tokens: int, like: torch.Tensor) -> torch.Tensor:
|
|
68
|
-
"""A boolean row mask over ``positions``, shaped to broadcast against ``like``.
|
|
69
|
-
|
|
70
|
-
One trailing singleton axis per axis ``like`` has after the token one, rather than exactly one.
|
|
71
|
-
On a hyper-connection trunk a point is ``[tokens, streams, width]``, and a ``[tokens, 1]`` mask
|
|
72
|
-
would line its token axis up against the *stream* axis: a shape error for most prompts, and --
|
|
73
|
-
worse -- a silent masking of the wrong thing for a prompt whose length happens to equal the
|
|
74
|
-
stream count.
|
|
75
|
-
"""
|
|
76
|
-
mask = torch.zeros(num_tokens, *([1] * (like.dim() - 1)), dtype=torch.bool, device=like.device)
|
|
77
|
-
for position in positions:
|
|
78
|
-
if 0 <= position < num_tokens:
|
|
79
|
-
mask[position] = True
|
|
80
|
-
return mask
|
|
81
|
-
|
|
82
|
-
|
|
83
67
|
def _process_point(demux: _Demux, site: Address, full: torch.Tensor) -> torch.Tensor:
|
|
84
68
|
"""Apply per-request steering/lens then capture at one hook ``site``.
|
|
85
69
|
|