feat: MoE expert cache — keep hot routed experts resident in VRAM - #5
Merged
Conversation
Phase 0 of the MoE expert-cache work: capture routed expert ids (ffn_moe_topk) per layer per decode step via the scheduler eval callback — backend-agnostic, so CPU-offloaded expert layers are traced too. Companion simulate.py replays traces against cache policies (static top-S, LRU, LFU-decay, periodic re-election) and prints hit-rate vs VRAM-budget curves plus routing-skew stats to decide whether expert caching is worth building at all.
…and index by nb Sizing the host buffer as n_used*n_tokens under-allocated: ffn_moe_topk is a view over the argsort output, so nb[1] spans n_expert ids per row and ggml_backend_tensor_get writes the whole strided range (heap corruption, malloc corrupted-top-size abort mid-prefill).
…ad cost per policy
Groundwork for the hot/cold expert-pack split: an id of -1 means the routed expert is not owned by this pack. The op contributes a zeroed output row for that slot so hot-pack and cold-pack results merge additively. CUDA fast paths (mmvq_id/mmq_id/mmf_id) still assert non-negative ids and are handled next.
mmid.cu gains a generic kernel that zeroes the dst rows of (token, slot) pairs whose expert id is -1; mmq and mmf launch it after the shared mm_ids_helper (which already drops unmatched ids from its compact mapping), and mmvq launches it before its kernels, which now early-return instead of using a negative id as a channel index (out of bounds read otherwise). test-backend-ops gains skip_ids variants of MUL_MAT_ID covering the mmvq/mmq/mmf/fallback dispatch corners.
Negative ids inflated nex_prev for every expert, opening a phantom region at the head of the compact arrays that nothing wrote; the mmq quantize preprocessing iterates the flat range and read garbage row indices (illegal memory access). Skipped ids now occupy no compact position, ids_src1 is sentinel-filled, and the quantize gather/scatter kernels skip negative entries.
Same treatment as mmvq: kernels early-return on a negative expert id instead of using it as a channel index, and the host launcher pre-zeroes the dst rows of skipped slots. Fixes wrong results (and a latent out-of-bounds channel read) for f16/f32 MUL_MAT_ID with -1 ids at small batch sizes.
… the compact map caused pool-layout-dependent illegal memory accesses when ids contain -1 (tile-padded reads consumed uninitialized row indices)
Ops whose ids may contain -1 announce it via op_params[0]; the CUDA dispatch then uses only the paths with verified skip support (mmvq for small batches, the general gather/scatter path otherwise). Large-batch mmq with sparse ids shows a pool-layout-dependent fault that resists isolation (all indexing audits pass; sanitizer and cuda-gdb unusable on the target driver); since the expert-pack design only decodes through packs, the safe route costs nothing in practice. mmq/mmf guards from earlier commits remain as hardening.
init_moe_expert_cache() runs post-load: reads a moe-trace profile (GGML_MOE_CACHE_PROFILE) and slot count (GGML_MOE_CACHE_SLOTS), and for each CPU-resident MoE layer builds GPU-resident hot packs of the top-S experts plus i32 remap tables (hot: global->slot or -1; cold: global id or -1). Graph wiring (build_moe_ffn dual path) is the next commit; these tensors are unused until then, so this is behavior-neutral.
…ld_moe_ffn When a layer carries hot packs, routed ids are remapped through moe_map_hot/cold (get_rows) and each expert projection runs as two skip-flagged mul_mat_id ops — hot pack on GPU, original tensors for cold — whose zero-padded outputs sum to the exact single-tensor result. Engaged for qwen35moe (separate gate/up/down, no expert scales); layers with merged gate_up, scales, or no packs use the unchanged path.
…y side Per-projection hot/cold interleave let the scheduler's split minimizer assign hot up/down matmuls to CPU, copying the GPU pack weights over PCIe every layer. Building one unbroken gate-up-swiglu-down chain per side keeps the hot pack on GPU; skipped (-1) rows stay zero through swiglu so a single add of the two chain outputs is exact.
Plumbed through llama_model_params; env vars remain as fallback so existing deployments keep working. Flags allow per-model settings in multi-model serving.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Caches the most-frequently-routed experts of CPU-offloaded MoE layers in VRAM. Decode runs hot experts on GPU and only the cold remainder on CPU, merged exactly. Opt-in, off by default, bit-exact output.
--moe-cache-profileis a routing profile CSV from the newllama-moe-tracetool (one-time capture per model:MOE_TRACE_OUT=trace.csv llama-moe-trace -m model.gguf -ncmoe 99 -p "<representative prompt>" -n 512). Merging workload traces iscat a.csv b.csv > merged.csv— a merged profile measures within 1% of per-workload specialists. Env fallbacks:GGML_MOE_CACHE_PROFILE/GGML_MOE_CACHE_SLOTS.Results (RTX 3060 12GB, ncmoe 99, fa 1)
Qwen3.6-35B-A3B Q4_K_M (21.1 GiB model, 256 experts/layer):
GLM-4.7-Flash Q4_K_M (17.1 GiB, deepseek2 arch, 64 experts/layer):
Stacks multiplicatively with
--spec-type draft-mtp(Qwen: 41.7 → 69.3 t/s generation, +66%; use 112 slots to leave room for the draft context).How
init_moe_expert_cache): ranks experts per layer by decode-routing frequency from the profile, packs the top-S (gate/up/down slabs) into a GPU buffer, and builds two id-remap tables per layer (global id -> pack slot/global id -> cold id, unused side = -1).build_moe_ffn): routed ids are remapped into hot and cold id sets; each side runs a complete gate→up→swiglu→down chain (ggml_mul_mat_idwith skip semantics), and one add merges them. Skipped (-1) rows produce zero rows andswiglu(0,0)=0, so the sum equals the single-tensor result exactly. Each chain is deliberately unbroken so the backend scheduler never splices CPU ops between GPU ops (interleaving lets the split minimizer migrate pack weights to CPU every layer).mul_mat_idlearns id=-1 "skip" semantics on CPU and CUDA (vec/general paths; routed away from mmq/mmf via an op flag). Covered by newtest-backend-opscases.Scope and fallbacks
Follow-ups