ggml: treat experts as cache residents during MoE offloading

[avifenesh]

Overview

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

• Patch commit: dc54faf7f
• 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 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 for the final patch were scanned for CUDA errors, OOMs, invalid arguments, aborts, and exceptions; none were found.

AI usage disclosure: YES - Codex assisted with local code review, benchmarking, cleanup, and preparing this fork-local staging PR. Before opening an upstream PR, I will manually review the patch and write/edit the final upstream PR text myself.

[gemini-code-assist]

Code Review

This pull request optimizes MoE weight offloading by tracking resident experts on the backend and only copying missing ones, reducing redundant data transfers. Changes include the introduction of the ggml_backend_sched_moe_loaded structure and updates to the scheduler's compute and lifecycle management logic. Feedback suggests explicitly invalidating the MoE resident state during a scheduler reset to maintain consistency when staging tensors are cleared.

I am having trouble creating individual review comments. Click here to see my feedback.

ggml/src/ggml-backend.cpp (1892-1898)

While the invalidation logic in ggml_backend_sched_compute_splits is robust enough to handle stale cache entries by checking src_data and dst_data pointers, it would be cleaner to explicitly invalidate the MoE resident state here when the scheduler is reset. Since hv_tensor_copies is cleared, the staging tensors are effectively deallocated, and the resident expert metadata should ideally reflect this state.

References

1. Defensive programming: ensure that invalid states (stale cache entries after a reset) are explicitly handled or cleared to maintain consistency.