pipeline-stages.md

Multi-GPU: layer pipeline stages

--devices A,B,... splits the model's layers into one pipeline stage per GPU. Each stage owns its layers whole: weights, the KV cache of its attention layers, the recurrent state of its GDN layers, and the scratch it runs in. The point is memory: a model that does not fit one card, or a context that does not, spreads across several with every card's memory usable for KV.

What it is and is not

It is a memory feature. The stages run in sequence: while one computes, the others wait. Decode is weight-bandwidth-bound and each stage reads only its own weights, so a single stream is about as fast as one GPU, minus the boundary hops. Splitting the batch into micro-batches would make each stage read its weights once per micro-batch and win nothing.

It is not tensor parallelism. No layer is divided across devices, so no collective runs inside a layer.

Multi-GPU is a Linux feature. Repeating one device id (--devices 0,0) exercises the whole stage path on a single card and is accepted everywhere, which is how the path is tested without a second GPU. Distinct ids are refused on Windows.

How it works

• StagePlan (core/stage_plan.h) says which contiguous layers each stage owns. --stage-layers 30,34 sets the counts. Without it default_stage_layers (models/qwen3_5/load.cpp) sizes each layer from the artifact and calls solve_stage_plan, an exact memory-balancing solver (maximum page groups every stage can hold, then least-full stage), verified against a brute-force partition oracle. It works from the bytes free on each device, the stored KV page-group and per-slot GDN state bytes of each layer (from the engine's KV storage format and slot count), and rough constants for what a stage needs beyond weights (context, workspace, graphs; more on rank 0 for the head), so it is a good default rather than exact; if it finds nothing fits it deals the layers evenly and lets the Program's planning report the device that runs out.
• Placement. bind_text places every layer's weights on its stage's device (Bindings::place, recorded for rank 0 as well so an object shared between a rank-0 layer and a later stage's layer is refused as a conflicting placement). The embedding, final norm and head stay on rank 0.
• The head stays on rank 0. The last stage sends the residual back; rank 0 then runs the final norm, head and sampling exactly as on one GPU. That costs one extra ~20 us hop per forward pass and keeps the round buffers, prefill buffers and every other head-side structure where they are.
• Per-rank state. DeviceKVPagePool is one page-group allocator over planes bound to several ranks' backings, so admission, prefix reuse and the context-cache policy never see how many devices hold the cache. Block tables are one host shadow with a device copy per attention-bearing rank (KVExecutionTablePool); GDN state is one LinearAttentionStatePool shard per stage (StateImageDevicePool); the host image of a state keeps the whole model's layers at their global offsets. Data movement takes RankStreams, so each plane's copy goes on its own rank's stream, and a single stream given to a multi-rank object fails instead of using the wrong device.
• Capacity. Each further device has its own affine reservation curve (RankCapacityCurve); the resolved page-group count is the smallest any device allows, and a failure names the device that could not hold the plan. Ranks sharing one physical device split its free memory.
• StageLink (core/stage_link.h) moves the residual between stages, and the small control tensors (positions, KV rows, slot indices, valid columns) the layers read from rank 0, through a ring of pinned-host slots. The copies are ordinary graph nodes. Decode captures the whole pass into one multi-device graph, so StageLink tracks, per event, which capture made its last record: a wait inside a capture may only target an event recorded in that capture, and an eager wait may not target one whose last record was in a capture (cudaErrorInvalidValue, even after the graph has launched).
• TextContext::run_staged packs the control tensors on rank 0, sends each stage its block up front, runs stage 0, and for each later stage rebinds the context's device, stream and workspace (ScopedDeviceRank, ScopedArenaRank), receives the residual and control block, runs the stage's layers and passes the residual on.

What is not covered yet

• Prefill does not overlap stages. A prefill chunk runs through the stages in turn and the engine synchronizes after each chunk, so at any moment one stage is busy. Overlapping stages needs micro-chunks inside a chunk (later stages start on micro-chunk 0 while stage 0 runs micro-chunk 1); that changes the chunk shapes the kernels see, and its benefit can only be measured on real cards.
• Vision works across stages with resident or CPU residency. The tower, the encode workspace and the scatter of the image's embeddings into the prompt stay on rank 0 with the embedding, before the layers run; the only thing later stages need is the multimodal RoPE positions, three rows per column, which travel in the control block (sized for them: eight int32 per boundary column). The default split counts a resident tower on the head stage. Overlay residency stays refused at load, before the weights: it borrows its encode window from the evict-ranked weights, and a later stage holds whole layers that nothing could lend.
• MTP works across stages: the replay records and their fold are per state shard, on the shard's own device, and the draft layer and head stay on rank 0. DFlash and DFlash2 work too; see below. The context cache works: its transactions copy each rank's planes and state shards on that rank's streams.

DFlash across stages

A masked drafter reads target features, the outputs of a few layers spread through the model: Qwen3.8-27B's DFlash2 adapter taps layers 5, 19, 33, 47 and 61 of 64 at width 5,120, the 35B-A3B DFlash adapter layers 1, 6, 11, 16, 22, 27, 32 and 37 of 40 at width 2,048. The drafter, its context KV and the feature buffers live on rank 0 with the head, so a tapped layer on a later stage sends its output there, the way the residual travels.

1. A link per tapped layer on a later stage (StageRuntime::features), from that stage to rank 0, built beside back in make_stage_runtime: a StageLink of two slots, each holding the widest pass (stage_boundary_columns) in BF16, plus a stream of the stage's own and the events that hand work to it and back.
2. Capture on the stage. As a tapped layer finishes, a stage-local tap copies its output into that tap's slab of a staging tensor in the stage's workspace and starts the transfer on the feature stream, so it crosses while the stage's later layers run. The stage's stream takes the feature stream back after the residual has left, which keeps the staging alive until every copy is done.
3. Re-capture on rank 0. Rank 0 has nothing to do until the residual returns, so run_staged queues the receives there first: each tapped output lands in one residual's room on rank 0's stream and goes to the sink's own capture_layer, so the prefill copy, the batch scatter and the check that every tapped layer was published are the ones one device runs. Stage 0's taps go to the sink directly.
4. Graphs and planning. The links' send and receive halves are capture-safe and the feature stream joins the capture through its events, so the DFlash decode graphs capture the extra copies with the residual's. The workspace plan reserves the staging slabs on each later stage and one receive slab on rank 0, sized for the stage with the most taps, and the default split counts the drafter's weights on the head stage, as it counts MTP's layer.

Cost, for a 30/34 split of the 27B, which puts the taps 33, 47 and 61 on the second stage: 30 KiB per column, so 240 KiB for an eight-column DFlash2 verify and 30 MiB for the default 1,024-token prefill chunk, sent while the stage's later layers run; only what is still in flight when the stage's last layer finishes adds to the pass. Each tap's two slots pin twice the widest pass in host memory, 60 MiB for those three taps at that chunk, beside the residual links' own slots.

Verification

• ninfer_multi_gpu_test, three suites in one executable (ranks share device 0, the staged protocol forced): the stage link (byte integrity across sizes and ring reuse, a two-hop chain, payload halves through separate graphs, the slot-reuse dependency read off the graph, eager use of a slot before and after a capture), the multi-rank KV pool and block-table copies against the single-rank pool as the oracle in both plane orders, and the stage solver against a brute-force partition enumeration.
• ninfer_kv_capacity_test: the per-device curve, including a device without KV and errors that name the device.
• ninfer_qwen3_5_stages_real_test (set NINFER_TEST_ARTIFACT): greedy output with --devices 0,0 and 0,0,0 equals --device 0 byte for byte, since the layers run the same kernels on the same data. Rows: graphs and eager, forced pinned-host transport, an uneven split (NINFER_TEST_STAGE_LAYERS, 20,44 by default), MTP, DFlash or DFlash2 (whichever drafter the artifact carries, each compared with the same backend on one device), Vision plain and with MTP (an image prompt spanning two prefill chunks, when the artifact has a tower), and a continuation that reuses the context cache. NINFER_TEST_DEVICE_IDS=0,1 puts the stages on real cards.
• ninfer_qwen3_5_loading_real_test: the default split covers the model, gives the head stage fewer layers, and gives a device with twice the memory more.

Ranks sharing a device cannot show a wrong-device pointer: a stale pointer into "another rank's" memory still works there. That class of bug needs a real second card.

Measured on two real GPUs (2026-09-21)

Rented Linux boxes, CUDA 12.8 build, Qwen3.6-27B groupwise-int (the pinned Hugging Face artifact), int8 KV. Neither box has peer access, so every boundary transfer goes through pinned host memory.

	2x RTX 3090 (PCIe 3.0 x16 each)	2x RTX A4000 (x8 + x16, cross-NUMA)
ninfer_qwen3_5_stages_real_test	every row byte-identical to one device: 2 stages, 3 stages (devices 0,1,0), graphs and eager, forced staged transport, uneven split, MTP, Vision (plain and MTP, an image over two prefill chunks), prefix reuse	the model does not fit one 16 GB card, so every row is compared with the first: identical across cut points, graphs and eager, staged transport, MTP, prefix reuse
CLI, prefill chunk 1024, splits 20/44, 44/20, 32/32 and the default	output identical to one device	output identical across splits, and identical to the 3090 output
Decode, plain	46.9 tok/s on one card, 48.5 split	24.1 tok/s
Decode, MTP3	100.2 tok/s on one card, 105.3 split	52.6 tok/s
Prefill	1.56k tok/s on one card and split	825 tok/s
Weights per card (default split)	15.3 GiB on one card, 8.35 GiB on the first split card	7.71 GiB
--kv-capacity auto at --max-context 262144	one card: refused (minimum reservation 9.2 GB, 8.76 GB free after weights); split: 262,144 tokens, 11.2 GiB free afterwards	262,144 tokens, 3.9 GiB free afterwards

A ninfer-serve check on the 3090 pair (four lanes, concurrent and repeated chat requests, prefix cache on) gave the same output split as on one card for every request run alone, run in parallel once, and repeated; a second concurrent round differed on one request, and one card differs from itself the same way, because lane batching changes the arithmetic.

What this says: the split costs nothing measurable at decode on PCIe 3.0 x16 (one 20 KB residual crossing per stage boundary per token is about 25 us against a ~21 ms step), and it is what makes the 27B's full 262,144-token context fit at all. It does not speed anything up: prefill is the same and decode is within noise. The A4000 pair runs the 27B with its full context and its decode rate matches the plan's estimate for a pipeline on 448 GB/s cards.

On a card with fewer SMs than the tuned route tables assume (the A4000 has 48), a GDN gating schedule whose cooperative grid would not be resident falls through to the first later route that fits the device; its speed there has not been tuned.

Later checks on real pairs, none with peer access: Ternary Bonsai 2 27B on two RTX 4090s generates byte for byte what one GPU does in every row of the stage test, and suspends and resumes with its prefix reuse intact; Qwen3.8-Flash-Next runs its experts over two RTX 3090 Ti, two RTX 4090 and two RTX 5090 (measurements).

DFlash across stages, measured (2026-10-06)

Rented Linux pairs, none with peer access, so every hop goes through pinned host memory; CUDA 12.8 builds for 80, 86 and 89, CUDA 13.1 for 120a; the pinned Qwen3.8-27B (DFlash2 bundle, 1cbd84e7) and Qwen3.6-35B-A3B (DFlash bundle, ee449580) artifacts; int8 KV.

ninfer_qwen3_5_stages_real_test with the stages on two cards (NINFER_TEST_DEVICE_IDS=0,1) passed all 22 rows, 11 per artifact, on 2x RTX 3090 (sm_86), 2x RTX 4090 (sm_89), 2x RTX 5090 (sm_120a) and 2x A100 PCIe 40 GB (sm_80): two and three stages, graphs and eager, forced staged transport, an uneven split, MTP, and DFlash2 or DFlash in two stages, staged transport, three stages eager and an uneven split, each over a 1,300-token prompt in three chunks and a continuation served from the context cache, byte-identical to the same backend on one card.

Speed with the CLI (--greedy --no-thinking, seven drafts): a 37-token chat prompt with 256 new tokens ("short") and a 3,322-token prompt with 64 ("long"), tok/s on one card and then split over the two by the default split, mean of two or three runs. Every split run generated the one-card text.

	2x RTX 3090	2x RTX 4090	2x RTX 5090	2x A100 PCIe
27B DFlash2, short, decode	206.1 / 223.2	264.8 / 264.8	416.4 / 418.4	242.5 / 233.8
27B DFlash2, long, prefill	1,247 / 1,270	2,267 / 2,227	2,965 / 2,990	1,357 / 1,297
27B DFlash2, long, decode	106.5 / 114.6	135.6 / 135.6	213.6 / 214.3	125.7 / 130.9
27B plain, short, decode	47.9 / 48.9	55.8 / 55.5	91.7 / 91.0	52.3 / 54.7
35B-A3B DFlash, short, decode	482.8 / 490.7	837.0 / 818.2	1,260 / 1,225	475.0 / 365.9
35B-A3B DFlash, long, prefill	6,320 / 6,233	12,033 / 11,667	11,300 / 11,750	7,327 / 5,173
35B-A3B DFlash, long, decode	195.8 / 205.5	343.0 / 337.3	471.7 / 462.2	199.2 / 197.2
35B-A3B plain, short, decode	206.1 / 207.9	274.0 / 268.7	not run	207.6 / 181.9

The 3090 column was measured with taps sent together after a stage's last layer; sending each tap as its layer finishes left decode equal within 1% on the 4090 and A100 pairs and raised the 35B-A3B's long prefill from 11,433 to 11,667 tok/s (4090) and from 4,777 to 5,173 (A100). The weights per card in the default split were 10.3 + 7.74 GiB for the 27B with DFlash2 (18.0 GiB on one card) and 10.6 + 9.41 GiB for the 35B-A3B (20.0 GiB); the 27B's workspace grew from 183.9 to 233.9 MiB per card for the boundary and the tap staging.

What this says: on the GeForce pairs the split costs DFlash and DFlash2 no more than it costs plain decoding, within about 3% either way (the 3090 pair is faster split, as it was on 2026-09-21), and the features ride along at no measurable price at decode. The A100 pair is the exception, and not because of DFlash: the 35B-A3B lost 12% split without speculation and 23% with DFlash. The same split with both stages on one A100 (--devices 0,0, local or forced staged links) cost nothing (466-473 tok/s on one card, 476-481 split), each A100 alone ran at the same speed, and an nsys profile of the split showed the kernels themselves running longer (the MoE path kernel +21%, the verify GDN layers +46%) rather than gaps opening between them. That points at the cards' clocks while they wait for each other (the two sit on different sockets of that host); GPU clock metrics are not permitted in the container, so the cause is not confirmed.

Tensor parallelism: status

Not built. What was established, so the next attempt starts from evidence rather than from the plan:

• Transport (tools/tp_probe.cu, rented boxes, 2026-09-21): neither a 2x RTX 3090 box (one slot PCIe 3.0 x4), a second 2x RTX 3090 box with both slots x16, nor a 2x RTX A4000 box had peer access or NVLink. The x16 pair measured 15 us (host-mapped kernel) to 24 us (staged, in a graph) for one column, 33-38 us for four, 75-78 us for sixteen, 130 us for thirty-two and 3.4 ms for a 20 MB prefill chunk, half the x4 box's time at every size. An all-reduce of one decode column (20 KB, FP32) costs 17-33 us whether done by a host-mapped kernel, a staged pinned-host copy or NCCL; at 640 KB it is 190-440 us and at 20 MB 5-13 ms. The transports are within about 20% of each other at every size, so a first TP should reuse StageLink-style staged copies (they already capture into graphs) and a local reduce, not write spin-wait kernels.
• Where it pays. Projected against this pipeline on two A4000s: 1.6-1.7x at one decode column, about 1.5x at four, roughly break-even at sixteen, a loss at thirty-two (MTP3 with eight lanes). TP prefill is bound by the link (~840 ms of all-reduce per 1,024-token chunk), so prefill speed belongs to the micro-chunk overlap above, not TP. Link width is per machine and decides everything; run the probe on the target box first. On the x16 3090 pair the 21 ms decode step would become about 16 ms with MLP-only TP (64 collectives at ~27 us against half the MLP bytes) and about 14 ms with full TP (128 collectives), i.e. roughly 1.3x and 1.5x at one column.
• Smallest useful first step: MLP-only TP on the dense 27B (about 64% of the weight bytes; projected ~1.35x decode). Each rank keeps attention and GDN whole and replicated (weights and KV), and takes half the MLP: gate_up split by rows (gate rows then up rows, so the SwiGLU pairing holds) and down split along K. A rank's down projection can reuse linear_add: rank 0 adds its partial into the real residual, the other rank into zeros, and one BF16 exchange plus an add gives every rank the same new residual (one extra rounding against the single-device path, so the oracle is a tolerance on logits and perplexity, not byte equality).
• What it needs before any speed can be claimed: loader-time slicing of row-split quantized parents (contiguous row ranges for gate_up, per-plane column ranges for down, group-aligned); the linear_swiglu Q4 kernel at 17,408 rows and linear_add Q5 at K=8,704, each registered for those exact shapes and qualified against the FP64 oracle and tuned on the target card (routes inherited from the full-width shapes are hypotheses); per-rank replicas of the KV pool and state pool with fan-out in the transactions; and the exchange itself. The decisive measurement is the probe on the machine that will run it.