[hash13]

Prefill: allow deeper overlap across multi-GPU pipeline stages

Problem

In the current multi-GPU prefill path, every stage launches the next stage asynchronously, but still waits for that stage to finish before returning.

With more than two GPUs this creates a recursive wait chain:

GPU0 waits for GPU1
GPU1 waits for GPU2
GPU2 waits for GPU3

As a result, intermediate stages cannot accept the next prompt chunk while later GPUs are still processing the previous one.

On a 4-GPU setup this can lead to visibly staggered GPU utilization during prompt processing instead of a fully filled pipeline.

Change

This PR changes the lifetime of the per-stage asynchronous hand-off.

Intermediate stages now return after handing their chunk to the direct successor instead of draining the complete remaining GPU chain.

The pipeline is drained only once, when the outer/root Prefill::run() has submitted all prompt chunks.

Conceptually, the current behavior is closer to:

GPU0: C0 C1 ---- C2 ----
GPU1:    C0 ---- C1 ----
GPU2:       C0 ---- C1
GPU3:          C0 ----

The intended behavior after this change is:

GPU0: C0 C1 C2 C3 C4
GPU1:    C0 C1 C2 C3
GPU2:       C0 C1 C2
GPU3:          C0 C1

This does not introduce tensor parallelism and does not change the CUDA kernels or model computation.

It only increases overlap between existing layer-split prefill stages.

Implementation

The change:

• moves the successor std::future from local Prefill::run() state to the Prefill instance
• keeps one outstanding job per direct successor
• prevents intermediate stages from recursively draining all later stages
• preserves the existing double-buffered host hand-off buffers
• drains the complete pipeline once at the end of the root prefill call
• keeps each individual GPU stage serial; two chunks are never executed concurrently on the same Prefill instance

Expected effect

The main benefit should be visible on configurations with 3 or more GPUs and prompts large enough to contain several prefill chunks.

For example, with four GPUs and a 2K prefill chunk size, a 32K prompt provides enough chunks to keep all stages active concurrently after pipeline warm-up.

Short prompts or a single prefill chunk should see little or no improvement.

Decode/token-generation performance is not expected to change.

Motivation / observed behavior

On a 4× RTX 3060 system, prompt processing currently shows GPU activity progressing largely stage by stage, while decode performance is normal.

This suggests that part of the low prefill throughput may come from the recursive synchronization between layer-split stages rather than raw GPU compute.

This PR is intended as a small, isolated change to test that hypothesis without changing expert streaming, tensor parallelism, CUDA kernels, or the model architecture.

The most useful validation is whether multiple GPUs stay active concurrently during long prompt prefill and whether prompt throughput improves without affecting model output or decode performance.

[hash13]

Test hardware

Component	48 GB test setup	64 GB test setup
CPU	Intel Xeon E5-2698 v4	Intel Xeon E5-2698 v4
GPUs	4× NVIDIA GeForce RTX 3060 12 GB	4× NVIDIA GeForce RTX 3060 12 GB
PCIe	PCIe 3.0 x8 per GPU	PCIe 3.0 x8 per GPU
Topology	All GPUs connected via CPU PCIe root ports / PHB	All GPUs connected via CPU PCIe root ports / PHB
NUMA	Single NUMA node	Single NUMA node
Host RAM	48 GB DDR4 (mmap-expert)	64 GB DDR4 (all experts in RAM)
OS	Debian 13	Debian 13
Measured H2D bandwidth	~5.9 GB/s	~5.9 GB/s

Prefill benchmark results

Prompt size	48 GB Before	48 GB After	48 GB Improvement	64 GB After, 130 W	64 GB After, 170 W
8K	525.7 tok/s	529.4 tok/s	+0.7%	692.5 tok/s	699.9 tok/s
16K	545.9 tok/s	600.7 tok/s	+10.0%	884.5 tok/s	896.6 tok/s
32K	698.5 tok/s	839.4 tok/s	+20.2%	1,417.7 tok/s	1,452.8 tok/s
64K	753.9 tok/s	964.2 tok/s	+27.9%	1,939.8 tok/s	1,965.9 tok/s

There is no "Before" benchmark for the 64 GB configuration, so no optimization percentage can be calculated for the 64 GB results.

For the 48 GB configuration, the optimization has little effect at 8K, but the benefit increases significantly with prompt size: about 10% at 16K, 20% at 32K, and 28% at 64K.

[hash13]

Tested with qwen3.8-flash-next-iq3_xxs

[hash13]

I tested the suggestions #559 on top of #598 on a 4× RTX 3060 system with layer_split auto.

I added two experimental changes:

• an optional dedicated CUDA copy stream for the inter-stage D2H handoff, to overlap the handoff copy with compute;
• a configurable handoff queue depth, allowing two chunks to be in flight between stages instead of one.

Test setup:

Qwen3.8-Flash-Next IQ3_XXS, 4× RTX 3060 12 GB, --expert-cache auto --prefill auto --max-context 131072 --kv int8, with unique nonces at the start of every prompt to avoid prefix-cache effects.

I tested all four combinations:

async handoff	depth	8K	16K	32K	64K
0	1	679	866	1399	1903
1	1	667	856	1386	1892
1	2	677	869	1404	1917
0	2	690	876	1413	1922

The dedicated async D2H stream did not improve performance. With queue depth 1 it was actually around 0.6–1.8% slower, and with depth 2 it was essentially neutral compared with the synchronous handoff.

Increasing the handoff depth from 1 to 2 without the async copy stream produced only a small improvement, roughly 1% in the latest back-to-back run.

So the main proposed async-copy optimization does not appear to address the remaining bottleneck on this setup.

For now I would therefore not add the dedicated copy-stream implementation to #598. The only potentially useful part is the slightly deeper handoff queue, but the gain is small enough that it probably needs more validation before being included.

[hash13]

Prefill chunk-size measurements

Test setup:

• Model: Qwen3.8-Flash-Next IQ3_XXS
• GPUs: 4× RTX 3060 12 GB
• layer_split auto
• --expert-cache auto
• --max-context 131072
• --kv int8
• 3 runs per size
• Metric: median prompt processing throughput in tok/s
• Unique nonce at the start of every prompt to avoid prefix-cache effects
• Based on the original refill: improve multi-GPU processing speed ~27% PP improvement on 4× RTX 3060 #598 changes, without the later experimental async-handoff patch

--prefill	2K	4K	8K	16K
auto	314.4	508.7	684.6	870.4
1024	286.6	430.0	—	—
3072	319.8	514.1	778.9	1040.1
4096	321.0	517.6	698.6	1109.8

Change relative to auto:

--prefill	2K	4K	8K	16K
1024	-8.8%	-15.5%	—	—
3072	+1.7%	+1.1%	+13.8%	+19.5%
4096	+2.1%	+1.7%	+2.0%	+27.5%

The results suggest that smaller prefill chunks can significantly improve the multi-GPU layer-split pipeline introduced by #598, but simply making the chunks as small as possible is not beneficial.

1024 is clearly too small and loses substantial throughput because of the additional per-chunk overhead.

The optimum also depends on prompt size. For the ~5.9K-token prompt, 3072 performs best, improving throughput from 684.6 to 778.9 tok/s (+13.8%). For the ~11.6K-token prompt, 4096 performs best, reaching 1109.8 tok/s instead of 870.4 tok/s with auto (+27.5%).

This suggests that the current auto chunk selection is not optimal for a multi-stage layer-split pipeline. A useful follow-up could be an adaptive chunk-size heuristic that considers both prompt length and the number of pipeline stages, rather than always preferring the largest efficient chunk.

[Niko1221]

Thanks! Taken into 0.1.39. With one or two GPUs the pipeline runs exactly as before: the last stage has no successor, so nothing waits differently. We checked one GPU here (output byte-identical); the 3+ GPU overlap and the +10-28% prompt figures are your 4x RTX 3060 measurements. We merged it with #663.

[hash13]

@Niko1221 Thanks for the merge!

I compared the current 0.1.39 version using several fixed prefill sizes (1024, 2048, 4096, 8192) and auto. I also compared 0.1.39 auto directly against 0.1.38 auto using the same hardware and benchmark setup.

Cache: unique nonce is placed at the START of every prompt

All values below are median prefill throughput in tok/s.

Actual tokens approx.	0.1.39 prefill 1024	0.1.39 prefill 2048	0.1.39 prefill 4096	0.1.39 prefill 8192	0.1.39 auto	0.1.38 auto	0.1.39 auto gain
257	157.8	157.5	157.9	158.3	157.6	159.6	−1.3%
437	199.8	200.5	199.8	199.6	201.0	200.4	+0.3%
801	228.1	228.5	227.5	228.3	229.3	227.9	+0.6%
2,970	435.8	560.4	558.9	506.5	502.1	510.7	−1.7%
5,862	518.0	716.8	711.0	714.5	718.0	696.8	+3.0%
11,648	567.3	914.0	1,124.1	916.4	918.1	728.7	+26.0%
23,223	613.6	1,473.4	1,411.6	1,479.3	1,483.5	980.7	+51.3%
46,369	638.7	2,008.7	1,674.0	2,005.8	2,010.3	1,084.1	+85.4%
92,609	661.4	2,467.5	1,945.2	2,470.8	2,472.6	1,128.6	+119.1%

The fixed prefill-size tests also show that 1024 is clearly too small for longer contexts on this 4-GPU setup, while 2048, 8192, and auto perform similarly at large prompt sizes. 4096 performs particularly well around 12–14k actual tokens, but falls behind again at larger contexts.

So on this system, 0.1.39 roughly doubles long-context prefill performance compared with 0.1.38. At the largest tested prompt, throughput increased from 1,128.6 tok/s to 2,472.6 tok/s, or about 2.19× the previous performance.