Skip to content

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

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 tp_size/tp_rank.

TYPE: Dict[str, Tuple[Module, str]]

enabled instance-attribute

enabled = False

rules instance-attribute

rules: Dict[str, Tuple[Module, str]] = {}

instrument

instrument(envoy: Any) -> None

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.

fragmented

fragmented(location: str) -> bool

whole

whole(location: str, value: Any) -> 'tuple[Any, Any]'

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

split(location: str, whole: Any) -> Any

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.