engine

engine

Carry saved values back out of the engine.

A worker's saved values live in the worker process, and nothing in vLLM's own output carries them. The engine knows when a request finishes, which is when its worker has nothing left to run, so that is where they are fetched and attached to the output the request produced.

Every finished request is collected for, not only the traced ones: a registered block (see nnsight.modeling.vllm.registration) runs on requests nnsight never created, and this is the one place their values can be handed back on the output they belong to.

NNsightLLMEngine

Bases: LLMEngine

An engine that attaches each finished request's saved values to its output.

step

step() -> Any

merge_collected

merge_collected(payloads: list) -> dict

Combine what each rank returned from collect_nnsight.

Not "take the first non-empty": a trace's values come from the rank holding the sampled output, while a registered block's come from whichever rank ran the layers it read — under pipeline parallelism those are different ranks, and taking one would silently drop the other.

Where two ranks did report the same name — tensor parallelism, where every rank runs the block and each gathers the same whole value — the earliest rank's wins. They are equal, but they are on different devices, and a value whose device depended on which rank answered last would be a confusing thing to hand back next to a traced one.

attach

attach(output: Any, entry: dict) -> None

Put a request's collected values on its output.

saves carries both kinds, since from the caller's side they are simply what was saved for this request. nnsight_saves keeps the trace's own apart, because that is what gets pushed back into the trace's variables and a registered value must not land there — see collect_nnsight.

A request that asked for several sampled sequences ran the block once per sequence, so each has values of its own. They go on the completion they belong to — output.outputs[i].saves, alongside that sequence's text and token ids — while output.saves stays the primary sequence's, which is all there is unless n was set.

acollect async

acollect(engine: Any, request_id: str, output: Any = None) -> Optional[dict]

One finished request's collected values, merged across the ranks.

The awaitable form, for the engines whose collective_rpc is a coroutine. output is the finished RequestOutput, which the worker cannot build and needs in order to serve tracer.result; without one this is simply the call that winds the request's workers up and frees what they held.