Repository navigation
common : --fit grows the ubatch of MoE models into spare VRAM, for up to 2.5x faster prompt processing - #321
Conversation
A MoE ubatch reads the weights of nearly every expert, and uploads them when the experts are in host memory, so prompt processing pays one pass over all experts per ubatch. Without -ub, MoE models now use the logical batch (up to 4096) instead of 512. -ub still overrides it. With a draft or MTP model the larger default applies only when the tensor overrides keep the experts in host memory. The draft context follows the main context's ubatch, so both compute buffers grow, and with the whole model on the GPU that can overrun VRAM.
|
Initial version ( Recording the first approach and its numbers here, because the body will change if the Approach: MoE models (GGUF
Results:
Why it's being reworked: a fixed default ignores how much VRAM is free. On a nearly full card:
The next version lets |
Replaces the fixed MoE default from the previous commit. Without -ub, --fit now places the weights and the context at the default ubatch as before, then doubles n_ubatch towards n_batch (up to 4096) while the larger compute buffers still fit in the memory left over on every device. The fit never gives up context or device weights for a bigger ubatch, and a draft or MTP context that follows the main ubatch is measured with it. -ub still overrides it, and -fit off keeps the default.
LeoBorcherding
changed the title
The growth stays on by default. -fitub off (or LLAMA_ARG_FIT_UBATCH=off) keeps the ubatch at its default while --fit still places the weights and the context, so the behaviour can be switched off without pinning -ub 512.
-fitub becomes -fitu (as -fitp, -fitt, -fitc), the option moves next to the other --fit options, and the log line names it as the way to turn the growth off.
|
@danielhanchen Reopening this because it looks like it fits with your open MoE PRs instead of overlapping them. If you don't think we need it, feel free to close it. What it does: MoE models still process prompts in chunks of 512 tokens unless you pass How it fits with your PRs. I merged this PR into each head and measured prompt speed with experts in RAM (
With the whole model in VRAM, the gain is smaller but still there: about 1.2x on the R9700 and 1.46x to 1.62x on the A100. On the L4 it stays at 512, because there's no room left over. Dense models are never touched. Correctness: greedy outputs are byte-identical with and without this PR on every pair, card and config, and answer speed doesn't change. Running each PR with The merges were local. Both had small conflicts, in |
…ubatch The growth inferred that the draft follows the target from equal ubatch values, so --spec-draft-n-ubatch set to the target's ubatch was measured at the grown size while the draft context is created at the user's value. Pass whether the draft follows from the caller instead.
|
Codex review on a mirror of |
What this does
llama.cpp reads a prompt in chunks called ubatches. The default chunk is 512 tokens. For MoE models (Qwen3-30B-A3B, gpt-oss, DeepSeek, Qwen3.8-Flash-Next and so on), every chunk has to run through all the experts' weights once. So a 4,600-token prompt means 10 trips through every expert. When the experts live in system RAM, each trip also copies them to the GPU again.
Bigger chunks mean fewer trips, but each chunk needs more GPU memory while it runs. So this PR lets
--fitdecide:--fitplaces the weights and the context exactly as it does today, at 512.n_batch, max 4096) as long as it still fits within the fit's safety margin.expert_count > 0). Dense models keep 512.-ub(orLLAMA_ARG_UBATCH) still wins, and-fit offkeeps 512.-fitu off(--fit-ubatch off,LLAMA_ARG_FIT_UBATCH=off) turns off just the ubatch growth. The rest of--fitworks as before.MoE model, n_ubatch = 2048 fits in free memory (set -ub or -fitu off to override).Why isn't this already the default?
As far as I can tell, nobody chose 512 for MoE on purpose. It's left over from three things:
--fitwas built to answer "will this fit?", not "how fast can this go?" It shrinks context and moves layers to RAM until the model fits, but it never touchedn_ubatch. This PR adds that.--cpu-moe,--n-cpu-moe,--fitspilling). Each 512-token chunk copies every expert to the GPU again. That setup only became common recently.Results
Prompt processing, same machine, with and without this PR:
-cmoe)--fitspills experts-cmoe)-cmoe)-cmoe)-cmoe)-cmoe)--fitspills experts-cmoe)-ub 1024by hand does the same--fitspills expertsGeneration speed (writing the answer) doesn't change. What changes is how long you wait for the first word:
Time saved before the answer starts
What you actually feel is the wait between sending the prompt and the first word of the answer (time to first token). On a 4.5k-token prompt:
These are 4,565 prompt tokens divided by the measured prompt speeds above, not separately timed. Longer prompts save proportionally more.
Accuracy
Perplexity on wikitext-2 (Qwen3.8-Flash-Next, experts in RAM): 5.365 at ubatch 512 and 5.370 at 2048. The error bars are ±0.08, 16x bigger than the change. This is the same tiny numeric difference you get today by passing
-ub 2048yourself. Greedy output is identical on every config, and when there's no room the output matches the base exactly.Where to look in the code
common.cpp:1288-1304: after a successful fit, MoE models with no-ubcall the growth step.fit.cpp:910-975:common_fit_n_ubatchdoubles the ubatch, re-measures memory with the fit's own probe, and keeps a step only if every device still has its margin free. Weights and context are never touched.fit.cpp:180: the draft/MTP memory measurement, factored out of the fit's existing code so the growth step can reuse it.common.cpp:1241-1244: readsexpert_countfrom the GGUF.arg.cpp:1683: records that-ubwas given, so an explicit value always wins.arg.cpp:2942-2955: the-fitu/--fit-ubatchswitch, next to the other--fitoptions.Tradeoffs
-ub 2048still works: on the spilled Q6_K it gives +110% on long prompts, at the cost of ~8% on short prompts and 316 MiB of experts moved to RAM.How it was tested
5c19a5a75), comparing the base and this PR.--fitversion. On the configs with spare VRAM, both versions are equally fast (within 0.5%).-fit off -c 49152on a nearly full L4 loads, same as the base.-np> 1.llama-fit-paramsdoesn't print the grown-ubyet.AI disclosure (per CONTRIBUTING): written and benchmarked with Claude Code. Every number above comes from a logged run, and the logs are available on request.