Support for Multi-GPU and --ssd-streaming#568
Draft
iSevenDays wants to merge 4 commits into
Draft
Conversation
iSevenDays
force-pushed
the
mgpu-ssd-merge-main
branch
2 times, most recently
from
July 17, 2026 07:50
c231111 to
2c5c235
Compare
iSevenDays
marked this pull request as draft
July 17, 2026 08:51
Brings multi-tier (2-GPU) SSD streaming to antirez main's tensor-parallel
mgpu layer. Previously SSD streaming was single-device only: one selected-
expert cache, one pinned staging pool, one upload stream, all on device 0.
antirez main refused the combination upfront:
ds4.c: if (e->ssd_streaming && e->multi_tier) {
"--ssd-streaming is not compatible with multi-GPU placement"
}
ds4_cuda.cu loader had a second guard:
if (g_n_gpus != 1) { "CUDA SSD streaming requires single-GPU placement" }
Both gates removed for the non-TP path so:
--ssd-streaming --gpu-vram 45,45 --ctx 131072
loads the 81GB q4k model across 2 GPUs (streaming, low resident VRAM/card).
Adapts the cchuter feat/mgpu-ssd port (commit 23245bb + 9522bac) to antirez
main's symbols/structure. antirez main's SSD is simpler than the cchuter
fork (the expert LRU cache is stubbed), so the per-tier struct only carries
the selected_cache + 4-slot staging pool + upload stream — no runtime caps.
ds4_cuda.cu
- New per-tier struct ds4_ssd_ctx { selected_cache, stage_raw/stage/
stage_event/stage_bytes, upload_stream } and static g_ssd[DS4_MAX_GPUS].
ssd_current() maps cudaGetDevice() -> g_gpu[t].device_id -> g_ssd[t]
(fallback g_ssd[0]).
- Converted all ~60 references to the old per-device globals
(g_stream_selected_cache, g_stream_selected_stage*, and
g_stream_selected_upload_stream) to fields on ssd_current(). The master
switch (g_ssd_streaming_mode) stays GLOBAL.
- ds4_gpu_set_ssd_streaming / _release_resident loop all tiers
(cudaSetDevice per tier) so each g_ssd[t] resets/frees on the device
that owns it.
- cuda_stream_selected_stage_pool_alloc creates the upload stream on the
CURRENT device (per-layer dispatch has cudaSetDevice'd).
- cuda_stream_selected_cache_begin_load derives logical_tier from
cudaGetDevice() (was hardcoded to 0); removed the g_n_gpus != 1 refuse.
- routed_moe_launch consumer re-affirms cudaSetDevice(g_gpu[logical_tier])
before ssd_current() so the ambient device matches the OUT tensor's tier.
- ds4_gpu_cleanup adds per-tier SSD teardown inside the existing g_gpu[]
loop (selected_cache + staging pool freed on each device).
- Embed-device fix (critical correctness, ported from 9522bac):
ds4_gpu_embed_token_hc_tensor and ds4_gpu_embed_tokens_hc_tensor wrap
their kernel launches in WITH_DEVICE(out_hc's device) so they never run
on a stale ambient device. Without this, a prefill after a decode faults
with cudaErrorIllegalAddress (embed runs on the stale head-tier device
while reading an emb-tier pointer).
ds4.c
- Gate relaxed: --ssd-streaming is now refused ONLY under CUDA tensor
parallelism (e->cuda_tensor_parallel). TP keeps half the routed experts
resident per tier, which genuinely conflicts with SSD streaming; plain
multi-tier (pipeline layer-split from --gpu-vram) is allowed.
- Multi-tier init branch wires ds4_gpu_set_ssd_streaming + auto-cache +
budget + slab bytes (mirrors the single-tier sequence, AFTER
ds4_gpu_init_multi populates g_gpu[]).
- Placement packer made SSD-aware:
* tensor_is_ssd_streamed_routed_expert() identifies the
.ffn_{gate,up,down}_exps.weight suffix.
* engine_compute_entry_bytes excludes routed-expert tensors under SSD
(non-TP) so they are not charged against the resident budget.
* engine_install_per_device_caches skips them so they are not registered
resident (would blow the budget and defeat streaming).
* engine_balance_placement_for_ssd splits the transformer layer range
across all requested GPUs budget-proportionally (contiguous) so
streaming work is balanced instead of collapsing onto tier 0.
- metal_graph_ssd_assert_tier_device(g, il) helper re-affirms the CUDA
device to the layer's home tier before each SSD producer stage (ambient
device drifts across attention/router kernels). Inserted before every
ds4_gpu_stream_expert_cache_{begin_selected_load,prepare_selected_batch,
seed_selected,seed_experts} call.
- Embed-device repoint: in metal_graph_prefill_layer_major (2 sites) and
metal_graph_verify_suffix_tops_impl, repoint g->active_tier = g->emb_tier
before metal_graph_upload_prompt_embeddings_hc so the embed kernel runs
on the device that owns prefill_tokens + token_embd. No-op single-tier.
- metal_graph_selected_async_load (SSD async staging worker): added
home_tier field to the job struct, threaded through _start_tensor.
Worker thread defaults to device 0; it now calls
ds4_gpu_set_current_device(job->home_tier) at run() entry so the per-tier
SSD cache is resolved correctly when staging from a service pthread.
Build (required gate): clean.
PATH=/usr/local/cuda/bin:$PATH make -j$(nproc) cuda CUDA_ARCH=sm_89
=> 0 errors, 0 warnings in ds4.c / ds4_cuda.cu (only the pre-existing
rax.c -Wuse-after-free warning remains).
Live 2-GPU verification pending (GPUs held by production service on :8002);
this commit is CODE + clean BUILD only.
Co-Authored-By: Claude <noreply@anthropic.com>
…al cache The CUDA fault: every short /v1/chat/completions (<=18 tokens by default, so most chat turns) returned HTTP 500 "cuda prefill failed" in ~0.02 s with NO CUDA error logged at default verbosity. Root cause (PRIMARY), in ds4.c metal_graph_eval_token_raw_swa_streaming: The per-token decode-style prefill path reads metal_graph_cur_hc(g), which expands via DS4_GPU_GRAPH_CLASS_P_ACCESSOR to cur_hc_by_tier[g->active_tier]. Multi-tier init leaves active_tier at the sentinel -1 (ds4.c:16747), and this function never repointed it before the embed. cur_hc_by_tier[-1] is an out-of-bounds slot (cur_hc_by_tier is the FIRST field of the struct), so ds4_gpu_embed_token_hc_tensor got out_hc=NULL, silently returned 0, ok went false, and the prefill returned false. Hence "cuda prefill failed" with zero CUDA state set (the silent log). After the first successful token, active_tier was left at the last layer tier (typically 1 = head tier), so the next token embed wrote cur_hc to the WRONG tier. Same class of bug as the earlier batch-embed device-mismatch (fixed by WITH_DEVICE on ds4_gpu_embed_token_hc_tensor plus active_tier=emb_tier repoint in metal_graph_prefill_layer_major), but on the per-token decode-streaming path that the earlier fix did not cover. Root cause (SECONDARY), in ds4_cuda.cu g_current_logical_tier: cudaSetDevice is per-thread, but g_current_logical_tier was a process-global cache. The async staging worker (metal_graph_selected_async_load_run in ds4.c) read the main thread just-set tier from the cache and SKIPPED its own cudaSetDevice, leaving the worker on device 0 while it believed it was on tier 1. It staged experts into g_ssd[0].selected_cache (wrong device) while the consumer ran on tier 1 and found g_ssd[1].selected_cache zero. So "CUDA streaming selected experts are unavailable for layer 21". Fix (PRIMARY), ds4.c metal_graph_eval_token_raw_swa_streaming: at function entry, call metal_graph_set_active_tier_no_copy(g, g->emb_tier). This sets the active_tier field AND calls cudaSetDevice(emb_tier), so the ambient device matches the field. The per-layer decode dispatcher early-return (tier == active_tier) then no-ops correctly because the ambient device is already right. No-op for single-tier (placement == NULL). Fix (SECONDARY), ds4_cuda.cu: declare g_current_logical_tier as thread_local so each thread tracks its own device. The worker ds4_gpu_set_current_device now misses the cache on first call (its own thread-local is -1) and actually calls cudaSetDevice. Also ports the reference fail-fast from feat/mgpu-ssd (commit 2c5c235): cuda_err_is_sticky plus _exit(1) inside cuda_ok on cudaErrorIllegalAddress and friends, gated by DS4_CUDA_NO_FAIL_FAST. Without this, a single transient illegal-access poisoned the context forever and the server kept returning HTTP 500 until manual restart; with it, a supervisor (systemd) restarts the process and service self-recovers. Observability (the bug was silent at default verbosity): gated behind DS4_SSD_DEBUG=1, off by default. - ds4_gpu_debug_probe(where, il, tier) in ds4_cuda.cu: logs the ambient CUDA device, the expected tier physical device, and any pending CUDA error after a forced cudaDeviceSynchronize on EVERY device (cross-device illegal-access errors hide from an ambient-only sync). - ds4_ssd_dbg() in ds4.c mirrors the gate. - Probes cover the entire multi-tier prefill path: reset_prefill_state, set_active_tier_batch entry/done, encode_layer_batch entry/post-attn/ post-ffn, the embed upload (both split_commands branches), the SSD expert staging (cuda_stream_prefill_batch_selected_load + decode_cuda_selected_load + ssd_assert_tier_device), the routed-MoE consumer, and the per-token eval_token_raw_swa_streaming path including the batched static decode loop. - moe_check FAIL dump in routed_moe_launch prints every field of the selected_cache comparison so the next reader of a sticky-error log sees exactly which condition failed. Verified: 3 chat requests on port 8012 with --ssd-streaming --ssd-streaming-cache-experts 40GB --gpu-vram 45,45 --ctx 262144 all return HTTP 200 with valid completions (was: HTTP 500 cuda prefill failed in 0.02 s). /v1/models reports context_length 262144. Zero CUDA errors in the log across all three requests. Decode is slow (~30-60 s/short reply) because the SSD-streaming expert cache is cold and only ~12 GiB per GPU is resident; that is the streaming-mode tradeoff, not a defect. Co-Authored-By: Claude <noreply@anthropic.com>
…SD streaming
The antirez-main SSD streaming port worked (correct output at 262k ctx) but
re-read every routed expert from SSD on every decoded token: ~0.45 tok/s with
VRAM stuck at ~12.6 GiB/card. Port the persistent resident LRU expert cache
from feat/mgpu-ssd so warm experts stay resident in VRAM.
ds4_cuda.cu — the LRU port (per-tier in g_ssd[t].expert_cache):
* Structs: cuda_stream_expert_cache_slot, cuda_stream_expert_cache
(capacity/count/tick + device slabs gate_ptr/up_ptr/down_ptr +
std::vector<slots>), added before struct ds4_ssd_ctx.
* Extended ds4_ssd_ctx with expert_cache + per-tier runtime_cap,
runtime_gate/down_bytes, memory_cap_notice (the OOM-driven cap state).
* Enum DS4_CUDA_STREAM_EXPERT_DEFAULT/MAX + g_stream_expert_budget_override.
* Ported the full LRU function set (verbatim from feat/mgpu-ssd, adapted to
resolve the active tier via ssd_current()): release_all, invalidate,
requested_budget, configured_budget, budget_visible_to_shared,
reserve_bytes, live_budget (the cudaMemGetInfo minus-reserve sizer that
emits the "capped from X to Y experts (available Z GiB, reserve W GiB)"
log), expert_bytes, note_size, shrunken_cap, note_oom_cap, try_alloc,
prepare, find, lru_slot, copy_to_compact, load_slot, seed_one.
* Wired cuda_stream_selected_cache_begin_load to consult the LRU: on hit,
D2D-copy gate/up/down from the LRU slab to the per-request compact cache
(no SSD read); on miss, stream from SSD into the LRU's LRU slot, then
D2D to compact; falls back to direct SSD-to-compact on LRU failure.
* Updated the extern SSD API: ds4_gpu_set_ssd_streaming (per-tier cap reset
+ free expert_cache on disable), ds4_gpu_set_streaming_expert_cache_budget
(sets g_stream_expert_budget_override + invalidates/releases every tier),
ds4_gpu_stream_expert_cache_{configured,current}_count,
release_resident, budget_for_expert_size, seed_selected.
* Added a no-op cuda_model_load_progress_finish stub (progress meter exists
in mgpu-ssd but not antirez main; lets the LRU logging port verbatim).
ds4.c — fix the multi-tier byte-budget conversion gap:
The multi-tier init path called ds4_engine_configure_streaming_auto_cache
but skipped ds4_engine_configure_streaming_cache_budget, so
--ssd-streaming-cache-experts NGB left ssd_streaming_cache_experts=0 and
the LRU's budget_visible_to_shared() gate kept the cache disabled. Added
the missing call (single-tier path has always done this).
VERIFIED on 2x RTX 4090 D (sm_89), ctx=262144:
Before: 0.45 tok/s warm, VRAM 12.6 GiB/card (no LRU; every expert re-read).
After: warm gen avg 10-14 t/s (peak chunk 38 t/s), VRAM 44.4 GiB/card.
Server log on the user's exact command
(DS4_CUDA_STREAMING_EXPERT_CACHE_RESERVE_GB=4 ... --ssd-streaming-cache-experts 40GB ...):
"CUDA streaming expert cache capped from 5556 to 4701/4717 experts
(available 34.99 GiB, reserve 4.00 GiB, 6.75 MiB/expert)"
/v1/models context_length=262144; correctness: 8+8 -> 16.
Build: PATH=/usr/local/cuda/bin:$PATH make -j$(nproc) cuda CUDA_ARCH=sm_89 clean.
Co-Authored-By: Claude <noreply@anthropic.com>
Under multi-tier SSD streaming the DSpark support pack and drafter kernels run on tier 1 (the configured executor tier), but metal_graph_configure_dspark_capture allocated every scratch tensor (dspark_target_hidden, dspark_hc_mean_weights, dspark_draft_hc, dspark_raw_cache, etc.) via ds4_gpu_tensor_alloc(bytes), which hard-codes tier 0 (ds4_cuda.cu:2592). The first decode token then ran metal_graph_dspark_capture_decode_layer(40) on tier 1: hc_weighted_sum_tensor read cur_hc (tier 1) against dspark_target_hidden and dspark_hc_mean_weights (tier 0), raising a sticky CUDA "illegal memory access" the moment layer 41's rms_norm_plain launched. The fail-fast gate restarted the service on every chat request. Also removes the upstream --ssd-streaming + --mtp incompatibility gate (ds4_engine_open_internal) so the support model can load resident while the main model streams. Fix: resolve exec_tier from g->dspark_exec_tier in metal_graph_configure_dspark_capture and use ds4_gpu_tensor_alloc_ptr_on(exec_tier, ...) for all DSpark scratch. Single-tier (Metal/CPU) builds are unchanged because alloc_ptr_on(0, ...) collapses to the legacy allocator there. Verified: 3/3 chat requests return content "13"; VRAM ~46.5 GiB/card; DSpark engaged (layers 40,41,42 capture clean); no CUDA errors in probes. Co-Authored-By: Claude <noreply@anthropic.com>
iSevenDays
force-pushed
the
mgpu-ssd-merge-main
branch
from
July 21, 2026 15:14
0f8fa24 to
e26e2f0
Compare
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.
Correctness Tests.
Test 1 — Single-GPU SSD:
Server (background):
(I actually used v4 flash official q2 quant, so you can just ln (link) it directly)
Request (separate call, once "listening on" appeared in logs):
Test 2 — 2-GPU SSD:
Server (background):
Request:
Both returned "content": "16".
context window of 131k is supported on just a single 4090d (48Gb VRAM)
Quickstart:
Info: thanks user @cchuter for initial multi-gpu support. This MR resolves issues with main branch, and adds ssd-streaming support.