Skip to content

feat(patch): multi-hook sweeps + source-run lifecycle#237

Merged
RhizoNymph merged 3 commits into
feat/patch-sweep-spansfrom
feat/patch-multihook-lifecycle
Jul 5, 2026
Merged

feat(patch): multi-hook sweeps + source-run lifecycle#237
RhizoNymph merged 3 commits into
feat/patch-sweep-spansfrom
feat/patch-multihook-lifecycle

Conversation

@RhizoNymph

@RhizoNymph RhizoNymph commented Jul 5, 2026

Copy link
Copy Markdown
Owner

Stacked on #235 (base feat/patch-sweep-spans).

Adds first-class multi-hook activation-patching sweeps and a source-run lifecycle so one-call sweeps stop leaking (and polluting the LRU with) fresh uuid runs.

Multi-hook sweeps

  • New hooks: list[str] | None on PatchSweepRequest. When given (non-empty, each injectable — validated, 400 otherwise) the (layers x positions) grid runs at every hook; hooks wins over the single hook field. The corrupt baseline and noise floor are computed once and shared across hooks; only patched cells fan out per hook.
  • Response gains hook_grids: list[HookGrid] (one {hook, grid, argmax} per hook). Top-level grid/hook/argmax mirror the first hook, so single-hook clients keep working; positions/layers/clean/corrupt/noise_floor/alignment/skipped stay shared (skipped cells carry their hook). hooks=["post_block"] matches hook="post_block" bar the extra hook_grids entry.
  • Client sweep_layers_positions(..., server_side=True, hooks=[..]) returns dict[str, SweepResult] keyed by hook; passing hooks on the per-cell path raises ValueError.

Hook-complete auto-capture

  • _auto_capture_clean now taps all three injectable hooks (pre_attn, post_attn, post_block) at the swept layers in one forward, so a kept auto-captured run is reusable for hook-comparison follow-up sweeps at a different hook.

Source-run lifecycle

  • Auto-drop: after a sweep that auto-captured completes, the fresh run is dropped (recent uuid runs sit at the LRU MRU end and would evict older deliberate captures first). keep_source: bool = False retains it; a pre-existing run is never dropped. Drop failures are log-warn.
  • DELETE /v1/patch_source/{run_id} -> collective_rpc("drop_patch_source_run") (new gpu_worker RPC mirroring lease_patch_source_runs), unioned across ranks: 200 {"dropped": true} or 404. An explicit owner drop succeeds even if the run is leased (store drop_run clears the lease); LRU eviction still never touches a leased run.
  • Both drop paths invalidate the admission-side _PatchSourceCache so a dropped run is not reported as still-existing (which would skip auto-capture and 400 the next sweep).
  • Client PatchStudy.drop_run wired to the DELETE route (returns bool); one-call keep_source passthrough.

Conflicts

Expected merge conflicts with sibling branch feat/patch-sweep-stream (SSE streaming) in api_router.py's fan-out section, protocol.py, and docs/features/activation_patching.md.

@RhizoNymph RhizoNymph merged commit d4d920b into feat/patch-sweep-spans Jul 5, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant