vllm¶
Bases: Remotable
A vLLM engine whose internals can be traced.
Interventions are written exactly as for any other model — the module tree
mirrors the architecture vLLM loaded — but they run inside the engine's worker
process. Sampling settings (temperature, max_tokens, top_p, ...)
are passed to trace/invoke rather than configured on the model, since
each invoke is its own vLLM request. Read generated tokens through
model.logits / model.samples under tracer.iter, or the whole
finished request through tracer.result.
Examples:
Single prompt, edit an activation, read the logits::
>>> model = VLLM("gpt2", dispatch=True)
>>> with model.trace("The Eiffel Tower is in", temperature=0.0):
... model.transformer.h[8].output[:] = 0
... logits = model.logits.save()
>>> model.tokenizer.decode(logits.argmax(dim=-1))
Several prompts is several invoke blocks (each is one request), not a
list — a shared save escapes each into its own name::
>>> with model.trace(temperature=0.0) as tracer:
... with tracer.invoke("The Eiffel Tower is in"):
... a = model.logits.save()
... with tracer.invoke("The capital of Japan is"):
... b = model.logits.save()
Streaming with mode="async" — saves arrive on the finished output::
>>> model = VLLM("gpt2", dispatch=True, mode="async")
>>> with model.trace("Hello", max_tokens=5) as tracer:
... logits = model.logits.save()
>>> async for output in tracer.backend: # doctest: +SKIP
... last = output
>>> last.saves["logits"] # doctest: +SKIP
CUDA graphs on, at the cost of declaring up front which locations a
trace may touch (see taps)::
>>> model = VLLM("gpt2", dispatch=True, taps=["transformer.h.*.output"])
>>> with model.trace("The Eiffel Tower is in", temperature=0.0):
... model.transformer.h[8].output[:] = 0 # in place
... hidden = model.transformer.h[6].output.clone().save()
A GPU-less client running a trace on a remote nnsight-serve engine — the client only builds a meta tree and never dispatches::
>>> model = VLLM("gpt2") # no GPU needed
>>> with model.trace("Hello", serve="http://host:8000"): # doctest: +SKIP
... logits = model.logits.save()
| ATTRIBUTE | DESCRIPTION |
|---|---|
vllm_entrypoint |
The underlying
|
tokenizer |
The tokenizer vLLM resolved for the checkpoint.
|
The logits for this request's step, before sampling.
A hookable run-level value like a module's .output — reading it parks the
worker until the engine produces this step's logits; writing it swaps them.
Under tracer.iter each pass sees the next decoded step's logits::
with model.trace("Hello", temperature=0.0) as tracer:
logits = model.logits.save()
The token ids the sampler drew from logits for this step.
Read or edit them inside a trace; setting them replaces the tokens the engine continues generation from — force a token::
with model.trace("Hello", temperature=0.0, max_tokens=3) as tracer:
for _ in tracer.iter[:3]:
model.samples = torch.zeros_like(model.samples) # feed token 0
Refuse: there is no local forward for a fake-tensor pass to run.
scan reads shapes by running the
model's own forward under a fake-tensor mode, on the meta module, with no
weights. Here that module is a client-side shell: the forward runs in the
engine's worker, on real weights, under torch.inference_mode. Refused
up front rather than at the engine, which would build itself and then ask
a fake mode to run a request it can't fake.
edit(*, name: str | None = None, inplace: bool = True, serve: str | None = None, api_key: str | None = None, backend: Any = None) -> Any
Install a block on the engine, to run for every request it handles.
The vLLM form of Envoy.edit.
An ordinary edit is replayed by the envoy that stores it, which here is
the client — where there are no weights, so it would never run. This
sends the block to the engine instead, where every request afterwards
gets its own copy: requests you trace, and requests submitted by
something that has never heard of nnsight.
What the block saves comes back on that request's output, as
output.saves — the same place a trace's values arrive. For a request
you are tracing, read it through tracer.result.saves.
The block is written like a trace body against the same envoy tree. It
belongs to no particular request, so there is nothing to
tracer.invoke(...).
| PARAMETER | DESCRIPTION |
|---|---|
name
|
What requests may address this edit by. A request that passes
TYPE:
|
inplace
|
Only
TYPE:
|
serve
|
An nnsight-serve URL to install the block on, the counterpart
of
TYPE:
|
api_key
|
Sent as the
TYPE:
|
backend
|
Optional backend for the underlying trace.
TYPE:
|
| RETURNS | DESCRIPTION |
|---|---|
Any
|
|
Any
|
the block follow a request across its generated tokens rather than |
Any
|
seeing only the prefill; and the handle to |
Any
|
|
Any
|
with. |
Any
|
clears every one still installed. |
Examples:
Read one layer out of everything the engine runs::
>>> with model.edit() as (tracer, edit): # doctest: +SKIP
... hidden = model.model.layers[16].output[0].save()
>>> outputs = model.generate(prompts, max_tokens=5) # doctest: +SKIP
>>> outputs[3].saves["hidden"] # doctest: +SKIP
>>> edit.clear() # doctest: +SKIP
Against a served engine, from a client with no GPU::
>>> with model.edit(serve="http://host:8000") as (tracer, edit):
... model.model.layers[16].output[0][:] = 0 # doctest: +SKIP
Named, so a request can choose::
>>> with model.edit(name="probe") as (tracer, edit): # doctest: +SKIP
... score = model.model.layers[16].output[0][-1].norm().save()
>>> with model.edit(name="steer") as (tracer, edit2): # doctest: +SKIP
... model.model.layers[8].output[0][:] += v
>>> outputs = model.generate(prompts, max_tokens=5, edits=["probe"]) # doctest: +SKIP
>>> outputs[0].saves["score"] # the probe ran; the steer did not
Clear every edit still installed on the engine.
The local form drops a list held on the envoy; here each edit lives on
the workers, so each is cleared through its own handle — which means this
is synchronous-engine only, like clear itself. Use aclear_edits on an
async engine, where it raises rather than half-clearing.
async
¶
clear_edits, awaited — the async engine's form.
A separate method rather than a clear_edits that returns something
awaitable on an async engine: a coroutine nobody awaits never runs at
all, so the sync-looking call would leave every edit installed and say
nothing. That is the failure Registration._rpc refuses, and it is worth
refusing here too — clear/aclear already come in this pair.
Generate — as a trace when used as a with block, plainly when not.
with model.generate(...) is trace: vLLM has no forward/generate
split, so the two are the same thing and generation length is
max_tokens (max_new_tokens is accepted and rewritten). Read the
generated tokens through model.logits/model.samples under
tracer.iter, or through tracer.result.
Called without a with block it just runs the engine and hands back
vLLM's request outputs, so an edit's values — which arrive on the
output — can be read without reaching past the model for
model.vllm_entrypoint::
>>> with model.edit() as (tracer, edit): # doctest: +SKIP
... hidden = model.model.layers[16].output[0].save()
>>> outputs = model.generate(prompts, max_tokens=5) # doctest: +SKIP
>>> outputs[5].saves["hidden"] # doctest: +SKIP
Which of the two it is comes from the call site — the same test
traceable makes for a method
used either way: capturing a block that is not there raises, and that is
the signal to run plainly.
Dispatch the trace to the engine instead of running it here.
Overrides interleave, which starts
each worker in this process alongside the model's forward. There is no
forward to run here — the weights live in the engine's worker — so the
workers are not started; they are serialized onto the requests by
_call and started by the model runner on the other side.