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.
Files changed (173) hide show
  1. {interp_engine-1.2.7 → interp_engine-1.2.8}/PKG-INFO +23 -34
  2. interp_engine-1.2.8/README.md +117 -0
  3. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/README.md +24 -6
  4. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/bench_spec.py +6 -4
  5. {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/PERFORMANCE.md +30 -9
  6. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/_hooks.py +25 -0
  7. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/freeze.py +22 -8
  8. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/requests.py +1 -17
  9. {interp_engine-1.2.7 → interp_engine-1.2.8}/pyproject.toml +1 -1
  10. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_hyper_connections.py +276 -0
  11. interp_engine-1.2.7/README.md +0 -128
  12. {interp_engine-1.2.7 → interp_engine-1.2.8}/.gitignore +0 -0
  13. {interp_engine-1.2.7 → interp_engine-1.2.8}/LICENSE +0 -0
  14. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/__init__.py +0 -0
  15. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/cells.py +0 -0
  16. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/probe.py +0 -0
  17. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/publish.py +0 -0
  18. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/report_bench.py +0 -0
  19. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__eager.json +0 -0
  20. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__vllm-cudagraph.json +0 -0
  21. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__vllm-dspark-cudagraph.json +0 -0
  22. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__vllm-dspark.json +0 -0
  23. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__vllm-freeze.json +0 -0
  24. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/deepseek-v4-flash-0731__vllm.json +0 -0
  25. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/gemma-2-2b__eager.json +0 -0
  26. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/gemma-2-2b__vllm-cudagraph.json +0 -0
  27. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/gemma-2-2b__vllm-freeze.json +0 -0
  28. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/gemma-2-2b__vllm.json +0 -0
  29. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/llama-3.1-8b__eager.json +0 -0
  30. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/llama-3.1-8b__vllm-cudagraph.json +0 -0
  31. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/llama-3.1-8b__vllm-freeze.json +0 -0
  32. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/llama-3.1-8b__vllm.json +0 -0
  33. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3-4b__eager.json +0 -0
  34. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3-4b__vllm-cudagraph.json +0 -0
  35. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3-4b__vllm-freeze.json +0 -0
  36. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3-4b__vllm.json +0 -0
  37. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3.8-27b__eager.json +0 -0
  38. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3.8-27b__vllm-cudagraph.json +0 -0
  39. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3.8-27b__vllm-freeze.json +0 -0
  40. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results/qwen3.8-27b__vllm.json +0 -0
  41. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/results-latest.md +0 -0
  42. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/run_all.sh +0 -0
  43. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/run_bench.py +0 -0
  44. {interp_engine-1.2.7 → interp_engine-1.2.8}/benchmarks/workloads.py +0 -0
  45. {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/AGENT_INTEGRATION.md +0 -0
  46. {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/ARCHITECTURE_QUIRKS.md +0 -0
  47. {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/COMPATIBILITY.md +0 -0
  48. {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/ENGINE_HOOK_MAPPINGS.md +0 -0
  49. {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/GRADIENTS.md +0 -0
  50. {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/INTERNALS.md +0 -0
  51. {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/PORTING.md +0 -0
  52. {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/README.md +0 -0
  53. {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/SUPPORTED_POINTS.md +0 -0
  54. {interp_engine-1.2.7 → interp_engine-1.2.8}/docs/USAGE.md +0 -0
  55. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/__init__.py +0 -0
  56. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/_loop.py +0 -0
  57. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/address.py +0 -0
  58. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/arch.py +0 -0
  59. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/attn_config.py +0 -0
  60. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/attn_scores.py +0 -0
  61. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/autograd_support.py +0 -0
  62. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/capture.py +0 -0
  63. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/chat_compose.py +0 -0
  64. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/chat_conventions.py +0 -0
  65. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/chat_formatters.py +0 -0
  66. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/cuda_preflight.py +0 -0
  67. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/dispatch.py +0 -0
  68. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/facts.py +0 -0
  69. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/hooks.py +0 -0
  70. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/lens.py +0 -0
  71. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/load.py +0 -0
  72. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/mappers.py +0 -0
  73. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/model.py +0 -0
  74. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/moe_routing.py +0 -0
  75. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/points.py +0 -0
  76. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/protocol.py +0 -0
  77. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/residual_basis.py +0 -0
  78. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/select.py +0 -0
  79. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/steer.py +0 -0
  80. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/steer_specs.py +0 -0
  81. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/sync.py +0 -0
  82. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/tokenize.py +0 -0
  83. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_backend.py +0 -0
  84. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/__init__.py +0 -0
  85. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/_demux.py +0 -0
  86. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/_payload.py +0 -0
  87. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/_tree.py +0 -0
  88. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/attn.py +0 -0
  89. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/capture.py +0 -0
  90. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/graphs.py +0 -0
  91. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/lens/__init__.py +0 -0
  92. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/lens/intervene.py +0 -0
  93. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/lens/readout.py +0 -0
  94. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/lens/unembed.py +0 -0
  95. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/mhc.py +0 -0
  96. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/native.py +0 -0
  97. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_capture/steering.py +0 -0
  98. {interp_engine-1.2.7 → interp_engine-1.2.8}/interp_engine/vllm_plugin.py +0 -0
  99. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/conftest.py +0 -0
  100. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/harness.py +0 -0
  101. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/model_expectations.yaml +0 -0
  102. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/synthetic_families.py +0 -0
  103. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_address.py +0 -0
  104. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_attn_config_tripwire.py +0 -0
  105. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_attn_probs_indexing.py +0 -0
  106. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_attn_scores.py +0 -0
  107. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_attn_z_gqa.py +0 -0
  108. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_autograd_support.py +0 -0
  109. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_bench_workloads.py +0 -0
  110. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_capability_refusals.py +0 -0
  111. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_capture_addressing.py +0 -0
  112. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_chat_compose.py +0 -0
  113. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_chat_formatters.py +0 -0
  114. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_chat_templates.py +0 -0
  115. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_core.py +0 -0
  116. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_cuda_preflight.py +0 -0
  117. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_doc_code_fences.py +0 -0
  118. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_eager_autograd.py +0 -0
  119. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_facts.py +0 -0
  120. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_family_points.py +0 -0
  121. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_freeze_dsv4_gpu.py +0 -0
  122. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_freeze_parity_gpu.py +0 -0
  123. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_freeze_set.py +0 -0
  124. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_freeze_warmup.py +0 -0
  125. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_gated_attn_out.py +0 -0
  126. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_head_contributions.py +0 -0
  127. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_hook_call_conventions.py +0 -0
  128. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_layer_kinds.py +0 -0
  129. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_load.py +0 -0
  130. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_logit_transform.py +0 -0
  131. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_mappers.py +0 -0
  132. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_mlp_internals.py +0 -0
  133. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_model_expectations.py +0 -0
  134. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_moe.py +0 -0
  135. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_multimodal_arch.py +0 -0
  136. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_new_models_gpu.py +0 -0
  137. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_no_chat_template.py +0 -0
  138. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_normalized_hook.py +0 -0
  139. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_packaging.py +0 -0
  140. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_parity_gpt2.py +0 -0
  141. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_per_layer_attn_dims.py +0 -0
  142. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_points_registry.py +0 -0
  143. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_protocol.py +0 -0
  144. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_published_benchmarks.py +0 -0
  145. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_qk_norm.py +0 -0
  146. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_qkv_layout.py +0 -0
  147. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_reasoning_spans.py +0 -0
  148. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_release.py +0 -0
  149. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_resid_mid.py +0 -0
  150. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_residual_basis.py +0 -0
  151. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_sandwich_norms.py +0 -0
  152. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_select.py +0 -0
  153. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_sliding_window_attn.py +0 -0
  154. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_small_models_gpu.py +0 -0
  155. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_steer_context.py +0 -0
  156. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_steer_math_parity.py +0 -0
  157. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_sync_loop.py +0 -0
  158. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_sync_parity.py +0 -0
  159. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_unified_free_functions.py +0 -0
  160. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_unresolved_families.py +0 -0
  161. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_capture_gpu.py +0 -0
  162. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_capture_scales.py +0 -0
  163. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_graph_path.py +0 -0
  164. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_graphs_on_gpu.py +0 -0
  165. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_hook_availability.py +0 -0
  166. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_kv_isolation.py +0 -0
  167. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_new_points.py +0 -0
  168. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_only_families.py +0 -0
  169. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_plugin.py +0 -0
  170. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vllm_wire_grammar.py +0 -0
  171. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_vocabulary_boundary.py +0 -0
  172. {interp_engine-1.2.7 → interp_engine-1.2.8}/tests/test_worker_lens_capture_readout.py +0 -0
  173. {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.7
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
- <p align="center">
35
- <img src="https://neuronpedia.s3.amazonaws.com/site-assets/ielogo.png" alt="interp-engine logo, a magnifying glass where the handle is a rocket" width="160">
36
- </p>
37
- <p align="center">
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
- <p align="center">
51
- <img src="https://neuronpedia.s3.amazonaws.com/site-assets/ie-perf.png" alt="Tokens per second while capturing and generating: eager against IE-vLLM and IE-vLLM-freeze on qwen3.8-27b and deepseek-v4-flash-0731, 8 requests in flight" width="100%">
52
- </p>
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. [`validator/`](validator/), which compares/validates it against TransformerLens, and nnsight/nnterp on real architectures.
60
- 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.
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
- model = load_model("Qwen/Qwen3-8B") # vLLM default, use backend='eager' to override
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 _without giving up capture or steering_. Every column below is capture-capable.
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
- Every multiplier is against eager on the same workload, and is the ratio of the _unrounded_ figures — tok/s is printed whole at 10 and above, so dividing two cells by hand can differ in the last place. Eager's generation loop is synchronous underneath, so eight requests serialize rather than batch — which is why it earns almost nothing from the second table and why the multipliers there are so much larger. `deepseek-v4-flash-0731` is block-quantized FP8 on both backends rather than bf16, so read its row as two backends serving the same quantized weights; it also keeps a tenth, because at 3 tok/s that digit is worth several percent. [`benchmarks/results-latest.md`](benchmarks/results-latest.md) has every figure at full precision.
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 [`validator/`](validator/).
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"` installs **reads** only, and a steering op needs a write tap to land in, so this row priced
229
- half the feature and reported the other half as `n/a` -- with a message that blamed graph replay for
230
- it, which is the thing freeze exists to work around. It now also passes `freeze_writes`, whose value
231
- is the sentinel `run_bench.STEER_WRITES` rather than a site: the `steer` workload writes mid-stack,
232
- that layer differs per model, and a freeze write is a `load_model` argument, so it has to be resolved
233
- from the config before a model exists to ask.
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` is what makes the `steer` cell a number instead of `n/a`. `"auto"` installs
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. The value is the sentinel `run_bench.STEER_WRITES`, resolved there to the mid-stack
271
- # `resid_post` the workload steers, because the layer differs per model and a freeze write has to
272
- # be named before the model exists.
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, via static `copy_` taps. Keeps graph
130
- replay and turns Dynamo off (`VLLM_USE_BREAKABLE_CUDAGRAPH=1`), so the win is replay without
131
- compile. Reads only: name write sites with `freeze_writes`, which wraps them for steering. Because
132
- breakable `add_eager` keeps the wrap as ordinary PyTorch on the live tensor, the write path is not
133
- limited to additive — `orthogonal`, `projection_cap` and the j-lens `steer`/`ablate`/`swap` ops all
134
- ride the same wrap, per-request and with `position_mask` honoured (`register_freeze_write`). An op
135
- outside that set is refused rather than silently skipped.
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 on the three conventional trunks above; its capture,
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 = torch.zeros(n, 1, dtype=torch.bool, device=delta.device)
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
- if op == "add":
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
 
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "interp-engine"
3
- version = "1.2.7"
3
+ version = "1.2.8"
4
4
  description = "A fast, standardized interpretability engine that supports most modern models and architectures. Powers Neuronpedia."
5
5
  readme = "README.md"
6
6
  license = "Apache-2.0"