Part of #32
Problem
Adding DMI hooks to a new model currently requires manually forking the model file (~187 lines of changes per model): adding HookPoints in __init__, inserting hook calls at specific positions in forward(), and hand-writing get_hook_specs() in forward execution order. This must be repeated per model × per framework.
Goal
A tool that lets users specify hook points at arbitrary positions in a model's forward pass (not just module boundaries), and automatically compiles a hooked model file with all boilerplate generated.
Key requirements
- Arbitrary insertion: hooks often go mid-forward (e.g.,
resid_mid after a residual add), not just at module boundaries. Must support inserting after any statement in forward()
- Frontend for config: users need a way to see the model's forward code and specify where to insert hooks — either a visual frontend, a declarative config with code pattern matching, or annotation-based marking
- Auto-generated boilerplate:
HookPoint() declarations in __init__, HookedRootModule inheritance, get_hook_specs() in correct execution order, hook name normalization — all generated by the compiler from the config
- Multi-framework: should work for both HF and vLLM model files
Feasibility
Python's ast module can parse forward functions and do source-to-source transformation. The compiler knows insertion order so get_hook_specs() can be auto-generated. Main challenge is designing a good UX for specifying arbitrary insertion points and handling naming variation across model architectures.
Tasks
Part of #32
Problem
Adding DMI hooks to a new model currently requires manually forking the model file (~187 lines of changes per model): adding HookPoints in
__init__, inserting hook calls at specific positions inforward(), and hand-writingget_hook_specs()in forward execution order. This must be repeated per model × per framework.Goal
A tool that lets users specify hook points at arbitrary positions in a model's forward pass (not just module boundaries), and automatically compiles a hooked model file with all boilerplate generated.
Key requirements
resid_midafter a residual add), not just at module boundaries. Must support inserting after any statement in forward()HookPoint()declarations in__init__,HookedRootModuleinheritance,get_hook_specs()in correct execution order, hook name normalization — all generated by the compiler from the configFeasibility
Python's
astmodule can parse forward functions and do source-to-source transformation. The compiler knows insertion order soget_hook_specs()can be auto-generated. Main challenge is designing a good UX for specifying arbitrary insertion points and handling naming variation across model architectures.Tasks