ggml: allow prefetching tensor overrides

[am17an]

Overview

This PR adds support to prefetch tensor overrides for each layer, overlapping with the compute for the current layer. Only the CUDA implementation is provided, guarded by the --prefetch-weights flag. At this stage I would consider this a PoC, so keeping as draft for now for comments and further tests.

Note that --no-mmap is necessary for this work, otherwise the operations are implicitly serialized due to the weights not being pinned.

Performance Analysis

Dense Models

For dense models, the optimization is relatively straightforward. Let's call $C_{n}$ the time taken for computing the $n^{th}$ layer on the GPU, and $T_{n+1}$ the time to transfer layer $n^{th}+1$ weights from the CPU to the GPU.

If you choose to override weights on the CPU, then this overlaps $T_{n+1}$ with $C_{n}$, whereas currently everything happens sequentially. So if $C_{n} > T_{n+1}$ we can "hide" the transfer latency and have it behave like a GPU. Although when $C_{n} >> T_{n+1}$, $C_{n}$ dominates so this is less useful. On the other hand, decreasing $T_{n+1}$ is only possible using newer hardware like PCIe Gen5 or using lower bpw.

There are two natural dimensions where we can increase $C_{n}$ without increasing $T_{n+1}$, those are ubatch size and the kv-cache depth. Here is a graph using a relatively recent model, since this is a linear attention model, the compute doesn't go up as fast a quadratic attention model with increasing depth. We override ffn_(gate|up|down).* to the CPU, which are the bulk of the weights in each layer. We can see it benefits at all batch sizes, but the gap is lesser at higher batch sizes (since $C_{n}$ starts to dominate)

llama-bench -m /opt/models/Qwen3.5-27B-Q4_K_M.gguf -fa 1 -p 2048 -ub 512,1024,2048 -d 0,10000,20000,30000,40000,50000 -n 0 -ot "ffn_(gate|up|down).*=CPU" -pw 0,1 --mmap 0

MoE models

For MoE models, the situation is different because of selective copying of experts (added in #15346). This is a massive improvement for smaller ubatch sizes, naturally prefetching cannot do this as it does not know which experts will be selected in the next layer. The situation is worse for larger MoE models with more experts and larger inner dims, so likely this will be slower for large MoE models unless they are deep (more layers) and not wide (larger expert dims).

The equation for prefetching to be beneficial becomes $C_{n} + T_{c}' > T_{c+1}$, where $T_{c}'$ is the time to transfer the selected experts for the $n^{th}$ layer. The same stuff applies, if we scale $C_{n}$ it becomes more beneficial, but we can also scale $T_{c}'$ by increasing the ubatch size, hence increasing expected number of used experts in a ubatch. In the graph below we see comparable performance at ubatch=512, but much better at 1024,2048. Since 50k context fully fits on this GPU, I also the added the theoretical maximum of fully offloading to the GPU

llama-bench -m /opt/models/gpt_oss-20b-mxfp4.gguf -fa 1 -p 2048 -ub 512,1024,2048 -d 0,10000,20000,30000,40000,50000 -n 0 -ncmoe 999 -pw 0,1 --mmap 0.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES, I used AI to write parts of the code, and I used it for debugging the ggml-backend code while struggling with the graph allocator. I have reviewed and tested it myself.

[am17an]

Simplified the code to not use any backend specific code, so any backend which supports events and streams should be able to use this. I think by default this should be on for dense models, and for MoE models it makes sense for smaller models.

[JohannesGaessler]

This looks very promising! I'm currently busy with tensor parallelism but after that I will look into whether I can help with this.

[pandruszkow]

What sort of work needs to be done to get this merged?

[am17an]

@pandruszkow someone needs to review it

[mholtgraewe]

Are there any plans to merge this PR any time soon? It's been in a limbo for a few months now, but it's sorely needed...

[ddh0]

I agree, it would be very nice if this was merged when ready. :)

[am17an]

@ggerganov is there some version of this which could be merged? Seems like it would a nice and unique feature for hybrid inference.

[JohannesGaessler]

Why are per-backend changes necessary?

[am17an]

It's to advertise the async stream availability, it can be done without it though

[pwilkin]

This would really help with mixed CPU/GPU inference, from my profiling, the data copy cost is currently the absolutely dominant cost of the entire process when running etc. on --cpu-moe.

[pwilkin]

@am17an BTW can you make --prefetch-weights automatically enable --no-mmap with a warning message? I'm absolutely sure most people won't read the fine print at the top :)

[0cc4m]

Why does this need a new flag and not just a check that events and tensor_set_async is available? You enabled it for CUDA only, but made no changes in the CUDA backend otherwise, so I assume the capability was already available.

[am17an]

Yes I can do that. It's already available in CUDA

[0cc4m]

So is there another requirement besides events and async set? copy_stream sounds like a separate transfer-only cuda stream, which e.g. Vulkan also provides, we call it transfer queue. But whether the transfer queue or the compute queue is used is decided by the backend, and currently I think we usually use the compute queue for async operations, because otherwise we have to synchronize the queues.

[am17an]

At this point there is no distinction between these things. From the ggml side it just creates another stream for the same device.

[ORippler]

1. To me, this looks like a formalization of how asynchonous copy/compute should behave on ggml backends. I would appreciate a write-up/simple sketch of this, especially for the cases where we do a copy from backend a -> backend b.
   • Linking a possibly related PR that is blocked from the currently missing, shared understanding of how async copy and compute should work on a ggml backend: #20793. CC @aendk
2. Event-based synchronization is currently bugged afaik, so we should ideally resolve this first before merging.

[pwilkin]

Bump, any updates on this? :)

[aendk]

Suggested change

[aendk]

I think, from a readability perspective, this should be moved to a separate function.
(The same goes for the long-standing MoE weight case, which puts ~80 LoC between the input copy logic and the activation copy logic between L1481 and L1564)
IMO, ggml_backend_sched_compute_splits() is barely maintainable as-is, so we should be careful when adding features onto it.

[0cc4m]

IMO this would need changes to the backend API to explicitely ask for a separate (transfer) stream/queue, otherwise it's just gonna end up in the main queue alongside other async calls. This is exactly the kind of thing where copy/DMA hardware should be used explicitely to handle the copies completely separated from the actual compute.

[am17an]

On CUDA it doesn't need that distinction, but perhaps we could instantiate a backend using a copy or compute enum, just thinking out loud.

[0cc4m]

Ah, right, I had missed the separate backend instance. Then it might not be important, if the GPU scheduler handles it well-enough. Though maybe some kind of separate data streaming abstraction could still help.

[noonghunna]

Some measurements from running this on a large CPU-offloaded MoE, plus a small gate that addresses the regression you predicted in the PR description.

Setup: Laguna-S-2.1 (117.6B, 256 experts, top-10, 48 layers) at IQ4_NL, 2× RTX 3090 (PCIe, no NVLink), 28 expert layers offloaded to CPU via --override-tensor, --no-mmap, 262K context. Within-binary A/B (-pw 0 reproduces the unpatched engine exactly), 2 boots per arm.

Prefill improves, first-token latency regresses hard

metric	-pw 0	-pw 1	delta
prefill 10K	902.5 t/s	1035.3 t/s	+14.7%
prefill 90K	734.2 t/s	835.7 t/s	+13.8%
TTFT, short prompt	1085 ms	1604 ms	+47.8%
decode	38.40 t/s	36.55 t/s	−4.8%

Mechanism: it's bytes, and it's the case you called out

Prefetch runs one split ahead of the router, so it cannot know which experts will be selected and must copy whole expert tensors — forfeiting the selective-expert copy from #15346 that the non-prefetch path uses. Bytes moved on the short-prompt graph: 13,668 MiB → 28,128 MiB (2.06×).

The control that settles it: disabling the selective copy on the -pw 0 path gives 30,774 MiB and TTFT 1725 ms — the regression reproduces with no prefetch machinery in the graph at all. So this is not event overhead, allocation, or stream sync.

Why it is asymmetric between short and long prompts: at 256 experts / top-10 the selective copy moves 44% of the weight set at 34 tokens but 84% at ubatch 2048. Prefetch therefore forfeits a 56% saving on short prompts and only 16% on long ones, while the overlap gain stays roughly constant.

A MoE-batch gate recovers TTFT and keeps the prefill win

Dense graphs are never gated — there is no selective copy to lose there, so prefetch is unconditionally profitable.

sched->prefetch_active = sched->prefetch_weights;
if (sched->prefetch_weights && min_batch > 0) {
    bool    has_moe   = false;
    int64_t moe_batch = 0;
    for (int i = 0; i < graph->n_nodes; i++) {
        const struct ggml_tensor * node = graph->nodes[i];
        if (node->op == GGML_OP_MUL_MAT_ID) {
            has_moe = true;
            if (node->ne[2] > moe_batch) {
                moe_batch = node->ne[2];
            }
        }
    }
    if (has_moe && moe_batch < min_batch) {
        sched->prefetch_active = false;
    }
}

Then consult prefetch_active at the runtime sites instead of prefetch_weights. One detail worth keeping: leave the nodes_per_input allocation bound on prefetch_weights, so buffer sizing does not vary per graph.

Measured crossover is between ~250 and ~1040 MoE batch; 512 works as a default.

	-pw 0	-pw 1 ungated	-pw 1 gated
TTFT	1084 ms	1613 ms	1096 ms
prefill ~5.2K	856 t/s	1026–1030	1026–1030
decode	38.34 t/s	36.60	37.00

One caveat the gate does not fix

The VRAM cost. The scheduler reserves for the worst-case graph at n_ubatch, where prefetch is active regardless of the gate, so the double-buffer claims that memory for the life of the process. On this setup it permanently costs ~15% of the VRAM that would otherwise hold resident expert weights — which is why decode only partially recovers (37.00 vs 38.34) rather than returning to baseline. Anyone tight on VRAM pays that whether the gate is in or not.

That is a property of worst-case graph reservation rather than anything in this PR, but it is worth knowing: on a memory-constrained MoE setup the steady-state decode cost can outweigh the prefill gain, so --prefetch-weights is probably best left opt-in rather than defaulted on for MoE.

Happy to run other configurations if that would help. The gate is yours to take or leave — it is a small change to your design, not a competing approach.

[yalun753]

Good to see this direction getting formalized. I can add measured data on the MoE side, since that's where you (correctly, I think) expect prefetching to struggle.

We run a 157B MoE (DeepSeek-V4-Flash, 256 experts/layer) the same way — experts mmap'd in host RAM, computed on GPU. Key observation for the prefetch equation: expert selection barely overlaps across adjacent layers. A 30-turn session touched ~29 GB of distinct experts and cross-layer expert-set overlap (Jaccard) is low — so "guess the next layer's experts" has very little prediction headroom on large scattered-routing models.

We also tested prefetching at the OS level (posix_madvise(WILLNEED) after DONTNEED, +5 ms wait): cudaHostRegister still took ~3 ms. Prefetch did not hide the cold-expert fault.

Where this landed for us: instead of prefetching, we pin the mmap'd expert tensors (cudaHostRegister, whole-tensor merged registration) so per-token transfers go DMA directly, and keep hot experts in a GPU-side LRU cache. Decode on Qwen3.6-35B-A3B: 46.8 -> 50.2 t/s at 2.9 GB VRAM (4090). On the 157B model the cold-expert first touch (~2.1 ms) can't be hidden and routing is scattered, so ~4-5 t/s there is the physical decode ceiling.

One data point that supports your C_n + T_c' > T_{c+1} analysis: at batch size 1 there's nothing to overlap — a single expert is ~1 MB (Gen4 copy ~60 us vs MMVQ kernel <10 us), and a double-buffered H2D pipeline measured no gain (5.0 vs 5.5 t/s). The overlap headroom only exists at large ubatch (prefill), which matches your dense-model graphs.

Shipped open-source if useful as reference: https://github.com/yalun753/moe-l2

[crusaderky]

@am17an am17an#25 merges a recent master branch into this PR
Working on benchmarks of my own...