Skip to content

interleaver

interleaver

T module-attribute

T = TypeVar('T')

NNSIGHT_PREFIX module-attribute

NNSIGHT_PREFIX = '__nnsight'

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:Interleaver managing execution flow.

TYPE: 'Interleaver'

path

Optional provider path prefix used to build requester/provider strings (e.g. "model.transformer.h.0"). May be None or empty — :meth:eproperty._build_requester falls back to the eproperty key alone in that case. This is how tracer-level eproperties such as :attr:InterleavingTracer.result work without a path prefix.

TYPE: Optional[str]

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.

interleaver instance-attribute

interleaver: 'Interleaver'

path instance-attribute

path: Optional[str]

eproperty

eproperty(key: str = None, description: str = None, iterate: bool = True)

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:

  1. Donate its __name__ and __doc__ to the descriptor (the name becomes the default key; the docstring is what users see in help(model.transformer.h[0].output)).
  2. 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 by Envoy).
  • :func:requires_operation_output / :func:requires_operation_input — operation-level hooks for .source tracing (used by OperationEnvoy).
  • 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 obj.path (<path>.<key>). Defaults to the stub function's name. Multiple eproperties can share a key (e.g. Envoy.input and Envoy.inputs both use "input") to provide different views on the same underlying value.

TYPE: str DEFAULT: None

description

A short label shown in the repr tree. Only eproperties with a description appear in the tree.

TYPE: str DEFAULT: None

iterate

Whether to append an iteration suffix (.i0, .i1, …). Defaults to True.

TYPE: bool DEFAULT: True

name instance-attribute

name: str = None

key instance-attribute

key = key

description instance-attribute

description = description

iterate instance-attribute

iterate = iterate

__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

postprocess(func: Callable) -> 'eproperty'

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

preprocess(func: Callable) -> 'eproperty'

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

transform(func: Callable) -> 'eproperty'

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: preprocess returns value.clone() so users can thing[:] = 0 without aliasing surprises; the transform returns the (mutated) clone so the model still sees their edits.
  • Per-head attention access: preprocess reshapes [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 transform and just use preprocess.

__get__

__get__(obj: IEnvoy, owner: Any) -> Any

__set__

__set__(obj: IEnvoy, value: Any)

provide

provide(obj: IEnvoy, value: Any) -> Any

Provide a value from the model side into the interleaving system.

Events

Bases: Enum

Enum for different types of events in the interleaving process.

VALUE class-attribute instance-attribute

VALUE = 'value'

SWAP class-attribute instance-attribute

SWAP = 'swap'

END class-attribute instance-attribute

END = 'end'

EXCEPTION class-attribute instance-attribute

EXCEPTION = 'exception'

SKIP class-attribute instance-attribute

SKIP = 'skip'

BARRIER class-attribute instance-attribute

BARRIER = 'barrier'

Cancelation

Bases: Exception

Exception raised when a request is canceled.

EarlyStopException

Bases: Exception

Exception raised to stop the execution of the model.

SkipException

SkipException(value: Any)

Bases: Exception

Exception raised to skip the execution of the model.

value instance-attribute

value = value

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: Dict[str, Mediator]

tracer

The tracer object that created this interleaver. Occationaly useful to know the tracer type for this interleaving.

TYPE: Optional[InterleavingTracer]

batcher

The batcher object that manages the slice of inputs associtated with each mediator.

TYPE: Batcher

current

The current mediator that is being processed. Must be update before resuming a given mediator.

TYPE: Mediator

PARAMETER DESCRIPTION
mediators

A list of mediator objects.

TYPE: List[Mediator] DEFAULT: []

tracer

The tracer object that created this interleaver.

TYPE: InterleavingTracer DEFAULT: None

batcher

The batcher object that manages the slice of inputs associtated with each mediator.

TYPE: Batcher DEFAULT: None

defer_exceptions instance-attribute

defer_exceptions = False

interleaving property

interleaving: bool

Check if the interleaver is currently interleaving.

RETURNS DESCRIPTION
bool

True if the interleaver is interleaving, False otherwise

TYPE: bool

initialize

initialize(mediators: List[Mediator], tracer: InterleavingTracer, batcher: Batcher = None)

cancel

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

iterate_requester(requester: str)

Append the current mediator's iteration index to a requester string.

The iteration is determined by one of two sources:

  • If mediator.iteration is set (user is inside an explicit tracer.iter[i] loop), use that value directly. This is how iterator tracers constrain requests to a specific generation step.
  • If mediator.iteration is None (user is not in an iter context, or a one-shot hook has just cleared it after matching), fall back to mediator.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. "model.layer.0.output").

TYPE: str

RETURNS DESCRIPTION

The requester with iteration suffix (e.g. "model.layer.0.output.i0").

wrap_module

wrap_module(module: Module)

Prepare a module for lazy hook execution.

Unlike previous versions that registered permanent input/output hooks on every module, this method only installs:

  1. A skippable forward wrapper — replaces module.forward with 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.
  2. A sentinel output hook — an empty register_forward_hook that returns output unchanged. This is required because PyTorch's Module.__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_hook and :func:hooks.output_hook can 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: Module

__enter__

__enter__()

__exit__

__exit__(exc_type, exc_val, exc_tb)

handle

handle(provider: Optional[str] = None, value: Optional[Any] = None, iterate: bool = False)

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: Optional[str] DEFAULT: None

value

The value being provided.

TYPE: Optional[Any] DEFAULT: None

iterate

Whether to append an iteration suffix to the provider and bump the per-mediator tracker.

TYPE: bool DEFAULT: False

RETURNS DESCRIPTION

The (potentially modified) value after the final mediator has

handled it. Used by operation_fn_hook for recursive

source tracing, where the injected function is returned from

meth:Mediator.handle via a SWAP event and flows back

through this broadcast.

check_dangling_mediators

check_dangling_mediators()

check_cache_full

check_cache_full()

Print a warning if a module to be cached was missed.

__deepcopy__

__deepcopy__(memo)

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: Interleaver

intervention

The intervention function to mediate

TYPE: Callable

info

Information about the tracing context associated with this mediator

TYPE: Info

name

Optional name for the mediator

TYPE: Optional[str]

batch_group

Optional batch group for the mediator to determine which slice of tensors are being intervened on

TYPE: Optional[List[int]]

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: SimpleQueue

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: SimpleQueue

worker

The thread that runs the intervention function

TYPE: Thread

history

A set of providers that have been seen by the mediator. Used to detect out of order interventions.

TYPE: Set[str]

iteration_tracker

Per-provider-path counter maintained by :func:IteratorTracer.register_iter_hooks (and by :meth:Interleaver.handle for provided values). One-shot intervention hooks read this to know which generation step is currently firing. Defaults to 0 for any path that has not yet been observed.

TYPE: Dict[str, int]

iteration

The target iteration this mediator is currently constrained to, or None. Set by :class:IteratorTracer before each yield so subsequent requests target the correct step. Cleared back to None by a one-shot hook after it matches a non-zero target, so later requests in the same intervention fall back to the tracker-based "current step" resolution.

TYPE: Optional[int]

user_cache

A list of caches to be used by the mediator

TYPE: List[Cache]

all_stop

Optional number of times to execute this mediator

TYPE: Optional[int]

PARAMETER DESCRIPTION
intervention

The intervention function

TYPE: Callable

info

Information about the tracing context

TYPE: 'Tracer.Info'

name

Optional name for the mediator

TYPE: Optional[str] DEFAULT: None

stop

Optional number of times to execute this mediator

TYPE: Optional[int] DEFAULT: None

intervention instance-attribute

intervention = intervention

name instance-attribute

name = name if name else f'Mediator{id(self)}'

idx instance-attribute

idx = None

info instance-attribute

info = info

batch_group instance-attribute

batch_group = batch_group

interleaver instance-attribute

interleaver = None

event_queue instance-attribute

event_queue = Value()

response_queue instance-attribute

response_queue = Value()

worker instance-attribute

worker = None

skip_container instance-attribute

skip_container = None

history instance-attribute

history = set()

user_cache instance-attribute

user_cache: List['Cache'] = list()

hooks instance-attribute

hooks: List[Any] = list()

iteration_tracker instance-attribute

iteration_tracker = defaultdict(int)

iteration instance-attribute

iteration = 0

all_stop instance-attribute

all_stop: Optional[int] = stop

args instance-attribute

args = list()

cross_invoker instance-attribute

cross_invoker = None

original_globals instance-attribute

original_globals = {}

deferred_exception instance-attribute

deferred_exception = None

transform instance-attribute

transform = None

lock instance-attribute

lock = 0

alive property

alive

frame property

frame: FrameType

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

Value()
value instance-attribute
value = None
lock instance-attribute
lock = allocate_lock()
has_value instance-attribute
has_value = False
get
get()
wait
wait()
put
put(value: Any)
restore
restore(value: Any)

__enter__

__enter__()

__exit__

__exit__(exc_type, exc_val, exc_tb)

start

start(interleaver: Interleaver)

Start the mediator's intervention thread.

PARAMETER DESCRIPTION
interleaver

The interleaver managing this mediator

TYPE: Interleaver

cancel

cancel()

Cancel the intervention thread and its ephemeral state.

handle

handle(provider: Optional[str] = None, value: Optional[Any] = None)

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. "model.layer.0.output.i0").

TYPE: Optional[str] DEFAULT: None

value

The value being provided (e.g. the module's output tensor).

TYPE: Optional[Any] DEFAULT: None

RETURNS DESCRIPTION

The (potentially modified) value from batcher.current_value.

handle_value_event

handle_value_event(requester: Any, provider: Any) -> bool

Handle a value event by providing the requested value or recording a missed provider.

PARAMETER DESCRIPTION
requester

The identifier of the requester

TYPE: str

provider

The identifier of the provider

TYPE: str

Returns: bool: Indicating whether the request was fulfilled by this processor, If so, continue processing events.

handle_swap_event

handle_swap_event(provider: Any, requester: Any, swap_value: Any)

Handle a swap event by swapping the value if the provider matches the requester.

PARAMETER DESCRIPTION
requester

The identifier of the requester

TYPE: str

provider

The identifier of the provider

TYPE: str

swap_value

The value to swap in

TYPE: Any

RETURNS DESCRIPTION
bool

Indicating whether the swap was fulfilled by this processor, If so, continue processing events.

handle_exception_event

handle_exception_event(exception: Exception)

Handle an exception event by raising the exception.

PARAMETER DESCRIPTION
exception

The exception to raise

TYPE: Exception

RETURNS DESCRIPTION
bool

Flag to stop processing events.

handle_barrier_event

handle_barrier_event(provider: Any, participants: Set[str])

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_end_event

handle_end_event()

Handle an end event by stopping the mediator.

handle_skip_event

handle_skip_event(provider: Any, requester: Any, value: Any)

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(value: Optional[Any] = None)

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: Optional[Any] DEFAULT: None

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: Events

requester

The identifier of the requester, plus any additional data for the event.

TYPE: Any

RETURNS DESCRIPTION
Any

The response from the provider

request

request(requester: str)

Request a value from a specific provider.

PARAMETER DESCRIPTION
requester

The identifier of the provider to request a value from

TYPE: str

RETURNS DESCRIPTION
Any

The requested value

swap

swap(requester: str, value: Any)

Send a swap event to replace the value of a provider.

PARAMETER DESCRIPTION
requester

The identifier of the requester

TYPE: str

value

The value to swap in

TYPE: Any

stop

stop()

Stop the execution of the model by raising an EarlyStopException.

skip

skip(requester: Any, value: Any)

end

end()

Signal that execution should continue without further intervention.

exception

exception(exception: Exception)

Signal that an exception occurred during intervention.

PARAMETER DESCRIPTION
exception

The exception that occurred

TYPE: Exception

push

push()

Push local variables to the interleaver state.

pull

pull()

Pull variables from the interleaver state to the frame globals.

set_user_cache

set_user_cache(cache: 'Cache')

Set the user cache for this mediator.

PARAMETER DESCRIPTION
cache

The cache to set

TYPE: 'Cache'

remove_hooks

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.

__getstate__

__getstate__()

Get the state of the mediator for serialization.

__setstate__

__setstate__(state)

Set the state of the mediator for deserialization.