CUDA SSD streaming: host-RAM tier for routed experts (working branch, 1.65-6X decode)
- Ngôn ngữ chính
- C
- Star
- 22.3k
- Fork
- 2.1k
- Merge trung bình
- 1 ngày 3 giờ
- Pull request đã merge (30 ngày)
- 4
Mô tả
Working branch, one commit, four files:
https://github.com/pmasala/ds4/commit/d6330c2
Posting this as an issue rather than a pull request: the design choices at
the end are yours.
## THE PROBLEM
On the CUDA SSD-streaming path a routed expert can only come from the small
per-layer VRAM buffer or from the SSD, and
`cuda_stream_selected_cache_begin_load` invalidates and recopies every
selected expert on every token. MoE routing keeps picking the same experts,
so most of those reads are re-reads, and decode ends up bound by the SSD
instead of compute. The seeding hooks Metal relies on (`seed_experts`,
`seed_selected`, `budget_for_expert_size`) are return-0/1 stubs on CUDA, so
the built-in hotlist never helped here.
The OS page cache is not an alternative here, because the streaming path
deliberately does not use it: expert reads go through a second fd opened
`O_RDONLY | O_DIRECT` (`cuda_model_stage_read`, ds4_cuda.cu:1873), and on the
buffered fallback each chunk is dropped again with
`posix_fadvise(POSIX_FADV_DONTNEED)`. That is the right call for an 80 GB
model on a 31 GB machine — the kernel would cache whatever was read last, not
what the router keeps selecting — but it leaves no reuse at all. A small
cache that follows the routing decisions holds only what is actually reused.
## WHAT THE CHANGE DOES
A pinned host-RAM slab becomes a second, faster source for the same VRAM
buffer:
- **Miss:** stream SSD -> VRAM as before, and copy the staged bytes into the
slab on the way. Filling the cache costs nothing extra.
- **Hit:** copy RAM -> VRAM on the non-blocking upload stream. No SSD read.
Eviction reuses Metal's route-hotness policy: experts the router keeps
selecting stay in the slab, cold ones are evicted. `ds4_metal.m` remains the
reference implementation.
The VRAM buffer and `slot_ids` are untouched, so the MoE kernel and its
numerics do not change. The budget planner sizes the slab, so
`--ssd-streaming-cache-experts` still applies — clamped by the planner as if
it were VRAM, see open question 1 — and the startup log matches what is
actually allocated. Layers whose experts do not fit the slab entry size
skip the cache and take the plain SSD path, which is what the startup log
already says. GLM and batched prefill also keep the plain SSD path.
One side effect: `ds4_gpu_stream_expert_cache_configured_count()` now returns
the planner budget where it used to return 0. This is harmless on CUDA: its
only reachable consumer is
`metal_graph_stream_prefill_batch_selected_addr_enabled()` (ds4.c:17707),
whose result is OR'd at ds4.c:33125 and ds4.c:57569 with the CUDA-specific
gate, and that gate's predicate is strictly weaker — so the combined value
cannot change.
Only `ds4_cuda.cu` changes. Pinning ~14.5 GiB costs about six seconds at
engine open. On a machine with less RAM the cache degrades rather than
turning off: if `cudaMallocHost` fails, the allocator halves the request and
retries, giving up only below 64 experts (ds4_cuda.cu:323-336). There are two
true off states, both logged — `disabled (no RAM budget)` when the budget
does not cover a single expert, and `unavailable (pinned alloc failed)` after
the retries are exhausted — plus the explicit opt-out
`DS4_CUDA_HOST_EXPERT_CACHE_GB=0`. In every off state the plain SSD path is
used unchanged; the new test covers exactly these cases.
`DS4_CUDA_HOST_EXPERT_CACHE_GB` overrides the slab size in GiB and disables
the tier when set to 0. It exists because the planner budget is VRAM-oriented
and gets clamped as such, and because the flag has no "off" value — see open
question 1.
## NUMBERS
Linux/WSL2, RTX 3500 Ada Laptop 12 GB, sm_89, CUDA 13.0, 31 GB RAM, NVMe.
DeepSeek V4 Flash IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8 imatrix, 80 GB. Each
figure is the mean of alternating runs; spread was a few percent.
decode t/s main branch gain
ctx 256 0.84 2.60 3.1x
ctx 2048 0.66 1.09 1.65x
mixed q2/q4 GGUF 0.80 1.63 2.0x
One honest caveat: the gain is the ratio between an SSD read and a RAM copy
on this machine — about 2.6 GB/s of effective O_DIRECT read during decode
against 12 GB/s host to VRAM. Where storage is faster the baseline is faster
and the gain will be smaller. The mechanism does not change: every hit
replaces an SSD read with a RAM copy, so the gain tracks the hit rate.
Prefill is unchanged, since the cache serves only decode: 9.28 vs 9.07 t/s at
ctx 2048 by ds4-bench, within the spread. Hit rates ran from 72% to 82%, the
highest on the 30474-token `--long-context` regression, so long context does
not hurt the cache.
## TESTING
make cuda CUDA_ARCH=sm_89 && make cpu && make cuda-regression && make test
make -C gguf-tools quality-score CUDA_ARCH=sm_89
DS4_TEST_SSD_STREAMING=1 ./ds4_test --logprob-vectors
DS4_TEST_SSD_STREAMING=1 ./ds4_test --server
DS4_TEST_SSD_STREAMING=1 ./ds4_test --long-context
All of the above pass with the cache active, `make cuda-regression` passes,
and `make cpu` builds warning-free.
Correctness check: sweeping the cache size so the hit rate moves from 33% to
74% leaves the greedy output byte-identical on both models. If the cache ever
served wrong bytes, a different hit rate would have produced different text.
Official-continuation scoring, cache on vs off
(`DS4_CUDA_HOST_EXPERT_CACHE_GB=0`, same binary), gives per-case avg_nll
equal to six decimals.
The new `tests/test_cuda_host_expert_cache.c` needs no GGUF and covers the
host-RAM and VRAM configurations where the cache cannot come up: as an
optimization it must either start and report what it pinned, or stay off and
say why. It has 30 checks; 11 of them encode the new behavior, so they fail
on main by design.
## PRE-EXISTING ISSUES, REPRODUCED ON MAIN
1. `./ds4_test --tool-call-quality` under `DS4_TEST_SSD_STREAMING=1` fails
four assertions, preceded by "CUDA streaming selected experts are
unavailable for layer 0": the `use_stream_selected_cache` guard at
ds4_cuda.cu:21194 is left unpopulated for the batched FFN path.
2. `DS4_CUDA_WEIGHT_CACHE_LIMIT_GB=1` crashes the streaming path with an
illegal memory access during the attn_sinks upload.
(A third pre-existing problem, greedy output not being reproducible run to
run on 12 GB cards, also occurs with this cache disabled, so it is not caused
by this branch; I will file it separately.)
## OPEN QUESTIONS
1. Should the planner know about the host tier? Today it plans a
VRAM-oriented budget that also happens to size a host-RAM allocation, and
clamps it as VRAM: a `--ssd-streaming-cache-experts 16GB` request is cut
to 9.0 GiB here by the graph working-set check (ds4.c:55418), which is why
the env override exists at all. Keeping the change out of ds4.c avoided
cross-backend risk, but the two are different resources. If the planner
becomes host-tier aware, the override stops being needed and only the
diagnostic `=0` remains.
2. The hotness policy is now duplicated between ds4_metal.m and ds4_cuda.cu,
with a comment asking that the two stay in step. Moving it to shared plain
C would be better; validating the Metal side of that refactor needs a Mac,
which I don't have.
3. The three per-expert copies could be merged into one. It measures only 3%
faster and would change the VRAM layout, which is what guarantees
bit-identical output today. I'd keep them separate; your call.
Hướng dẫn đóng góp
Đánh giá
Issue này chưa được đánh giá.