Skip to content

llama : MoE cache seeds a long first prompt chunk from the last prompt's routing - #345

Open
LeoBorcherding wants to merge 2 commits into
unslothai:moe-offload-prefetch-cachefrom
LeoBorcherding:moe-cache-seed-last-prompt
Open

LeoBorcherding wants to merge 2 commits into
unslothai:moe-offload-prefetch-cachefrom
LeoBorcherding:moe-cache-seed-last-prompt

Conversation

@LeoBorcherding

@LeoBorcherding LeoBorcherding commented Oct 31, 2026 •

Copy link
Copy Markdown

Stacks on #268. With experts in host memory and the MoE cache, the first big chunk of every prompt read its own routing before uploading experts, so those uploads ran on the compute stream instead of ahead on #268's copy queue. This seeds that chunk from the previous prompt's routing, so its copies overlap compute like every later chunk's do.

What changes

  • Routing counts are kept across prompts instead of cleared at each one.
  • A seeding ubatch of at least 4 x prefetch_min_batch tokens (1024) with no routing of its own copies every expert ahead on the copy queue and seeds from the last prompt's routing. Later chunks of the same prompt still seed from their own prompt.
  • The server's small tail ubatch reads the routing of the big ubatch before it, so its fallback follows the last prompt too.
  • Short chunks keep the routed-only copy: they leave many experts unrouted and have little compute to hide the copies under.
  • HIP: ggml_backend_cuda_copy_regions_async now works there too (the copy-regions kernel for device copies, one hipMemcpyAsync per upload), so a prefetched copy takes its cached experts from the MoE cache and uploads only the misses. Before, HIP re-uploaded every expert of a seeded chunk, cached or not. CUDA >= 12.8 already had this.

Only the prompt path with --moe-cache-mib and experts on the host is affected. Outputs and the cache hit rate are unchanged: greedy decode hashes match with and without this on every server.

Results

prompt speed

Qwen3.8-Flash-Next coder IQ1_M, -cmoe --moe-cache-mib 16000 -c 8192 -fa on, -b = -ub, MTP off. Warm 4.58k-token prompts, two servers per build, builds alternating on the same machine.

card ubatch #268 #268 + this PR change
Radeon AI PRO R9700, Windows 11, ROCm 7.1 default (512) 547 tok/s 547 tok/s none, chunks under 1024 tokens keep the old path
Radeon AI PRO R9700, Windows 11, ROCm 7.1 2048 1024 tok/s 1461 tok/s +43%
Radeon AI PRO R9700, Windows 11, ROCm 7.1 8192 1364 tok/s 1740 tok/s +28%
A100 40 GB, Colab default (512) 625 tok/s 592 tok/s none (one cold first prompt at 484 pulls the mean down)
A100 40 GB, Colab 2048 1144 tok/s 1222 tok/s +7%
A100 40 GB, Colab 8192 1577 tok/s 1793 tok/s +14%
RTX PRO 6000 Blackwell 96 GB, Colab 2048 3485 tok/s 3507 tok/s +1%
RTX PRO 6000 Blackwell 96 GB, Colab 8192 4484 tok/s 4782 tok/s +7%

R9700: the second card in that box, which also drives the display (decode there runs ~15% below the first card). #268 fd7489efb measured Oct 30; this PR at 153050ec5 on Nov 1. Each day ran an alternating 1ce1b6d3a control (the seeding commit alone: 1121 / 1689 on Oct 30, 1122 / 1679 on Nov 1), which matched within 1%. A100 measured Oct 30 at 1ce1b6d3a against #268; 153050ec5 changes only HIP, so the A100 rows still apply. RTX PRO 6000 is an Oct 23 run of the seeding change. An earlier A100 run agreed (+3% at 2048, +11% at 8192), as did an L4 (+6%, +11%).

Stacked with #321 and no flags (#321 picks ub 2048): A100 1199 tok/s against 625 for #268 alone. The R9700 stack was measured before 153050ec5 (1122 against 547); with it, the ub 2048 row above is what #321's default gets.

The gain grows with the ubatch, because the seeding chunk is a bigger share of the prompt: at 8192 this 4.58k prompt is one chunk, so the whole synchronous upload pass comes off the compute stream. Decode is unchanged on the R9700 and the A100 (38.4 vs 38.4 and 47.4 vs 47.4 tok/s at ub 8192), and the greedy decode hashes match between builds on each card, 153050ec5 included. On the R9700 at ub 8192 the HIP copy change cuts uploads from 23.5 GB to 12.4 GB per prompt and the per-layer waits for them from ~327 ms to ~22 ms.

Accuracy

No change. On a Colab L4, llama-perplexity on wikitext-2 (100 chunks of 512, 25.5k scored tokens) gives bit-identical logits with and without this change at ub 2048: PPL 5.3695 for both, mean KLD 0.000000, top-1 agreement 99.996% (the base file's 16-bit logit storage, same as a build compared with itself). With MTP (#240, n-max 2) draft acceptance is unchanged (0.801 vs 0.803 decode, 0.677 vs 0.677 on long prompts, inside run-to-run noise), and the greedy outputs hash the same. That run was the same change on an earlier stack of #268 + #264; the cards above also give matching greedy hashes per card.

Known cost: decode on Blackwell

On the RTX PRO 6000 Blackwell (sm_120), decode after a prompt drops from ~89 to ~83 tok/s with this at ub 8192. It isn't the cache: hits and misses are the same. Once #268's copy queue runs a kernel without the PDL launch attribute beside the compute stream, later CUDA graphs lose their PDL overlap. Before this change the ub-8192 seeding chunk was the one prompt chunk that never used the copy queue; now it does, and that's where the prompt gain comes from. The same drop already happens without this PR at ub 2048, where every chunk is prefetched.

Ampere and HIP have no PDL, so the A100 and AMD cards don't see it. The fix belongs in #268's copy queue (nothing non-PDL beside compute on sm_90+), and making the copy kernel PDL-aware halves the low-state rate in testing. I'd rather fix it there than gate this off.

Not run yet

  • H100 (sm_90): whether Hopper shows the same decode drop.
  • -np > 1 and multi-GPU: the counts are per context, so not expected to matter, but not run.

…t's routing

With experts in host memory and the MoE cache, the last large ubatch of each
llama_decode call seeds the cache. When that ubatch had no earlier routing in
the prompt (the first server chunk, or a whole prompt in one chunk) it read its
own ids first, so its expert uploads ran on the compute stream after a
per-layer readback instead of ahead on the prefetch copy queue.

Keep the routing counts across prompts. A long seeding ubatch with no routing
of its own now copies all experts ahead like any other large ubatch and seeds
from the last prompt's routing. Later chunks of the same prompt still seed from
the prompt's own routing. A small ubatch with outputs (the server's last prompt
tokens) also reads the routing of the large ubatch before it, so the fallback
follows the last prompt. Short chunks keep the routed-only copy: they leave
many experts unrouted and have little compute to hide the copies under.
…cache

ggml_backend_cuda_copy_regions_async returned false on HIP, so the prefetch
copy of a large ubatch could only take long runs of cached experts from the
MoE cache and uploaded the rest from host memory again. With the cache seeded
from the last prompt that was every expert of every layer: 23.5 GB per
4.58k-token prompt on an R9700 at ub 8192, 16 GB of it already in VRAM.

On HIP the device copies now go through the copy-regions kernel and the
uploads are one hipMemcpyAsync each (there is no batched copy API), so a
cached expert is copied from the cache and only the misses are uploaded.
@LeoBorcherding

Copy link
Copy Markdown
Author

Pushed 153050ec5. On HIP, ggml_backend_cuda_copy_regions_async returned false, so the prefetch copy of a large ubatch could take only long runs of cached experts from the MoE cache and uploaded everything else from host memory again. With the first chunk now seeded from the last prompt, that was every expert of every layer: 23.5 GB per 4.58k-token prompt on an R9700 at ub 8192, 16 GB of it already in VRAM. HIP now uses the copy-regions kernel for the cache copies and one hipMemcpyAsync per upload (there is no batched copy API), so only the misses are uploaded. CUDA >= 12.8 already took this path, so the NVIDIA rows are unchanged.

R9700 (Windows 11, ROCm 7.1, the second card in the box), warm 4.58k-token prompts, two servers per build alternating, against the previous head 1ce1b6d3a:

ub before after change
2048 (what #321 picks) 1122 tok/s 1461 tok/s +30%
8192 1679 tok/s 1740 tok/s +3.6%

Uploads at ub 8192 drop from 23.5 GB to 12.4 GB per prompt, and the compute stream's per-layer waits for them from ~327 ms to ~22 ms. Decode is unchanged, and greedy decode and long-answer hashes match the previous head in all 12 servers.

One ub 8192 server of the new head read 1561 tok/s with decode at 13.6-20 tok/s while another process was using the other card in the box; two reruns were clean and it is left out of the mean.

The body's R9700 rows and chart are updated to this head.

@LeoBorcherding

Copy link
Copy Markdown
Author

Codex review on a mirror of 153050ec5 is clean. CI is next.

This branch has not been deployed

No deployments
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