ggml: treat experts as cache residents during MoE offloading

[avifenesh]

Overview

In short, this PR improves TPS and latency while using MoE models by making experts cache residents and avoiding copy and overwrite when an expert is already cached.

When MoE weights are offloaded from host memory, the scheduler copies the experts used by the current GGML_OP_MUL_MAT_ID split into a backend-side staging tensor. In long prompt-cache workloads, the same staging tensor can be reused across turns, so experts copied for earlier tokens may still be resident.

This change tracks resident MoE experts per scheduler tensor copy and treats the staging tensor as an expert cache:

• Tracks resident expert IDs with a ggml_bitset_t
• computes missing_ids = used_ids & ~loaded_ids
• copies only missing contiguous expert ranges
• Skips the split-input overwrite wait when no copy is needed
• Invalidates resident state when the expert count, expert size, source data pointer, or destination staging pointer changes

The change is internal to ggml/src/ggml-backend.cpp. It does not add a user-facing flag, public API, or backend-specific implementation.

Benchmark Setup

• Base: upstream/master at b64739ea3
• Patch commit: 0fcbeeaa6
• Hardware: NVIDIA GeForce RTX 5090 Laptop GPU, 23981 MiB VRAM
• Driver: 595.58.03
• Build: CUDA 13.1, build-cuda13-clean/bin/llama-server, CUDA graphs off
• Server mode: --cpu-moe, -c 65536, -ngl auto, -np 1
• Expert staging path: default op-offload enabled
• Prompt cache: disabled with --no-cache-prompt except for the medium-turns context-fill suite
• Requests: sequential

Only CPU-MoE op-offload measurements are included below because that is the path changed by this PR.

Results

64k Multi-Prompt Op-Offload

This is the short steady-state control: 60s measured duration after ~10s warmup, comparing the resident-expert cache with a local cache-disabled control build on the same op-offload path.

Model	Cache	Duration	Warmup	Requests	Prompt tok/s	Completion tok/s	Total tok/s	Avg latency
Qwen3.6-35B-A3B UD-Q4_K_M	disabled control	63.97s	11.31s	10	13.46	28.52	41.98	6.397s
Qwen3.6-35B-A3B UD-Q4_K_M	enabled	60.71s	11.19s	10	14.18	30.04	44.23	6.071s
Gemma 4 26B A4B Q4_K_M	disabled control	63.06s	10.44s	8	10.78	23.76	34.54	7.882s
Gemma 4 26B A4B Q4_K_M	enabled	60.85s	10.50s	8	11.17	25.24	36.41	7.607s

Deltas:

• Qwen3.6-35B-A3B: total tok/s +5.36%, avg latency -5.09%
• Gemma 4 26B A4B: total tok/s +5.43%, avg latency -3.49%

Long-Context Prefill Pressure

Single cold request, no warmup, prompt-cache disabled, generated prompt calibrated through /tokenize, max_tokens=128.

Model	Cache	Prompt tokens	Latency	Prompt tok/s	Total tok/s
Qwen3.6-35B-A3B UD-Q4_K_M	unpatched baseline	60028	176.31s	340.47	341.20
Qwen3.6-35B-A3B UD-Q4_K_M	enabled	60028	25.08s	2393.31	2398.42
Gemma 4 26B A4B Q4_K_M	unpatched baseline	60044	120.57s	498.01	499.07
Gemma 4 26B A4B Q4_K_M	enabled	60044	30.81s	1948.92	1953.08

Deltas:

• Qwen3.6-35B-A3B: 7.03x faster, latency -85.77%
• Gemma 4 26B A4B: 3.91x faster, latency -74.45%

Medium-Turn Context Fill

Sequential ~5k-token turns with prompt-cache enabled, stopping around 60k prompt tokens. This simulates a long chat where the earlier prompt state is reused and reports processed prompt tokens separately from cached tokens.

Model	Cache	Turns	Max prompt tokens	Duration	Avg latency	Processed prompt tok/s
Qwen3.6-35B-A3B UD-Q4_K_M	unpatched baseline	12	61141	198.76s	16.491s	307.83
Qwen3.6-35B-A3B UD-Q4_K_M	enabled	12	60525	45.84s	3.745s	1321.33
Gemma 4 26B A4B Q4_K_M	unpatched baseline	12	61001	152.69s	12.662s	399.80
Gemma 4 26B A4B Q4_K_M	enabled	12	60300	57.55s	4.734s	1048.54

Deltas:

• Qwen3.6-35B-A3B: 4.34x faster, processed prompt tok/s 4.29x, latency -77.29%
• Gemma 4 26B A4B: 2.65x faster, processed prompt tok/s 2.62x, latency -62.61%

Validation

• git diff --check
• cmake --build build-cuda13-clean --target llama-server test-backend-ops -j 10
• build-cuda13-clean/bin/test-backend-ops test -o MUL_MAT_ID: 764/764 tests passed
• Relevant rerun logs were scanned for CUDA errors, OOMs, invalid arguments, aborts, and exceptions; none were found.

---

AI usage disclosure: YES - Codex gpt 5.5 xhigh was used to write this code. It assisted with code writing, local code review, benchmarking, cleanup, and preparing the MD-styled parts of this PR, like the tables.

While an LLM wrote the code, I did the design, the research, and the codebase reading, I directed, reviewed and steered during development, designed the benchmarks, reviewed the code more than once and instructed changes, and I'm doing the final signing on this code and PR.
I own this code, and I'm responsible for the output of the tools I use.

[ssakar]

How do the perplexity numbers look like?

[avifenesh]

@ssakar
Thanks for asking, I didn't check it before and missed an issue. With --cpu-moe --op-offload, PPL went to nan after the first chunk.

The problem was that the resident-expert state could outlive the scheduler's temporary buffer contents. I pushed a fix that scopes the resident bitset to the current scheduler compute epoch.

After the fix, PPL matches the parent commit exactly on my run:

• model: unsloth/Qwen3.6-35B-A3B-GGUF, Qwen3.6-35B-A3B-UD-Q4_K_M.gguf
• command: llama-perplexity -ngl 999 --cpu-moe --op-offload -fa on -c 512 -b 512 -ub 512 --chunks 16 --no-warmup
• parent: PPL = 7.1051 +/- 0.27827
• patched: PPL = 7.1051 +/- 0.27827

The chunk values also match exactly:
[1]4.0972,[2]5.7571,[3]5.1944,[4]5.1812,[5]5.3006,[6]5.4650,[7]5.7343,[8]6.2433,[9]6.7152,[10]7.1059,[11]7.0291,[12]7.1499,[13]7.5467,[14]7.1223,[15]7.0790,[16]7.1051

[avifenesh]

@ssakar I checked the docs and didn't see other checks that I missed. Anything else that you think I should run to validate no other blasts?

[ernestuz-b]

@avifenesh, Hi,

I was thinking on something similar to this, and then found this PR, as it looks like a very good starting point.

I had a look at your patch. I rebased it locally onto the current upstream master and added some simple telemetry around the MoE resident-copy path.

tl;dr: with the NaN fix in place, I am not seeing the copy-skipping gains in my tests. The resident bitsets are populated, and many later experts would have hit, indeed, but the epoch invalidation clears the state before it can be used, so you lose any speed gain.

Long Explanation

I do not have the branch online at the moment, but I can share the local telemetry diff or logs if useful.

The telemetry is inserted in ggml/src/ggml-backend.cpp, inside ggml_backend_sched_compute_splits(), around the existing GGML_OP_MUL_MAT_ID CPU-MoE/offload copy path.

The key location is just before/around the existing computation of:

missing_ids = used_ids & ~loaded_ids;

I added counters roughly like this:

// after used_ids has been built, and before copy_experts() mutates loaded_ids

bool has_missing_ids = false;
size_t resident_hits = 0;
size_t resident_misses = 0;

missing_ids.resize(loaded_id_size);
ggml_bitset_t * missing_ids_data = missing_ids.data();
ggml_bitset_t * used_ids_data = used_ids.data();

for (size_t i = 0; i < loaded_id_size; ++i) {
    missing_ids_data[i] = used_ids_data[i] & ~loaded_ids[i];

    if (sched->moe_cache_trace_enabled) {
        const ggml_bitset_t hit_ids = used_ids_data[i] & loaded_ids[i];
        resident_hits   += ggml_backend_sched_bitset_popcount(&hit_ids, 1);
        resident_misses += ggml_backend_sched_bitset_popcount(&missing_ids_data[i], 1);
    }

    has_missing_ids = has_missing_ids || missing_ids_data[i] != 0;
}

if (sched->moe_cache_trace_enabled) {
    sched->moe_cache_trace.ops_seen++;
    sched->moe_cache_trace.used_experts_unique +=
        ggml_backend_sched_bitset_popcount(used_ids_data, loaded_id_size);
    sched->moe_cache_trace.loaded_experts_before +=
        ggml_backend_sched_bitset_popcount(loaded_ids, loaded_id_size);
    sched->moe_cache_trace.resident_hits += resident_hits;
    sched->moe_cache_trace.resident_misses += resident_misses;
    sched->moe_cache_trace.skipped_experts += resident_hits;
    sched->moe_cache_trace.skipped_bytes_estimate += resident_hits * expert_size;

    if (!has_missing_ids) {
        sched->moe_cache_trace.empty_copy_ops++;
    }
}

I added diagnostic counters around the invalidation path as well, to check if resident state existed just before being cleared:

const bool invalid_epoch =
    loaded->epoch != moe_loaded_epoch;
const bool invalid_n =
    loaded->n_expert != n_expert;
const bool invalid_size =
    loaded->expert_size != expert_size;
const bool invalid_src =
    loaded->src_data != input_data;
const bool invalid_dst =
    loaded->dst_data != input_cpy_data;

if (sched->moe_cache_trace_enabled &&
        (invalid_epoch || invalid_n || invalid_size || invalid_src || invalid_dst)) {
    const size_t pre_invalid_loaded =
        ggml_backend_sched_bitset_popcount(loaded_ids, loaded_id_size);
    const size_t pre_invalid_hit =
        ggml_backend_sched_bitset_intersection_popcount(used_ids_data, loaded_ids, loaded_id_size);

    sched->moe_cache_trace.pre_invalid_loaded += pre_invalid_loaded;
    sched->moe_cache_trace.pre_invalid_hit += pre_invalid_hit;

    if (pre_invalid_loaded > 0) {
        sched->moe_cache_trace.invalid_loaded++;
    }

    if (invalid_epoch) sched->moe_cache_trace.invalid_epoch++;
    if (invalid_n)     sched->moe_cache_trace.invalid_n++;
    if (invalid_size)  sched->moe_cache_trace.invalid_size++;
    if (invalid_src)   sched->moe_cache_trace.invalid_src++;
    if (invalid_dst)   sched->moe_cache_trace.invalid_dst++;
}

and inside copy_experts(first_expert, last_expert):

if (sched->moe_cache_trace_enabled) {
    sched->moe_cache_trace.copy_ranges++;
    sched->moe_cache_trace.copied_experts += last_expert - first_expert + 1;
    sched->moe_cache_trace.copied_bytes += expert_size_copy + padding_end;
}

The telemetry is enabled with:

LLAMA_MOE_CACHE_TRACE=1
LLAMA_MOE_CACHE_TRACE_EVERY=100

In my tests, with the NaN fix, the resident bitset is invalidated on every MoE copy opportunity. So the safe version appears to be correct, but unfortunately, I am not seeing any actual resident hits.

Example result:

ops=1440
invalid=1440
hit=0
miss=263871
loaded_before=0
copied=156.1GB
skipped_est=0.0GB
empty=0

After I added the invalidation diagnostics above, they show that the cache state does exist before invalidation:

pre_invalid_loaded=241788
pre_invalid_hit=219222
invalid_loaded=1320
invalid_epoch=1440

So the patch populates the resident bitsets as expected, and many later selected experts would have hit, but the epoch invalidation clears the state before missing_ids = used_ids & ~loaded_ids can benefit from it.

I tried ignoring epoch-only invalidation while still invalidating on structural changes (n_expert, expert_size, src_data, dst_data). That produced real hits and a large apparent prefill speedup:

hit=17640
miss=22482
copied=13.3GB
skipped_est=10.7GB
empty=1212

But, again, as expected, it is not correct: llama-perplexity produced NaNs. So the epoch invalidation is protecting against real stale temporary-buffer contents.

My current understanding is:

• the scheduler split-input copy is allocated from graph-temporary memory;
• the same dst_data address can be reused across scheduler computes;
• pointer stability is therefore not sufficient to prove the expert bytes are still valid;
• the epoch check is necessary for correctness;
• but with that check, the cache lifetime is too short to produce reuse in this benchmark.

So I am not yet seeing evidence that the safe version of the patch still supports the performance claims in this configuration. There may be another workload or configuration where reuse happens within the same scheduler compute epoch, but in this CPU-MoE/op-offload run I get zero actual skipped copies.

It may be useful to add resident-hit / skipped-copy telemetry to the PR itself, or rerun the benchmarks against the commit that includes the NaN fix.

My test command was roughly:

LLAMA_MOE_CACHE_TRACE=1 \
LLAMA_MOE_CACHE_TRACE_EVERY=100 \
../llama-build/bin/llama-bench \
  -v \
  -m "$MODEL" \
  -ngl 999 \
  -ncmoe 999 \
  -nopo 0 \
  -fa 1 \
  -ctk q8_0 \
  -ctv q8_0 \
  -sm layer \
  -p 2048 \
  -n 128 \
  -b 1024 \
  -ub 1024

Model:

Qwen_Qwen3.6-35B-A3B-Q4_K_M.gguf (by Bartowski)

My guess is that the NaN problem comes from loaded_ids persisting longer than the allocation lifetime of the input_cpy contents.

The resident bitset can survive and still say “expert X is loaded”, and the dst_data pointer can even be the same, but input_cpy is still a scheduler/graph temporary. If ggml_gallocr has reused that memory between scheduler computes, then that the pointers are equal does not imply that the old expert bytes are still there. With the epoch invalidation disabled, the code can skip copying an expert range and then MUL_MAT_ID reads stale temporary-buffer contents as if they were quantized expert weights, which explains the immediate NaNs.

So my guess is that cross-epoch reuse needs one of two designs (please, correct me if I am wrong):

• Make the staged expert tensor allocation genuinely persistent for the intended cache lifetime, so gallocr cannot reuse/clobber it.

OR

• Use a separate persistent expert-cache buffer owned by the MoE cache logic, with explicit expert_id -> slot_id bookkeeping.

I was tempted to implement option 1, because it might be closer to your current patch, but it has memory ownership / allocator-lifetime implications that I do not understand well enough. Option 2 looks cleaner architecturally, though it is a bigger change because compact slots require ID remapping or an op-level/cache-aware path.

In any case, I think the immediate useful next step for this PR would be adding resident-hit / skipped-copy telemetry after the NaN fix, because correctness alone does not show whether the optimization is still active.

Happy to share the local telemetry diff if it helps.

[avifenesh]

@ernestuz-b Hi, yes, I checked the numbers, and your observation is correct.
I assessed the options, and there is some direction to get gains from tracking residency, but it comes from a fairly different angle and not by reshaping this PR.
I'm closing this PR, and if my number will be proved as correct i will open a separate PR from the correct angle.
Thanks for pushing, I'm actually coming with strong caching-oriented background, thats litteraly my job, and I still missed the invalidation cost.