Support pipeline parallelism with MoE experts in host RAM

[drrros]

Support pipeline parallelism with MoE experts in host RAM

TL\DR: x1.2-x2 prompt processing of long prompts on multiple GPUs with with MoE experts in host RAM

Problem: after the discussion in #29761 i decided to look deeper into the pipeline parallelism thing. Basically it is this: when part of the MoE weights sits in host RAM, llama.cpp silently switches off pipeline parallelism between GPUs (src/llama-context.cpp, !has_tensor_overrides()). The check itself comes from #11397, where copies of the host weights were reserved n_copies (4) times per device. So with overrides the scheduler goes fully serial - any op with host weights is offloaded to the first device that can do it, which on CUDA is always CUDA0. Result: all experts are uploaded through one GPU, and the other GPUs are just waiting.

slaren wrote in #11397 that "the pipeline parallelism implementation doesn't work at all if a part of the model needs to be evaluated on the CPU". With fit it is a bit subtle actually: the experts are really evaluated on CPU only on small batches (generation), on prompt batches the ops go to GPU through op offload. But the check does not care about that, pipeline is off anyway.

There is also #20252, people asking there why pipeline does nothing for them. ggerganov answers that PP scales nearly linear and it is either a regression or something specific to the reporter's system. But that thread is about models fully in VRAM. Once part of the weights is in RAM there is nothing to suspect, it just is not supported today.

What i have

A branch, 6 files, +272/-7, where pipeline parallelism keeps working with the weights in host RAM:

• one copy of a host weight per device instead of 4 copies held for the whole graph. The allocator reuses it between splits;
• user inputs (KQ mask and friends) are uploaded once per device, on the stream of the split that uses them, instead of 4 copies per device. With pipeline this alone cuts the compute buffers from 9133 to 2988 MiB per GPU (Qwen3.5-35B in VRAM, c=262144). The mask grows as ubatch x context, so on long contexts the saving is bigger;
• ops with host weights run on the device of their own layer, not always on GPU0. That is the main point actually - every GPU uploads experts of its own layers while the others are computing;
• the MoE ids read (sched : copy only the used experts when offloading prompt processing #15346, used to copy only the selected experts) makes the host synchronize with the split backend, so the host stands until that one finishes and cannot queue the next graphs - the pipeline stalls. On big batches all experts of the layer are uploaded instead, since most of them are used anyway;
• graph topology is kept constant with pipeline. Since CUDA: fuse shared experts into MMVQ #29184 a graph with 1 output row does not match the reserved worst-case graph anymore for models with shared experts on CUDA (the fusion of shared experts into MMVQ adds alloc deps on the single-row path), so every ubatch reallocates the compute buffers and that synchronizes all devices. The regression bisects to CUDA: fuse shared experts into MMVQ #29184 exactly: Qwen3.5-35B fully in VRAM, 3 GPUs, server, 52421-token prompt, c=65536 - 7131 / 7632 t/s on the parent 4ebdf2c, 3740 / 3747 on bed0a85 (CUDA: fuse shared experts into MMVQ #29184), 3701 / 3716 on master cb7934c, and 7148 / 7675 with only this fix on top of master. Models without shared experts (gpt-oss for example) are not affected;
• fit gives every GPU an equal slice of layers in this case, instead of putting all host weights on the last device (-ts 8,8,33 becomes 16,16,17);
• two small fixes that came along the way: ggml_backend_sched_reserve_size drops the backend assignments made by the callback (it resets before split_graph), and CUDA top-k. About top-k separately below, the short version is that with the pinned CCCL >= 3.4.3 it runs as a per-row DeviceTopK loop (DeviceSegmentedTopK is still a TODO upstream), which is worst exactly on many rows of moderate length. The radix select from ROCm: add radix TOP_K for long rows #27466 is up to x125 faster there ([1536, 4096]), x45-51 on 512 rows of 1536-8192 columns, and only x3.8-4.2 on the longest rows (65536 columns). Against the argsort path of CCCL < 3.4.3 it is x1.5-5.5.

Some of these are useful without the host weights thing at all, i will try to split them so.

Benchmarks (3x RTX PRO 4000 blackwell, PCIe 5.0x16, EPYC 9274F, 384 GB DDR5 (12), CUDA 13.3)

One note about CCCL first: all llama-bench and server numbers below are built against the 3.3.3 that ships with CUDA 13.3 (the top-k section is the exception, there i measured both). Master built like the release builds do (CCCL 3.4.3, #29792) has a much slower top-k for the Qwen3.8 and DeepSeek indexers, so master numbers below are probably too good to it - i did not measure that. And from the CI workflow comments, CUDA >= 13.5 ships CCCL >= 3.5 on its own, so local builds on a new toolkit fall into the same slow loop, not only the release builds.

llama-bench:

with the default fit layout, -b 4096 -ub 4096 -fa on -lm none -n 128, CUDA_SCALE_LAUNCH_QUEUES=4x, master -> patch t/s (MiMo also -ctk q8_0 -ctv q8_0, Qwen3.8 next also --lazy-mode off):

Model	pp512	pp2048	pp8192	pp32768	tg128
Qwen3.8-Flash-Next Q8_0 (151 GB)	749 -> 736	1236 -> 1215	1377 -> 1681	1249 -> 2369	46.3 -> 47.7
DeepSeek-V4-Flash Q8_K_XL (151 GB)	332 -> 316	687 -> 667	740 -> 1003	711 -> 1484	24.0 -> 24.2
MiMo-V2.6-Flash MXFP4 (167 GB)	371 -> 355	876 -> 705	1119 -> 1454	1048 -> 2083	34.8 -> 34.7

llama-server:

(RFC 7540 prompt, 52421 tokens, MiMo 51341 because of its tokenizer, no cache), PP t/s of the 1st / 2nd request:

Model	master	patch	TG
Qwen3.8 (fit, c=262144)	952 / 961	1756 / 1982	40.8 -> 41.6
DeepSeek-V4-Flash (fit)	625 / 629	1128 / 1321	23.0 -> 23.1
MiMo-V2.6-Flash (fit, c=131072)	883 / 886	1741 / 1834	33.6 -> 33.3

Generation stays about the same, that is not what pipeline is for: TG in the tables is +2..3% on Qwen3.8, +0.4..1% on DeepSeek, -0.3% on MiMo. The Qwen3.8 bump is above the noise, but it is probably the other fit layout, not the pipeline.

2 GPUs with fit (Qwen3.8, which does not fit in VRAM either way): pp16384 1254 -> 1496.

All experts in RAM (Qwen3.5-35B-A3B, -ncmoe 999, no fit): pp8192 3078 -> 4393, pp32768 2814 -> 6274. Over the ubatch size (pp8192, -b 8192) it looks like: parity up to ub 256, then +45% at 512, +51% at 1024, +66% at 2048, and +46% at 4096. Not monotone, the peak is at 2048 and i do not have a clean explanation for that yet.

Models fully in VRAM: the topology fix alone gives back the pre-#29184 speed on Qwen3.5-35B (3 GPUs, 52k prompt, the bisect above): 3701 -> 7148. The rest of the branch does not regress anything: llama-bench pp8192 gpt-oss-20b 17812 -> 17746, Llama 3.3 8B 12615 -> 12583, both are noise.

CUDA top-k (test-backend-ops perf -o TOP_K, us; master picks the path by the CCCL version: with >= 3.4.3 it is the per-row DeviceTopK loop, that is the releases and CI since #29792, before it the argsort path):

ne, k	master 3.3.3	master 3.4.3	patch
[2048, 4096], 512	1703	37656	318
[65536, 512], 512	5928	6170	1470
[65536, 8192], 2048	95203	99531	26500

First two rows are the real shapes of the indexers (QSA in qwen4exp does top-k over n_kv/4 columns, lightning in deepseek4 up to 65536 columns, both k=512), number of rows is number of tokens in the ubatch. The last row is just a stress case, no model gives it to me. Peak GPU memory is on par with master, on the stress case patch 2418 MiB, master 2734 MiB on 3.3.3 and 2350 MiB on 3.4.3.

Correctness. test-backend-ops -o TOP_K and -o ARGSORT pass on CCCL 3.3.3 and 3.4.3, ctest 6/6 (alloc, backend-sampler, llama-archs, save-load-state, sampling, batch-alloc). In the CI style -DGGML_SCHED_NO_REALLOC=ON build the patch passes a 52k prompt on the server while master aborts on its second ubatch. Not gonna hide it: on that model master runs with pipeline off there (sched copies = 1) and the failing graph has 7641 nodes against 6856 reserved, so it is not the #29184 case. Cause is not analyzed, probably qwen4exp specific - in this mode generation aborts near the 128th token both on master and with the patch. The #28739 scenario (GGML_OP_OFFLOAD_MIN_BATCH=0 -ncmoe 999 -ub 64) does not regress (158 -> 160).

And about PCIe - pinned data goes up at 57.6 GB/s on one GPU alone and at 106.8 GB/s when all three go at once, 1.85x. So the links themselves are symmetric and fine. The problem on master is that the single upload path sits on GPU0 - nsys shows expert uploads taking 18 of 52 s on GPU0 on a 52k prompt (profiled with the master preset parameters, ub 6144, not the settings from the tables above). The pipeline does not make uploads faster, it just spreads the same traffic over three links.

Notes on the numbers, so you do not have to guess where the patch is worse:

• single-ubatch prompts with fit are -2..-5%. There is nothing to overlap there, and the balanced layout keeps slightly fewer experts in VRAM;
• MiMo pp2048 (-20% in the table) is a llama-bench artifact, not a real regression. llama-bench makes a context per test, so a 2048-token prompt becomes a "full ubatch" and triggers the upload-all-experts path even though the graph runs alone. On the server the same model is at parity (643 -> 656 on a 1746-token prompt);
• both sides use CUDA_SCALE_LAUNCH_QUEUES=4x. Without it the patch loses 22..27% on Qwen3.8 (pp8192 1685 -> 1314). Master without it i did not measure, and i probably should, will add.

Anyway, I've tested this locally and it works really good for me, the plan is to upstream this either as a single PR (if bit less than 300 lines are in reviewable range\limit), or as a few small PRs as needed to simplify review. Currently I have a draft PR (#29963) with all changes included, the draft PR is split by commits like this (and so can be split by PRs accordingly):

#	PR	Files	+/-	Needs	What it gives alone
1	cuda : use radix select for TOP_K with multiple long rows	ggml-cuda/top-k.cu	+15/-2	-	the top-k table above, up to x125 on CCCL 3.4.3
2	ggml : upload user inputs once per device before the splits	ggml-backend.cpp	+54	-	compute buffers 9133 -> 2988 MiB per GPU, PP unchanged
3	llama : keep graph topology constant with pipeline parallelism	llama-context.cpp/.h, llama-graph.cpp	+33/-1	-	PP qwen 3.8 35b 3701 -> 7148, the #29184 bisect above
4	ggml : pipeline parallelism with weights in host memory	ggml-backend.cpp, llama-context.cpp	+94/-4	2, 3	manual balanced layout, pp32768 1095 -> 1995; on master's fit layout it stays at master speed, the spread check does not enable pipeline there
5	common : balanced layer slices for fit with pipeline parallelism	fit.cpp	+76	4	the llama-bench fit tables above, Qwen3.8 1249 -> 2369

1 to 3 are independent and useful on their own, 4 needs 2 and 3, 5 goes after 4. The sched code in 2 and 4 (host-weight predicate, input upload loop, MoE path) touches the same text as #21067 (weight prefetch on a copy stream), it also has a few lines in src/llama-context.cpp but does not go into fit.cpp.

If something above looks wrong, most likely it is my measurement - say where and i will re-check.

[eapache]

I have a multi-GPU system that is quite asymmetric:

• an RTX PRO 4500 over PCIe 5x16
• an RTX 3060 over PCIe 3x4

I won’t have time to test this for a bit, but I wonder if the limited PCI bandwidth to the 3060 will make this approach actually slower than the version currently on master which runs everything through the 4500?

[drrros]

if the limited PCI bandwidth to the 3060 will make this approach actually slower than the version currently on master which runs everything through the 4500

I think it could, but it's better to test on hardware ofc. Perhaps some optional --balance-layers could be better approach?

[drrros]

--balance-layers

Or, even better, an option to fit param: on \ off \ balance

[eapache]

Yeah I did a quick local test with gemma-4-26B-A4B-it-UD-Q4_K_XL.gguf , -ncmoe 999, 1:1 layer split. On master the 4500 peaks at RX of ~10GiB/s, the 3060 is mostly idle, and I get about 1100 t/s pp. With #29963 the 3060 peaks at RX of about 3GiB/s, the 4500 is mostly idle, and I get only 600 t/s pp.

[eapache]

Explicitly hard-code-disabling pipeline_parallel on the PR branch recovers full performance, so it is definitely that vs any of the other changes.

[eapache]

I tried disabling the asymmetry guard in your MR and benchmarking some more natural layer splits (given the relative speeds of the two cards) like 4:1, 6:1, 10:1 - clearly better than 1:1, but ultimately nothing was as fast as just disabling the parallelism and using the 4500 exclusively. Even with all transformer layers on the 4500 and only the output layer on the 3060, enabling the pipeline still imposed a roughly 14% penalty.