vllm¶
The base envoys of a vLLM family: the same values as nnterp.components, read where vLLM keeps them.
vLLM runs a family through its own implementation, which differs from transformers' in three ways a value has to absorb:
- No batch axis. A request's activations are
[tokens, hidden_size]: every prompt token on the prefill, one row on each decode step. The values keep nnterp's layouts, so they are served as[1, tokens, hidden_size]and code written against[:, -1]runs on both engines. - Live buffers. A served tensor is the model's own, and the next fused kernel rewrites it after the block has read it, so a saved reference comes back holding later data. Every value here is a private copy, handed back to the model when the block moves on; in-place edits to the copy and assignment both reach the model.
- The residual stream in two halves. Most blocks fuse the residual add
into the next norm and are called with, and return,
(hidden_states, residual): the stream is their sum.FusedLayeris that block,Layerone that is called with the stream (alone, or after the positions) and returns it (alone, or with something beside it).
The attention module hands its queries, keys and values to vLLM's attention
layer (its attn child) and gets the per-head outputs back, each
[tokens, num_heads * head_dim]; they are that child's inputs and output,
served in nnterp's layouts with the heads split out. What the layer computes
between them (the scores, the pattern) is inside its kernel, not Python on
this engine, so those two are recomputed from the queries and keys:
read-only, and there on the prefill only, where a step holds every key.
A family vLLM runs through its transformers backend (transformers' own
modules inside the engine) builds on nnterp.components.vllm.transformers_backend
instead, where the batch axis is already there and the engine's attention
layer is mounted as the attention module's attn child.
module-attribute
¶
DECODE_STEP = "recomputed from this step's queries and keys, and on a decode step the earlier tokens' keys are in vLLM's cache, which only its kernel reads; read it on the prefill"
module-attribute
¶
__all__ = ['Attention', 'DECODE_STEP', 'Flat', 'FusedLayer', 'Layer', 'Mlp', 'argument', 'batched', 'on_decode_step', 'on_transformers_backend', 'project', 'summed', 'unbatched', 'with_argument']
Bases: Attention
A vLLM attention module: its contribution, what goes into and comes out of the engine's attention layer, and the pattern.
The module projects (and, on a rotary family, rotates) its queries, keys
and values and calls vLLM's attention layer on them, self.attn(q, k,
v), which returns the per-head outputs the output projection then
mixes. Those four are the attn child's inputs and output, each
[tokens, num_heads * head_dim], served with the heads split out by the
module's own head_dim. A family whose attention module names that
child, or its head width, another way points the values at it.
They are this step's rows: on a decode step the keys and values are the new token's alone, and the earlier ones are in the engine's cache, which the kernel reads for itself.
The scores and the pattern are computed inside the kernel, which serves
nothing between its inputs and its output, so here they are recomputed
from the queries and keys, with the layer's own scale, softcap and
window. That makes them read-only (the kernel never takes a pattern: edit
the queries or keys to change what a head attends to, or its head outputs
to change what it wrote) and there on the prefill only, where a step
holds every key; a decode step raises Unavailable (on_decode_step).
class-attribute
instance-attribute
¶
attention_scores = DerivedEProperty(scores, description='The scores entering the softmax, [1, num_heads, query, key], recomputed from the queries and keys; prefill only, read-only', unavailable=on_decode_step)
class-attribute
instance-attribute
¶
attention_probabilities = DerivedEProperty(probabilities, description='The attention pattern, [1, num_heads, query, key], recomputed from the queries and keys; prefill only, read-only', unavailable=on_decode_step)
attention_queries(value: Tensor) -> Queries
The queries the attention layer receives, after the projection, any query norm and the rotary embedding.
attention_keys(value: Tensor) -> Keys
This step's keys as the attention layer receives them; num_kv_heads wide under grouped-query attention.
attention_values(value: Tensor) -> Values
This step's values as the attention layer receives them; num_kv_heads wide under grouped-query attention.
attention_head_outputs(value: Tensor) -> HeadOutputs
Each head's output as the attention layer returns it, before the output projection mixes them.
attention_premix(value: Tensor) -> HeadOutputs
The output projection's input split by head, after any gate the attention applies; the head outputs where there is none.
Flat(*args: Any, heads: str | None = None, batch: bool | Callable[[Envoy], bool] = False, factor: Callable[[Envoy], float] | None = None, **kwargs: Any)
Bases: EProperty
An EProperty over a [tokens, ...] tensor, served as a private [1, tokens, ...] copy.
Takes what EProperty takes (a path for a key, select for one
argument of a call). The decorated function receives the batched copy and
returns the value; its annotation is the layout. In-place edits are handed
back to the model when the block moves on, and an assignment is checked
against the rows this request has before it is swapped in. An output
that is a tuple serves its first element unless select says another:
a norm vLLM calls with the residual beside the stream returns
(normed, residual), and the same norm called alone returns the tensor.
| PARAMETER | DESCRIPTION |
|---|---|
heads
|
For a
TYPE:
|
batch
|
Whether the tensor already carries the batch axis,
TYPE:
|
factor
|
For a value that is the tensor times a scalar of the model
(Granite's
TYPE:
|
Bases: Layer
A vLLM block with the residual add fused into the next norm (Llama, Qwen, Gemma, most families).
It is called forward(positions, hidden_states, residual) and returns
(hidden_states, residual): hidden_states is what the last sublayer
produced and residual the stream before it, so the stream is their
sum, entering and leaving. The first block is called with residual
None and the embeddings as hidden_states.
The sum is computed for the read, so it is nobody's buffer. An edit goes
back as a change to hidden_states (the next norm adds the two), and
a read that edits nothing changes nothing: the difference it hands back is
exactly zero.
class-attribute
instance-attribute
¶
attention_norm_output = Layer.attention_norm_output
class-attribute
instance-attribute
¶
layer_mid = summed('post_attention_layernorm.inputs', description='The residual stream after the attention add: hidden_states + residual, [1, tokens, hidden_size]')
Skip this block, handing hidden ([1, tokens, hidden_size]) on as its residual stream.
Safe on a request that runs alone. The engine batches whatever is in flight into one step, and a skipped block has to answer for all of it; see nnsight's vLLM guide before skipping on a shared engine.
Bases: Layer
A vLLM block that is called with the residual stream and returns it.
Two things vary between the blocks that work this way, and a family
states them: STREAM, where the stream sits in the block's call (GPT-2's
block takes it alone, most take the positions first), and the base's
returns_tuple, for a block that returns the stream with something
beside it (Exaone4's and Cohere's return (hidden_states, residual)
with the whole stream first, where a FusedLayer returns its two halves).
Skip this block, handing hidden ([1, tokens, hidden_size]) on as its residual stream.
Safe on a request that runs alone. The engine batches whatever is in flight into one step, and a skipped block has to answer for all of it; see nnsight's vLLM guide before skipping on a shared engine.
Bases: Mlp
A vLLM feed-forward module: its contribution is its output.
The neurons are the down projection's input, as on transformers. On a
plain MLP (GPT-2's GELU) that is also the activation's output, and it is
read there for both: vLLM builds an activation by name once and shares the
one module between every block, so act_fn's own output cannot say
which block called it. A gated MLP is different on vLLM: its gate and up
projections are one gate_up_proj and its act_fn (SiluAndMul,
GeluAndMul) computes act(gate) * up in one kernel, so what it
returns is already the neurons and act(gate) alone is never a tensor.
There mlp_activation is unavailable (GATED) and mlp_neurons is
act_fn's output. A mixture of experts has neither (ROUTED).
mlp_activation(value: Tensor) -> Neurons
act(up_proj(x)) on a plain MLP, [1, tokens, intermediate_size]: read as the down projection's input, which it is.
DECODE_STEP when a read made now would be served a decode step of the request, else None.
A request's first visit to its attention layer is the prefill, every
prompt token at once, and each later one a decode step of one token. The
worker knows which visit a read would ask for before anything is read,
so the value refuses there, by name. Outside a trace there is no step to
ask about and the answer is the prefill's: support lists it available.
[tokens, ...] as a private [1, tokens, ...] copy; with head_dim, the last axis split into heads.
[tokens, num_heads * head_dim] (or [tokens, num_heads, head_dim], where a
module keeps its heads apart) becomes [1, tokens, num_heads, head_dim], or
[1, num_heads, tokens, head_dim] with heads_first.
Whether envoy's model runs on vLLM's transformers backend, where every activation carries the batch axis.
unbatched(value: Tensor, rows: Tensor, name: str, head_dim: int | None = None, heads_first: bool = False) -> Tensor
What batched served, back as the [tokens, ...] the model holds, refusing another shape.
A written value is spliced into the step the engine is running, among other requests' rows. One of the wrong height would reach the next kernel as it is, where the mismatch can end the engine and every request in it, so it is refused here, in the block that wrote it.
One argument of a call served as (args, kwargs), however the caller passed it; None when absent.
summed(key: str, hidden: tuple = (0, 'x'), residual: tuple = (1, 'residual'), description: str | None = None) -> EProperty
The stream at a call that takes it in two halves, hidden + residual, [1, tokens, hidden_size].
A fused block hands its stream to the next norm as norm(hidden_states,
residual), and the norm adds the two before it normalizes (in place):
the stream is their sum, computed for the read. key is the call's
inputs, hidden and residual where the halves sit in it, as
(index, name). An edit goes back as a change to the hidden half, and
a read that edits nothing hands the difference back as exactly zero; a
call without the residual (a block's first norm) is its stream alone.
inputs with that argument replaced, where the caller passed it.