Repository navigation
llama : MoE cache seeds a long first prompt chunk from the last prompt's routing - #345
LeoBorcherding wants to merge 2 commits into
Conversation
…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.
|
Pushed 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
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. |
|
Codex review on a mirror of |
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
prefetch_min_batchtokens (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.ggml_backend_cuda_copy_regions_asyncnow works there too (the copy-regions kernel for device copies, onehipMemcpyAsyncper 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-miband 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
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.R9700: the second card in that box, which also drives the display (decode there runs ~15% below the first card). #268
fd7489efbmeasured Oct 30; this PR at153050ec5on Nov 1. Each day ran an alternating1ce1b6d3acontrol (the seeding commit alone: 1121 / 1689 on Oct 30, 1122 / 1679 on Nov 1), which matched within 1%. A100 measured Oct 30 at1ce1b6d3aagainst #268;153050ec5changes 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,
153050ec5included. 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-perplexityon 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
-np > 1and multi-GPU: the counts are per context, so not expected to matter, but not run.