cuda: MoE expert cache, adaptive VRAM caching of CPU-resident experts - #24524
cuda: MoE expert cache, adaptive VRAM caching of CPU-resident experts #24524leloch wants to merge 3 commits into
Conversation
…tion Adds a backend-agnostic function table (ggml_moe_cache, zero-initialized here, populated by a GPU backend at registry time) through which the CPU mul_mat_id kernel can offload cached expert rows: - ggml-cpu mul_mat_id: thread 0 plans hits/misses against the cache and dispatches all hit rows in one batched GPU launch while the remaining threads compute the miss rows; results are collected into dst before the node ends, so outputs are identical to the pure-CPU path. - ggml-cpu swiglu: skips dst rows the cache computed fused on the GPU (glu_hits mask), enabling a fused gate+up+GLU fast path. - sched: offers the GPU-side copy tensor of a CPU MUL_MAT_ID dst to the cache before the CPU split runs (redirect_offer) and lets the cache populate it directly, skipping the host round-trip copy (redirect_finalize); host weight-buffer teardown notifies the cache (invalidate) so async fills never read freed memory. With no GPU backend registered the table stays empty and every hook is a null-check no-op.
Caches hot CPU-resident MoE expert weights in spare VRAM and computes their mul_mat_id rows on the GPU during single-token decode: - per-(expert size, type) slot pools with exact source-tensor stride, LRU eviction with at-capacity admission throttling, async pinned-staging insert workers, idle-time prefetch backfill, hot-set persistence across runs - paired gate/up pools with a fused gate+up+SwiGLU batched matvec (engaged per layer only after the gate->up->GLU wiring is observed) - GPU-resident dst handoff for down projections via a pinned relay image and stream-ordered events (no P2P requirement) - decode-only fill, stable shape census before any allocation, role-group budgeting for mixed-quant models - baseline-sampled bail-out: the cache measures itself against the pure CPU path and trims its VRAM + disables itself if it is not winning; CUDA errors degrade to the CPU path, and the allocator can reclaim the cache (ggml_moe_cache_trim) on OOM before retrying - numerical self-test (GGML_CUDA_MOE_CACHE_SELFTEST=1) of the batched dispatch path, no model required Disabled with GGML_CUDA_MOE_CACHE=0; engagement requires experts >= 256 KiB of a supported quant type (see docs/moe-cache.md).
When the cache is available and the model spills heavily (model bytes >= 1.8x usable VRAM), the fit prefers keeping all experts on the CPU -- leaving maximum VRAM for the expert cache -- over statically placing a few layers. Measured: wins on heavily-spilling MoE models, keeps the stock placement near the fit boundary, and GGML_CUDA_MOE_CACHE=0 restores the previous behavior exactly. --moe-cache N: 0 disables the cache, N caps its per-device VRAM budget in MiB, absent = auto.
|
Hi @leloch, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
|
You can open a RFC in discussions page, as such is PR is just too large for any maintainer to review. Ask the AI to perhaps create it so that it clearly mentions the benefits vs the maintenance burden of such a change. |
|
Makes sense, thank you, opened the discussion here: #24528 |
…Erkenntnissen aktualisiert AtomicBot-ai#58 KV Cache Size Limiting: - Status ☐→⏭️ (verschoben) - PR ggml-org#18747 ist noch OPEN (nicht gemerged) - TurboQuant (bereits vorhanden) ist komplementär und höhere Priorität - PR liefert nur Infrastruktur, keine echten PagedAttention-Benefits AtomicBot-ai#61 Persistent VRAM Expert Cache: - PR ggml-org#23170 ist 'no-op when made correct' (RFC ggml-org#24528) - Eigentliche Lösung: PR ggml-org#24524 (closed, 2222 Zeilen, invertiertes execution model, MUL_MAT_ID auf CPU, GPU cached rows parallel) - 10% Experten → 80% cache hits, Top 30% → 95% hits - arXiv: ProMoE, MoE-Infinity, DuoServe-MoE, Caching Analysis - Empfehlung: PR ggml-org#24524 Design manuell portieren (1-2 Wochen)
…at_id / scheduler integration Portiert PR ggml-org#24524 (leloch) Commit 1/3: Backend-agnostische Function-Table (ggml_moe_cache, zero-initialized) als Brücke zwischen ggml-cpu und CUDA-Backend. - ggml-backend-moe-cache.h (neu, 95 Zeilen): struct ggml_moe_cache_api mit begin/plan/dispatch/collect/stats/redirect_offer/redirect_finalize/ glu_hits/invalidate/node_time. Zero-initialized global, populated by CUDA backend at registry time. - ggml-cpu.c (+71 Zeilen): Thread 0 in mul_mat_id partitioniert hits/misses gegen cache, dispatcht hit-rows als batched GPU launch, andere Threads berechnen miss-rows. Collect am Ende synchronisiert GPU results in dst. Sentinel-id Handling (i02 < 0 → zero dst row). - ops.cpp (+13 Zeilen): swiglu_f32 glu_hits mask — überspringt dst rows die cache fused auf GPU berechnet hat. - ggml-backend.cpp (+59 Zeilen): invalidate hook in buffer_free, redirect_offer vor CPU split, redirect_finalize an consumer copy site. Ohne CUDA-Backend: alle Hooks sind null-check no-ops (zero-initialized). Build grün auf Hydra (CUDA+CPU). Keine Verhaltensänderung ohne Registration. ROADMAP AtomicBot-ai#61 Status: ☐ → 🔄. Session-Plan hinzugefügt.
Portiert PR ggml-org#24524 (leloch) Commit 2/3: Vollständige CUDA-Implementierung des persistenten VRAM Expert Cache. - moe-cache.cu (neu, 1771 Zeilen): Per-(expert_size,type) slot pools mit LRU-Eviction, async pinned-staging insert workers, idle-time prefetch backfill, hot-set persistence across runs. Paired gate+up pools mit fused gate+up+SwiGLU batched matvec. GPU-resident dst handoff für down projections via pinned relay image + stream-ordered events. Decode-only fill, stable shape census, role-group budgeting. Baseline-sampled bail-out judge. CUDA errors degradieren zu CPU path. - moe-cache.cuh (neu, 20 Zeilen): Registration entry point. - mmvq.cu (+33 Zeilen): ggml_cuda_moe_cache_mmv — batched matvec über slot-pool experts mit device-side index array, optional gate fusion. - mmvq.cuh (+15 Zeilen): Declaration. - ggml-cuda.cu (+13 Zeilen): Include moe-cache.cuh, pool OOM trim hook (ggml_moe_cache_trim), registration in ggml_backend_cuda_reg. Aktiviert via GGML_CUDA_MOE_CACHE=1 env var. Build grün auf Hydra (CUDA+CPU, RTX 3070 Ampere). Ohne env var: Cache inaktiv, kein Verhaltensänderung.
…oval Portiert PR ggml-org#24524 Commit 3/3 + Fork-spezifisches Cleanup: - arg.cpp (+27 Zeilen): --moe-cache N CLI flag (0=off, N=VRAM budget MiB, absent=auto). Setzt GGML_CUDA_MOE_CACHE / GGML_CUDA_MOE_CACHE_BUDGET_MB env vars vor Backend-Registration. - fit.cpp (+109 Zeilen): MoE-cache-aware fit placement. Bei heavily-spilling MoE Modellen (model_bytes >= 1.8x usable VRAM) bevorzugt fit alle Experten auf der CPU zu lassen — der dynamische Cache nutzt die freie VRAM besser als statisches partial placement. Expert-size gate verhindert falsche Entscheidung bei kleinen Experten. - ggml-backend.cpp (-124 Zeilen): LFRU Expert Cache Code entfernt (expert_cache_enabled, expert_cache_hits/misses, expert_valid, tensor_copied, valid_bitset, GGML_EXPERT_CACHE env var). Dead code mit 0% hit-rate. thecodacus Prefetch (prefetch_experts) BLEIBT unangetastet. Build grün auf Hydra (CUDA+CPU). test-backend-ops läuft durch ohne Crash.
Overview
This is an interesting one: it adds a CUDA-side adaptive cache for CPU-resident MoE expert weights. With the core idea that some experts are 'hot' while others are not used much. Empirically: 10% of experts account for 80–81% of cache hits, and the top 30% account for 95–96%. Disclaimer: it is AI-generated/assisted. I have been trying to do this for months, with Opus 4.X, not much luck.
However, Fable 5 took all that research and experiment and was able to produce something that actually gives me decent gains on my hardware. I appreciate the policy around all of this and do not expect it to be merged, but perhaps a good 'source of inspiration' for future work.
When MoE model spills its experts to system RAM (
--cpu-moe,--n-cpu-moe, or the auto-fit), token generation is dominated by the CPU reading expert weights at RAM bandwidth. This cache keeps the hottest experts in otherwise-idle VRAM , and, unlike every prior attempt at this, the MUL_MAT_ID node stays on the CPU: inside the CPUmul_mat_idkernel, thread 0 dispatches one batched matvec over the cached (hit) rows on the GPU while the remaining CPU threads compute the miss rows concurrently, exactly as they would have anyway. Misses cost nothing extra, so the worst case degrades to the vanilla CPU path.Zero-config and on by default;
--moe-cache 0disables,--moe-cache Ncaps the per-device budget at N MiB. Vanilla path untouched when disabled; non-CUDA builds compile with the API table as a zero-initialized no-op.Measured on 4× RTX 3090 / EPYC 7R13 48c / 8-ch DDR4 256GB, llama-bench tg300, identical command lines:
-ngl 99 -ncmoe 99Quality: decode-path perplexity matches the pure-CPU path within a fraction of one standard error; prompt/batch processing is bit-untouched (PPL exactly equal). As with any
-nglchange, GPU rounding can flip a near-tie token under greedy decoding; nine architectures were validated this way.Safety rails: decode-only fill (prompt routing measurably thrashes the cache), stable shape census before any VRAM is spent, paired gate+up pools with fused gate+up+SwiGLU dispatch, GPU-resident handoff of down-projection results to GPU consumer splits, on-disk hot-set persistence across runs, baseline-sampled bail-out (sustained slower-than-CPU → trim and disable for the run), VRAM surrender on allocator OOM, multi-GPU sharding.
test-backend-opsMUL_MAT_ID 789/789 with cache on and off; built-in numerical self-test; CPU-only build clean; 3-commit bisectable series.Additional information
This problem has history: #20757 (feature request for exactly this), prior attempts #21609 #21614 #21620 #23170 (all closed), #17044, and the prefetch-based direction #21067. All prior designs move the MUL_MAT_ID node to the GPU and try to make the weight copies cheaper, which puts every cache miss on the critical path as a synchronous PCIe transfer , batot1 measured ~3× decode regression from that on #20757, and the Metal slot-pool branch in the same thread ran 2× slower even at 97–99% hit rate, purely from per-layer sync points. The hybrid hit/miss split used here avoids that structurally, which is why the results are uniformly non-negative.
Two specific connections: #23170's post-mortem concluded the fix for its no-op-when-correct problem is "a separate persistent expert-cache buffer owned by the MoE cache logic, with explicit expert_id -> slot_id bookkeeping" , that is this design. And #21067 is complementary, not competing: prefetch hides cold transfers at large ubatch (prompt), this cache eliminates hot re-transfers at ubatch 1 (decode); a build with both would cover both phases.
I am opening this as a draft and do not expect it to be merged, given the AI policy below , I know how the April expert-cache PRs ended, and I'm not trying to relitigate that. The results seemed too good not to share, and the design is in genuinely new territory versus the prior attempts. If the maintainers prefer this live as a comment on #20757 or a discussion , or nowhere , I'll close it immediately. I'm happy to run any benchmark, ablation, or A/B against #21067 on this hardware that would help evaluate the idea.
Requirements