fragments¶
fragments
¶
Which values a vLLM engine splits across ranks, and how to reassemble them.
vLLM shards its linear layers under tensor parallelism, so the value at a
ColumnParallelLinear or RowParallelLinear is one rank's piece of the real
tensor. A user asked for the layer, not a piece of it, so those are gathered
before a worker sees them and re-split before vLLM's own forward carries on.
A FusedMoE layer needs the same correction for a different reason: an MoE
block that defers the combine (reduce_results=False — Qwen-MoE, DeepSeek)
returns per-rank partial sums that the outer block all-reduces afterwards, so
the value at the experts module is a partial too. It is gathered and re-split on
the same terms, with the group size taken from the expert layout rather than
tp_size alone.
Everything about when — once per visit, only when something is waiting, put
back on the way out — belongs to
Interleaver.handle and is
shared with every other distributed runtime; see
nnsight.intervention.fragments.
That sharing is what this module is. Doing the same job from
VLLMBatcher instead costs two
extra pairs of forward hooks per parallel layer, bracketing the interleaver's
own, plus a memo and an explicit release: Batcher.narrow runs once per
parked worker, so several workers reading one value would otherwise run several
collectives and deadlock the ranks. On the interleaver the bracket is already
once-per-visit, so none of that is needed — no memo, no watch/release, no
extra hooks.
VLLMFragments
¶
Bases: Fragments
Records which of an engine's locations hold one rank's piece.
| ATTRIBUTE | DESCRIPTION |
|---|---|
enabled |
Whether this engine is sharded at all. False on one rank, and then this costs one attribute check per handled location.
|
rules |
Location -> the module that produced it and which side it is.
The module is kept because the collective to run and the arithmetic
to undo it both depend on its type and its
TYPE:
|
instrument
¶
Record whether either side of this envoy's module is a piece.
Called for every envoy as the tree is built, which is the one moment both
the module and its path are in hand. Holding onto the module is safe for
the same reason the previous hook-based version was: vLLM builds the tree
once, after load_model, and does not swap modules under it — a hook
registered on a swapped-out module would have gone just as dead.
whole
¶
The real tensor behind value, and how to cut it back down.
vLLM's rules describe a location and nothing else — no value here carries
its own layout — so the way back is split with this location bound to it.
split
¶
This rank's piece of whole, as vLLM's own forward expects it.
Applied to whatever intervention code left behind, so an edit made to the
assembled tensor is carried back into the model rather than dropped — and
to a value that was never gathered at all (a .skip replacement, or the
argument of an ad-hoc call), which is already the real tensor.