CUDA streaming: per-size-class expert caches, keep cache warm across prefills#504
Open
iCreil wants to merge 1 commit into
Open
CUDA streaming: per-size-class expert caches, keep cache warm across prefills#504iCreil wants to merge 1 commit into
iCreil wants to merge 1 commit into
Conversation
…prefills Mixed-quant GGUFs (e.g. the official Flash q2-q4-imatrix) interleave layers whose routed experts have different byte sizes. The CUDA expert cache was a single global slab keyed on the last-seen size: every q2<->q4 layer transition reset the runtime cap and released the whole cache, so decode on mixed models was stuck at ~1 t/s with a permanent cap-notice flood. Additionally, prefill batch loads released the entire resident expert cache unconditionally before allocating their staging buffers, forcing a full cache re-warm on every request (this also affects uniform models). Changes: - keep up to 4 expert caches, one per (gate,down) byte-size class; runtime caps and cap notices are tracked per class - begin_compact_load: try the staging allocation first and only sacrifice the expert cache when VRAM is actually short (retry once) - let the prefill batch path read the global cache (hits are device-to- device copies instead of host uploads), appending on miss only while free capacity remains, never evicting: a long prompt cannot cycle out the decode working set On an RTX PRO 6000 Blackwell 96GB with the official Flash q2-q4-imatrix GGUF, a 200-word completion goes from not finishing within 400s (~0.9 t/s) to ~25 t/s sustained decode with --ssd-streaming-cache-experts 9472 and DS4_CUDA_STREAMING_EXPERT_CACHE_RESERVE_GB=12 (97% cache hit rate per the VERBOSE stats). Uniform-quant GGUFs keep their existing behaviour.
giannisanni
pushed a commit
to giannisanni/neutronstar
that referenced
this pull request
Jul 7, 2026
…fills Port of iCreil's PR against upstream main, applied to the glm-local CUDA streaming stack (comments translated to English): 1. Per-size-class expert caches (4 classes keyed on gate/down expert bytes). Mixed-quant models no longer thrash the cache on layer-size transitions. For this fork's GLM merge it means the Q2_K MTP draft layer (blk.78) can be cached alongside the IQ2_XXS layers instead of falling back to model views, and the slab-budget gate now sees a per-class budget. 2. Prefill batch loads try their staging allocation before releasing the resident expert cache, and only sacrifice it when VRAM is truly short. Short chat prompts no longer throw away the warm decode working set. 3. Prefill loads may append to free cache capacity but never evict valid slots, so long prompts cannot cycle out the decode working set. Prefill hits become device-to-device copies. Tested on the 4060 Ti / GLM-5.2 rig: one-shot regression clean (0.33 t/s, correct output), two-turn chat clean (turn 2 resume 0.35/0.32 t/s, memory low point 7.7GB, no OOM).
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
1 participant
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.
Fixes #503.
See the issue for the full analysis and benchmarks. Summary of the three changes, all in
ds4_cuda.cu:1. Per-size-class expert caches.
g_stream_expert_cachebecomes a fixed array of 4 caches keyed by(gate_expert_bytes, down_expert_bytes).note_size()is replaced byclass_index()(assigns a free slot on first sight; recycles class 0 in the — never observed — case of more than 4 distinct sizes). Runtime caps and cap-notice dedup move to per-class arrays.prepare()operates on the class cache; cross-class transitions no longer release anything.2. Conditional release in
begin_compact_load. The unconditionalrelease_all()on theallow_global_cache=0path becomes: tryensure_slots()first; on allocation failure release the expert caches and retry once. Long-prompt VRAM pressure is still handled, short prompts keep the warm cache.3. Prefill batch reads the cache.
prepare_selected_batchnow passesallow_global_cache=1plus a newallow_cache_evict=0mode: batch loads copy hits device-to-device and may append to free capacity, but never evict valid slots (an LRU cache smaller than a long prompt's expert union would otherwise be fully cycled, destroying the decode working set). The single-token decode path keeps full LRU behaviour (allow_cache_evict=1).Notes:
memsetover the struct containing astd::vectoris replaced by explicit member resets (vector().swap()), which also fixes a latent leak of the slots buffer.make cuda-genericwith Flash q2-imatrix (uniform: behaviour unchanged) and Flash q2-q4-imatrix (fixed: ~0.9 → ~25 t/s sustained decode, 97% cache hit rate perDS4_CUDA_STREAMING_EXPERT_CACHE_VERBOSE).