--prefetch-experts-slots: lookahead H2D prefetch of host-resident MoE expert weights

[leshchukandrej]

PR body — --prefetch-experts-slots: lookahead H2D prefetch of host-resident MoE expert weights

Branch: feat/prefetch-experts-slots — commit e502dd002 on top of v0.4.5 (99be96e2e)
Scope: 8 files, +327 / −0, this feature only (no other changes ride along)

What it does

When a MoE model's expert weights are not GPU-resident (--n-cpu-moe N, or
auto-offload with a fit margin such that part of the model streams from host RAM),
every prefill batch forces the scheduler to upload each used expert tensor
host→device right before the MUL_MAT_ID split that consumes it. That H2D
transfer sits on the critical path and serializes behind the compute of the
previous split, so large-batch prefill (prompt processing) is dominated by
waiting on PCIe/NVLink instead of by kernels.

This PR makes those uploads overlap compute. During prefill it fires full-tensor
H2D copies ahead of need through a second backend instance on the same
device into rotating staging buffers, and the consuming split performs a
per-split cross-stream event wait that — because the copy was fired one split
earlier — is already satisfied by launch time.

split i        :  [ MUL_MAT_ID compute on expert weights ]
                    ^ uploads for split i+2 fired here (1-deep lookahead)
prefetch stream:     |======= H2D expert(i+2) into staging slot =======|
split i+2      :  [ wait ready[i+2] (no-op) ]  [ MUL_MAT_ID compute ]

User-visible surface

• --prefetch-experts-slots N — default 0 = off, zero memory overhead.
  N >= 2 enables the pipeline (needs at least 2 slots to rotate); 3 is
  recommended, 4 is the cap.
• GPU memory cost: N × max_expert_tensor, lazy-allocated on first fire.
• Optionally: llama_context_params.prefetch_experts_slots /
  ggml_backend_sched_set_prefetch_experts_slots() for library users.

How it works (mechanism)

ggml_backend_sched gains a small prefetch state machine
(ggml/src/ggml-backend.cpp):

1. Gate — only fires for splits whose first node is GGML_OP_MUL_MAT_ID,
   whose expert weight input is host-resident with
   GGML_BACKEND_BUFFER_USAGE_WEIGHTS, and whose routed batch is prefill-scale
   (ids ≥ 2 × n_expert). In callback_eval mode (decode, MTP draft) it never
   fires — decode is unaffected by construction.
2. Staging — a second backend instance (ggml_backend_dev_init on the same
   device, only if the device advertises async + events caps) uploads the
   full expert tensor into slot buffers via ggml_backend_tensor_set_async.
   With a large batch essentially every expert is used, so routing ids carry no
   information worth waiting for — full-tensor prefetch is the right trade.
3. Lookahead — while split i computes, the copy for split i + 1 + LOOKAHEAD
   is already in flight (LOOKAHEAD = 1).
4. Sync — each staged split does one ggml_backend_event_wait(ready[slot])
   on its backend right before graph launch. Because of the lookahead this is a
   no-op in steady state; it is the only synchronization point that preserves
   tool-call/tool_choice semantics (an earlier one-wait-per-graph variant was
   dropped for that reason).
5. Safety — the staged copy is only re-pointed at the slot buffer for the
   duration of its split and restored right after launch (kernels have already
   captured the address), so any fallback path can never observe a dangling
   slot. On any allocation/cap failure prefetch disables itself and the regular
   copy path is used — correctness never depends on prefetch.

The only user-visible knob is the slot count: lookahead depth and wait mode are
hardcoded to their measured-optimal values.

Lossless

Prefetch changes when bytes arrive on device, never what is computed:

• slots = 0 (default): the setter fully disables the feature — no state, no
  allocations, and the scheduler's compute loop behaves exactly as before.
• slots ≥ 2: the copied bytes are the same host weights; the consuming graph
  is unchanged (same tensors, same kernels, same order); the ready-event wait
  guarantees the data is on device before launch, making the overlap invisible
  to the numerics. The only mutation is a temporary re-point of the input copy's
  buffer/data pointer, restored before the next eval.
• No sampler, graph, or scheduling-order changes of any kind.

Empirical evidence (mindport build, greedy temperature 0, identical
prompt/seed, --n-cpu-moe 20 on a 24B A3B, q4_0 KV cache):

run	TTFT (prompt_ms)	output
prefetch OFF	3501	682 chars
prefetch ON (slots 3)	3786	682 chars
byte-identical:		True

Measured effect (host-expert configs)

Large prefill batches are where the feature pays; decode is untouched by design.

config	prompt	TTFT OFF	TTFT ON	decode OFF→ON	recall
24B A3B, -ncmoe 20	~42k tok (ctx 60k)	15.71 s	13.97 s (−11%)	45.8 → 46.0 t/s	20/20 both
24B A3B, -ncmoe 20	200 tok (ctx 4k)	0.51 s	0.41 s (−20%)	73.8 → 74.6 t/s	19/20 both
21.8 GB 35B A3B Q4_K_M, auto-offload	~42k tok	21.19 s	16.46 s (−22%)	41.8 → 41.5 t/s	20/20 both
21.8 GB 35B A3B Q4_K_M, auto-offload	200 tok	0.76 s	1.64 s (regression)	61.5 → 62.6 t/s	17/20 both

Notes: on the 35B no -ncmoe was used — the auto-offload fit margin
(-fitt 700) alone leaves ~half the expert weights host-resident, which is the
realistic deployment this feature targets. At 200 tokens the ON arm's 0.76→1.64 s
TTFT is the staging fixed cost dominating a tiny prefill — the feature is aimed
at large batches and is simply off by default for everyone else. Recall/outputs
are identical across arms in every row.

Files

file	change
common/arg.cpp	--prefetch-experts-slots N flag
common/common.{h,cpp}	common_params.prefetch_experts_slots plumbing
include/llama.h, src/llama-cparams.h	cparams field + default (0)
src/llama-context.cpp	wire cparams → sched_reserve()
ggml/include/ggml-backend.h, ggml/src/ggml-backend.cpp	scheduler prefetch state machine (+297)

[Anbeeld]

This is a pretty complicated addition that's not directly adjacent to BeeLLama's features. I'll prefer it to go through upstream's review process via your PR there: ggml-org#28414. Thanks.