Repository navigation
Conversation
| # `level` keeps this to the hooks added above — hooks the caller | ||
| # attached before the call survive. | ||
| self.remove_all_hook_fns(level=context_level) | ||
| for hp, direction, handle in added_hooks: |
There was a problem hiding this comment.
With stop_at_layer set and reset_hooks_end=False, the stop hook is kept along with the caller's hooks. Before this change, the next run_with_cache cleared it as a side effect. Now nothing does, so every later forward (including a plain bridge(tokens)) silently returns the residual stream at that layer instead of logits. Could the stop hook be removed unconditionally, since it's internal rather than something the caller asked to keep, with the reset_hooks_end=False case added to the stop-at-layer test?
There was a problem hiding this comment.
Thanks for catching this — you're right that retaining requested hooks must not retain the internal stop hook. run_with_cache() now tracks its stop handle separately and always removes that exact handle in finally, including when reset_hooks_end=False or the forward raises. Requested cache hooks and previously retained caller hooks still survive according to the retention setting.
I expanded the stop-at-layer regression to cover both reset settings and both execution helpers. It checks that the stopped call returns the residual stream, the internal stop hook is gone, requested hooks remain, and both a subsequent plain forward and a subsequent cache call return the baseline logits. Additional exception-path tests cover a forward failing before the stop point.
Validation: 294 passed, 1 skipped, 3 xfailed across the focused hook orchestration/reset/alias/cache/stop-layer/gating tests; no new skips or xfails. Scoped pycln/isort/Black checks and full mypy (388 source files) also pass. The complete test suite was not rerun.
Description
Preserve hooks explicitly retained with
reset_hooks_end=Falsewhen subsequent Bridge helpers clean up their temporary hooks.Retained hooks keep their lexical context level after a helper exits. A later sibling
hooks(),run_with_hooks(), orrun_with_cache()call reuses that level, so level-based cleanup can silently remove those retained hooks. Inrun_with_cache(), this can also remove retained hooks on unrelated hook points. This is an uncovered lifecycle case related to #1636 and #1638: directly registered hooks are already protected, but hooks intentionally retained from an earlier helper are not.Each helper now records the exact
LensHandleinstances it registers and removes only its own temporary handles. Registration-failure cleanup and early-stop hooks use the same ownership rule. Public hook APIs, lexical context levels, explicit level-based resets, permanent-hook behavior, and context-clearing behavior remain unchanged.Regression tests cover all three helpers, forward/backward hooks, same/unrelated hook points, nested retained scopes, exceptions, partial registration failures, caller-owned same-level hooks, and early-stop cleanup. A tiny Native model test also verifies that a retained intervention continues to affect logits after an unrelated cache call and that an explicit reset restores the baseline.
Type of change
Validation
uv run --no-sync mypy .: success, no issues in 388 source files.git diff --check. These are the Makefile-equivalent tools used on native Windows without GNU Make.The complete unit-test suite and GPU tests were not run. The focused run emitted existing stub-adapter alias and dependency deprecation warnings; this patch adds no warning sites. No documentation update is needed because this restores the existing retained-hook contract without changing the public API.
Checklist