Skip to content

feat: overlap CPU cold-expert and GPU hot-expert computation in the MoE cache - #9

Merged
thecodacus merged 5 commits into
perffrom
fable5/moe-cache-overlap
Aug 22, 2026
Merged

feat: overlap CPU cold-expert and GPU hot-expert computation in the MoE cache#9
thecodacus merged 5 commits into
perffrom
fable5/moe-cache-overlap

Conversation

@thecodacus

Copy link
Copy Markdown
Owner

What

Runs the MoE expert cache's two chains concurrently: the cold-expert chain executes on a CPU worker thread while the hot-pack chain runs on the GPU, joined at the merge. On by default; disable with --no-sched-async-cpu (llama-server, llama-cli) or bench both sides in one run with llama-bench --sched-async-cpu 0,1 (column sac).

Measured (RTX 3060, Qwen3.6-35B, 88 slots, merged profile)

off on
generation + MTP (llama-cli, 256 tok) 70.0 72.7–74.2 (+4–5%)
generation, cache only 59.9 61.0 (+1.8%)
pp256 (after batch gate) 309.1 297.7 (−3.7%)

Outputs byte-identical on vs off, with and without MTP. Serving stress test at 32k ctx / 96 slots: 509 t/s prefill on a 1.9k-token prompt, 69.8 t/s over 400 generated tokens.

How

  • ggml sched: ggml_backend_sched_set_async_cpu() — CPU splits execute on a persistent worker; the main loop keeps launching later splits and joins lazily at the first split that is CPU or reads a CPU-resident tensor. Leftover jobs are drained on eval entry, synchronize, and free.
  • graph: the cold chain is built before the hot chain so the hot split follows it with no cross-backend inputs (it launches while the worker computes), and the merge add is pinned to CPU for decode-size batches (n_tokens <= 8) so it cannot pull the hot chain into a split that waits on the cold result. Large prefill batches keep the previous graph — the CPU merge and its activation copies cost more than the overlap hides there.
  • fix en route: the expert cache pack buffer is now marked GGML_BACKEND_BUFFER_USAGE_WEIGHTS. Without it, the scheduler's assignment pass does not pin pack tensors to their backend, and a CPU-assigned consumer propagates backward and drags the hot matmuls (plus a per-layer weight copy) onto the CPU.
  • plumbing: llama_context_params.sched_async_cpu (default true) → cparams → sched setter; common args flag pair; llama-bench test dimension.

Notes

  • Plain decode without speculative decoding measures ~1–2% slower with the overlap at bench (46.7 vs 47.5 tg64) — the split overhead is not fully repaid when the chains are short. With --spec-type draft-mtp the verify batches lengthen both chains and the gain is consistent.
  • Deployment sizing note: keep ~900 MB VRAM free beyond the pack — runtime CUDA pool growth on large prompts allocates beyond what the load-time check sees.

…YNC_CPU)

The dual-path MoE cache ran its two chains serially: the scheduler
blocks on CPU splits, and the cold chain's input copies drain the whole
CUDA stream. Three changes make them concurrent, gated behind
GGML_SCHED_ASYNC_CPU=1:

- sched: CPU splits execute on a persistent worker thread; the main
  loop keeps launching later splits and joins lazily at the first split
  that is CPU or reads a CPU-resident input
- graph: cold chain is built before the hot chain so the hot split
  follows it and carries no cross-backend inputs (launches while the
  worker computes)
- graph: the merge add is pinned to CPU so it does not pull the hot
  chain into a split that waits on the cold result
Without the weights usage flag the scheduler's assignment pass does not
pin the pack tensors to their backend, so a CPU-assigned consumer (the
pinned merge in async mode) propagates backward and pulls the hot
matmuls onto the CPU, copying pack weights over PCIe every layer.
Replaces the GGML_SCHED_ASYNC_CPU env var: the overlap is now on by
default and controlled through llama_context_params.sched_async_cpu.
- ggml: ggml_backend_sched_set_async_cpu() setter, env read removed
- llama: context param + cparams plumbing, graph gate reads cparams
- common: --sched-async-cpu / --no-sched-async-cpu (server, cli)
- llama-bench: sched_async_cpu test dimension (--sched-async-cpu 0,1
  benches both in one run, column shown as 'sac' when non-default)
For prefill batches the CPU merge and its per-layer activation copies
cost more than the overlap hides (pp256 regressed 307 -> 257 t/s).
Decode and MTP verify batches (n_tokens <= 8) keep the overlap.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant