interleaver¶
interleaver
¶
IEnvoy
¶
Bases: Protocol
Interface for objects that participate in the interleaving system.
Any object that uses :class:eproperty descriptors must satisfy this
protocol by providing:
| ATTRIBUTE | DESCRIPTION |
|---|---|
interleaver |
The :class:
TYPE:
|
path |
Optional provider path prefix used to build requester/provider
strings (e.g.
TYPE:
|
Notes
Implementors that have no meaningful path (e.g. tracers) do not
need to declare a path attribute — :meth:eproperty._build_requester
uses getattr(obj, "path", "") so a missing attribute is treated
the same as None / "". The attribute is declared
Optional[str] here for type clarity.
eproperty
¶
A descriptor for defining hookable properties on :class:IEnvoy objects.
eproperty exposes values through the interleaving request/swap
mechanism. During a trace, reading an eproperty issues a blocking
request to the interleaver; writing to it schedules a swap.
The decorated stub¶
Every eproperty is defined by decorating a stub method. The body
of that stub is never executed for its return value — it is a
placeholder whose only jobs are:
- Donate its
__name__and__doc__to the descriptor (the name becomes the defaultkey; the docstring is what users see inhelp(model.transformer.h[0].output)). - Carry the decorators stacked on top of it that perform the real
work — registering the PyTorch hook (or operation hook) that will
eventually deliver the value the user is about to
request().
Concretely::
@eproperty()
@requires_output # ← does the work: registers a one-shot
def output(self): ... # forward hook on self._module before
# the request blocks the worker thread.
On every __get__ the descriptor calls self._hook(obj) (the
decorated stub). Because the decorator wraps the empty stub, that
invocation runs the decorator's pre-setup — typically registering the
appropriate hook so the value will arrive when the model executes —
and then calls the (no-op) stub. The descriptor then issues the
actual request(requester) call, which blocks until the hook fires.
The pre-setup decorators live in :mod:nnsight.intervention.hooks:
- :func:
requires_output/ :func:requires_input— module-level one-shot forward / pre-forward hooks (used byEnvoy). - :func:
requires_operation_output/ :func:requires_operation_input— operation-level hooks for.sourcetracing (used byOperationEnvoy). - Custom backends (e.g. vLLM) supply their own decorators in the same
pattern; the contract is "make sure a provider for this requester
string will fire before
request()blocks".
A bare eproperty with no pre-setup decorator is also valid for
things that are provided externally — e.g. InterleavingTracer.result
is fed by Envoy.interleave calling self.interleaver.handle("result", ...),
so no per-access hook setup is needed.
Path / key resolution¶
The interleaver is obtained from obj.interleaver. The path prefix
is obtained from obj.path if the attribute exists and is truthy;
if absent or empty, the key alone is used as the requester string.
This is how tracer-level eproperties like InterleavingTracer.result
work without a path prefix.
Supported implementors¶
Any class satisfying the :class:IEnvoy protocol can host
eproperties. In this codebase that includes:
- :class:
Envoy— module-level.output,.input,.inputs - :class:
OperationEnvoy— operation-level.output,.input,.inputs(source tracing) - :class:
InterleavingTracer— tracer-level.result - vLLM :class:
VLLM—.logits,.samples
| PARAMETER | DESCRIPTION |
|---|---|
key
|
The interleaving key appended to
TYPE:
|
description
|
A short label shown in the repr tree. Only eproperties with a description appear in the tree.
TYPE:
|
iterate
|
Whether to append an iteration suffix (
TYPE:
|
__call__
¶
__call__(hook: Callable[..., T]) -> 'T | eproperty'
Register the decorated stub.
hook is the user's stub method (e.g. def output(self): ...)
with any pre-setup decorators from :mod:nnsight.intervention.hooks
already applied. The body is treated as a no-op; what matters is
what the decorators do when hook(obj) is invoked from
:meth:__get__ — typically registering a one-shot PyTorch hook so
the value will be produced by the time the request blocks. The
stub's __name__ becomes the default key.
postprocess
¶
Register a post-processing function called on __set__.
Runs on the user-supplied value just before it is swapped into the
running model. Used by :class:Envoy.input to repack a single value
back into the (args, kwargs) shape the model's hook expects.
preprocess
¶
Register a pre-processing function called on __get__.
Runs on the raw value pulled from the interleaver before it is
returned to the user. Used by :class:Envoy.input to extract the
first positional argument from (args, kwargs).
When a corresponding :meth:transform is also registered, the value
returned by preprocess is captured by the transform's closure —
in-place mutations the user makes are visible inside the transform.
transform
¶
Register a one-shot __get__ -> swap-back transform.
transform complements :meth:preprocess. When preprocess
returns a new object (a clone, a reshape, a view onto a slice),
in-place edits the user makes to that object are invisible to the
running model — the model still holds the original value. transform
closes that loop: at request time the preprocessed value is bound into
the callable via functools.partial and parked on the current
mediator; once the user is done with their edits and the worker yields
control, the mediator invokes the transform and batcher.swaps the
return value back into the model.
The function signature is transform() -> Any (no args — the value
is captured by the closure). Whatever the transform returns replaces
the original model-side value for the rest of the forward pass.
Use cases¶
- Safe mutable view:
preprocessreturnsvalue.clone()so users canthing[:] = 0without aliasing surprises; the transform returns the (mutated) clone so the model still sees their edits. - Per-head attention access:
preprocessreshapes[B, S, H]into[B, n_heads, S, head_dim]; the transform reshapes back to[B, S, H]so the model continues with the user-edited heads.
Example::
class MyEnvoy(Envoy):
@eproperty(key="output")
@requires_output
def heads(self): ...
@heads.preprocess
def heads(self, value):
# Expose attention heads as a separate dim.
B, S, H = value.shape
return value.view(B, S, self.n_heads, H // self.n_heads) .transpose(1, 2)
@heads.transform
@staticmethod
def heads(value):
# Reshape back to the model's [B, S, H] layout.
return value.transpose(1, 2).reshape(value.shape[0], value.shape[2], -1)
Notes¶
- Transform is one-shot per access — it fires once when the value event for that access is processed and is then cleared.
- The transform can be a plain function (commonly decorated with
@staticmethod) since the preprocessed value already carries the context via the closure. - If you only want to view the value and don't intend to swap it
back, omit
transformand just usepreprocess.
Events
¶
Bases: Enum
Enum for different types of events in the interleaving process.
Cancelation
¶
Bases: Exception
Exception raised when a request is canceled.
EarlyStopException
¶
Bases: Exception
Exception raised to stop the execution of the model.
SkipException
¶
Bases: Exception
Exception raised to skip the execution of the model.
Interleaver
¶
Interleaver(mediators: List[Mediator] = [], tracer: InterleavingTracer = None, batcher: Batcher = None)
Manages the interleaving of model execution and interventions.
This class coordinates the flow between the model's forward pass and user-defined intervention functions, allowing for inspection and modification of intermediate values.
| ATTRIBUTE | DESCRIPTION |
|---|---|
mediators |
A dictionary of mediator names to mediator objects. Each meidator is responsible for a single invoke, or intervention function.
TYPE:
|
tracer |
The tracer object that created this interleaver. Occationaly useful to know the tracer type for this interleaving.
TYPE:
|
batcher |
The batcher object that manages the slice of inputs associtated with each mediator.
TYPE:
|
current |
The current mediator that is being processed. Must be update before resuming a given mediator.
TYPE:
|
| PARAMETER | DESCRIPTION |
|---|---|
mediators
|
A list of mediator objects.
TYPE:
|
tracer
|
The tracer object that created this interleaver.
TYPE:
|
batcher
|
The batcher object that manages the slice of inputs associtated with each mediator.
TYPE:
|
interleaving
property
¶
Check if the interleaver is currently interleaving.
| RETURNS | DESCRIPTION |
|---|---|
bool
|
True if the interleaver is interleaving, False otherwise
TYPE:
|
initialize
¶
initialize(mediators: List[Mediator], tracer: InterleavingTracer, batcher: Batcher = None)
cancel
¶
Cancel all mediators / intervention threads.
After each mediator's worker thread is torn down, every hook it
registered (module one-shot, cache, operation, gradient, iter
tracker) is removed via :meth:Mediator.remove_hooks. This is
the single cleanup path for all dynamic hooks registered during
the session.
iterate_requester
¶
Append the current mediator's iteration index to a requester string.
The iteration is determined by one of two sources:
- If
mediator.iterationis set (user is inside an explicittracer.iter[i]loop), use that value directly. This is how iterator tracers constrain requests to a specific generation step. - If
mediator.iterationisNone(user is not in an iter context, or a one-shot hook has just cleared it after matching), fall back tomediator.iteration_tracker[requester]. The tracker is maintained by persistent hooks registered by :class:IteratorTracer(see :func:register_iter_hooks), which increment it after every forward pass for each module path.
This dual-mode behavior lets the same requester syntax work both inside and outside an iter loop.
| PARAMETER | DESCRIPTION |
|---|---|
requester
|
The base requester string (e.g.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
|
The requester with iteration suffix (e.g. |
wrap_module
¶
Prepare a module for lazy hook execution.
Unlike previous versions that registered permanent input/output hooks on every module, this method only installs:
- A skippable forward wrapper — replaces
module.forwardwith a thin wrapper that checks for__nnsight_skip__in kwargs. If present, the module's original forward is bypassed and the skip value is returned directly. - A sentinel output hook — an empty
register_forward_hookthat returnsoutputunchanged. This is required because PyTorch'sModule.__call__fast-paths when no hooks are registered: if a module has zero hooks at call time, dynamically added hooks during the forward pass will never fire. The sentinel ensures PyTorch always goes through the hook dispatch path so that one-shot hooks registered by :func:hooks.input_hookand :func:hooks.output_hookcan be picked up mid-forward.
Actual interception of inputs/outputs is handled lazily by one-shot
hooks registered on-demand by each mediator (see hooks.py).
| PARAMETER | DESCRIPTION |
|---|---|
module
|
The PyTorch module to prepare.
TYPE:
|
handle
¶
Broadcast a provider value to all mediators.
Used by :meth:eproperty.provide and :meth:Envoy.interleave to
push values (e.g. vLLM logits, generation results) into the
interleaving system. Unlike module hooks, these values are not
produced by the normal PyTorch forward pass — they are pushed
from the model runner side, so this method acts as a fan-out to
every mediator and bumps the per-mediator iteration counter for
the provider path when iterate=True.
When iterate=True, the per-mediator iteration_tracker for
this provider path is bumped after each mediator processes the
value. This mirrors the behavior of the persistent iter hooks
registered by :class:IteratorTracer, but for values that flow
through provide() instead of a PyTorch forward hook.
| PARAMETER | DESCRIPTION |
|---|---|
provider
|
The provider string identifying this value.
TYPE:
|
value
|
The value being provided.
TYPE:
|
iterate
|
Whether to append an iteration suffix to the provider and bump the per-mediator tracker.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
|
The (potentially modified) value after the final mediator has |
|
|
handled it. Used by |
|
|
source tracing, where the injected function is returned from |
|
|
meth: |
|
|
through this broadcast. |
Mediator
¶
Mediator(intervention: Callable, info: 'Tracer.Info', name: Optional[str] = None, batch_group: Optional[List[int]] = None, stop: Optional[int] = None)
Mediates between the model execution and a single intervention function.
This class handles the communication between the model's forward pass and a user-defined intervention function, allowing for inspection and modification of intermediate values.
| ATTRIBUTE | DESCRIPTION |
|---|---|
interleaver |
The interleaver that this mediator is currently running in
TYPE:
|
intervention |
The intervention function to mediate
TYPE:
|
info |
Information about the tracing context associated with this mediator
TYPE:
|
name |
Optional name for the mediator
TYPE:
|
batch_group |
Optional batch group for the mediator to determine which slice of tensors are being intervened on
TYPE:
|
event_queue |
Where the mediator (worker thread) puts events to be processed by the interleaver (main thread). Will only ever have 1 or 0 items in the queue.
TYPE:
|
response_queue |
Where the interleaver (main thread) puts responses to events, to then be processed by the mediator (worker thread). Will only ever have 1 or 0 items in the queue.
TYPE:
|
worker |
The thread that runs the intervention function
TYPE:
|
history |
A set of providers that have been seen by the mediator. Used to detect out of order interventions.
TYPE:
|
iteration_tracker |
Per-provider-path counter maintained by
:func:
TYPE:
|
iteration |
The target iteration this mediator is
currently constrained to, or
TYPE:
|
user_cache |
A list of caches to be used by the mediator
TYPE:
|
all_stop |
Optional number of times to execute this mediator
TYPE:
|
| PARAMETER | DESCRIPTION |
|---|---|
intervention
|
The intervention function
TYPE:
|
info
|
Information about the tracing context
TYPE:
|
name
|
Optional name for the mediator
TYPE:
|
stop
|
Optional number of times to execute this mediator
TYPE:
|
frame
property
¶
Get the frame of the intervention function.
| RETURNS | DESCRIPTION |
|---|---|
FrameType
|
The frame of the intervention function |
MissedProviderError
¶
Bases: Exception
Exception raised when a provider is missed.
OutOfOrderError
¶
Bases: MissedProviderError
Exception raised when interventions are defined out of order.
Value
¶
start
¶
start(interleaver: Interleaver)
Start the mediator's intervention thread.
| PARAMETER | DESCRIPTION |
|---|---|
interleaver
|
The interleaver managing this mediator
TYPE:
|
handle
¶
Process a provided value against this mediator's pending event.
Called directly by one-shot hooks (see hooks.py) or by
:meth:Interleaver.handle for source operations. This method
saves and restores the interleaver's current mediator and the
batcher's current_value / current_provider, so it is safe
to call from within any hook without corrupting shared state.
Depending on the pending event, this will:
- VALUE: Deliver the value to the worker thread.
- SWAP: Replace the value in the batcher.
- SKIP: Inject __nnsight_skip__ into kwargs so
nnsight_forward bypasses the module.
- BARRIER: Synchronize participating mediators.
- END: Cancel the mediator (intervention finished).
- EXCEPTION: Re-raise the exception from the worker thread.
| PARAMETER | DESCRIPTION |
|---|---|
provider
|
The provider string (e.g.
TYPE:
|
value
|
The value being provided (e.g. the module's output tensor).
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
|
The (potentially modified) value from |
handle_value_event
¶
Handle a value event by providing the requested value or recording a missed provider.
| PARAMETER | DESCRIPTION |
|---|---|
requester
|
The identifier of the requester
TYPE:
|
provider
|
The identifier of the provider
TYPE:
|
Returns: bool: Indicating whether the request was fulfilled by this processor, If so, continue processing events.
handle_swap_event
¶
Handle a swap event by swapping the value if the provider matches the requester.
| PARAMETER | DESCRIPTION |
|---|---|
requester
|
The identifier of the requester
TYPE:
|
provider
|
The identifier of the provider
TYPE:
|
swap_value
|
The value to swap in
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
bool
|
Indicating whether the swap was fulfilled by this processor, If so, continue processing events. |
handle_exception_event
¶
Handle an exception event by raising the exception.
| PARAMETER | DESCRIPTION |
|---|---|
exception
|
The exception to raise
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
bool
|
Flag to stop processing events. |
handle_barrier_event
¶
Handle a barrier event by setting a barrier.
Propagates each participant's nested handle return value back
into batcher.current_value. Without this, a SWAP fired in
a participant's body during the barrier walk produces a new
tensor (concat path) that Mediator.handle's prev_value
restore immediately discards — making cross-invoke transfers
of swapped values silently no-op. The nested handle's return
value is the post-restore value (captured before the restore
runs in :meth:Mediator.handle), so re-assigning it here
carries the swap forward to the outer handle context.
handle_skip_event
¶
Handle a skip event by appending the replacement value into kwargs.
Instead of raising a SkipException (as in the old permanent-hook
approach), each mediator's replacement value is appended to a list
under kwargs["__nnsight_skip__"] along with that mediator's
batch_group. The nnsight_forward wrapper installed by
:meth:Interleaver.wrap_module consumes the list, sorts by batch
start, and concats along dim 0 so multi-invoke skips produce a
full-batch tensor that downstream modules can consume.
This approach works with the one-shot hook system because the input
hook has access to (args, kwargs) via the batcher and can
modify kwargs in-place before the forward call.
respond
¶
Respond from the interleaver (main thread) to the mediator (worker thread) the value for a pending event.
| PARAMETER | DESCRIPTION |
|---|---|
value
|
The value to provide
TYPE:
|
send
¶
send(event: Events, requester: Any)
Send an event to interleaver (main thread) from this mediator (worker thread), and wait for it to be processed by the interleaver.
| PARAMETER | DESCRIPTION |
|---|---|
event
|
The event to send
TYPE:
|
requester
|
The identifier of the requester, plus any additional data for the event.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Any
|
The response from the provider |
request
¶
Request a value from a specific provider.
| PARAMETER | DESCRIPTION |
|---|---|
requester
|
The identifier of the provider to request a value from
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Any
|
The requested value |
swap
¶
Send a swap event to replace the value of a provider.
| PARAMETER | DESCRIPTION |
|---|---|
requester
|
The identifier of the requester
TYPE:
|
value
|
The value to swap in
TYPE:
|
exception
¶
Signal that an exception occurred during intervention.
| PARAMETER | DESCRIPTION |
|---|---|
exception
|
The exception that occurred
TYPE:
|
set_user_cache
¶
Set the user cache for this mediator.
| PARAMETER | DESCRIPTION |
|---|---|
cache
|
The cache to set
TYPE:
|
remove_hooks
¶
Remove every hook registered on behalf of this mediator.
Drains self.hooks — the single list that tracks module one-shot
hooks, cache hooks, operation hooks, gradient hooks, and iter-tracker
hooks. .remove() is idempotent on every handle type used here, so
calling this is safe even if some hooks have already self-removed
after firing.