batching¶
Scope a worker to its own request's tokens inside a scheduled step.
A stacked-tensor model gives each invoke a row range that is fixed the moment the
trace is written. vLLM gives neither. It packs every request the scheduler picked
for a step into one flat [total_tokens, hidden] slab — a whole prompt's tokens
on prefill, a single token per decode step — and which requests are in that slab
changes from step to step as they arrive and finish.
So a worker's group is a token span, recomputed every step by
NNsightGPUModelRunner
rather than assigned once up front. The row math itself is the base
Batcher's, moved off dim 0 onto
whichever axis carries tokens: dim 0 for a model vLLM has its own definition
for, dim 1 for one served through vLLM's Transformers backend (see
VLLMBatcher._token_dim).
What a block is served is a view into that slab, so an in-place edit lands
where the model will read it. The cost is that a read aliases live memory: vLLM's
fused kernels overwrite activation buffers in place, so a tensor kept past the
point it was read holds whatever was written over it (see NNSIGHT_VLLM_CLONE_READS
and the "clone what you keep" rule in docs/models/vllm.md).
Gathering a sharded value is not here — see
fragments. The split is the point: narrowing
happens once per parked worker, while a collective must happen once per value
however many workers read it.
Bases: Batcher
A Batcher over vLLM's flat token axis.
property
¶
Whether narrowing applies — always.
The base skips narrowing for a lone invoke, because one invoke is the whole batch. That never holds here: the engine fills a step with whatever it has, so a request's tokens sit alongside other requests' — another trace's, another tenant's, or a decode of a request whose own block already finished. A worker is only ever entitled to its own span, so there is no case in which handing it the whole slab is right.
narrow(value: Any, group: BatchGroup) -> Any
Serve a worker its token span — a view, or a copy under CLONE_READS.
A view is the default because it is what makes an in-place edit work: the
block writes into the slab the model goes on to read. It is also what makes
a read unreliable, since vLLM's fused kernels (fused_add_rms_norm,
MLA's in-place rotation of the q_proj output) overwrite those buffers a
few ops later — a value kept past its read point comes back holding a later
layer's data, with nothing to indicate it. The documented answer is to
.clone() in the block, which is easy to forget and silent when forgotten.
CLONE_READS trades the other way for a whole engine: every value handed
to a block is a private copy, so anything kept is what was computed, and
.save(), tracer.cache() and appends under tracer.iter are all safe
without a clone at each site.
In-place edits do not survive it. With a copy served there is nothing
aliasing the slab, so layers[10].output[0][:] += v writes to the copy
and the model never sees it — silently, the same way it is silent today
when a save aliases. Assign instead: read, edit the copy, write it back
with layers[10].output = out, which goes through widen and lands.
An eproperty that registered a transform write-back is unaffected —
that path already splices its edited value back through widen.
Off by default. The copy is of the narrowed span, not the whole slab, but it is still one allocation per module read per step.
Whether CLONE_READS is set to anything but a falsy spelling.
Read from the environment rather than CONFIG because the batcher that does
the narrowing lives in the engine's worker process, built when the engine
loads the model. A CONFIG.APP field set in the client process would never
reach it; an env var set before VLLM(...) is inherited when the worker is
spawned.