Skip to content

cuda : MoE look-ahead expert prefetch (PR-A consumer) + LFRU eviction guard - #126

Open
ddvnguyen wants to merge 3 commits into
feat/763-reconcile-qwen4exp-mtpfrom
feat/moe-lookahead-pra
Open

ddvnguyen wants to merge 3 commits into
feat/763-reconcile-qwen4exp-mtpfrom
feat/moe-lookahead-pra

Conversation

@ddvnguyen

Copy link
Copy Markdown
Owner

What

PR-A of the MoE look-ahead expert prefetch: the consumer half plus the --moe-lookahead
gate, 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:

Consumer

ggml_backend_cuda_moe_prefetch_experts() (plus a tensor variant) asks the CUDA MoE cache to
prefetch 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 synchronize
and 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 through cudaMemcpyBatchAsync (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_experts signature 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.

ec23a599f ports colibri's PILOT_EVICT_GUARD: a speculative fill may displace a resident only
when 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 the
low 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 and
speculative 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).
  • 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.
  • CLI: --help lists --moe-lookahead N; default and --moe-lookahead 0 load stock;
    --moe-lookahead 8 --moe-expert-cache-size 84 loads; --moe-lookahead 8 with cache-size 0 fails
    loud at load (--moe-lookahead requires --moe-expert-cache-size > 0).
  • Build note: /opt/software/cuda/13.2 does 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 gs qwen4exp-mtp lineage, so it targets feat/763-reconcile-qwen4exp-mtp
(PR #120) - the same stacking PR #123 uses. baseline-flash-next is currently cdd11021b, a parent
of that lineage, so targeting it directly would re-show PR #120's whole diff. GitHub retargets this
to baseline-flash-next automatically once #120 merges.

Open for PR-P1 (producer), not addressed here

  1. prefetch_legacy_layer calls acquire_legacy_cache(experts) with compute_stream=nullptr; if
    layer 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 > 0 and fail at load if the budget does not
    fit, so the no-op is unreachable rather than silent.
  2. A failed cudaMemcpyBatchAsync clears the new booking but does not restore the resident it
    displaced; the miss falls to the demand path. Safe now that the guard keeps the victim cold.
  3. phase_prefetch_dropped is exposed only through the test accessor; the runtime telemetry line
    does not report it yet.

ddvnguyen and others added 3 commits September 13, 2026 23:03
…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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant