Skip to content

Feature Request: Two-tier GPU+RAM expert cache for MoE offload (pluggable eviction policy) #20757

Description

@e1n00r

Prerequisites

  • I am running the latest code.
  • I carefully followed the README.md.
  • I searched for existing issues and found none covering a persistent VRAM expert slot cache.
  • I reviewed the Discussions, and have a new and useful enhancement to share.

Feature Description

Problem

Running GPT-OSS-120B (128 experts/layer, 36 layers, ~57 GB expert weights) on an 8 GB GPU works today via --cpu-moe with MoE tensors kept in CPU RAM. The dense/attention layers run on GPU; expert weights are copied CPU→GPU on each forward pass by ggml_backend_sched_compute_splits().

The bottleneck: every decode step copies the same ~4 hot experts per layer from RAM to GPU, uses them, and discards the GPU copy. For a 36-layer model with top-4 routing that's 144 RAM→GPU copies per token — most of them identical to the previous token. The GPU-side expert data is thrown away between passes even though the same experts will be needed again immediately.

MoE routing is skewed. In practice ~15–20% of experts handle ~80% of tokens. Keeping those hot experts resident in VRAM between passes — in a fixed-address slot buffer — reduces per-token PCIe traffic to near zero after a short warm-up period.

Proof of concept

I built a Python PoC (tinyserve) on HuggingFace transformers for GPT-OSS-120B implementing a two-tier expert cache:

Tier 1: GPU VRAM  — persistent slot buffer, N fixed-address expert slots, SLRU eviction  [proposed]
Tier 2: CPU RAM   — pinned memory, backing store for all local experts                    [proposed]
Tier 3: SSD/mmap  — full weight tensor, demand-paged by OS                               (existing)

On a cache hit (Tier 1), the expert is already on GPU — zero copies. On a Tier 1 miss, the expert is copied from pinned RAM (Tier 2) into the freed slot. On a Tier 2 miss (cold start), the expert is demand-paged from the mmap'd file.

Results on RTX PRO 2000 (8 GB VRAM), GPT-OSS-120B:

Phase tok/s Cache hit rate
Cold start (0–80 t) 1.9–2.5 48–56%
Warming (80–160 t) 2.6–5.1 53–78%
Steady state (160+) 12–14 ~98–100%

Pure CPU offload with no cache on the same hardware: ~0.5–1 tok/s. At steady state the model is GPU-compute bound — the PCIe pipe is mostly idle.

What llama.cpp already has

  1. --cpu-moe and --n-cpu-moe in all tools (common/arg.cpp:2284,2291) — already route MoE expert tensors to CPU RAM for all tools including llama-cli, llama-server, and llama-bench. Also available as env vars LLAMA_ARG_CPU_MOE / LLAMA_ARG_N_CPU_MOE. This part is done; the gap is what happens after the weights are on CPU.

  2. Selective expert copy in ggml_backend_sched_compute_splits() (ggml/src/ggml-backend.cpp:1445–1564) — already reads the ids tensor at runtime and copies only the used expert sub-rows CPU→GPU via ggml_backend_tensor_set_async(). Crucially, it copies by byte offset into a GPU tensor that mirrors the full CPU tensor layout (expert_offset = first_id * expert_size, line 1529) — no slot remapping, no persistence across passes. This is the hook point for the cache.

  3. POSIX_MADV_WILLNEED in llama_mmap (src/llama-mmap.cpp:436) — used at model load time for prefetch. POSIX_MADV_DONTNEED does not currently exist in the codebase and would need to be added for page release.

What is new

Tier 1: persistent GPU slot buffer

Today the GPU-side expert allocation (input_cpy in ggml_backend_sched_compute_splits()) mirrors the full CPU tensor shape — it has n_expert slots even though only top_k are used. Each forward pass, the selected experts are written into their original offset positions, then discarded at pass end. Expert IDs sent to the kernel are global IDs that index this full-size tensor.

The proposed change: allocate a GPU buffer of N slots at model load (N << n_expert, controlled by --moe-expert-cache-size N). A persistent mapping (expert_id → slot_idx) carries over between passes. On a hit, the slot is already populated — zero copy. On a miss, a slot is evicted per policy, the new expert is copied from Tier 2 into that slot, and the mapping is updated. The kernel receives remapped slot indices rather than global expert IDs, requiring a small change to the GGML_OP_MUL_MAT_ID path in ggml_backend_sched_compute_splits().

This is the primary performance lever. The 14 tok/s in the PoC comes from Tier 1 hits.

Tier 2: pinned RAM backing store

Today, CPU-side expert data lives in the mmap'd weight file — demand-paged from SSD on first access, held in OS page cache thereafter. No pinned allocation is made.

The proposed change: at model load, allocate a CPU pinned memory region for all local experts. Hot experts are available for fast non-blocking H2D transfer on Tier 1 misses. Without pinned memory, every Tier 1 miss copies from pageable memory, which is slower and cannot overlap with GPU compute.

Proposed design

1. Pluggable eviction policy interface

To maximize flexibility and allow further advancements without touching cache logic, the policy is plug-and-play. Swapping policies requires no cache code changes:

// Sketch — naming should follow ggml conventions
struct ggml_moe_cache_policy {
    const char * name;
    void    (* on_hit)      (void * ctx, int32_t expert_id, int32_t slot);
    int32_t (* select_evict)(void * ctx);
    void    (* on_load)     (void * ctx, int32_t expert_id, int32_t slot);
    void *  ctx;
};

// Selected via --moe-expert-cache-policy <name>
void ggml_moe_cache_register_policy(struct ggml_moe_cache_policy policy);

// Built-in policies
struct ggml_moe_cache_policy ggml_moe_cache_policy_lru (int capacity);
struct ggml_moe_cache_policy ggml_moe_cache_policy_slru(int capacity);
struct ggml_moe_cache_policy ggml_moe_cache_policy_lfu (int capacity);
struct ggml_moe_cache_policy ggml_moe_cache_policy_fifo(int capacity);

2. SLRU as the recommended default

SLRU is the recommended default for the Tier 1 slot buffer. The reason is specific to how MoE routing interacts with the two-tier design.

Plain LRU has a failure mode here: prefill sequences activate all experts in a roughly flat distribution, wiping the GPU slot buffer. When decode resumes, Tier 1 is cold — every token falls through to a Tier 2 copy (CPU→GPU) until the cache re-warms. In a two-tier system this is particularly costly because even Tier 2 hits involve a PCIe copy, eliminating the primary performance advantage.

SLRU splits the slot buffer into a probationary segment (~20% of slots) for new admissions and a protected segment (~80%) for experts accessed more than once. Eviction always targets probationary first. A prefill burst fills and cycles through the probationary segment but cannot displace protected entries. Hot decode experts, having been seen multiple times, occupy protected slots and survive the burst intact.

The PoC also uses a frequency-gated admission filter: an expert is only admitted to the Tier 1 slot buffer on its second miss. Cold experts seen once during a prefill never enter either segment, protecting both from pollution at the point of admission rather than only at eviction. In the PoC, SLRU + admission filter improved steady-state hit rate by 8–15 percentage points over plain LRU on mixed prefill+decode traffic.

3. SSD-tier management — PoC tricks (~60 LOC)

Two optimisations from the PoC transfer directly:

  • Readahead (POSIX_MADV_WILLNEED): after each decode step, fire a prefetch hint for the experts just used — they're likely needed next token. The pattern exists at src/llama-mmap.cpp:436 for model load; the same call applied per-expert after each pass costs nothing and gives the OS time to warm the pages.
  • Page release (POSIX_MADV_DONTNEED): when an expert is evicted from Tier 2, release its OS page cache entry. This call does not currently exist in the codebase and would need to be added to src/llama-mmap.cpp.

A further optimisation from the PoC — double-buffering the PCIe transfer (two alternating staging buffers that overlap Tier 2→Tier 1 copies with GPU compute) — is out of scope for a first implementation but is a natural follow-on once the slot buffer is in place.

Help wanted

I'm primarily a Python developer. The PoC demonstrates the approach works and the numbers are real, but I'm not equipped to write production-quality C++ for the ggml backend scheduler internals.

Looking for:

  • A C++ contributor familiar with ggml_backend_sched_compute_splits() to implement the slot buffer and two-tier cache
  • Feedback on whether the policy interface fits ggml conventions — specifically whether the slot-index remapping belongs inside ggml_backend_sched_compute_splits() or as a new op/wrapper
  • Any existing work in this direction I might have missed

Happy to provide benchmark data, answer design questions, and test any implementation.

Files relevant to a C++ implementer

File Relevance
ggml/src/ggml-backend.cpp:1445–1564 Expert copy loop — both tiers insert here; slot remap needed at line 1529
common/arg.cpp:2284,2291 Existing --cpu-moe / --n-cpu-moe parameter wiring (already in all tools)
src/llama.cpp:481–509 MoE tensor CPU placement logic for auto-fit
src/llama-mmap.cpp:436 POSIX_MADV_WILLNEED pattern; DONTNEED needs adding

AI disclosure: I used Claude Code to help analyse the llama.cpp codebase and structure this RFC. The Python PoC, benchmark numbers, and design rationale are my own. I reviewed this text myself and can explain and defend every technical claim here.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions