interleaver

interleaver

An interleaver whose hooks survive CUDA-graph replay, at declared taps.

A replayed CUDA graph runs no Python, so the forward hooks Interleaver.instrument installs never fire under enforce_eager=False. vLLM's breakable graphs (VLLM_USE_BREAKABLE_CUDAGRAPH=1) leave a seam: a callable handed to the recording's add_eager while the graph is being captured is run at that point of every replay, against the recording's own tensors. A tap is Interleaver.handle registered as one of those callables — the same handoff a hook makes, replayed.

Two things follow from replay that an eager hook never meets. The value at a location lives at a fixed address the graph rewrites every step, so a kept reference — .save(), a list appended under tracer.iter — aliases memory the next step overwrites: clone what you keep. And the callable's return is discarded, so an edit has to land in place; in-place edits already do, and a replacement swap is copied back into the live tensor, shape-checked.

Taps are the only locations a replayed step reaches. Every other module's handoff lives in its forward, which a replayed graph never runs; a block parked on one is told so when its request ends (see Requests.finish_dangling).

VLLMInterleaver

VLLMInterleaver(taps: Iterable[str] = (), **kwargs: Any)

Bases: Interleaver

An Interleaver that records its handoff into CUDA graphs at taps.

ATTRIBUTE	DESCRIPTION
taps	Full locations ("model.model.layers.3.output") a replayed step serves. Empty on an enforce_eager engine, where this class behaves exactly as its base.

taps instance-attribute

taps = frozenset(taps)

instrument

instrument(envoy: Any) -> None

Route the module through this interleaver, and open its forward if a tap reaches inside.

A tap on an operation — {path}.source.{op}.output — is served by the module's source-instrumented forward, which hands each operation to handle as it runs. That forward is installed lazily, on the first .source access, and the graph recording runs before any trace could ask for it; so install it here, as the tree is built, for exactly the modules the taps name — and check the op exists while a typo is still a load error rather than a request that parks forever.

handle

handle(provider: str, value: Any) -> Any

replay

replay(location: str, live: Any) -> None

Serve location from a replayed graph: handle, then land any edit in place.

A graph is recorded for a batch size, so its tensors carry that many rows and this step fills the first batcher.total of them; the rest is padding no worker owns. Trimmed to a view before the handoff, so a worker is narrowed out of the rows that exist — and an edit copied back into the view lands in the graph's memory all the same.