Info

Last Execution: 2026-09-02

Package	Version
nnsight	0.8.0
Python	3.12.13
torch	2.13.0+cu130
transformers	5.15.0
vllm	0.27.1

vLLM¶

VLLM("repo") is an nnsight model whose forward pass is a vLLM engine. You write the same with model.trace(...) block you would write against a HuggingFace model; nnsight serializes it onto the request, runs it inside vLLM's worker interleaved with the forward, and hands the saved values back. PagedAttention, continuous batching and tensor parallelism are unchanged.

What follows is a first pass: loading an engine, reading a value, writing one, and sending several prompts. The vLLM section is the reference, with a page per job and every snippet run against Qwen3-8B.

vLLM runs on CUDA only, so these cells need a GPU. The integration is an optional extra:

pip install "nnsight[vllm]"

Load an engine¶

The constructor takes a HuggingFace repo id and anything vllm.LLM accepts. dispatch=True builds the engine now; without it the constructor builds only a meta-device copy of the module tree, and the first trace builds the engine.

In [1]:

from nnsight.modeling.vllm import VLLM

model = VLLM(
    "openai-community/gpt2",
    gpu_memory_utilization=0.15,
    dispatch=True,
)

print(type(model.vllm_entrypoint).__name__, model.dispatched)
assert model.dispatched

from nnsight.modeling.vllm import VLLM model = VLLM( "openai-community/gpt2", gpu_memory_utilization=0.15, dispatch=True, ) print(type(model.vllm_entrypoint).__name__, model.dispatched) assert model.dispatched

LLM True

Read¶

Any module on the tree is a location. Two more sit past the model: model.logits, this step's pre-sampling logits, and tracer.result, the finished RequestOutput.

In [2]:

prompt = "The Eiffel Tower is in the city of"

with model.trace(prompt, temperature=0.0, max_tokens=1) as tracer:
    mlp = model.transformer.h[6].mlp.output.clone().save()
    logits = model.logits.save()
    result = tracer.result.save()        # served last: nothing can be read after it

n_tokens = len(model.tokenizer(prompt)["input_ids"])
print("mlp.output  ", tuple(mlp.shape), f"({n_tokens} prompt tokens, no batch axis)")
print("model.logits", tuple(logits.shape))
print("next token  ", repr(model.tokenizer.decode(logits.argmax(-1))))
print("generated   ", repr(result.outputs[0].text))

assert mlp.shape[0] == n_tokens
assert logits.shape[0] == 1
assert model.tokenizer.decode(logits.argmax(-1)) == " Paris"

prompt = "The Eiffel Tower is in the city of" with model.trace(prompt, temperature=0.0, max_tokens=1) as tracer: mlp = model.transformer.h[6].mlp.output.clone().save() logits = model.logits.save() result = tracer.result.save() # served last: nothing can be read after it n_tokens = len(model.tokenizer(prompt)["input_ids"]) print("mlp.output ", tuple(mlp.shape), f"({n_tokens} prompt tokens, no batch axis)") print("model.logits", tuple(logits.shape)) print("next token ", repr(model.tokenizer.decode(logits.argmax(-1)))) print("generated ", repr(result.outputs[0].text)) assert mlp.shape[0] == n_tokens assert logits.shape[0] == 1 assert model.tokenizer.decode(logits.argmax(-1)) == " Paris"

mlp.output   (10, 768) (10 prompt tokens, no batch axis)
model.logits (1, 50257)
next token   ' Paris'
generated    ' Paris'

Write¶

Assignment into a served value edits the running model. The value is vLLM's own buffer, so an in-place write is the direct form.

In [3]:

with model.trace(prompt, temperature=0.0, max_tokens=1):
    model.transformer.h[8].output[:] = 0
    zeroed = model.logits.argmax(-1).save()

print("layer 8 zeroed ->", repr(model.tokenizer.decode(zeroed)))
assert model.tokenizer.decode(zeroed) != " Paris"

with model.trace(prompt, temperature=0.0, max_tokens=1): model.transformer.h[8].output[:] = 0 zeroed = model.logits.argmax(-1).save() print("layer 8 zeroed ->", repr(model.tokenizer.decode(zeroed))) assert model.tokenizer.decode(zeroed) != " Paris"

layer 8 zeroed -> '\n'

Several prompts¶

One tracer.invoke(...) is one vLLM request, and one request takes one prompt. Several prompts means several invokes in the same trace; vLLM's continuous batcher runs them together. A name saved in every invoke comes back as a list, in invoke order, and each entry is sized to its own prompt.

In [4]:

prompts = ["The Eiffel Tower is in",
           "The capital of Japan is the city of",
           "Water boils at"]

with model.trace(temperature=0.0, max_tokens=1) as tracer:
    for p in prompts:
        with tracer.invoke(p):
            hidden = model.transformer.h[6].output.clone().save()   # one name, three invokes

for p, h in zip(prompts, hidden):
    print(f"{p!r:<38} {tuple(h.shape)}")

assert len(hidden) == len(prompts)
assert [h.shape[0] for h in hidden] == [len(model.tokenizer(p)["input_ids"]) for p in prompts]

prompts = ["The Eiffel Tower is in", "The capital of Japan is the city of", "Water boils at"] with model.trace(temperature=0.0, max_tokens=1) as tracer: for p in prompts: with tracer.invoke(p): hidden = model.transformer.h[6].output.clone().save() # one name, three invokes for p, h in zip(prompts, hidden): print(f"{p!r:<38} {tuple(h.shape)}") assert len(hidden) == len(prompts) assert [h.shape[0] for h in hidden] == [len(model.tokenizer(p)["input_ids"]) for p in prompts]

'The Eiffel Tower is in'               (7, 768)
'The capital of Japan is the city of'  (8, 768)
'Water boils at'                       (3, 768)

What changes about writing a block¶

• No batch axis. vLLM packs every in-flight request's tokens into one [total_tokens, hidden] slab and nnsight narrows your block to its own request's rows, so a value is [pos, width] on the prefill step and [1, width] on each decode step. The last position is [-1], never [:, -1, :].
• Clone what you keep. A served value is the model's live buffer, and a later layer rewrites it in place. mlp.output.clone().save() above is a copy; without the .clone() it would come back holding whatever ran through that memory next. Reducing (.mean(0), .norm()) is a copy too.
• A decoder layer's .output is a pair. On Llama, Qwen and Mistral trunks vLLM fuses the residual add into the next layer's norm, so layers[i].output is (hidden, residual) and the residual stream after the layer is their sum. GPT-2, above, returns a plain tensor.

One more thing to expect: a request's numbers depend on the batch it was scheduled with. Reduction order inside a fused kernel follows the batch, and in bf16 that moves the last digit, so a near-tied argmax can land differently between runs. On the second prompt above GPT-2 puts 0.124 on ' Kyoto' and 0.124 on ' Tokyo'. Greedy decoding pins the sampler, not the arithmetic, so compare distributions or effect sizes rather than expecting an exact token or an exact float back.

Where to go next¶

The rest of the section:

Locations	What can be asked for, and how it is named
Loading models	Eager, CUDA-graph and async engines; what nnsight forces
Capabilities and limits	What refuses, and how errors surface
Capture	One location, every layer, every step, many prompts
Steering, Patching, Ablation	Writing into the running model
Generation	Sampling, per-step values, n > 1
Editing the engine	model.edit(): install a block once, for every request
Async and servers	mode="async", nnsight-serve, GPU-less clients
Tensor parallelism	Sharded models read as whole tensors
Performance	Eager against CUDA-graph taps, measured

The design is written up in NNsight × vLLM: Interpretability at Production Scale.