Conversation
…ok-ahead design
…kahead Wire the consumer half of docs/moe-lookahead-design.md: a new --moe-lookahead N gate (default 0, env LLAMA_ARG_MOE_LOOKAHEAD) that asks the CUDA MoE cache to prefetch the next layer's expert weights during the current layer's compute. The prefetch path mirrors acquire with is_prefetch=true, wait_for_compute=false and pin=false: no stream synchronize, no compute-stream wait, pinned slots are skipped, expert ids outside [0, n_experts) are ignored, and prefetch telemetry is accounted separately from demand. Host-to-device copies are batched through cudaMemcpyBatchAsync when CUDART >= 12.80, with a per-copy fallback otherwise. Targets whose buffer is not MoE-cached are skipped with a one-shot warning. --moe-lookahead 0 leaves the decode step unchanged. A non-zero width without --moe-expert-cache-size > 0 fails loudly at model load. Adds test-moe-cache --lookahead-prefetch-only and the lookahead prefetch legacy layer case. Assisted-by: Oh My Pi (deepseek-v4.1-flash)
The speculative prefetch path picked its victim with plain LRU, so a prediction could displace an expert that demand had already proven hot. The eviction guard the design requires (docs/moe-lookahead-design.md, "Invariants to preserve" #2 and the colibri safety invariant at c:1462-1478) was missing: only the pin guard - never evict a slot a running GEMM is reading - was implemented, on both the demand and speculative paths. Port colibri's PILOT_EVICT_GUARD. A speculative fill may displace a resident only when that resident is not genuinely warm; a resident is protected when it has at least 2 demand accesses AND is clearly hotter than the prediction, by the 25% + 4-frequency hysteresis in LFRU score units (score ported from colibri c/tier.h:40-43: frequency in the high bits, recency saturating in the low byte). A blocked prediction is dropped rather than forced in, so it can never thrash a demand-loaded expert. Heat is recorded on demand accesses only. Counting predictions would let a speculation inflate the very score that decides whether it may displace a resident, which is the failure colibri hit in ggml-org#490. Refactor the eviction core so the demand and speculative paths cannot drift: ggml_cuda_moe_cache_select_victim_locked picks the LRU unpinned slot and applies the guard for speculative fills, and ggml_cuda_moe_cache_install_fill_locked records the eviction telemetry and publishes the new entry. The copy-batching split in the prefetch path is unchanged, and the free-slot, pin and unknown-eid behaviour is identical. Add phase_prefetch_dropped so a guard that drops everything (colibri ggml-org#490) is visible rather than silent, and extend the test prefetch accessor with it. Add the lookahead prefetch eviction guard case: a 3-slot pool, a warm resident survives a prediction for a never-seen expert (dropped, no copy, no eviction, still a demand hit), while a once-demanded resident is evicted for it. Assisted-by: Oh My Pi (deepseek-v4.1-flash)
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
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
PR-A of the MoE look-ahead expert prefetch: the consumer half plus the
--moe-lookaheadgate, on top of the design doc. The producer (
GGML_OP_MOE_PREFETCH/llama-graph.cpp,"PR-P1" in the doc) is out of scope here.
Three commits over
feat/763-reconcile-qwen4exp-mtp:137fa17bbdocs: look-ahead design doc (folds measured evidence from upstream RFC RFC: MoE expert cache, VRAM caching of hot CPU-resident experts with hybrid hit/miss execution ggml-org/llama.cpp#24528)cbb19d802feat(moe-cache): the consumer behind--moe-lookaheadec23a599ffix(moe-cache): the LFRU eviction guard (see below)Consumer
ggml_backend_cuda_moe_prefetch_experts()(plus a tensor variant) asks the CUDA MoE cache toprefetch the next layer's expert rows during the current layer's compute. Gate:
--moe-lookahead N(env
LLAMA_ARG_MOE_LOOKAHEAD), default 0.Invariants kept:
is_prefetch=true,wait_for_compute=false,pin=false; no stream synchronizeand no compute-stream wait on the prefetch path; pinned slots are never chosen; expert ids outside
[0, n_experts)are ignored; H2D copies are batched throughcudaMemcpyBatchAsync(CUDART >=12.80) instead of one copy per expert id; a target whose buffer is not MoE-cached is skipped with a
one-shot warning; the exported
ggml_backend_cuda_moe_prefetch_expertssignature is unchanged;width 0 leaves the decode step untouched.
Eviction guard (the review finding)
The first consumer revision picked its victim by plain LRU on the prefetch path, so a prediction
could displace an expert that demand had already proven hot. Only the pin guard (never evict a slot
a running GEMM is reading) existed, on both the demand and speculative paths.
ec23a599fports colibri'sPILOT_EVICT_GUARD: a speculative fill may displace a resident onlywhen that resident is not genuinely warm - protected when it has at least 2 demand accesses AND
is clearly hotter than the prediction, by the 25% + 4-frequency hysteresis in LFRU score units
(score ported from colibri
c/tier.h:40-43: frequency in the high bits, recency saturating in thelow byte). A blocked prediction is dropped, never forced in. Heat counts demand accesses only;
counting predictions would let a speculation inflate the score that decides whether it may displace
a resident (the failure colibri hit in ggml-org#490).
The eviction core is now shared (
select_victim_locked/install_fill_locked) so the demand andspeculative paths cannot drift; the prefetch path's copy batching is unchanged.
Verification (RTX 5060 Ti CUDA0 + RTX 3060 CUDA1, CUDA 13.2.1)
tests/test-moe-cache: 19/19 cases pass,test-moe-cache: OK, exit 0 (was 18 before the new case).lookahead prefetch eviction guard: 3-slot pool. A warm resident (3 demand accesses)survives a prediction for a never-seen expert - dropped, no copy, no miss, no eviction, still a
demand hit. A once-demanded resident is evicted for the prediction.
--lookahead-prefetch-only:lookahead prefetch legacy layer OK+ the guard case, exit 0.--helplists--moe-lookahead N; default and--moe-lookahead 0load stock;--moe-lookahead 8 --moe-expert-cache-size 84loads;--moe-lookahead 8with cache-size 0 failsloud at load (
--moe-lookahead requires --moe-expert-cache-size > 0)./opt/software/cuda/13.2does not exist on this box (only 12.9, 13.2.1, 13.2.2,13.3.1), so the build used
-DCUDAToolkit_ROOT=/opt/software/cuda/13.2.1.Base
This branch sits on the
gsqwen4exp-mtplineage, so it targetsfeat/763-reconcile-qwen4exp-mtp(PR #120) - the same stacking PR #123 uses.
baseline-flash-nextis currentlycdd11021b, a parentof that lineage, so targeting it directly would re-show PR #120's whole diff. GitHub retargets this
to
baseline-flash-nextautomatically once #120 merges.Open for PR-P1 (producer), not addressed here
prefetch_legacy_layercallsacquire_legacy_cache(experts)withcompute_stream=nullptr; iflayer L+1's pool is not installed yet the lease is empty and the prefetch silently no-ops. The
producer must guarantee the target pool is installed first (spec Q1). Preferred fix: install all
MoE-cached pools at model load when
--moe-lookahead > 0and fail at load if the budget does notfit, so the no-op is unreachable rather than silent.
cudaMemcpyBatchAsyncclears the new booking but does not restore the resident itdisplaced; the miss falls to the demand path. Safe now that the guard keeps the victim cold.
phase_prefetch_droppedis exposed only through the test accessor; the runtime telemetry linedoes not report it yet.