vllm: gather a fused projection into the layout tp=1 has

[JadenFiotto-Kaufman]

The bug

Under tensor parallelism VLLMFragments.whole ran a plain all-gather for anything that is a ColumnParallelLinear. QKVParallelLinear and MergedColumnParallelLinear are subclasses whose per-rank shard is itself fused, so the ranks' concatenation is [q0 k0 v0 | q1 k1 v1], not [q | k | v]. split was the exact inverse, so reads and writes were self-consistent — just not in the layout tp=1 has and every recipe assumes. Per-head slicing and gate, up = value.chunk(2, -1) were silently wrong from tp=2 up.

From the stress sweep (Llama-3.2-1B, tp=2, against an in-process transformers reference):

qkv_proj.output vs HF [q|k|v]:                    max|d| = 18.188
qkv_proj.output reordered [q0q1|k0k1|v0v1] vs HF: max|d| =  0.031
after zeroing gate_up[:, 8192:] (the 'up' half):  mlp.output max|.| = 0.492   (expected 0.0)

docs/models/vllm-parallelism.md promised the opposite ("the same complete tensor ... read at their full width"), and tests/vllm/test_tensor_parallel.py compared column-parallel reads as a per-row multiset, which passes on either layout — that is what let it through.

What changed

All of the code change is in src/nnsight/modeling/vllm/fragments.py.

• whole un-interleaves the last dim after the gather; split cuts one rank's fused piece straight back out of the whole. Both dispatch on self.rules[location], as before; nothing outside the file changes.
• Where the widths come from. module.output_partition_sizes — vLLM's own, and already per rank: ColumnParallelLinear.__init__ does output_partition_sizes = [divide(size, tp_size) for size in self.output_sizes] for exactly the two subclasses that set output_sizes (0.27.1 linear.py:479-485), and leaves a single entry — the whole shard — for every other column-parallel layer. QKVParallelLinear inflates its own output_sizes by tp_size (linear.py:1083-1088) precisely so that division lands on the per-rank width, replication included. Read against the 0.27.1 in /disk/.../v27 on hakone; I am confident about this one.
• KV replication. With fewer KV heads than ranks vLLM sets num_kv_heads = 1 and replicates K and V across num_kv_head_replicas adjacent ranks — its weight loader takes shard_rank = tp_rank // num_kv_head_replicas (linear.py:1379) — so a naive un-interleave would hand back a tensor several KV heads too wide. One per group is taken, the way _one_per_group already does for DCP, and split gives every rank in a group the same K/V slice back.
• Unrecognised fused subclasses warn rather than being assumed unfused. This is not hypothetical: MinimaxM3QKVParallelLinearWithIndexer is a QKVParallelLinear that packs five sub-shards, one of them replicated across every rank rather than across a KV group (linear.py:1492-1503, 1565), so the QKVParallelLinear rule is applied only at exactly three sub-shards and anything else warns and stays in rank order.
• Unfused ColumnParallelLinear, RowParallelLinear, DCP-group layers and the MoE paths take exactly the code they took before.

Docs: docs/models/vllm-parallelism.md keeps the "same complete tensor" sentence (now true) and gains what the fused layout is, that nnsight normalises it, and a note on replicated KV.

Tests: tests/vllm/test_tensor_parallel.py compares column-parallel reads element-wise against the one-rank engine instead of as a sorted multiset (same for the ad-hoc-call test), its module docstring is rewritten, and it gains

• TestShardedEdit::test_half_edit_means_the_same_half — zero the upper half of gate_up_proj.output; since the activation is silu(gate) * up the MLP must go to exactly zero, which it does only if that half really is up. This is the edit _zero_all cannot catch.
• TestFusedLayout — the reordering arithmetic on its own, parametrized over real layouts including the replicated-KV cases (8q/2kv at tp=4, MQA at tp=8) that no checkpoint these fixtures can run reaches: Qwen2.5-0.5B has 14 query heads and cannot shard past 2.

Verified on 2x A100, tp=2

This was opened as a draft claiming only arithmetic: I have no vLLM and no multi-GPU machine, so I could not run a line of it. It has since been run on real hardware, and the fix is confirmed.

	unpatched tp=2	patched tp=2
qkv_proj.output vs HF [q|k|v] (Llama-3.2-1B)	max|d| 18.188	0.031
the same comparison, reordered	0.031	18.188
gate_up_proj.output vs HF [gate|up]	4.505	0.008
mlp.output after zeroing the up half	max|.| 0.4922	exactly 0.0000
attn.c_attn vs HF, gpt2 in fp32	13.0416	exactly 0.0000
Qwen2.5-0.5B, element-wise vs tp=1	—	min row cosine 1.0000

The reordered comparison swapping to precisely the two numbers the direct one had is the signature of a pure permutation, which is what this was. Plain unfused column-parallel, row-parallel, the row-parallel input gather and five edits are byte-identical across unpatched-tp2, patched-tp2 and patched-tp1 — the change reaches the fused layers and nothing else. output_partition_sizes was confirmed per rank, at runtime, on both ranks: Llama [1024,256,256] / [4096,4096], gpt2 [384,384,384] / [1536], with the single-entry plain layer correctly returning None.

Verified without a GPU as well: TestFusedLayout's 12 arithmetic cases, and a scratch harness stubbing the vLLM linear classes and the collectives that drives whole/split end to end on every rank for QKV tp=2, GQA-replicated tp=4, MQA tp=8, gate_up tp=2, a four-way merged column and a plain unfused column.

The tests had a bug of their own — fixed in e3b1b9f

The first hardware run was 33 passed, 2 failed, and neither failure was about sharding. _read and the ad-hoc tests saved .output[0] by reference, so what they compared afterwards was whatever the forward had since written over it — the documented vLLM live-buffer trap. It bit the one-rank reference exactly as hard as the sharded run, which is why it surfaced on qkv_proj and not on gate_up_proj, and why test_row_parallel_call_takes_and_returns_the_whole (cosine 0.1256) reproduced identically at tp=1 and on unpatched code. With the clones: 35 passed, qkv tp1-vs-tp2 cosine 0.9163 → 1.0000, ad-hoc 0.1256 → 1.0000 (max|d| 18.875 → 0.0000).

Cloned at every save whose value is compared later, logits included. The reads asserted only for their shape are left as they were.

Still untested

• KV replication on a real engine. Both test models have num_kv_head_replicas == 1 at tp=2, so nothing that ran reaches the dedupe — it needs tp > 8 or a multi-query model. Covered only by TestFusedLayout and the stub harness.
• Any tp > 2.
• The warning path against a real MinimaxM3QKVParallelLinearWithIndexer, exercised only against a stand-in class.
• FusedMoE and DCP-group layers — untouched by this change, and unexercised by this run.

TestFusedLayout needs no GPU at all but sits in tests/vllm/, which CI ignores. Say the word and I will move it somewhere CI runs it; fragments.py imports fine without vLLM.

Known gap: _KimiGDNMergedColumnParallelLinear

_fused_sub_shards reads replication only off num_kv_head_replicas, so a merged column that replicates a sub-shard by any other route is not covered and does not warn. _KimiGDNMergedColumnParallelLinear is exactly that: it gives every rank the same copy of one projection (output_sizes[i] *= tp_size, loaded with tp_rank forced to 0), which reads here as an ordinary merged column, so _unfuse emits tp_size copies of it and the whole comes back too wide.

Nothing regresses — the plain gather was equally wrong there before this branch — but it is a real gap, recorded here rather than left to be discovered later, and noted in the function's docstring. The fix is a replicated_shard_id check in _fused_sub_shards; it is left out because no checkpoint available here can test it.

🚨 The nnsight:vllm skill is now wrong and is not mine to fix

plugins/nnsight/skills/vllm/references/parallel-and-architectures.md:45 says:

A fused projection (qkv_proj, gate_up_proj) gathers in rank order — [q₀ k₀ v₀ | q₁ k₁ v₁] — so slice it by head, not by [:q_size].

That was correct before this change and is wrong after it — which is exactly why the skills agent never questioned the doc. It must be updated in the same batch as this PR, not after: the layout is now [q | k | v] at every tensor_parallel_size, and slicing [:q_size] is the right thing to do. I do not own the skills repo, so I have not touched it.

Deliberately left alone

• The transformers TP path (DistributedConfig, src/nnsight/modeling/tp/) does its own gather and was never exercised — tensor-parallel-torchrun did not run in the sweep. It is outside this file and I have not looked at whether it has the same defect. TP is not fixed generally until someone checks it.
• MinimaxM3QKVParallelLinearWithIndexer is not supported, only warned about. Its five-way layout with a wholly-replicated index_k is straightforward to add once someone can test it on a MiniMax checkpoint.
• split builds only the calling rank's piece rather than re-interleaving the whole tensor and then discarding the other ranks' columns. Same columns either way; one fewer full-width allocation per rank per write-back.
• Older vLLM: if a fused layer somehow lacks output_partition_sizes, the layer is treated as unfused and behaves exactly as it does on 0.8 today. I could only read 0.27.1.

🤖 Generated with Claude Code