feat: overlap CPU cold-expert and GPU hot-expert computation in the MoE cache - #9
Merged
Conversation
…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.
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
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 withllama-bench --sched-async-cpu 0,1(columnsac).Measured (RTX 3060, Qwen3.6-35B, 88 slots, merged profile)
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_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.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.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.llama_context_params.sched_async_cpu(default true) → cparams → sched setter; common args flag pair; llama-bench test dimension.Notes
--spec-type draft-mtpthe verify batches lengthen both chains and the gain is consistent.