Skip to content

CUDA: resident expert cache and faster selected-expert uploads for --ssd-streaming decode#605

Open
iCreil wants to merge 4 commits into
antirez:mainfrom
iCreil:feat/stream-miss-upload
Open

CUDA: resident expert cache and faster selected-expert uploads for --ssd-streaming decode#605
iCreil wants to merge 4 commits into
antirez:mainfrom
iCreil:feat/stream-miss-upload

Conversation

@iCreil

@iCreil iCreil commented Jul 25, 2026

Copy link
Copy Markdown

While testing --ssd-streaming decode on CUDA (RTX PRO 6000, 96 GB, the
q4 Flash GGUF served from tmpfs) I measured a flat ~2.3 tok/s no matter
how warm the page cache is: every decode step re-uploads all selected
experts of every layer from the host map. The backend rewrite left the
CUDA resident expert cache API as stubs (configured_count() returns 0,
the seed hooks are no-ops), so only the per-layer staging buffers exist.
This PR restores the resident cache and then chips away at what the
timings said was left. Warm single-stream decode goes from ~2.3 to
29-32 tok/s; the four commits are independent steps and each one is
A/B-able on the same binary via env switches. Happy to rework any part.

1. Restore the resident expert cache (~2.3 -> 16.7 tok/s)

Brings back a VRAM LRU cache for streamed experts in front of the
staging buffers: hits become device-to-device copies, misses load the
expert into the cache once. Compared to the pre-rewrite cache this
version also:

  • keeps one cache per (gate,down) byte-size class (4 classes):
    mixed-quant GGUFs interleave layers with different expert byte sizes,
    and a single cache keyed on the last-seen size is released and
    rebuilt on every size transition, never accumulating hits;
  • lets prefill batch loads read the cache (hits become d2d copies) and
    append into free capacity, but never evict: a long prompt cannot
    cycle out the warm decode working set. Decode loads evict normally;
  • allocates the staging buffers with the cache still in place and only
    sacrifices it (once, with a retry) when VRAM is genuinely short;
  • implements the existing prefill-seed/hotlist hooks on top of the same
    cache, and only reports a budget to the shared working-set estimates
    when the operator opted in (budget flag or env), so default runs keep
    the current planner behavior.

Budget comes from the existing --ssd-streaming-cache-experts /
DS4_CUDA_STREAMING_EXPERT_CACHE_N, live-capped against free VRAM with
a reserve (DS4_CUDA_STREAMING_EXPERT_CACHE_RESERVE_GB, default 16).
DS4_CUDA_STREAMING_EXPERT_CACHE_N=0 disables the cache (used for the
A/B numbers).

2. Skip the per-load VRAM query once the cache has settled (16.7 -> 23.5)

Host-side phase timing showed ~90% of the steady-state selected load
(~920us of ~1010us) inside cache prepare: when the VRAM-capped capacity
lands below the requested budget (e.g. 7000 requested, 5892 fitting),
the capacity >= target fast path can never hit and every load of
every layer falls through to cudaMemGetInfo (~0.9ms each — ~55ms of
the ~60ms decode step at 61 layers). Remember the request each class
settled at and return immediately while dims and budget are unchanged;
the mark is cleared on release and re-evaluated on any budget change.

3. Pipeline the streamed expert upload path (23.5 -> ~24.5)

A cache miss paid ~2.3ms for a ~13.5 MiB expert: each of the three
tensor copies pread()s the whole tensor into pinned staging as a single
64 MiB-chunk pass and then fully synchronizes the upload stream, so
reads and uploads never overlap. Split the copy into 2 MiB chunks
(DS4_CUDA_STREAM_COPY_CHUNK_MB) through the existing 4-buffer ring
and issue the three tensors of an expert as one async group with a
single flush.

4. Plain pageable uploads when there is no O_DIRECT path (~24.5 -> 29-32)

With the model map backed by page cache/tmpfs, uploading straight from
the (pageable) map runs ~3x faster than the staged pread bounce on this
box (~2.3 -> ~0.8ms per expert). Default to the plain copy only when no
O_DIRECT read path is active — real direct-I/O SSD streaming keeps the
staged path and its page-cache bypass. DS4_CUDA_STREAM_MISS_PLAIN_COPY=1/0
forces either way.

Numbers

RTX PRO 6000 (96 GB), DeepSeek V4 Flash q4 GGUF in tmpfs,
--ssd-streaming --ssd-streaming-cache-experts 7000 --ctx 32768 --prefill-chunk 1024, 200 generated tokens per request, single stream,
steady state (2-3 requests of warm-up):

configuration tok/s
main (no resident cache) ~2.3 flat
+ resident cache 16.7
+ settled budget 23.5
+ pipelined uploads ~24.5
+ plain pageable uploads 29-32

Cache hit rate at steady state is ~96%; what remains of the step is
roughly one third miss uploads, the rest compute and the per-layer
router readback. Greedy output is unchanged in every A/B above
(cache on vs off, staged vs plain uploads — the copies are byte-exact).

With --batched-session N the resident cache is shared by the
sessions, so the per-session pipelining scales aggregates too: on top
of the #604 branch I measured 30 / 34 / 36.5 tok/s aggregate at 1/2/4
concurrent clients with greedy streams identical to the solo run.

Notes and limitations

  • CUDA only: Metal keeps its own (already resident-capable) cache
    paths untouched, and the ROCm layer-cache path is untouched too.
  • The cache lives on the single-GPU streaming placement (same
    constraint as the selected staging buffers).
  • Exact configs and per-step measurements are in the commit messages;
    I can share the bench harness if useful. If you had a different shape
    in mind for streaming caching, happy to adapt or split this.

iCreil and others added 4 commits July 25, 2026 11:22
The GPU-backend rewrite reduced the CUDA streaming expert cache API to
stubs: configured_count() returns 0, seed_selected()/seed_experts() are
no-ops, and cuda_stream_selected_cache_begin_load() re-uploads every
selected expert from the host map on every layer of every decode step.
With --ssd-streaming the decode is PCIe-bound and stays flat no matter
how warm the OS page cache is.

Restore the resident VRAM expert cache in front of the selected staging
buffers:

- Cache hits become device-to-device copies into the staging buffers;
  misses are loaded once into the cache (LRU) and then copied.
- Caches are kept per (gate,down) byte-size class (4 classes):
  mixed-quant GGUFs interleave layers with different expert byte sizes,
  and a single cache keyed on the last-seen size would be released and
  rebuilt on every size transition, never accumulating hits.
- Decode loads may evict (LRU); prefill batch loads
  (prepare_selected_batch) may append into free capacity but never
  evict, so a long prompt cannot cycle out the warm decode working set.
- The staging buffers are allocated with the resident cache still in
  place; only if that allocation fails is the cache released and the
  allocation retried once (very long prompts on tight VRAM).
- Budget: DS4_CUDA_STREAMING_EXPERT_CACHE_N (default 512, engine
  override via ds4_gpu_set_streaming_expert_cache_budget), live-capped
  against cudaMemGetInfo with a reserve
  (DS4_CUDA_STREAMING_EXPERT_CACHE_RESERVE_GB, default 16), shrinking
  ~10% per allocation failure with a per-class runtime cap.
- configured_count()/budget_for_expert_size() only report a budget when
  the operator opted in (cache-N env, engine budget, or hotlist env), so
  the shared working-set estimates and addr fast paths keep the current
  default behavior.
- seed_selected()/seed_experts() implement the existing prefill-seed and
  hotlist-seed hooks on top of the same cache.
- DS4_CUDA_STREAMING_EXPERT_CACHE_VERBOSE logs per-layer
  hit/miss/direct counters.

The resident cache is only used on the single-GPU streaming placement
(same constraint as the selected staging cache).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NCDczkCWtiwuTeS9mmDei7
Host-side phase timing showed ~90% of the steady-state selected load
(~920us of ~1010us) going to cuda_stream_expert_cache_prepare: when
the VRAM-capped capacity ends up below the requested budget (e.g.
7000 requested, 5892 fitting), the "capacity >= target" fast path can
never hit, so every load of every layer falls through to
cuda_stream_expert_cache_live_budget and its cudaMemGetInfo call
(~0.9ms each — ~55ms of the ~60ms decode step at 61 layers).

Remember the requested cap each class settled at; steady-state loads
with unchanged dims and budget return the cache immediately. The mark
is cleared on release and naturally re-evaluated when the requested
budget changes (env, engine budget, or an OOM runtime cap).

Measured (RTX PRO 6000, Flash q4 tmpfs, --ssd-streaming, cache
budget 7000, 200-token requests): warm decode 16.7 -> 23.5 tok/s;
remaining per-load cost is dominated by the ~3.6% resident-cache
misses (PCIe uploads). Greedy output unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NCDczkCWtiwuTeS9mmDei7
A resident-cache miss costs ~2.3ms for a ~13.5 MiB expert (~6 GB/s on
a PCIe gen5 box): each of the three tensor copies pread()s the whole
tensor into pinned staging as a single 64 MiB-chunk pass and then
fully synchronizes the upload stream, so reads and device uploads
never overlap.

- Split the copy into small chunks (2 MiB default,
  DS4_CUDA_STREAM_COPY_CHUNK_MB) so the host read of chunk i+1
  overlaps the device upload of chunk i through the 4-buffer ring.
- Make the per-copy sync optional and share the staging-ring cursor
  across calls: the three tensors of one cached expert are issued as
  one async group with a single flush, pipelining across tensors too.
- DS4_CUDA_STREAM_MISS_PLAIN_COPY=1 uploads straight from the
  (pageable) model map instead, as a measurement switch for
  tmpfs/page-cache maps.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NCDczkCWtiwuTeS9mmDei7
… O_DIRECT

A/B on the target box (Flash q4 in tmpfs, resident cache active):
uploading a missed expert straight from the pageable model map runs
~3x faster than the staged pread bounce (~2.3ms -> ~0.8ms per 13.5 MiB
expert; warm decode 23.5 -> 29.3-31.8 tok/s, greedy unchanged). When
an O_DIRECT read path is active (real SSD streaming, page cache must
be bypassed) the staged path remains the default.
DS4_CUDA_STREAM_MISS_PLAIN_COPY=1/0 forces either way.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01NCDczkCWtiwuTeS9mmDei7
iSevenDays added a commit to iSevenDays/ds4 that referenced this pull request Jul 25, 2026
…tirez#605 race fix

Cherry-picked onto feat/mtp-native (which has upstream main + MTP/DSpark + expert LRU):
- antirez#604: 10 commits (session-batch decode, fused routed experts, VRAM reserve for resident sessions)
- antirez#605: 1 commit (batched server session recovery race fix)
- antirez#605 streaming improvements (4 commits): deferred — conflict with per-tier LRU, adapt later

MTP/DSpark: present (from upstream fc9efd1 via feat/mtp-native).
--batched-session --prefill-chunk 1024 now available (antirez#604).
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