cache

cache

Record module activations during a trace with tracer.cache().

A cache observes the values flowing past the interleaver and keeps the ones for the modules you asked for. Unlike reading a single location with .output, a cache captures every selected module across the whole run — every layer, and (in a generation loop) every step::

with model.trace(prompt) as tracer:
    cache = tracer.cache()                       # every module's output
cache["model.transformer.h.0"].output            # by path
cache.transformer.h[0].output                    # or by navigation

Because the interleaver already funnels every module input/output through handle (applying interventions first), a cache is just a post-intervention observer: it needs no per-module controllers of its own. Navigation and alias/index resolution are delegated to the model's Envoy tree, so a CacheView stays thin and aliases / ModuleList indexing work for free.

Entry dataclass

Entry(output: Any = None, inputs: tuple[tuple[Any, ...], dict[str, Any]] | None = None)

One visit to a module: its output and/or its inputs.

A module visited more than once in a run (e.g. a generation loop) produces one Entry per visit; see CacheView.

output class-attribute instance-attribute

output: Any = None

inputs class-attribute instance-attribute

inputs: tuple[tuple[Any, ...], dict[str, Any]] | None = None

input property

input: Any

The first input — the first positional argument, else the first keyword.

Cache

Cache(model: Envoy, modules: list[Envoy | str] | None = None, device: device | None = device('cpu'), dtype: dtype | None = None, detach: bool = True, include_output: bool = True, include_inputs: bool = False, non_blocking: bool = False)

Records selected modules' activations as the run reaches them.

Created by cache and registered on the calling mediator, so observe is fed every location the run reaches (post-intervention). Values for the selected module paths are stored in entries — one list of Entry per module path, an entry appended per visit.

ATTRIBUTE	DESCRIPTION
model	The root envoy, used to resolve paths / aliases for CacheView.
targets	The module paths to keep, or None to keep every module. TYPE: set[str] | None
entries	Recorded values, {module_path: [Entry, ...]}. TYPE: dict[str, list[Entry]]

model instance-attribute

model = model

device instance-attribute

device = device

dtype instance-attribute

dtype = dtype

detach instance-attribute

detach = detach

non_blocking instance-attribute

non_blocking = non_blocking

include_output instance-attribute

include_output = include_output

include_inputs instance-attribute

include_inputs = include_inputs

targets instance-attribute

targets: set[str] | None = None if modules is None else {m if isinstance(m, str) else m.path for m in modules}

entries instance-attribute

entries: dict[str, list[Entry]] = {}

__getstate__

__getstate__() -> dict

wants

wants(location: str) -> bool

Whether this cache would record location.

The same question observe answers by recording, asked without recording — for a caller that has to do work before the value can be offered, and only wants to do it for locations some cache actually keeps (see TPFragments, which must run a collective first). Shares _select with observe so the two can't drift.

subscriptions

subscriptions() -> 'dict[str, tuple[str, str]]'

Every location this cache keeps, named before the model starts.

The interleaver routes a value straight to the caches subscribed to its location, so a cache has to know its locations up front: an explicit module list names them, and modules=None means every module in the tree, read off the model here.

observe

observe(location: str, value: Any) -> None

Record value if location is a selected module's input/output.

observe_selected

observe_selected(selected: tuple[str, str], value: Any) -> None

Record a location the interleaver has already selected.

This is the fast counterpart to observe: an explicit-target cache receives its (path, slot) subscription from the interleaver, so there is no second provider-string parse on the hot path.

CacheView

CacheView(cache: Cache, envoy: Envoy | None, path: str | None = None)

Path- and attribute-addressable view over a Cache's entries.

The object returned by tracer.cache(). Read a module's captured value with .output / .inputs / .input after selecting it, either by path (cache["model.transformer.h.0"]) or by navigating the tree (cache.model.transformer.h[0] — or the short cache.transformer.h[0]). Navigation is resolved against the model's envoy tree, so renamed modules and ModuleList indices resolve the same way they do on the model (cache.model.transformer.h["second_layer"] works when 1 is renamed second_layer). When a module was visited multiple times (a generation loop), len(view) is the visit count and .output returns the list.

The object handed back is a virtual root above the model: its child is the model itself, reached by the model's name (cache.model), mirroring the full paths used as keys. _envoy is None there and a real envoy at every node below.

output property

output: Any

inputs property

inputs: Any

input property

input: Any

keys

keys() -> list[str]

The cached module paths at or below this view's node (all, at the root).

__getattr__

__getattr__(name: str) -> CacheView

__getitem__

__getitem__(key: str | int) -> CacheView

__contains__

__contains__(key: str) -> bool

__len__

__len__() -> int

__repr__

__repr__() -> str