Skip to content

Hook Interface & Model Compiler — Auto-inject HookPoints into model files #36

Description

@Samfisheryu

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

  • Design the hook config schema and frontend for specifying arbitrary insertion points
  • Prototype AST-based compiler (parse forward → insert hooks → generate boilerplate)
  • Validate: auto-compiled Qwen3 should match the existing hand-written hooked version
  • Extend to vLLM model variant generation

No activity

Activity on this issue will appear here.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions