GPUModelRunner

GPUModelRunner

Run interventions inside vLLM's worker, against the real weights.

This is where a trace written in another process actually happens. The runner builds its own VLLM over the module vLLM loaded, so the module tree here has the same paths the client wrote against; a worker arriving on a request then resolves straight onto the real modules.

Three points on vLLM's own path carry it:

• _update_states — the scheduler has just decided what runs this step, so new requests hand over their workers and every worker's token span is recomputed.
• execute_model — the forward, run with the interleaver open so hooks serve the parked workers.
• sample_tokens / _sample — logits and sampled ids never pass through a module, so they are offered to workers directly by location.

Workers run as greenlets on this thread. _update_states is called from execute_model, and hooks fire on whichever thread runs the forward, so the worker and the model take strict turns on one thread — there is nothing to synchronize during the forward. Collection (collect_nnsight) is the exception: under Ray it lands on a different thread than the forward, which is why saves and errors are snapshotted onto the mediator on the worker thread (record_saves, finish_dangling) rather than read live at collect time.

Request

Request(request_id: str)

What this engine carries for one of its requests.

vLLM names a request "{external}-{8 hex}" on the way in and an n > 1 child "{index}_{parent}", while the engine asks about — and returns — the external id; both spellings are parsed here once. mediator is the traced worker, if the request is an nnsight trace; copies are the registered blocks' per-request workers, moved to harvested once the request is over; error is a payload that failed to deserialize.

__slots__ class-attribute instance-attribute

__slots__ = ('id', 'stem', 'engine_id', 'index', 'mediator', 'copies', 'harvested', 'error')

id instance-attribute

id = request_id

stem instance-attribute

stem = request_id[:suffix.start()] if suffix else request_id

mediator instance-attribute

mediator: Any = None

copies instance-attribute

copies: dict[str, Any] = {}

harvested instance-attribute

harvested: dict[str, dict] = {}

error instance-attribute

error: Optional[dict] = None

named

named(ids) -> bool

Whether ids (the engine's, or the scheduler's) name this request.

key

key(ids) -> tuple[str, int]

(engine_id, sequence index) as the engine that asked about ids knows it.

The child reading is taken only when the parent it names is one the engine asked about, so an id that merely starts with digits and an underscore is not mistaken for somebody's second sequence.

workers

workers() -> list

Registered copies first, so a trace reads what they left behind.

saves

saves() -> dict

The trace's block-scope names that were marked with .save().

deferred

deferred() -> Optional[dict]

The request's deferred error, captured for the client, or None.

The request's own error first — a block that failed to deserialize, or an edits= name nothing is installed under — then the traced block's.

Requests

Requests()

The workers riding this engine's in-flight requests.

Two kinds of worker run here. A traced one arrives on its request, one per request, and goes home when that request finishes. A registered one is a block left on the engine ahead of time (see nnsight.modeling.vllm.registration): the template is deserialized once, and every request that arrives afterwards — whether or not it is an nnsight trace at all — gets a fresh copy with a scope of its own, whose saves are kept here until they are collected.

ATTRIBUTE	DESCRIPTION
requests	Request id -> Request, for as long as anything of it is still wanted. TYPE: dict[str, Request]
templates	Registration id -> the deserialized block each request's copy is built from, so the source is compiled once rather than per request. TYPE: dict[str, Any]
names	Registration id -> the name it was installed under, or None. A request that names the edits it wants (extra_args["nnsight_edits"]) gets copies of those and of every unnamed one; a request that names none gets copies of them all. TYPE: dict[str, str | None]

requests instance-attribute

requests: dict[str, Request] = {}

templates instance-attribute

templates: dict[str, Any] = {}

names instance-attribute

names: dict[str, str | None] = {}

out instance-attribute

out: set = set()

counts instance-attribute

counts: dict[str, int] = {}

nrows instance-attribute

nrows = 0

register

register(registration_id: str, payload: bytes, persistent_objects: dict, name: str | None = None) -> None

Take a block the engine should run for every request from now on.

unregister

unregister(registration_id: str) -> None

Stop running a block and forget anything it has not handed back.

add

add(new_requests: list['NewRequestData'], persistent_objects: dict) -> None

Take the worker off each new request that carries one.

A request with no nnsight payload is another tenant of the same engine and runs only the registered blocks — it still occupies tokens in the batch, so scope counts it. A payload that fails to deserialize is recorded as that request's error and surfaced at collect, rather than raised inside execute_model where it would take the engine every tenant shares down.

A request this engine already carries is one vLLM preempted and is resuming; its workers continue (the engine replays the tokens they already saw inside one recompute step, so a fresh block would be short by exactly those steps — scope keeps the steps they sat out off their count).

refuse_chunked

refuse_chunked(spans: list[tuple[str, int]], states: dict) -> None

Give a request whose prompt this step only partly prefills its error, not a worker.

Chunked prefill (off by default on an nnsight engine) splits a prompt across steps: a block would see a slice of its prompt now and the rest later, and the sample of every chunk but the last is one vLLM discards. states is the runner's per-request state, which knows how much of the prompt is computed before this step.

workers

workers() -> list

Every worker this engine is carrying, traced and registered alike.

scope

scope(model: 'VLLM', spans: list[tuple[str, int]]) -> None

Point every worker at its own tokens within this step's batch.

spans is each scheduled request and its token count, in the order the forward's tensors follow. A worker's span is only meaningful for the step it was computed in, so every worker's is recomputed: workers whose request isn't running now report no group and are dropped from the interleaver, and a block that already ran to completion is dropped too, since the interleaver starts anything not alive and a finished block must not run a second time.

unflatten

unflatten(model: 'VLLM') -> None

Re-point each worker from its tokens to its row.

Logits and sampled ids carry one row per request, not per token, so the spans that scoped the forward would select the wrong thing.

record_saves

record_saves() -> None

Note, on each scheduled worker, which of its values were saved.

Read from the thread-local save-set while still on the thread the workers ran on. Collection can happen on another thread — Ray dispatches it through its own RPC worker — where that thread-local is empty, so the answer is captured here and carried on the mediator instead. The set only grows across a request's steps, so re-recording each step keeps the latest superset.

record

record(mediator: Any) -> None

Snapshot one worker's saved names, and its error if it has one.

Split out because the snapshot is not only taken per step: a worker served at collect time — one parked on tracer.result, which only exists once the engine has assembled the output — binds its name after the last record_saves of the run, and would otherwise be recorded as having saved nothing.

harvest

harvest(finished: set[str]) -> None

Shelve finished requests' registered values until they are collected.

Driven by the scheduler's own finished set, and again by a collect that finds a request not yet harvested — the scheduler's pass happens at the top of the next step, which on the async path may come after the collect, or (for the last request in flight) never. Off the scheduler rather than only off a collect so a registration works for requests nobody is tracing — the OpenAI server's, say — where nothing would otherwise come asking. A copy still parked when its request ends was waiting for a location this request never reached, which is ordinary for a registration, so it is unwound quietly and whatever it saved is kept.

serve_result

serve_result(mediator: Any, output: Any) -> None

Hand a finished request's output to a worker parked on tracer.result.

A worker parked anywhere else is left alone — its own location is what it is waiting for, and finish_dangling reports it. Runs on the workers' own thread where the greenlet can be resumed; where that differs (Ray's collect) the switch is refused and the read is left unserved, exactly as it was before.

finish_dangling

finish_dangling(mediator: Any, taps: frozenset = frozenset(), quiet: bool = False) -> None

Surface a worker still parked when its request has finished.

A worker still alive at the end was waiting on a location the model never reached — the interleaver's check_dangling_mediators, but for a single request as it retires here rather than after a whole local run. Which of those a parked worker is, and what it should be told, is dangling_unwind's to decide, so a trace behaves the same on this engine as it does locally: a read past the model's point is an error, and a tracer.iter loop the request could not supply — bounded and open alike — is cut short with a warning, keeping the values from the steps that ran. Surfacing differs, since the client is in another process: an error becomes the request's deferred error and is raised there.

Runs on the workers' own thread, where the greenlet can be resumed — the throw is skipped where that thread differs (e.g. Ray's collect), leaving the worker to be dropped without a surfaced error.

taps is the engine's tap set when it replays CUDA graphs: a module location outside it is never reached by a replayed step, and the error says so rather than reporting the model ran past it. quiet only unwinds (a registered copy has no request to report to).

NNsightGPUModelRunner

Bases: GPUModelRunner

A vLLM model runner that interleaves interventions with the forward.

load_model

load_model(*args: Any, **kwargs: Any) -> None

capture_model

capture_model() -> Any

Record vLLM's CUDA graphs with the interleaver open.

A module hands off only while interleaving, and recording the graphs runs the forward outside any step — so open the interleaver (with no workers) for the recording, which is what lets VLLMInterleaver.handle register each tap's replay into the graph.

execute_model

execute_model(scheduler_output: 'SchedulerOutput', intermediate_tensors: Optional['IntermediateTensors'] = None) -> Any

sample_tokens

sample_tokens(*args: Any, **kwargs: Any) -> Any

collect_nnsight

collect_nnsight(request_ids: list[str], finished_request_ids: Optional[list[str]] = None, outputs: Optional[dict] = None) -> Optional[bytes]

Return the saved values and any deferred error of the named requests.

Keyed per request rather than merged, so two traces that happen to name a variable the same don't overwrite each other on the way home. Each entry is {"saves": {...}, "error": ..., "registered": {...}}.

A registered block's values are kept apart from the trace's own because they are not the same kind of thing: a name a trace saved on several requests is one shared container the invokes were writing into, and merge_shared_saves reassembles it on that assumption. A registration saving the same name on every request it runs on looks identical from the outside and is not — it is one value per request — so merging them together would quietly fold a thousand separate activations into one.

PARAMETER	DESCRIPTION
request_ids	Requests to collect saved values from. TYPE: list[str]
finished_request_ids	Those that are done, whose workers are wound up and forgotten afterwards. TYPE: Optional[list[str]] DEFAULT: None
outputs	Engine request id -> the RequestOutput that request produced, for serving tracer.result. The engine has it and the worker does not, so it has to be handed back across; a caller that does not pass it simply leaves result unserved. TYPE: Optional[dict] DEFAULT: None

nnsight_register

nnsight_register(registration_id: str, payload: bytes, name: str | None = None) -> None

Keep a block and run it for every request from now on.

See nnsight.modeling.vllm.registration. Runs on all ranks, so every rank builds the same per-request copies and their reads stay in lockstep — which is what a sharded model's gathers need. name is what requests may address it by (edits=[...]).

nnsight_clear_registered

nnsight_clear_registered(registration_id: str) -> None

Stop running a block and drop what it has not handed back.

nnsight_request_count

nnsight_request_count() -> int

How many requests' workers this runner still tracks.

A leak gauge: it should return to zero once every request has finished or been aborted. A number that only grows across requests means workers are outliving their requests — a finished one is freed in collect_nnsight, an aborted one when its stream is closed (see the async and serve backends).