From 66df47ae33a645134691d86e0e492cd3bbc1f94a Mon Sep 17 00:00:00 2001 From: Anton Sokolchenko Date: Tue, 21 Jul 2026 10:58:12 +0000 Subject: [PATCH 1/4] feat(ssd): per-tier SSD weight-streaming across multiple GPUs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- ds4.c | 212 ++++++++++++++++++++++++++++- ds4_cuda.cu | 380 ++++++++++++++++++++++++++++++++++------------------ 2 files changed, 458 insertions(+), 134 deletions(-) diff --git a/ds4.c b/ds4.c index 876cfe27f..d578dc8c9 100644 --- a/ds4.c +++ b/ds4.c @@ -15479,6 +15479,23 @@ static bool metal_graph_set_active_tier_no_copy(ds4_gpu_graph *g, int tier) { return true; } +/* SSD streaming producer/consumer device re-affirmation. The per-layer + * dispatcher (metal_graph_set_active_tier_*) sets the CUDA device to the + * layer's home tier at layer entry, but the ambient device can drift back + * to tier 0 across the attention/router kernels that run before the SSD + * producer (the resident path tolerates drift because each kernel re-sets + * its own device; SSD allocation does not). Call this right before any + * SSD producer stages data for layer `il` so ssd_current() resolves to the + * layer's home tier and the staging lands on the right device. No-op for + * single-tier (placement == NULL). */ +static inline void metal_graph_ssd_assert_tier_device(const ds4_gpu_graph *g, + uint32_t il) { + if (g && g->placement) { + const int tier = g->placement[il + 1]; + if (tier >= 0) (void)ds4_gpu_set_current_device(tier); + } +} + /* Upstream: --power N GPU duty-cycle throttling helpers. The single-tier * --power=100 path is a no-op; multi-tier inherits the same helpers via * graph_power_note_prefill_layer / graph_power_note_decode_token which we @@ -20507,6 +20524,7 @@ static bool metal_graph_decode_set_hash_selected_override( } const uint64_t gate_expert_bytes = gate_tensor_bytes / DS4_N_EXPERT; const uint64_t down_expert_bytes = down_tensor_bytes / DS4_N_EXPERT; + metal_graph_ssd_assert_tier_device(g, il); const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, @@ -20687,6 +20705,7 @@ static bool metal_graph_decode_selected_readahead_override( DS4_N_EXPERT_USED) == 0) { return false; } + metal_graph_ssd_assert_tier_device(g, il); const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, @@ -20750,6 +20769,7 @@ static bool metal_graph_decode_cuda_selected_load( const double t_read = profile ? now_sec() : 0.0; if (ok) { + metal_graph_ssd_assert_tier_device(g, il); const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, @@ -20831,6 +20851,7 @@ static bool metal_graph_cuda_stream_prefill_batch_selected_load( n_ids64 * sizeof(selected_ids[0])) != 0; const double t_read = profile ? now_sec() : 0.0; if (ok) { + metal_graph_ssd_assert_tier_device(g, il); const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, @@ -20887,6 +20908,12 @@ typedef struct metal_graph_selected_async_load { uint64_t gate_expert_bytes; uint64_t down_expert_bytes; int32_t selected_ids[DS4_MAX_EXPERT_USED]; + /* Home tier (logical) of the layer this job is staging for. The worker + * thread is a long-lived service pthread that defaults to device 0; SSD + * streaming requires cudaSetDevice to the tier that owns the layer + * BEFORE the load runs, so ssd_current() in ds4_cuda.cu resolves to the + * right per-tier cache. -1 means "no preference" (single-tier path). */ + int home_tier; } metal_graph_selected_async_load; static pthread_mutex_t g_metal_graph_selected_async_load_mutex = @@ -20904,6 +20931,12 @@ static metal_graph_selected_async_load g_metal_graph_selected_async_load_job; static void metal_graph_selected_async_load_run( metal_graph_selected_async_load *job) { job->ok = false; + /* Worker thread defaults to device 0; SSD streaming needs the device + * set to the tier that owns this layer so ssd_current() resolves to the + * right per-tier cache. home_tier is -1 in single-tier (no-op). */ + if (job->home_tier >= 0) { + (void)ds4_gpu_set_current_device(job->home_tier); + } if (!job->router_selected || !job->model || !job->layer || DS4_N_EXPERT == 0 || DS4_N_EXPERT > DS4_MAX_EXPERT || @@ -20948,6 +20981,7 @@ static void metal_graph_selected_async_load_run( } } job->ids_ok = true; + /* home_tier was applied at run() entry; no graph pointer in the job. */ const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(job->model, job->layer, @@ -21024,7 +21058,8 @@ static DS4_MAYBE_UNUSED bool metal_graph_selected_async_load_start_tensor( uint32_t il, uint64_t event_value, uint64_t gate_expert_bytes, - uint64_t down_expert_bytes) { + uint64_t down_expert_bytes, + int home_tier) { if (!job || !router_selected || event_value == 0) return false; if (!metal_graph_selected_async_load_ensure_worker()) return false; memset(job, 0, sizeof(*job)); @@ -21035,6 +21070,7 @@ static DS4_MAYBE_UNUSED bool metal_graph_selected_async_load_start_tensor( job->event_value = event_value; job->gate_expert_bytes = gate_expert_bytes; job->down_expert_bytes = down_expert_bytes; + job->home_tier = home_tier; pthread_mutex_lock(&g_metal_graph_selected_async_load_mutex); if (g_metal_graph_selected_async_load_has_job || @@ -21060,6 +21096,14 @@ static DS4_MAYBE_UNUSED bool metal_graph_selected_async_load_start( uint64_t event_value, uint64_t gate_expert_bytes, uint64_t down_expert_bytes) { + /* Set home_tier on the caller's job struct BEFORE _start_tensor copies it + * into the shared slot (under the mutex). Setting it after would race with + * the worker waking up and copying the job back out. */ + int home_tier = -1; + if (g && g->placement) { + const int tier = g->placement[il + 1]; + home_tier = (tier >= 0) ? tier : 0; + } return metal_graph_selected_async_load_start_tensor( job, g ? metal_graph_router_selected(g) : NULL, @@ -21068,7 +21112,8 @@ static DS4_MAYBE_UNUSED bool metal_graph_selected_async_load_start( il, event_value, gate_expert_bytes, - down_expert_bytes); + down_expert_bytes, + home_tier); } static bool metal_graph_selected_async_load_finish( @@ -21152,6 +21197,7 @@ static void rocm_graph_batch_selected_async_load_run( return; } } + /* home_tier was applied at run() entry; no graph pointer in the job. */ const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(job->model, job->layer, @@ -23433,6 +23479,7 @@ static bool metal_graph_encode_decode_layer_phase( /* The worker read valid ids but could not stage the load * (it is not allowed to wait on in-flight cache entries). * This thread is, so retry the same load synchronously. */ + metal_graph_ssd_assert_tier_device(g, il); const ds4_gpu_stream_expert_table retry_table = graph_stream_expert_table_make(model, layer, @@ -23462,6 +23509,7 @@ static bool metal_graph_encode_decode_layer_phase( ds4_gpu_routed_moe_set_selected_override(selected_ids, DS4_N_EXPERT_USED) != 0; if (ok) { + metal_graph_ssd_assert_tier_device(g, il); const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, @@ -29533,6 +29581,7 @@ static bool metal_graph_seed_streaming_expert_cache_from_prefill( } const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; + metal_graph_ssd_assert_tier_device(g, il); const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, @@ -29666,6 +29715,7 @@ static bool metal_graph_seed_streaming_expert_cache_from_hotlist( } const uint64_t gate_expert_bytes = layer->ffn_gate_exps->dim[1] * gate_row_bytes; const uint64_t down_expert_bytes = layer->ffn_down_exps->dim[1] * down_row_bytes; + metal_graph_ssd_assert_tier_device(g, il); const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, @@ -33004,6 +33054,12 @@ static bool metal_graph_prefill_layer_major( } if (!split_commands) { + /* Multi-tier: the batch embedding kernel reads prefill_tokens (staged + * on emb_tier) and the token_embd weight (staged on emb_tier); it + * must run on emb_tier's device. After a decode, active_tier is left + * on the head tier, so repoint at emb_tier (no boundary-hop copy is + * needed: the embed overwrites batch_cur_hc in full). No-op single-tier. */ + if (g->placement && g->active_tier != g->emb_tier) g->active_tier = g->emb_tier; ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), metal_graph_prefill_tokens(g), model, @@ -33161,6 +33217,10 @@ static bool metal_graph_prefill_layer_major( #endif double t_layer0 = (profile || throttle) ? now_sec() : 0.0; + /* Multi-tier: see comment at the !split_commands branch above — repoint + * active_tier at emb_tier before the embed upload so the kernel runs on + * the device that owns prefill_tokens + token_embd. No-op single-tier. */ + if (g->placement && g->active_tier != g->emb_tier) g->active_tier = g->emb_tier; ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), metal_graph_prefill_tokens(g), model, @@ -33836,6 +33896,9 @@ static bool metal_graph_verify_suffix_tops_impl( const double upload_t0 = timing ? now_sec() : 0.0; bool ok = metal_graph_upload_prompt_tokens(metal_graph_prefill_tokens(g), prompt, start, n_tokens); + /* Multi-tier: repoint active_tier at emb_tier so the embed kernel runs on + * the device that owns prefill_tokens + token_embd. No-op single-tier. */ + if (g->placement && g->active_tier != g->emb_tier) g->active_tier = g->emb_tier; if (ok) ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), metal_graph_prefill_tokens(g), model, @@ -40070,7 +40133,8 @@ static bool glm_graph_encode_sparse_ffn_one( il, selected_event, gate_out * gate_row_bytes, - down_out * down_row_bytes); + down_out * down_row_bytes, + -1); async_path_profiled = async_profile && async_load_started; } if (async_profile) { @@ -53981,6 +54045,31 @@ static void model_warm_weights_sharded(const ds4_model *m, * * Tensor names live in ds4_str slices (ptr+len), NOT NUL-terminated. * We bound every comparison by name.len to avoid out-of-buffer reads. */ + +/* Returns true if the tensor is a routed-expert weight that SSD streaming + * serves from host/disk instead of resident VRAM. When SSD streaming is + * enabled these tensors MUST be excluded from the multi-tier placement + * budget (they are not loaded resident — they stream per-tier at runtime + * via the per-tier SSD cache). The names are the flat ds4_tensor slice form + * of "blk..{ffn_gate_exps,ffn_up_exps,ffn_down_exps}.weight" (and the + * mtp.* mirrors). Suffix-matched so the "blk.." / "mtp.0." prefix does + * not matter. */ +static bool tensor_is_ssd_streamed_routed_expert(const ds4_tensor *t) { + const char *p = t->name.ptr; + int n = (int)t->name.len; + if (n <= 0 || !p) return false; + static const char k_gate[] = ".ffn_gate_exps.weight"; + static const char k_up[] = ".ffn_up_exps.weight"; + static const char k_down[] = ".ffn_down_exps.weight"; + const size_t ng = sizeof(k_gate) - 1; + const size_t nu = sizeof(k_up) - 1; + const size_t nd = sizeof(k_down) - 1; + if ((size_t)n >= ng && memcmp(p + n - ng, k_gate, ng) == 0) return true; + if ((size_t)n >= nu && memcmp(p + n - nu, k_up, nu) == 0) return true; + if ((size_t)n >= nd && memcmp(p + n - nd, k_down, nd) == 0) return true; + return false; +} + static int tensor_to_entry(const ds4_tensor *t, int n_layer) { const char *p = t->name.ptr; int n = (int)t->name.len; @@ -54067,11 +54156,25 @@ static bool engine_cuda_tp_output_env_requested(void); static int engine_compute_entry_bytes(const ds4_engine *e, size_t *out) { const int n_entries = DS4_N_LAYER + 2; const bool cuda_tp_ep = engine_cuda_tp_ep_requested(e); + /* When SSD streaming is on AND we are NOT using CUDA tensor parallelism, + * the routed-expert weights (the bulk of an MoE model) are served from + * host/disk by the per-tier SSD caches, not loaded resident. Exclude + * them from the placement budget so the packer sizes only dense weights + * + KV; the per-tier expert cache is grown lazily at runtime out of each + * tier's free VRAM. Without this, the packer refuses CPU-spill for a + * model whose routed experts were never going to be resident anyway. + * Under CUDA TP the experts are sharded resident, so they must stay in + * the budget. */ + const bool ssd_excludes_routed = + e->ssd_streaming && !cuda_tp_ep; for (int i = 0; i < n_entries; i++) out[i] = 0; for (uint64_t i = 0; i < e->model.n_tensors; i++) { const ds4_tensor *t = &e->model.tensors[i]; if (t->bytes == 0) continue; + if (ssd_excludes_routed && tensor_is_ssd_streamed_routed_expert(t)) { + continue; + } int entry = tensor_to_entry(t, DS4_N_LAYER); if (entry < 0 || entry >= n_entries) entry = 0; if (cuda_tp_ep && engine_cuda_tp_output_env_requested() && @@ -54352,6 +54455,60 @@ static int engine_compute_cuda_ep_placement( return 0; } +/* Rewrite e->placement[] for SSD streaming: split the transformer layer + * range (entries 1..DS4_N_LAYER) into CONTIGUOUS ranges across the + * cfg->n_gpus GPUs in proportion to each GPU's vram_bytes share. Token + * embedding (entry 0) stays on tier 0; the output head (entry n_layers+1) + * follows the last tier that owns a transformer layer. + * + * Why: with SSD streaming the routed-expert weights (the bulk of an MoE + * model) are excluded from the resident budget by engine_compute_entry_ + * bytes, so the small dense footprint collapses onto tier 0 and the other + * requested GPUs sit idle. But routed experts stream per-tier from SSD, + * so the layer RANGE — and thus the streaming work — must be balanced + * across every GPU the user requested. Balancing the layer range shares + * the streaming work evenly. + * + * Dense weights (attention + shared expert + norms) are small relative to + * any reasonable --gpu-vram budget, so this contiguous split never spills + * to CPU. e->n_placement_entries and e->multi_tier are updated in place. */ +static void engine_balance_placement_for_ssd(ds4_engine *e, + const ds4_gpu_config *cfg) { + const int n_gpus = cfg->n_gpus; + const int n_layers = DS4_N_LAYER; + if (n_gpus < 2 || n_layers < 1) return; + + double total = 0.0; + for (int d = 0; d < n_gpus; d++) total += (double)cfg->vram_bytes[d]; + if (total <= 0) return; + + /* Entry 0 (token embedding) on tier 0. */ + e->placement[0] = 0; + + /* Walk transformer layers in order, assigning each to the tier whose + * cumulative budget share covers this layer's fractional position. This + * yields contiguous, budget-proportional ranges. */ + int t = 0; + double cumulative = (double)cfg->vram_bytes[0]; + for (int i = 1; i <= n_layers; i++) { + const double layer_frac = (double)i / (double)n_layers; + while (t + 1 < n_gpus && layer_frac > cumulative / total) { + t++; + cumulative += (double)cfg->vram_bytes[t]; + } + e->placement[i] = t; + } + + /* Output head (entry n_layers+1) on the highest tier that owns a layer. */ + int last_tier = 0; + for (int i = 1; i <= n_layers; i++) { + if (e->placement[i] > last_tier) last_tier = e->placement[i]; + } + e->placement[n_layers + 1] = last_tier; + e->n_placement_entries = n_layers + 2; + e->multi_tier = last_tier > 0 ? 1 : 0; +} + /* Phase A: classify multi-tier on a freshly-opened engine (model loaded, * weights bound). Pure CPU — no GPU init required. Sets e->multi_tier, * e->n_placement_entries, e->placement[], and e->gpu_cfg. Returns 0 on @@ -54436,6 +54593,20 @@ static int engine_classify_multi_tier(ds4_engine *e, const ds4_gpu_config *cfg) e->n_placement_entries = DS4_N_LAYER + 2; engine_adjust_output_head_for_cuda_tp(e, entry_bytes); + /* SSD streaming rebalance (non-TP only): the greedy packer fills tier 0 + * first, so with routed experts excluded from the resident budget + * (engine_compute_entry_bytes) the small dense footprint collapses onto + * tier 0 and the other requested GPUs sit idle. But routed experts + * stream per-tier from SSD, so the layer RANGE — and thus the streaming + * work — must be balanced across every GPU the user requested. Rewrite + * placement[] with a contiguous, budget-proportional split. Dense + * weights are tiny relative to the per-tier budget, so this never + * causes CPU-spill. Under CUDA TP this is skipped: TP needs the + * sharded placement produced by engine_compute_cuda_ep_placement. */ + if (e->ssd_streaming && !cuda_tp_ep && e->gpu_cfg.n_gpus >= 2) { + engine_balance_placement_for_ssd(e, &e->gpu_cfg); + } + int first_tier = e->placement[0]; int multi_tier = 0; for (int i = 1; i < e->n_placement_entries; i++) { @@ -54555,9 +54726,18 @@ static int engine_install_per_device_caches(ds4_engine *e) { : "ds4: CUDA output TP cache duplication enabled for output head (%u ways)\n", output_tp_ways); } + /* SSD streaming (non-TP): routed-expert weights stream per-tier from + * the host mmap via the per-tier SSD caches; do NOT register them + * resident on each device (would blow the budget and defeat streaming). + * Only the dense weights (attention, shared expert, norms) are loaded + * resident. Under CUDA TP the experts are sharded resident, so keep + * them. */ + const bool ssd_skip_routed = + e->ssd_streaming && !cuda_tp_ep; for (uint64_t i = 0; i < e->model.n_tensors; i++) { const ds4_tensor *t = &e->model.tensors[i]; if (t->bytes == 0) continue; + if (ssd_skip_routed && tensor_is_ssd_streamed_routed_expert(t)) continue; int entry = tensor_to_entry(t, DS4_N_LAYER); if (entry < 0 || entry >= e->n_placement_entries) entry = 0; int logical_tier = e->placement[entry]; @@ -55447,9 +55627,10 @@ static int ds4_engine_open_internal(ds4_engine **out, *out = NULL; return 1; } - if (e->ssd_streaming && e->multi_tier) { + if (e->ssd_streaming && e->multi_tier && e->cuda_tensor_parallel) { fprintf(stderr, - "ds4: --ssd-streaming is not compatible with multi-GPU placement\n"); + "ds4: --ssd-streaming is not compatible with CUDA tensor parallelism " + "(TP keeps half the routed experts resident per tier; SSD streams them)\n"); ds4_engine_close(e); *out = NULL; return 1; @@ -55635,6 +55816,27 @@ static int ds4_engine_open_internal(ds4_engine **out, *out = NULL; return 1; } + /* SSD streaming: now that ds4_gpu_init_multi has populated g_gpu[], + * enable SSD per tier. ds4_gpu_set_ssd_streaming loops every tier, + * cudaSetDevice()ing to each so ssd_current() in the release/reset + * helpers resolves to g_ssd[t]; the per-tier caches size lazily + * on first use. Mirrors the single-tier sequence below. */ + ds4_gpu_set_glm_model(DS4_MODEL_FAMILY == DS4_MODEL_FAMILY_GLM_DSA); + ds4_gpu_set_ssd_streaming(e->ssd_streaming); + if (!ds4_engine_configure_streaming_auto_cache(e)) { + ds4_engine_close(e); + *out = NULL; + return 1; + } + ds4_gpu_set_streaming_expert_cache_budget(e->ssd_streaming_cache_experts); + if (e->ssd_streaming) { + /* Pin the expert cache's slab size class to the model's uniform + * per-expert bytes (mirrors the single-tier branch). */ + uint64_t slab_expert_bytes = 0; + if (ds4_streaming_routed_expert_bytes(&e->weights, &slab_expert_bytes)) { + ds4_gpu_set_streaming_expert_cache_expert_bytes(slab_expert_bytes); + } + } /* GPU-only multi-tier execution is now wired up * (B2-B6: per-tier graph allocation, dispatch loops, boundary * copies). CPU-spill placements were rejected by diff --git a/ds4_cuda.cu b/ds4_cuda.cu index aaa4df113..3f8293dc1 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -147,31 +147,56 @@ typedef struct { ds4_gpu_tensor slot_selected_tensor; } cuda_stream_selected_cache; -static cuda_stream_selected_cache g_stream_selected_cache; +/* Per-tier SSD streaming state. SSD streaming was originally single-device: + * one selected-expert cache, one pinned staging pool, one upload stream, all + * on device 0. In multi-GPU mode q4k's per-layer dispatch routes each layer + * to a tier via cudaSetDevice(g_gpu[tier].device_id) (ds4.c helpers around + * metal_graph_set_active_tier_*), and each tier may own a contiguous layer + * range. A single SSD cache therefore cannot serve two devices — CUDA + * streams are device-bound and cudaFree must run on the allocating device. + * + * Style B resolution (no ds4.c call-site signature changes): every SSD + * producer/consumer resolves the active tier itself via cudaGetDevice() + * (see ssd_current()). Per-layer dispatch has already cudaSetDevice()'d to + * the tier that owns the current layer before any SSD function runs, so + * ssd_current() returns the tier that owns the layer about to execute. */ +struct ds4_ssd_ctx { + cuda_stream_selected_cache selected_cache; + void *stage_raw[4]; + void *stage[4]; + cudaEvent_t stage_event[4]; + uint64_t stage_bytes; + cudaStream_t upload_stream; +}; +static ds4_ssd_ctx g_ssd[DS4_MAX_GPUS]; + +/* Forward declaration: body is defined after g_gpu[] / g_n_gpus below. */ +static ds4_ssd_ctx *ssd_current(void); static void cuda_stream_selected_cache_invalidate(void) { - g_stream_selected_cache.valid = 0; + ssd_current()->selected_cache.valid = 0; } static void cuda_stream_selected_cache_release(void) { - const int tier = g_stream_selected_cache.logical_tier; + ds4_ssd_ctx *ssd = ssd_current(); + const int tier = ssd->selected_cache.logical_tier; if (tier >= 0 && tier < g_n_gpus) { (void)ds4_gpu_set_current_device(tier); } - if (g_stream_selected_cache.gate_ptr) { - (void)cudaFree(g_stream_selected_cache.gate_ptr); + if (ssd->selected_cache.gate_ptr) { + (void)cudaFree(ssd->selected_cache.gate_ptr); } - if (g_stream_selected_cache.up_ptr) { - (void)cudaFree(g_stream_selected_cache.up_ptr); + if (ssd->selected_cache.up_ptr) { + (void)cudaFree(ssd->selected_cache.up_ptr); } - if (g_stream_selected_cache.down_ptr) { - (void)cudaFree(g_stream_selected_cache.down_ptr); + if (ssd->selected_cache.down_ptr) { + (void)cudaFree(ssd->selected_cache.down_ptr); } - if (g_stream_selected_cache.slot_selected_ptr) { - (void)cudaFree(g_stream_selected_cache.slot_selected_ptr); + if (ssd->selected_cache.slot_selected_ptr) { + (void)cudaFree(ssd->selected_cache.slot_selected_ptr); } - memset(&g_stream_selected_cache, 0, sizeof(g_stream_selected_cache)); - g_stream_selected_cache.logical_tier = -1; + memset(&ssd->selected_cache, 0, sizeof(ssd->selected_cache)); + ssd->selected_cache.logical_tier = -1; } typedef struct { @@ -260,6 +285,22 @@ static inline int ds4_tensor_device_idx(const ds4_gpu_tensor *t) { return d; } +/* Resolve the SSD context for the currently-active CUDA device. Per-layer + * dispatch (ds4.c) calls ds4_gpu_set_current_device(tier) — which maps tier + * to g_gpu[tier].device_id and cudaSetDevice()s — BEFORE any SSD + * producer/consumer runs, so cudaGetDevice() returns the device that owns + * the layer about to execute. Falls back to tier 0 when the runtime is + * unavailable (e.g. CPU-only build path) or the device is not in g_gpu[] + * (e.g. the default context before init). */ +static ds4_ssd_ctx *ssd_current(void) { + int dev = 0; + if (cudaGetDevice(&dev) != cudaSuccess) return &g_ssd[0]; + for (int t = 0; t < g_n_gpus; t++) { + if (g_gpu[t].device_id == dev) return &g_ssd[t]; + } + return &g_ssd[0]; +} + /* Debug/override flags are read once per CUDA init. The hot decode path calls * the xdev helpers many times per token, so they must not re-enter getenv(). */ static void cuda_xdev_env_refresh(void) { @@ -407,11 +448,8 @@ static void *g_model_stage_raw[4]; static void *g_model_stage[4]; static cudaEvent_t g_model_stage_event[4]; static uint64_t g_model_stage_bytes; -static void *g_stream_selected_stage_raw[4]; -static void *g_stream_selected_stage[4]; -static cudaEvent_t g_stream_selected_stage_event[4]; -static uint64_t g_stream_selected_stage_bytes; -static cudaStream_t g_stream_selected_upload_stream; +/* The per-tier SSD staging pool (stage_raw/stage/stage_event/stage_bytes + + * upload_stream) lives in g_ssd[tier] (see struct ds4_ssd_ctx above). */ static int cuda_ok(cudaError_t err, const char *what); static const char *cuda_model_range_ptr_from_fd( @@ -1596,29 +1634,35 @@ static int cuda_model_stage_read(void *stage, uint64_t stage_bytes, } static void cuda_stream_selected_stage_release(void) { + ds4_ssd_ctx *ssd = ssd_current(); for (size_t i = 0; i < 4; i++) { - if (g_stream_selected_stage_event[i]) { - (void)cudaEventDestroy(g_stream_selected_stage_event[i]); - g_stream_selected_stage_event[i] = NULL; + if (ssd->stage_event[i]) { + (void)cudaEventDestroy(ssd->stage_event[i]); + ssd->stage_event[i] = NULL; } - if (g_stream_selected_stage_raw[i]) { - (void)cudaFreeHost(g_stream_selected_stage_raw[i]); - g_stream_selected_stage_raw[i] = NULL; - g_stream_selected_stage[i] = NULL; + if (ssd->stage_raw[i]) { + (void)cudaFreeHost(ssd->stage_raw[i]); + ssd->stage_raw[i] = NULL; + ssd->stage[i] = NULL; } } - g_stream_selected_stage_bytes = 0; - if (g_stream_selected_upload_stream) { - (void)cudaStreamDestroy(g_stream_selected_upload_stream); - g_stream_selected_upload_stream = NULL; + ssd->stage_bytes = 0; + if (ssd->upload_stream) { + (void)cudaStreamDestroy(ssd->upload_stream); + ssd->upload_stream = NULL; } } static int cuda_stream_selected_stage_pool_alloc(uint64_t bytes) { - if (g_stream_selected_stage_bytes >= bytes) return 1; + ds4_ssd_ctx *ssd = ssd_current(); + if (ssd->stage_bytes >= bytes) return 1; cuda_stream_selected_stage_release(); + /* Upload stream is device-bound: create it on the CURRENT device (the + * per-layer dispatcher has already cudaSetDevice()'d to the tier that + * owns this layer). Stash it in this tier's ssd ctx. */ + ssd = ssd_current(); cudaError_t err = cudaStreamCreateWithFlags( - &g_stream_selected_upload_stream, cudaStreamNonBlocking); + &ssd->upload_stream, cudaStreamNonBlocking); if (err != cudaSuccess) { fprintf(stderr, "ds4: CUDA streaming selected upload stream creation failed: %s\n", @@ -1627,7 +1671,7 @@ static int cuda_stream_selected_stage_pool_alloc(uint64_t bytes) { return 0; } for (size_t i = 0; i < 4; i++) { - err = cudaMallocHost(&g_stream_selected_stage_raw[i], (size_t)bytes); + err = cudaMallocHost(&ssd->stage_raw[i], (size_t)bytes); if (err != cudaSuccess) { fprintf(stderr, "ds4: CUDA streaming selected staging allocation failed: %s\n", @@ -1636,9 +1680,9 @@ static int cuda_stream_selected_stage_pool_alloc(uint64_t bytes) { cuda_stream_selected_stage_release(); return 0; } - g_stream_selected_stage[i] = cuda_align_ptr( - g_stream_selected_stage_raw[i], g_model_direct_align); - err = cudaEventCreateWithFlags(&g_stream_selected_stage_event[i], + ssd->stage[i] = cuda_align_ptr( + ssd->stage_raw[i], g_model_direct_align); + err = cudaEventCreateWithFlags(&ssd->stage_event[i], cudaEventDisableTiming); if (err != cudaSuccess) { fprintf(stderr, @@ -1649,7 +1693,7 @@ static int cuda_stream_selected_stage_pool_alloc(uint64_t bytes) { return 0; } } - g_stream_selected_stage_bytes = bytes; + ssd->stage_bytes = bytes; return 1; } @@ -1679,6 +1723,9 @@ static int cuda_model_copy_to_device_streamed( chunk + (g_model_direct_align > 1 ? g_model_direct_align : 1); if (!cuda_stream_selected_stage_pool_alloc(stage_bytes)) return 0; + /* pool_alloc just sized this tier's stage pool + upload stream. Cache the + * tier ctx once: device does not change for the lifetime of this copy. */ + ds4_ssd_ctx *ssd = ssd_current(); uint64_t copied = 0; uint64_t chunk_idx = 0; while (copied < bytes) { @@ -1686,7 +1733,7 @@ static int cuda_model_copy_to_device_streamed( const uint64_t bi = chunk_idx % 4u; cudaError_t err; if (chunk_idx >= 4u) { - err = cudaEventSynchronize(g_stream_selected_stage_event[bi]); + err = cudaEventSynchronize(ssd->stage_event[bi]); if (err != cudaSuccess) { fprintf(stderr, "ds4: CUDA streaming selected staging wait failed for %s: %s\n", @@ -1696,8 +1743,8 @@ static int cuda_model_copy_to_device_streamed( } } const char *payload = NULL; - if (!cuda_model_stage_read(g_stream_selected_stage[bi], - g_stream_selected_stage_bytes, + if (!cuda_model_stage_read(ssd->stage[bi], + ssd->stage_bytes, offset + copied, n, &payload)) { fprintf(stderr, "ds4: CUDA streaming selected read failed for %s at %.2f MiB: %s\n", @@ -1707,7 +1754,7 @@ static int cuda_model_copy_to_device_streamed( } err = cudaMemcpyAsync(dst + copied, payload, (size_t)n, cudaMemcpyHostToDevice, - g_stream_selected_upload_stream); + ssd->upload_stream); if (err != cudaSuccess) { fprintf(stderr, "ds4: CUDA streaming selected copy failed for %s at %.2f MiB: %s\n", @@ -1716,8 +1763,8 @@ static int cuda_model_copy_to_device_streamed( (void)cudaGetLastError(); return 0; } - err = cudaEventRecord(g_stream_selected_stage_event[bi], - g_stream_selected_upload_stream); + err = cudaEventRecord(ssd->stage_event[bi], + ssd->upload_stream); if (err != cudaSuccess) { fprintf(stderr, "ds4: CUDA streaming selected staging record failed for %s: %s\n", @@ -1733,7 +1780,7 @@ static int cuda_model_copy_to_device_streamed( } const cudaError_t err = - cudaStreamSynchronize(g_stream_selected_upload_stream); + cudaStreamSynchronize(ssd->upload_stream); if (err != cudaSuccess) { fprintf(stderr, "ds4: CUDA streaming selected upload sync failed for %s: %s\n", @@ -2250,6 +2297,11 @@ extern "C" void ds4_gpu_cleanup(void) { c->scratch = NULL; c->scratch_bytes = 0; } + /* Per-tier SSD teardown: free this tier's selected-expert cache and + * staging pool on the device that owns them. cudaSetDevice(c->device_id) + * above makes ssd_current() resolve to g_ssd[i]. */ + cuda_stream_selected_cache_release(); + cuda_stream_selected_stage_release(); } for (int i = 0; i < DS4_MAX_GPUS; i++) { for (int j = 0; j < DS4_MAX_GPUS; j++) { @@ -2260,8 +2312,6 @@ extern "C" void ds4_gpu_cleanup(void) { } } } - cuda_stream_selected_cache_release(); - cuda_stream_selected_stage_release(); g_n_gpus = 0; g_cublas_ready = 0; @@ -11532,11 +11582,23 @@ extern "C" int ds4_gpu_embed_token_hc_tensor(ds4_gpu_tensor *out_hc, const void uint64_t weight_bytes = (uint64_t)n_vocab * n_embd * sizeof(uint16_t); if (weight_offset > model_size || weight_bytes > model_size - weight_offset) return 0; const int logical_tier = ds4_tensor_device_idx(out_hc); - const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "token_embd"); - if (!wptr) return 0; - uint32_t n = n_embd * n_hc; - embed_token_hc_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr, (const unsigned short *)wptr, token, n_embd, n_hc); - return cuda_ok(cudaGetLastError(), "embed token launch"); + /* Pin the kernel to out_hc's home device: cuda_resolve_weight_ptr returns + * a device-local pointer valid only on logical_tier's physical device. + * Without this the kernel can launch on a stale ambient device left by a + * prior decode on another tier and take an illegal memory access. */ + const int target_dev = (logical_tier >= 0 && logical_tier < g_n_gpus) + ? g_gpu[logical_tier].device_id : 0; + int ok = 0; + WITH_DEVICE(target_dev) { + const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, weight_bytes, logical_tier, "token_embd"); + if (!wptr) { ok = 0; } + else { + uint32_t n = n_embd * n_hc; + embed_token_hc_kernel<<<(n + 255) / 256, 256>>>((float *)out_hc->ptr, (const unsigned short *)wptr, token, n_embd, n_hc); + ok = cuda_ok(cudaGetLastError(), "embed token launch"); + } + } + return ok; } extern "C" int ds4_gpu_embed_tokens_hc_tensor( @@ -11557,18 +11619,31 @@ extern "C" int ds4_gpu_embed_tokens_hc_tensor( return 0; } const int logical_tier = ds4_tensor_device_idx(out_hc); - const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, - (uint64_t)n_vocab * n_embd * sizeof(uint16_t), - logical_tier, - "token_embd"); - if (!wptr) return 0; - uint64_t n = (uint64_t)n_tokens * n_hc * n_embd; - embed_tokens_hc_kernel<<<(n + 255) / 256, 256>>>( - (float *)out_hc->ptr, - (const int32_t *)tokens_t->ptr, - (const __half *)wptr, - n_vocab, n_tokens, n_embd, n_hc); - return cuda_ok(cudaGetLastError(), "embed tokens launch"); + /* Pin the kernel to out_hc's home device: cuda_resolve_weight_ptr returns + * a device-local pointer valid only on logical_tier's physical device, + * and out_hc/tokens_t live there too. Without this the kernel launched on + * the ambient device (left there by a prior decode on another tier) and + * the first prefill chunk took an illegal memory access. */ + const int target_dev = (logical_tier >= 0 && logical_tier < g_n_gpus) + ? g_gpu[logical_tier].device_id : 0; + int ok = 0; + WITH_DEVICE(target_dev) { + const char *wptr = cuda_resolve_weight_ptr(model_map, weight_offset, + (uint64_t)n_vocab * n_embd * sizeof(uint16_t), + logical_tier, + "token_embd"); + if (!wptr) { ok = 0; } + else { + uint64_t n = (uint64_t)n_tokens * n_hc * n_embd; + embed_tokens_hc_kernel<<<(n + 255) / 256, 256>>>( + (float *)out_hc->ptr, + (const int32_t *)tokens_t->ptr, + (const __half *)wptr, + n_vocab, n_tokens, n_embd, n_hc); + ok = cuda_ok(cudaGetLastError(), "embed tokens launch"); + } + } + return ok; } static int indexer_scores_launch( @@ -20859,25 +20934,37 @@ static int routed_moe_launch( } const uint64_t required_slot_count = (uint64_t)n_tokens * n_expert; const int logical_tier = ds4_tensor_device_idx(out); + /* SSD selected-expert cache: when SSD streaming has staged a compact + * (gate, up, down) slab for THIS layer's routed experts, read directly + * from the staged device pointers instead of going through the model + * range / selective-cache resolver. Multi-tier aware: ssd_current() + * resolves the tier whose layer is executing. The per-layer dispatcher + * sets the device, but ambient device can drift across the attention/ + * router kernels that ran before this consumer; re-affirm the OUT + * tensor's home tier so ssd_current() reads THIS layer's slab. */ + if (g_n_gpus > 1 && logical_tier >= 0 && logical_tier < g_n_gpus) { + (void)cudaSetDevice(g_gpu[logical_tier].device_id); + } + cuda_stream_selected_cache *sc = &ssd_current()->selected_cache; const int use_stream_selected_cache = allow_streaming && g_ssd_streaming_mode && - g_stream_selected_cache.valid && - g_stream_selected_cache.logical_tier == logical_tier && - g_stream_selected_cache.model_map == model_map && - g_stream_selected_cache.layer == layer_index && - g_stream_selected_cache.n_total_expert == n_total_expert && - g_stream_selected_cache.slot_count >= required_slot_count && - g_stream_selected_cache.gate_offset == gate_offset && - g_stream_selected_cache.up_offset == up_offset && - g_stream_selected_cache.down_offset == down_offset && - g_stream_selected_cache.gate_expert_bytes == gate_expert_bytes && - g_stream_selected_cache.down_expert_bytes == down_expert_bytes && - g_stream_selected_cache.gate_ptr && - g_stream_selected_cache.up_ptr && - g_stream_selected_cache.down_ptr && - g_stream_selected_cache.slot_selected_tensor.ptr && - g_stream_selected_cache.slot_selected_tensor.bytes >= + sc->valid && + sc->logical_tier == logical_tier && + sc->model_map == model_map && + sc->layer == layer_index && + sc->n_total_expert == n_total_expert && + sc->slot_count >= required_slot_count && + sc->gate_offset == gate_offset && + sc->up_offset == up_offset && + sc->down_offset == down_offset && + sc->gate_expert_bytes == gate_expert_bytes && + sc->down_expert_bytes == down_expert_bytes && + sc->gate_ptr && + sc->up_ptr && + sc->down_ptr && + sc->slot_selected_tensor.ptr && + sc->slot_selected_tensor.bytes >= required_slot_count * sizeof(int32_t); if (g_ssd_streaming_mode && allow_streaming && !use_stream_selected_cache) { @@ -20887,18 +20974,18 @@ static int routed_moe_launch( return 0; } if (use_stream_selected_cache) { - selected = &g_stream_selected_cache.slot_selected_tensor; + selected = &sc->slot_selected_tensor; } const char *gate_w = use_stream_selected_cache ? - g_stream_selected_cache.gate_ptr : + sc->gate_ptr : cuda_resolve_weight_ptr(model_map, gate_offset, gate_bytes, logical_tier, "moe_gate"); const char *up_w = use_stream_selected_cache ? - g_stream_selected_cache.up_ptr : + sc->up_ptr : cuda_resolve_weight_ptr(model_map, up_offset, gate_bytes, logical_tier, "moe_up"); const char *down_w = use_stream_selected_cache ? - g_stream_selected_cache.down_ptr : + sc->down_ptr : cuda_resolve_weight_ptr(model_map, down_offset, down_bytes, logical_tier, "moe_down"); if (!gate_w || !up_w || !down_w) return 0; @@ -22918,9 +23005,10 @@ static int cuda_stream_selected_ensure_bytes( static int cuda_stream_selected_ensure_i32(uint64_t count) { if (count == 0 || count > UINT64_MAX / sizeof(int32_t)) return 0; const uint64_t bytes = count * sizeof(int32_t); + cuda_stream_selected_cache *sc = &ssd_current()->selected_cache; return cuda_stream_selected_ensure_bytes( - (char **)&g_stream_selected_cache.slot_selected_ptr, - &g_stream_selected_cache.slot_selected_capacity, + (char **)&sc->slot_selected_ptr, + &sc->slot_selected_capacity, bytes, "selected-id remap"); } @@ -22960,11 +23048,6 @@ static int cuda_stream_selected_cache_begin_load( slot_count == 0) { return 0; } - if (g_n_gpus != 1) { - fprintf(stderr, - "ds4: CUDA SSD streaming requires single-GPU placement\n"); - return 0; - } std::vector expert_to_slot; std::vector compact_ids; @@ -23001,26 +23084,39 @@ static int cuda_stream_selected_cache_begin_load( } const uint64_t gate_bytes = compact_count * table->gate_expert_bytes; const uint64_t down_bytes = compact_count * table->down_expert_bytes; - const int logical_tier = 0; - if (g_stream_selected_cache.logical_tier != logical_tier && - (g_stream_selected_cache.gate_ptr || - g_stream_selected_cache.up_ptr || - g_stream_selected_cache.down_ptr || - g_stream_selected_cache.slot_selected_ptr)) { + /* Per-tier selected-expert cache. The per-layer dispatcher has already + * cudaSetDevice()'d to the tier that owns this layer, so ssd_current() + * returns that tier's selected_cache. All ensures/copies below land on + * the SAME device — cudaMalloc inside ensure_bytes targets the current + * device, and the stage/upload path uses this tier's stream. Resolve + * logical_tier from the active device so the cache tags land correctly + * for whatever tier owns this layer. */ + int ambient_dev = 0; + (void)cudaGetDevice(&ambient_dev); + int logical_tier = 0; + for (int t = 0; t < g_n_gpus; t++) { + if (g_gpu[t].device_id == ambient_dev) { logical_tier = t; break; } + } + cuda_stream_selected_cache *sc = &ssd_current()->selected_cache; + if (sc->logical_tier != logical_tier && + (sc->gate_ptr || + sc->up_ptr || + sc->down_ptr || + sc->slot_selected_ptr)) { cuda_stream_selected_cache_release(); + sc = &ssd_current()->selected_cache; } - if (ds4_gpu_set_current_device(logical_tier) != 0 || - !cuda_stream_selected_ensure_bytes( - &g_stream_selected_cache.gate_ptr, - &g_stream_selected_cache.gate_capacity, + if (!cuda_stream_selected_ensure_bytes( + &sc->gate_ptr, + &sc->gate_capacity, gate_bytes, "gate experts") || !cuda_stream_selected_ensure_bytes( - &g_stream_selected_cache.up_ptr, - &g_stream_selected_cache.up_capacity, + &sc->up_ptr, + &sc->up_capacity, gate_bytes, "up experts") || !cuda_stream_selected_ensure_bytes( - &g_stream_selected_cache.down_ptr, - &g_stream_selected_cache.down_capacity, + &sc->down_ptr, + &sc->down_capacity, down_bytes, "down experts") || !cuda_stream_selected_ensure_i32(slot_count)) { cuda_stream_selected_cache_invalidate(); @@ -23038,17 +23134,17 @@ static int cuda_stream_selected_cache_begin_load( const uint64_t gate_dst = (uint64_t)i * table->gate_expert_bytes; const uint64_t down_dst = (uint64_t)i * table->down_expert_bytes; if (!cuda_model_copy_to_device_streamed( - g_stream_selected_cache.gate_ptr + gate_dst, + sc->gate_ptr + gate_dst, table->model_map, table->model_size, gate_src, table->gate_expert_bytes, "stream gate expert copy") || !cuda_model_copy_to_device_streamed( - g_stream_selected_cache.up_ptr + gate_dst, + sc->up_ptr + gate_dst, table->model_map, table->model_size, up_src, table->gate_expert_bytes, "stream up expert copy") || !cuda_model_copy_to_device_streamed( - g_stream_selected_cache.down_ptr + down_dst, + sc->down_ptr + down_dst, table->model_map, table->model_size, down_src, table->down_expert_bytes, "stream down expert copy")) { @@ -23056,7 +23152,7 @@ static int cuda_stream_selected_cache_begin_load( return 0; } } - if (!cuda_ok(cudaMemcpy(g_stream_selected_cache.slot_selected_ptr, + if (!cuda_ok(cudaMemcpy(sc->slot_selected_ptr, slot_ids.data(), (size_t)slot_count * sizeof(int32_t), cudaMemcpyHostToDevice), @@ -23065,24 +23161,24 @@ static int cuda_stream_selected_cache_begin_load( return 0; } - g_stream_selected_cache.logical_tier = logical_tier; - g_stream_selected_cache.model_map = table->model_map; - g_stream_selected_cache.layer = table->layer; - g_stream_selected_cache.n_total_expert = table->n_total_expert; - g_stream_selected_cache.slot_count = slot_count; - g_stream_selected_cache.compact_count = (uint32_t)compact_count; - g_stream_selected_cache.gate_offset = table->gate_offset; - g_stream_selected_cache.up_offset = table->up_offset; - g_stream_selected_cache.down_offset = table->down_offset; - g_stream_selected_cache.gate_expert_bytes = table->gate_expert_bytes; - g_stream_selected_cache.down_expert_bytes = table->down_expert_bytes; - g_stream_selected_cache.slot_selected_tensor.ptr = - g_stream_selected_cache.slot_selected_ptr; - g_stream_selected_cache.slot_selected_tensor.bytes = + sc->logical_tier = logical_tier; + sc->model_map = table->model_map; + sc->layer = table->layer; + sc->n_total_expert = table->n_total_expert; + sc->slot_count = slot_count; + sc->compact_count = (uint32_t)compact_count; + sc->gate_offset = table->gate_offset; + sc->up_offset = table->up_offset; + sc->down_offset = table->down_offset; + sc->gate_expert_bytes = table->gate_expert_bytes; + sc->down_expert_bytes = table->down_expert_bytes; + sc->slot_selected_tensor.ptr = + sc->slot_selected_ptr; + sc->slot_selected_tensor.bytes = (uint64_t)slot_count * sizeof(int32_t); - g_stream_selected_cache.slot_selected_tensor.owner = 0; - g_stream_selected_cache.slot_selected_tensor.device_id = logical_tier; - g_stream_selected_cache.valid = 1; + sc->slot_selected_tensor.owner = 0; + sc->slot_selected_tensor.device_id = logical_tier; + sc->valid = 1; return 1; } @@ -27373,8 +27469,24 @@ extern "C" void ds4_gpu_set_glm_model(bool enabled) { extern "C" void ds4_gpu_set_ssd_streaming(bool enabled) { g_ssd_streaming_mode = enabled ? 1 : 0; - cuda_stream_selected_cache_invalidate(); - if (!g_ssd_streaming_mode) cuda_stream_selected_cache_release(); + /* Multi-tier: invalidate every tier's selected_cache, and when disabling + * free each tier's caches + staging pool on the device that owns them. + * cudaSetDevice before the release helpers so ssd_current() resolves to + * g_ssd[t]. Single-tier path (g_n_gpus == 0 at the time of the first + * call) is handled by the fallback loop iteration over n == 1, which + * runs on whatever device is current — matching the legacy behavior. */ + const int n = g_n_gpus > 0 ? g_n_gpus : 1; + int prev = -1; + (void)cudaGetDevice(&prev); + for (int t = 0; t < n; t++) { + if (g_n_gpus > 0) (void)cudaSetDevice(g_gpu[t].device_id); + cuda_stream_selected_cache_invalidate(); + if (!g_ssd_streaming_mode) { + cuda_stream_selected_cache_release(); + cuda_stream_selected_stage_release(); + } + } + if (prev >= 0) (void)cudaSetDevice(prev); } extern "C" void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) { @@ -27390,15 +27502,25 @@ extern "C" uint32_t ds4_gpu_stream_expert_cache_configured_count(void) { } extern "C" uint32_t ds4_gpu_stream_expert_cache_current_count(void) { - return g_stream_selected_cache.valid ? - g_stream_selected_cache.compact_count : 0; + /* Status readout: report the active tier's resident count. Per-layer + * dispatch has set the device; ssd_current() returns the right tier. */ + cuda_stream_selected_cache *sc = &ssd_current()->selected_cache; + return sc->valid ? sc->compact_count : 0; } extern "C" void ds4_gpu_stream_expert_cache_reset_route_hotness(void) { } extern "C" void ds4_gpu_stream_expert_cache_release_resident(void) { - cuda_stream_selected_cache_release(); + /* Release resident caches on every tier (each on its own device). */ + const int n = g_n_gpus > 0 ? g_n_gpus : 1; + int prev = -1; + (void)cudaGetDevice(&prev); + for (int t = 0; t < n; t++) { + if (g_n_gpus > 0) (void)cudaSetDevice(g_gpu[t].device_id); + cuda_stream_selected_cache_release(); + } + if (prev >= 0) (void)cudaSetDevice(prev); } extern "C" int ds4_gpu_stream_expert_cache_seed_selected( From d2479cafcbe357135dba0d757843b929b32a4501 Mon Sep 17 00:00:00 2001 From: Anton Sokolchenko Date: Tue, 21 Jul 2026 12:48:07 +0000 Subject: [PATCH 2/4] Fix multi-tier SSD decode-streaming prefill: tier-device + thread-local 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 --- ds4.c | 103 ++++++++++++++++++++++++++++++++++++++++- ds4_cuda.cu | 131 +++++++++++++++++++++++++++++++++++++++++++++++++++- ds4_gpu.h | 4 ++ 3 files changed, 235 insertions(+), 3 deletions(-) diff --git a/ds4.c b/ds4.c index d578dc8c9..a1409d143 100644 --- a/ds4.c +++ b/ds4.c @@ -15452,6 +15452,7 @@ static bool metal_graph_set_active_tier_batch(ds4_gpu_graph *g, int tier, uint32 return true; } if (tier < 0 || tier >= DS4_MAX_GPUS) return false; + (void)ds4_gpu_debug_probe("set_active_tier_batch entry", -1, tier); if (tier == g->active_tier) return true; if (ds4_gpu_set_current_device(tier) != 0) return false; if (g->active_tier >= 0) { @@ -15464,6 +15465,7 @@ static bool metal_graph_set_active_tier_batch(ds4_gpu_graph *g, int tier, uint32 } } g->active_tier = tier; + (void)ds4_gpu_debug_probe("set_active_tier_batch done", -1, tier); return true; } @@ -15492,8 +15494,23 @@ static inline void metal_graph_ssd_assert_tier_device(const ds4_gpu_graph *g, uint32_t il) { if (g && g->placement) { const int tier = g->placement[il + 1]; - if (tier >= 0) (void)ds4_gpu_set_current_device(tier); + if (tier >= 0) { + (void)ds4_gpu_set_current_device(tier); + (void)ds4_gpu_debug_probe("ssd_assert_tier_device", (int)il, tier); + } + } +} + +/* Multi-tier + SSD-staging debug trace gate (matches ds4_cuda.cu). Default + * off in production; set DS4_SSD_DEBUG=1 to enable the per-step prefill + * trace. */ +static int ds4_ssd_dbg(void) { + static int e = -1; + if (e < 0) { + const char *s = getenv("DS4_SSD_DEBUG"); + e = (s && s[0] == '1') ? 1 : 0; } + return e; } /* Upstream: --power N GPU duty-cycle throttling helpers. The single-tier @@ -20757,6 +20774,8 @@ static bool metal_graph_decode_cuda_selected_load( getenv("DS4_CUDA_STREAMING_EXPERT_CACHE_PROFILE") != NULL; const double t0 = profile ? now_sec() : 0.0; + int stage_tier = g->placement ? g->placement[il + 1] : 0; + (void)ds4_gpu_debug_probe("decode_cuda_sel_load entry", (int)il, stage_tier); if (ds4_gpu_end_commands() == 0) return false; const double t_sync = profile ? now_sec() : 0.0; @@ -20766,10 +20785,12 @@ static bool metal_graph_decode_cuda_selected_load( selected_ids, (uint64_t)DS4_N_EXPERT_USED * sizeof(selected_ids[0])) != 0; + (void)ds4_gpu_debug_probe("decode_cuda_sel_load post-read", (int)il, stage_tier); const double t_read = profile ? now_sec() : 0.0; if (ok) { metal_graph_ssd_assert_tier_device(g, il); + (void)ds4_gpu_debug_probe("decode_cuda_sel_load pre-stage", (int)il, stage_tier); const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, @@ -20780,6 +20801,7 @@ static bool metal_graph_decode_cuda_selected_load( &table, selected_ids, DS4_N_EXPERT_USED) != 0; + (void)ds4_gpu_debug_probe("decode_cuda_sel_load post-stage", (int)il, stage_tier); } const double t_load = profile ? now_sec() : 0.0; @@ -20841,6 +20863,8 @@ static bool metal_graph_cuda_stream_prefill_batch_selected_load( getenv("DS4_CUDA_STREAMING_PREFILL_BATCH_SELECTED_PROFILE") != NULL; const double t0 = profile ? now_sec() : 0.0; + int home_tier = g->placement ? g->placement[il + 1] : 0; + (void)ds4_gpu_debug_probe("batch_selected_load entry", (int)il, home_tier); if (ds4_gpu_end_commands() == 0) return false; const double t_sync = profile ? now_sec() : 0.0; @@ -20849,9 +20873,11 @@ static bool metal_graph_cuda_stream_prefill_batch_selected_load( 0, selected_ids, n_ids64 * sizeof(selected_ids[0])) != 0; + (void)ds4_gpu_debug_probe("batch_selected_load post-read", (int)il, home_tier); const double t_read = profile ? now_sec() : 0.0; if (ok) { metal_graph_ssd_assert_tier_device(g, il); + (void)ds4_gpu_debug_probe("batch_selected_load pre-stage", (int)il, home_tier); const ds4_gpu_stream_expert_table table = graph_stream_expert_table_make(model, layer, @@ -20863,6 +20889,7 @@ static bool metal_graph_cuda_stream_prefill_batch_selected_load( selected_ids, n_tokens, DS4_N_EXPERT_USED) != 0; + (void)ds4_gpu_debug_probe("batch_selected_load post-stage", (int)il, home_tier); } free(selected_ids); const double t_load = profile ? now_sec() : 0.0; @@ -28624,7 +28651,9 @@ static bool metal_graph_encode_layer_ffn_batch( } DS4_METAL_PROFILE_FFN_STAGE("router"); + int ffn_tier = g->placement ? g->placement[il + 1] : 0; if (ok) { + (void)ds4_gpu_debug_probe("ffn_batch pre-staging", (int)il, ffn_tier); ok = metal_graph_cuda_stream_prefill_batch_selected_load(g, model, layer, @@ -28632,6 +28661,7 @@ static bool metal_graph_encode_layer_ffn_batch( n_tokens, gate_expert_bytes, down_expert_bytes); + (void)ds4_gpu_debug_probe("ffn_batch post-staging", (int)il, ffn_tier); } #ifdef DS4_ROCM_BUILD @@ -28857,6 +28887,7 @@ static bool metal_graph_encode_layer_ffn_batch( g->tp_batch_out && g->tp_batch_in; const bool cuda_tp_owned_batch_moe = g->cuda_tp_ep && g->cuda_tp_prefill_ffn; + (void)ds4_gpu_debug_probe("ffn_batch pre-routed-moe", (int)il, ffn_tier); if (ok && cuda_tp_owned_batch_moe) { ok = metal_graph_encode_mixed_routed_rows( g, decode_items, decode_count, model, layer, il, n_tokens); @@ -28975,6 +29006,7 @@ static bool metal_graph_encode_layer_ffn_batch( (uint64_t)n_tokens * DS4_N_EMBD, il, pos0); } DS4_METAL_PROFILE_FFN_STAGE("routed_moe"); + (void)ds4_gpu_debug_probe("ffn_batch post-routed-moe", (int)il, ffn_tier); if (!shared_done) { DS4_METAL_ENCODE_PREFILL_SHARED_EXPERT(); } @@ -29090,8 +29122,9 @@ static bool metal_graph_encode_layer_batch( uint32_t il, uint32_t pos0, uint32_t n_tokens) { + int this_tier = g->placement ? g->placement[il + 1u] : 0; + (void)ds4_gpu_debug_probe("encode_layer_batch entry", (int)il, this_tier); if (g->placement) { - const int this_tier = g->placement[il + 1u]; if (!metal_graph_set_active_tier_batch(g, this_tier, n_tokens)) { return false; } @@ -29103,6 +29136,7 @@ static bool metal_graph_encode_layer_batch( if (!ok) { fprintf(stderr, "ds4: gpu layer %u attention batch encode failed\n", il); } + (void)ds4_gpu_debug_probe("encode_layer_batch post-attn", (int)il, this_tier); if (ok) { ok = metal_graph_encode_layer_ffn_batch(g, model, layer, il, pos0, n_tokens, NULL, 0); @@ -29110,6 +29144,7 @@ static bool metal_graph_encode_layer_batch( fprintf(stderr, "ds4: gpu layer %u ffn batch encode failed\n", il); } } + (void)ds4_gpu_debug_probe("encode_layer_batch post-ffn", (int)il, this_tier); if (ok) { ds4_gpu_tensor *tmp = metal_graph_batch_cur_hc(g); g->batch_cur_hc_by_tier[g->active_tier] = metal_graph_batch_next_hc(g); @@ -29145,19 +29180,50 @@ static bool metal_graph_eval_token_raw_swa_streaming( const bool batch_static_decode = static_decode_map && metal_graph_stream_decode_layer_batch_enabled(g); bool ok = true; + /* Multi-tier: the per-token embed kernel reads token_embd (staged on + * emb_tier) and writes to metal_graph_cur_hc(g) = cur_hc_by_tier[active_tier]. + * After init (active_tier=-1) OR after a prior decode/prefill left + * active_tier on a different tier (e.g. the head tier / tier 1, where the + * last decode layer ran), the embed would resolve cur_hc to the wrong + * tier or to an out-of-bounds slot (cur_hc_by_tier[-1]). The embed then + * silently no-ops (NULL out_hc) and the whole prefill returns false + * ("cuda prefill failed") with NO CUDA error logged. + * + * The per-layer decode dispatcher (metal_graph_set_active_tier_decode) + * early-returns when active_tier already equals the requested tier + * WITHOUT calling cudaSetDevice, so we MUST also switch the AMBIENT + * device here — otherwise a stale ambient device left by the previous + * token's last layer (e.g. tier 1) survives into il=0 of this token + * (tier 0), where encode_decode_layer_phase's set_active_tier_decode(0) + * no-ops and the layer's attention kernels run on device 1 reading + * device-0 raw_kv / norm tensors -> illegal memory access. + * + * metal_graph_set_active_tier_no_copy does both: cudaSetDevice(emb_tier) + * AND active_tier=emb_tier. The embed kernel itself uses WITH_DEVICE for + * defense-in-depth. No boundary-hop copy is needed: the embed overwrites + * cur_hc in full. No-op single-tier (placement == NULL). */ + if (g->placement && + !metal_graph_set_active_tier_no_copy(g, g->emb_tier)) { + return false; + } + (void)ds4_gpu_debug_probe("eval_token_raw_swa entry", -1, g->active_tier); if (static_decode_map) { if (!static_map_state_cache || !g->streaming_static_decode_map_current) { ok = metal_graph_stream_map_decode_static_all(model, weights); + (void)ds4_gpu_debug_probe("eval_token_raw_swa post-map_static_all", -1, g->active_tier); if (ok) g->streaming_static_decode_map_current = static_map_state_cache; } } else { g->streaming_static_decode_map_current = false; ok = metal_graph_stream_map_token(model, weights); + (void)ds4_gpu_debug_probe("eval_token_raw_swa post-map_token", -1, g->active_tier); } if (ok && !static_decode_map && DS4_N_LAYER > 0) { metal_graph_stream_readahead_layer_decode(model, weights, 0); + (void)ds4_gpu_debug_probe("eval_token_raw_swa post-readahead_l0", -1, g->active_tier); } if (ok) ok = ds4_gpu_begin_commands() != 0; + (void)ds4_gpu_debug_probe("eval_token_raw_swa pre-embed", -1, g->active_tier); if (ok) { ok = ds4_gpu_embed_token_hc_tensor(metal_graph_cur_hc(g), model->map, @@ -29167,9 +29233,12 @@ static bool metal_graph_eval_token_raw_swa_streaming( (uint32_t)token, DS4_N_EMBD, DS4_N_HC) != 0; + (void)ds4_gpu_debug_probe("eval_token_raw_swa post-embed", -1, g->active_tier); } if (batch_static_decode) { for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { + int bsd_tier = g->placement ? g->placement[il + 1] : 0; + (void)ds4_gpu_debug_probe("eval_token bsd pre-encode", (int)il, bsd_tier); ok = metal_graph_encode_decode_layer(g, model, &weights->layer[il], @@ -29180,6 +29249,7 @@ static bool metal_graph_eval_token_raw_swa_streaming( raw_row, n_raw, token); + (void)ds4_gpu_debug_probe("eval_token bsd post-encode", (int)il, bsd_tier); if (ok) { ds4_gpu_tensor *tmp = metal_graph_cur_hc(g); g->cur_hc_by_tier[g->active_tier] = metal_graph_after_ffn_hc(g); @@ -29188,13 +29258,17 @@ static bool metal_graph_eval_token_raw_swa_streaming( } } if (ok && logits) { + (void)ds4_gpu_debug_probe("eval_token bsd pre-output_head", -1, g->active_tier); ok = metal_graph_encode_output_head(g, model, weights, weights->output->dim[1]); + (void)ds4_gpu_debug_probe("eval_token bsd post-output_head", -1, g->active_tier); } const double t_encoded = (profile || throttle) ? now_sec() : 0.0; if (ok) ok = ds4_gpu_end_commands() != 0; + (void)ds4_gpu_debug_probe("eval_token bsd post-end_commands", -1, g->active_tier); const double t_done = (profile || throttle) ? now_sec() : 0.0; if (ok && logits) { ok = ds4_gpu_tensor_read(metal_graph_logits(g), 0, logits, (uint64_t)DS4_N_VOCAB * sizeof(float)) != 0; + (void)ds4_gpu_debug_probe("eval_token bsd post-logits_read", -1, g->active_tier); } const double t_read = (profile || throttle) ? now_sec() : 0.0; if (profile) { @@ -29218,11 +29292,14 @@ static bool metal_graph_eval_token_raw_swa_streaming( return ok; } if (ok) ok = ds4_gpu_end_commands() != 0; + (void)ds4_gpu_debug_probe("eval_token post-embed end_commands", -1, g->active_tier); double encode_s = 0.0; double execute_s = 0.0; for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { const double tl0 = profile ? now_sec() : 0.0; + int decode_tier = g->placement ? g->placement[il + 1] : 0; + (void)ds4_gpu_debug_probe("eval_token decode_loop pre-map", (int)il, decode_tier); if (!static_decode_map && !metal_graph_stream_map_layer_decode(model, weights, il)) { ok = false; break; @@ -29232,6 +29309,7 @@ static bool metal_graph_eval_token_raw_swa_streaming( } else if (!static_decode_map && logits) { metal_graph_stream_readahead_output(model, weights); } + (void)ds4_gpu_debug_probe("eval_token decode_loop pre-encode", (int)il, decode_tier); if (ok) ok = ds4_gpu_begin_commands() != 0; bool encoded_layer = false; if (ok) { @@ -29246,6 +29324,7 @@ static bool metal_graph_eval_token_raw_swa_streaming( n_raw, token); encoded_layer = true; + (void)ds4_gpu_debug_probe("eval_token decode_loop post-encode", (int)il, decode_tier); } if (encoded_layer) { ds4_gpu_tensor *tmp = metal_graph_cur_hc(g); @@ -32628,12 +32707,21 @@ static bool metal_graph_reset_prefill_state(ds4_gpu_graph *g) { g->mtp_n_raw = 0; metal_graph_dspark_cache_reset(g); metal_graph_dspark_capture_invalidate(g); + int prev_active = g->active_tier; + if (ds4_ssd_dbg()) { + fprintf(stderr, + "ds4[probe] reset_prefill_state ENTER active_tier=%d emb_tier=%d placement=%p\n", + prev_active, g->emb_tier, (void*)g->placement); + fflush(stderr); + } for (uint32_t il = 0; il < DS4_N_LAYER; il++) { const uint32_t ratio = ds4_layer_compress_ratio(il); if (ratio == 0) continue; const uint32_t coff = ratio == 4 ? 2u : 1u; const uint64_t attn_width = (uint64_t)coff * DS4_N_HEAD_DIM; const uint64_t attn_rows = (uint64_t)coff * ratio; + int home_tier = g->placement ? g->placement[il + 1] : 0; + (void)ds4_gpu_debug_probe("reset_prefill pre-fill", (int)il, home_tier); if (!metal_tensor_fill_f32(g->layer_attn_state_kv[il], 0.0f, attn_width * attn_rows)) return false; if (!metal_tensor_fill_f32(g->layer_attn_state_score[il], DS4_NEG_INF, attn_width * attn_rows)) return false; if (ratio == 4) { @@ -32643,6 +32731,7 @@ static bool metal_graph_reset_prefill_state(ds4_gpu_graph *g) { if (!metal_tensor_fill_f32(g->layer_index_state_score[il], DS4_NEG_INF, index_width * index_rows)) return false; } } + (void)ds4_gpu_debug_probe("reset_prefill post-all", -1, prev_active); return true; } @@ -33060,6 +33149,7 @@ static bool metal_graph_prefill_layer_major( * on the head tier, so repoint at emb_tier (no boundary-hop copy is * needed: the embed overwrites batch_cur_hc in full). No-op single-tier. */ if (g->placement && g->active_tier != g->emb_tier) g->active_tier = g->emb_tier; + (void)ds4_gpu_debug_probe("prefill_nosplit pre-embed", -1, g->active_tier); ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), metal_graph_prefill_tokens(g), model, @@ -33067,6 +33157,7 @@ static bool metal_graph_prefill_layer_major( prompt, start, n_tokens); + (void)ds4_gpu_debug_probe("prefill_nosplit post-embed", -1, g->active_tier); if (ok) ok = ds4_gpu_begin_commands() != 0; for (uint32_t il = 0; ok && il < DS4_N_LAYER; il++) { ok = metal_graph_encode_layer_batch(g, @@ -33221,6 +33312,13 @@ static bool metal_graph_prefill_layer_major( * active_tier at emb_tier before the embed upload so the kernel runs on * the device that owns prefill_tokens + token_embd. No-op single-tier. */ if (g->placement && g->active_tier != g->emb_tier) g->active_tier = g->emb_tier; + (void)ds4_gpu_debug_probe("prefill_split pre-embed", -1, g->active_tier); + if (ds4_ssd_dbg()) { + fprintf(stderr, + "ds4[probe] prefill_split embed: active_tier=%d emb_tier=%d n_tokens=%u start=%u ssd_streaming=%d\n", + g->active_tier, g->emb_tier, n_tokens, start, (int)g->ssd_streaming); + fflush(stderr); + } ok = metal_graph_upload_prompt_embeddings_hc(metal_graph_batch_cur_hc(g), metal_graph_prefill_tokens(g), model, @@ -33228,6 +33326,7 @@ static bool metal_graph_prefill_layer_major( prompt, start, n_tokens); + (void)ds4_gpu_debug_probe("prefill_split post-embed", -1, g->active_tier); const double t_embed_encoded = (profile || throttle) ? now_sec() : 0.0; const double t_embed_done = (profile || throttle) ? now_sec() : 0.0; if (profile) { diff --git a/ds4_cuda.cu b/ds4_cuda.cu index 3f8293dc1..36fba989a 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -120,7 +120,17 @@ static int g_cuda_exact_score_split_vec4_plain; static int g_cuda_exact_score_split_dim2; static int g_cuda_exact_score_split_fuse_inv_rope; static int g_cuda_moe_decode_graph; -static int g_current_logical_tier = -1; +/* Per-thread cache of the current logical tier. MUST be thread-local: + * cudaSetDevice is per-thread, so a process-global cache would let one + * thread's set-device answer a different thread's "am I already there?" + * check. Concretely: the async staging worker (metal_graph_selected_async_* + * in ds4.c) reads the main thread's just-set tier from a global cache and + * skips its own cudaSetDevice, leaving the worker on device 0 while it + * believes it is on tier 1. The worker then stages experts on device 0, + * the per-tier SSD selected_cache on tier 1 stays zero, and the consumer's + * routed_moe_launch fails with "streaming selected experts are unavailable". + * Thread-localizing the cache makes each thread track its own device. */ +static thread_local int g_current_logical_tier = -1; static int g_ssd_streaming_mode; typedef struct { @@ -301,6 +311,65 @@ static ds4_ssd_ctx *ssd_current(void) { return &g_ssd[0]; } +/* ========================================================================= + * Debug probe for the multi-tier SSD-staging fault investigation. + * + * Gate with DS4_SSD_DEBUG=1 (default OFF). Logs the ambient CUDA device, the + * expected physical device for `tier`, the logical tier, the layer, and any + * pending CUDA error. To catch *lazy* async illegal-memory-access errors + * promptly, it forces a cudaDeviceSynchronize() before sampling + * cudaGetLastError(); this is expensive, so it is opt-in. + * + * NB: cudaGetLastError() *clears* the sticky error, so a subsequent cuda_ok() + * will not _exit() on it — i.e. enabling DS4_SSD_DEBUG effectively suspends + * fail-fast so the full prefill trace can be captured. The CUDA context is + * still poisoned, so later ops will re-report and the trace is still + * terminal; we just get to see every step instead of dying on the first one. + * Returns the (cleared) cudaError_t ordinal, or 0 on success (or when off). */ +static int ds4_ssd_debug_on(void) { + static int e = -1; + if (e < 0) { + const char *s = getenv("DS4_SSD_DEBUG"); + e = (s && s[0] == '1') ? 1 : 0; + } + return e; +} + +extern "C" int ds4_gpu_debug_probe(const char *where, int il, int tier) { + /* Resolve the ambient device even when the probe is off so the call still + * has a well-defined side-effect-free return; the work happens inside the + * gated branch. */ + int dev = -1; + (void)cudaGetDevice(&dev); + int phys = (tier >= 0 && tier < g_n_gpus) ? g_gpu[tier].device_id : -99; + if (!ds4_ssd_debug_on()) return 0; + /* Sync EVERY device (not just ambient) and report the first non-zero + * error. Cross-device illegal-memory-access errors otherwise hide from + * an ambient-only cudaDeviceSynchronize. */ + cudaError_t err = cudaSuccess; + int err_dev = -1; + int prev_dev = -1; + (void)cudaGetDevice(&prev_dev); + for (int t = 0; t < g_n_gpus; t++) { + if (cudaSetDevice(g_gpu[t].device_id) != cudaSuccess) continue; + cudaError_t sync_err = cudaDeviceSynchronize(); + cudaError_t peek = cudaGetLastError(); + if (sync_err != cudaSuccess && peek == cudaSuccess) peek = sync_err; + if (peek != cudaSuccess && err == cudaSuccess) { + err = peek; + err_dev = g_gpu[t].device_id; + } + } + if (prev_dev >= 0) (void)cudaSetDevice(prev_dev); + fprintf(stderr, + "ds4[probe] %-30s il=%-3u tier=%d phys_dev=%d ambient_dev=%d n_gpus=%d err=%d%s%s%s\n", + where ? where : "?", (unsigned)il, tier, phys, dev, g_n_gpus, + (int)err, err ? " :: " : "", err ? cudaGetErrorString(err) : "", + err ? (err_dev >= 0 ? "" : "") : ""); + fflush(stderr); + return (int)err; +} + /* Debug/override flags are read once per CUDA init. The hot decode path calls * the xdev helpers many times per token, so they must not re-enter getenv(). */ static void cuda_xdev_env_refresh(void) { @@ -1330,9 +1399,42 @@ static float *cuda_q8_f32_ptr( return dev; } +/* CUDA "sticky" errors permanently invalidate the context: after one of them + * every subsequent CUDA call returns cudaErrorIllegalAddress, so the engine + * would keep surfacing HTTP 500 ("cuda prefill failed") on every request + * until manually restarted. Fail fast on these so a supervisor + * (systemd Restart=on-failure) can restart the process and restore service. + * Override for debugging with DS4_CUDA_NO_FAIL_FAST=1. */ +static bool cuda_err_is_sticky(cudaError_t err) { + switch (err) { + case cudaErrorIllegalAddress: /* "an illegal memory access was encountered" */ + case cudaErrorLaunchFailure: + case cudaErrorHardwareStackError: + case cudaErrorIllegalInstruction: + case cudaErrorInvalidPc: + case cudaErrorMisalignedAddress: + case cudaErrorAssert: + case cudaErrorCooperativeLaunchTooLarge: + return true; + default: + return false; + } +} + static int cuda_ok(cudaError_t err, const char *what) { if (err == cudaSuccess) return 1; fprintf(stderr, "ds4: CUDA %s failed: %s\n", what, cudaGetErrorString(err)); + if (cuda_err_is_sticky(err) && getenv("DS4_CUDA_NO_FAIL_FAST") == NULL) { + fprintf(stderr, + "ds4: CUDA %s failed: %s is a sticky unrecoverable context error; " + "exiting so a supervisor can restart the process " + "(set DS4_CUDA_NO_FAIL_FAST=1 to disable).\n", + what, cudaGetErrorString(err)); + fflush(stderr); + /* _exit, not exit: skip atexit handlers that could themselves touch + * the broken CUDA context and hang. The OS reclaims all resources. */ + _exit(1); + } return 0; } @@ -20971,6 +21073,33 @@ static int routed_moe_launch( fprintf(stderr, "ds4: CUDA streaming selected experts are unavailable for layer %u\n", layer_index); + if (ds4_ssd_debug_on()) { + int cur_dev = -1; + (void)cudaGetDevice(&cur_dev); + /* Find which tier the ambient device corresponds to. */ + int ambient_tier = -1; + for (int t = 0; t < g_n_gpus; t++) { + if (g_gpu[t].device_id == cur_dev) { ambient_tier = t; break; } + } + fprintf(stderr, + "ds4[probe] moe_check FAIL layer=%u logical_tier=%d cur_dev=%d ambient_tier=%d " + "sc(valid=%d logical=%d model_map=%p layer=%u n_total=%u slot_count=%u req=%u " + "gate_off=%llu up_off=%llu down_off=%llu gate_bytes=%llu down_bytes=%llu " + "gate_ptr=%p up_ptr=%p down_ptr=%p slot_ptr=%p slot_bytes=%llu)\n", + layer_index, logical_tier, cur_dev, ambient_tier, + (int)sc->valid, sc->logical_tier, (void*)sc->model_map, + (unsigned)sc->layer, (unsigned)sc->n_total_expert, + (unsigned)sc->slot_count, (unsigned)required_slot_count, + (unsigned long long)sc->gate_offset, + (unsigned long long)sc->up_offset, + (unsigned long long)sc->down_offset, + (unsigned long long)sc->gate_expert_bytes, + (unsigned long long)sc->down_expert_bytes, + (void*)sc->gate_ptr, (void*)sc->up_ptr, (void*)sc->down_ptr, + (void*)sc->slot_selected_tensor.ptr, + (unsigned long long)sc->slot_selected_tensor.bytes); + fflush(stderr); + } return 0; } if (use_stream_selected_cache) { diff --git a/ds4_gpu.h b/ds4_gpu.h index 2000bba8b..4e79255df 100644 --- a/ds4_gpu.h +++ b/ds4_gpu.h @@ -190,6 +190,10 @@ int ds4_gpu_stream_expert_cache_prepare_selected_batch( const int32_t *selected_ids, uint32_t n_tokens, uint32_t n_selected); +/* Multi-tier SSD debug probe. No-op when DS4_SSD_DEBUG!=1; otherwise logs the + * ambient device, expected tier device, and any pending CUDA error after a + * forced sync. Defined in ds4_cuda.cu / ds4_rocm.cu. */ +int ds4_gpu_debug_probe(const char *where, int il, int tier); #endif #ifdef DS4_ROCM_BUILD int ds4_gpu_stream_expert_cache_load_layer( From f836788e7d5506fbac48931be68a290691d52be2 Mon Sep 17 00:00:00 2001 From: Anton Sokolchenko Date: Tue, 21 Jul 2026 13:40:12 +0000 Subject: [PATCH 3/4] feat(ssd): port expert LRU cache from feat/mgpu-ssd to antirez-main SSD streaming MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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), 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 --- ds4.c | 10 + ds4_cuda.cu | 860 +++++++++++++++++++++++++++++++++++++++++++++++++--- 2 files changed, 830 insertions(+), 40 deletions(-) diff --git a/ds4.c b/ds4.c index a1409d143..0a8803f7f 100644 --- a/ds4.c +++ b/ds4.c @@ -55927,6 +55927,16 @@ static int ds4_engine_open_internal(ds4_engine **out, *out = NULL; return 1; } + /* Convert a byte budget (--ssd-streaming-cache-experts NGB) into + * an expert count; the LRU's budget_visible_to_shared() gates on + * ssd_streaming_cache_experts != 0. Single-tier path has always + * done this; multi-tier must too or 'NGB' silently leaves the + * LRU disabled. */ + if (!ds4_engine_configure_streaming_cache_budget(e)) { + ds4_engine_close(e); + *out = NULL; + return 1; + } ds4_gpu_set_streaming_expert_cache_budget(e->ssd_streaming_cache_experts); if (e->ssd_streaming) { /* Pin the expert cache's slab size class to the model's uniform diff --git a/ds4_cuda.cu b/ds4_cuda.cu index 36fba989a..4966c1d0f 100644 --- a/ds4_cuda.cu +++ b/ds4_cuda.cu @@ -39,7 +39,13 @@ enum { DS4_CUDA_SPLITKV_CHUNK = 512u, DS4_CUDA_SPLITKV_SCORE_CAP = 512u, DS4_CUDA_SPLITKV_S_MAX = 16u, - DS4_CUDA_SPLITKV_S_FLOOR = 4u + DS4_CUDA_SPLITKV_S_FLOOR = 4u, + /* SSD weight-streaming resident expert-cache sizing (LRU port from + * feat/mgpu-ssd). Cap is the user-facing expert count (--ssd-streaming- + * cache-experts); the runtime may shrink it further per-tier based on + * cudaMemGetInfo minus DS4_CUDA_STREAMING_EXPERT_CACHE_RESERVE_GB. */ + DS4_CUDA_STREAM_EXPERT_DEFAULT = 8u * 64u, + DS4_CUDA_STREAM_EXPERT_MAX = 61u * 384u }; /* struct ds4_gpu_tensor is defined in ds4_gpu.h (no longer opaque as of @@ -157,6 +163,48 @@ typedef struct { ds4_gpu_tensor slot_selected_tensor; } cuda_stream_selected_cache; +/* Resident LRU expert cache (port from feat/mgpu-ssd). One slab per tier + * holding the most-recently-used routed experts (gate/up/down weight tiles) + * resident in VRAM, so that decode tokens hitting already-cached experts + * skip the SSD read. The cache is sized lazily from cudaMemGetInfo() minus + * the reserve (DS4_CUDA_STREAMING_EXPERT_CACHE_RESERVE_GB), and seeded by + * cuda_stream_selected_cache_begin_load below. + * + * g_stream_expert_budget_override is GLOBAL — the user-facing cap (from + * --ssd-streaming-cache-experts) applied uniformly to every tier. */ +static uint32_t g_stream_expert_budget_override; + +struct cuda_stream_expert_cache_slot { + int valid; + const void *model_map; + uint64_t model_size; + uint32_t layer; + uint32_t n_total_expert; + uint32_t expert; + uint64_t gate_offset; + uint64_t up_offset; + uint64_t down_offset; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; + uint64_t age; +}; + +struct cuda_stream_expert_cache { + int valid; + uint32_t capacity; + uint32_t count; + uint64_t tick; + uint64_t gate_expert_bytes; + uint64_t down_expert_bytes; + char *gate_ptr; + char *up_ptr; + char *down_ptr; + uint64_t gate_capacity; + uint64_t up_capacity; + uint64_t down_capacity; + std::vector slots; +}; + /* Per-tier SSD streaming state. SSD streaming was originally single-device: * one selected-expert cache, one pinned staging pool, one upload stream, all * on device 0. In multi-GPU mode q4k's per-layer dispatch routes each layer @@ -172,11 +220,16 @@ typedef struct { * ssd_current() returns the tier that owns the layer about to execute. */ struct ds4_ssd_ctx { cuda_stream_selected_cache selected_cache; + cuda_stream_expert_cache expert_cache; void *stage_raw[4]; void *stage[4]; cudaEvent_t stage_event[4]; uint64_t stage_bytes; cudaStream_t upload_stream; + uint32_t runtime_cap; /* OOM-driven cap for this tier */ + uint32_t memory_cap_notice; /* last cap reported via stderr */ + uint64_t runtime_gate_bytes; /* expert slab class last seen */ + uint64_t runtime_down_bytes; }; static ds4_ssd_ctx g_ssd[DS4_MAX_GPUS]; @@ -23167,6 +23220,586 @@ static int cuda_stream_selected_ranges_valid( down_bytes <= table->model_size - table->down_offset; } +/* ========================================================================= + * Resident LRU expert cache (port from feat/mgpu-ssd). + * + * The cache lives per-tier in g_ssd[t].expert_cache; all functions here + * resolve the active tier via ssd_current() (per-layer dispatch has already + * cudaSetDevice()'d). The slab is sized lazily by cuda_stream_expert_cache_prepare + * using cudaMemGetInfo() on the current device minus the configured reserve. + * ========================================================================= */ + +/* No-op stub for the progress meter that exists in the mgpu-ssd build but + * not in antirez main. Keeping the call sites identical lets us backport + * the LRU logging verbatim. */ +static void cuda_model_load_progress_finish(void) { +} + +static void cuda_stream_expert_cache_release_all(void) { + ds4_ssd_ctx *ssd = ssd_current(); + if (ssd->expert_cache.gate_ptr) { + (void)cudaFree(ssd->expert_cache.gate_ptr); + } + if (ssd->expert_cache.up_ptr) { + (void)cudaFree(ssd->expert_cache.up_ptr); + } + if (ssd->expert_cache.down_ptr) { + (void)cudaFree(ssd->expert_cache.down_ptr); + } + ssd->expert_cache.slots.clear(); + memset(&ssd->expert_cache, 0, sizeof(ssd->expert_cache)); +} + +static void cuda_stream_expert_cache_invalidate(void) { + ds4_ssd_ctx *ssd = ssd_current(); + for (cuda_stream_expert_cache_slot &slot : ssd->expert_cache.slots) { + slot.valid = 0; + } + ssd->expert_cache.valid = 0; + ssd->expert_cache.count = 0; + ssd->expert_cache.tick = 0; +} + +static uint32_t cuda_stream_expert_cache_requested_budget(void) { + uint32_t cap = g_stream_expert_budget_override != 0 ? + g_stream_expert_budget_override : DS4_CUDA_STREAM_EXPERT_DEFAULT; + const char *env = getenv("DS4_CUDA_STREAMING_EXPERT_CACHE_N"); + if (env && env[0]) { + char *end = NULL; + errno = 0; + unsigned long v = strtoul(env, &end, 10); + while (end && (*end == ' ' || *end == '\t')) end++; + if (end != env && errno == 0 && end && *end == '\0') { + cap = v > DS4_CUDA_STREAM_EXPERT_MAX ? + DS4_CUDA_STREAM_EXPERT_MAX : (uint32_t)v; + } + } + if (cap > DS4_CUDA_STREAM_EXPERT_MAX) { + cap = DS4_CUDA_STREAM_EXPERT_MAX; + } + return cap; +} + +static uint32_t cuda_stream_expert_cache_configured_budget(void) { + uint32_t cap = cuda_stream_expert_cache_requested_budget(); + ds4_ssd_ctx *ssd = ssd_current(); + if (ssd->runtime_cap != 0 && cap > ssd->runtime_cap) { + cap = ssd->runtime_cap; + } + return cap; +} + +static int cuda_stream_expert_cache_budget_visible_to_shared(void) { + if (!g_ssd_streaming_mode) return 0; + if (g_stream_expert_budget_override != 0) return 1; + const char *env = getenv("DS4_CUDA_STREAMING_EXPERT_CACHE_N"); + if (env && env[0]) return 1; + env = getenv("DS4_CUDA_ENABLE_STREAMING_EXPERT_HOTLIST"); + if (!env || !env[0]) { + env = getenv("DS4_CUDA_STREAMING_EXPERT_HOTLIST"); + } + return env && env[0] && strcmp(env, "0") != 0; +} + +static uint64_t cuda_stream_expert_cache_reserve_bytes(void) { + uint64_t gb = 16; + const char *env = getenv("DS4_CUDA_STREAMING_EXPERT_CACHE_RESERVE_GB"); + if (env && env[0]) { + char *end = NULL; + errno = 0; + unsigned long long v = strtoull(env, &end, 10); + while (end && (*end == ' ' || *end == '\t')) end++; + if (end != env && errno == 0 && end && *end == '\0') { + gb = (uint64_t)v; + } + } + if (gb > UINT64_MAX / 1073741824ull) return UINT64_MAX; + return gb * 1073741824ull; +} + +static uint32_t cuda_stream_expert_cache_live_budget( + uint32_t requested, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + uint64_t reclaim_bytes, + int report) { + ds4_ssd_ctx *ssd = ssd_current(); + if (requested == 0 || + gate_expert_bytes == 0 || + down_expert_bytes == 0 || + gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { + return 0; + } + const uint64_t per_expert_bytes = + gate_expert_bytes * 2ull + down_expert_bytes; + if (per_expert_bytes == 0) return 0; + + size_t free_b = 0; + size_t total_b = 0; + cudaError_t err = cudaMemGetInfo(&free_b, &total_b); + if (err != cudaSuccess) { + fprintf(stderr, + "ds4: CUDA streaming expert cache memory query failed: %s; " + "using direct selected loads\n", + cudaGetErrorString(err)); + (void)cudaGetLastError(); + return 0; + } + + uint64_t free_bytes = (uint64_t)free_b; + if (reclaim_bytes > UINT64_MAX - free_bytes) { + free_bytes = UINT64_MAX; + } else { + free_bytes += reclaim_bytes; + } + uint64_t reserve = cuda_stream_expert_cache_reserve_bytes(); + const uint64_t total_bytes = (uint64_t)total_b; + if (total_bytes != 0 && reserve > total_bytes / 2ull) { + reserve = total_bytes / 2ull; + } + if (free_bytes <= reserve) { + if (report && ssd->memory_cap_notice != requested) { + cuda_model_load_progress_finish(); + fprintf(stderr, + "ds4: CUDA streaming expert cache disabled: available %.2f GiB <= reserve %.2f GiB\n", + (double)free_bytes / 1073741824.0, + (double)reserve / 1073741824.0); + ssd->memory_cap_notice = requested; + } + return 0; + } + + uint64_t usable = free_bytes - reserve; + uint64_t max_slots64 = usable / per_expert_bytes; + if (max_slots64 > UINT32_MAX) max_slots64 = UINT32_MAX; + uint32_t capped = requested; + if ((uint64_t)capped > max_slots64) capped = (uint32_t)max_slots64; + if (report && capped != requested && ssd->memory_cap_notice != capped) { + cuda_model_load_progress_finish(); + fprintf(stderr, + "ds4: CUDA streaming expert cache capped from %u to %u experts " + "(available %.2f GiB, reserve %.2f GiB, %.2f MiB/expert)\n", + requested, + capped, + (double)free_bytes / 1073741824.0, + (double)reserve / 1073741824.0, + (double)per_expert_bytes / 1048576.0); + ssd->memory_cap_notice = capped; + } + return capped; +} + +static uint64_t cuda_stream_expert_cache_expert_bytes( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (gate_expert_bytes == 0 || + down_expert_bytes == 0 || + gate_expert_bytes > (UINT64_MAX - down_expert_bytes) / 2ull) { + return 0; + } + return gate_expert_bytes * 2ull + down_expert_bytes; +} + +static void cuda_stream_expert_cache_note_size( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + ds4_ssd_ctx *ssd = ssd_current(); + if (ssd->runtime_gate_bytes == gate_expert_bytes && + ssd->runtime_down_bytes == down_expert_bytes) { + return; + } + ssd->runtime_gate_bytes = gate_expert_bytes; + ssd->runtime_down_bytes = down_expert_bytes; + ssd->runtime_cap = 0; + ssd->memory_cap_notice = 0; +} + +static uint32_t cuda_stream_expert_cache_shrunken_cap(uint32_t cap) { + if (cap == 0) return 0; + const uint32_t release = (cap + 9u) / 10u; + return cap > release ? cap - release : 0; +} + +static void cuda_stream_expert_cache_note_oom_cap( + uint32_t failed_cap, + uint32_t new_cap, + uint64_t expert_bytes, + const char *errstr) { + ds4_ssd_ctx *ssd = ssd_current(); + if (ssd->runtime_cap != 0 && + ssd->runtime_cap <= new_cap) { + return; + } + ssd->runtime_cap = new_cap; + const uint32_t released = + failed_cap > new_cap ? failed_cap - new_cap : 0; + cuda_model_load_progress_finish(); + fprintf(stderr, + "ds4: CUDA streaming expert cache allocation failed at %u experts " + "/ %.2f GiB%s%s\n", + failed_cap, + expert_bytes != 0 ? + (double)((uint64_t)failed_cap * expert_bytes) / 1073741824.0 : + 0.0, + errstr && errstr[0] ? ": " : "", + errstr && errstr[0] ? errstr : ""); + if (new_cap != 0) { + fprintf(stderr, + "ds4: shrinking resident cache margin by %u experts / %.2f GiB; " + "runtime cache cap now %u experts\n", + released, + expert_bytes != 0 ? + (double)((uint64_t)released * expert_bytes) / 1073741824.0 : + 0.0, + new_cap); + } else { + fprintf(stderr, + "ds4: disabling resident expert cache after OOM; using direct selected loads\n"); + } +} + +static int cuda_stream_expert_cache_try_alloc( + uint32_t cap, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + char **gate_ptr, + char **up_ptr, + char **down_ptr, + const char **errstr) { + *gate_ptr = NULL; + *up_ptr = NULL; + *down_ptr = NULL; + if (errstr) *errstr = NULL; + if (cap == 0 || + (uint64_t)cap > UINT64_MAX / gate_expert_bytes || + (uint64_t)cap > UINT64_MAX / down_expert_bytes) { + return 0; + } + const uint64_t gate_bytes = (uint64_t)cap * gate_expert_bytes; + const uint64_t down_bytes = (uint64_t)cap * down_expert_bytes; + + void *gate = NULL; + void *up = NULL; + void *down = NULL; + cudaError_t err = cudaMalloc(&gate, (size_t)gate_bytes); + if (err != cudaSuccess) goto fail; + err = cudaMalloc(&up, (size_t)gate_bytes); + if (err != cudaSuccess) goto fail; + err = cudaMalloc(&down, (size_t)down_bytes); + if (err != cudaSuccess) goto fail; + + *gate_ptr = (char *)gate; + *up_ptr = (char *)up; + *down_ptr = (char *)down; + return 1; + +fail: + if (errstr) *errstr = cudaGetErrorString(err); + (void)cudaGetLastError(); + if (gate) (void)cudaFree(gate); + if (up) (void)cudaFree(up); + if (down) (void)cudaFree(down); + return 0; +} + +static cuda_stream_expert_cache *cuda_stream_expert_cache_prepare( + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes, + uint32_t target_cap) { + const uint64_t expert_bytes = + cuda_stream_expert_cache_expert_bytes(gate_expert_bytes, + down_expert_bytes); + if (expert_bytes == 0) return NULL; + cuda_stream_expert_cache_note_size(gate_expert_bytes, down_expert_bytes); + + const uint32_t requested_cap = cuda_stream_expert_cache_configured_budget(); + if (requested_cap == 0) return NULL; + if (target_cap == 0 || target_cap > requested_cap) target_cap = requested_cap; + if (target_cap == 0) return NULL; + cuda_stream_expert_cache *ec = &ssd_current()->expert_cache; + const int same_dims = + ec->valid && + ec->gate_expert_bytes == gate_expert_bytes && + ec->down_expert_bytes == down_expert_bytes; + if (!same_dims && ec->valid) { + cuda_stream_expert_cache_release_all(); + } + if (same_dims && + ec->capacity != 0 && + ec->capacity >= target_cap && + ec->slots.size() == ec->capacity) { + return ec; + } + + uint64_t reclaim_bytes = 0; + if (same_dims && + ec->capacity != 0 && + (uint64_t)ec->capacity <= UINT64_MAX / expert_bytes) { + reclaim_bytes = (uint64_t)ec->capacity * expert_bytes; + } + uint32_t cap = + cuda_stream_expert_cache_live_budget(target_cap, + gate_expert_bytes, + down_expert_bytes, + reclaim_bytes, + reclaim_bytes == 0); + if (cap == 0) return NULL; + if (same_dims && + ec->capacity != 0 && + ec->capacity >= cap && + ec->slots.size() == ec->capacity) { + return ec; + } + + cuda_stream_expert_cache_release_all(); + ec = &ssd_current()->expert_cache; + while (cap != 0) { + if ((uint64_t)cap > UINT64_MAX / gate_expert_bytes || + (uint64_t)cap > UINT64_MAX / down_expert_bytes) { + fprintf(stderr, "ds4: CUDA streaming expert cache size overflow\n"); + return NULL; + } + + char *gate_ptr = NULL; + char *up_ptr = NULL; + char *down_ptr = NULL; + const char *alloc_error = NULL; + if (!cuda_stream_expert_cache_try_alloc(cap, + gate_expert_bytes, + down_expert_bytes, + &gate_ptr, + &up_ptr, + &down_ptr, + &alloc_error)) { + const uint32_t new_cap = + cuda_stream_expert_cache_shrunken_cap(cap); + cuda_stream_expert_cache_note_oom_cap(cap, + new_cap, + expert_bytes, + alloc_error); + cap = new_cap; + if (cap != 0) { + cap = cuda_stream_expert_cache_live_budget(cap, + gate_expert_bytes, + down_expert_bytes, + 0, + 1); + } + continue; + } + + try { + ec->slots.resize(cap); + } catch (...) { + fprintf(stderr, "ds4: CUDA streaming expert cache metadata allocation failed\n"); + (void)cudaFree(gate_ptr); + (void)cudaFree(up_ptr); + (void)cudaFree(down_ptr); + cuda_stream_expert_cache_release_all(); + return NULL; + } + + ec->valid = 1; + ec->capacity = cap; + ec->count = 0; + ec->tick = 0; + ec->gate_expert_bytes = gate_expert_bytes; + ec->down_expert_bytes = down_expert_bytes; + ec->gate_ptr = gate_ptr; + ec->up_ptr = up_ptr; + ec->down_ptr = down_ptr; + ec->gate_capacity = + (uint64_t)cap * gate_expert_bytes; + ec->up_capacity = + (uint64_t)cap * gate_expert_bytes; + ec->down_capacity = + (uint64_t)cap * down_expert_bytes; + return ec; + } + return NULL; +} + +static int cuda_stream_expert_cache_find( + cuda_stream_expert_cache *cache, + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t n_total_expert, + uint32_t expert, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + if (!cache || !cache->valid) return -1; + for (uint32_t i = 0; i < cache->capacity; i++) { + const cuda_stream_expert_cache_slot &slot = cache->slots[i]; + if (slot.valid && + slot.model_map == model_map && + slot.model_size == model_size && + slot.layer == layer && + slot.n_total_expert == n_total_expert && + slot.expert == expert && + slot.gate_offset == gate_offset && + slot.up_offset == up_offset && + slot.down_offset == down_offset && + slot.gate_expert_bytes == gate_expert_bytes && + slot.down_expert_bytes == down_expert_bytes) { + return (int)i; + } + } + return -1; +} + +static uint32_t cuda_stream_expert_cache_lru_slot( + cuda_stream_expert_cache *cache) { + for (uint32_t i = 0; i < cache->capacity; i++) { + if (!cache->slots[i].valid) return i; + } + uint32_t slot = 0; + uint64_t best_age = cache->slots[0].age; + for (uint32_t i = 1; i < cache->capacity; i++) { + if (cache->slots[i].age < best_age) { + best_age = cache->slots[i].age; + slot = i; + } + } + return slot; +} + +static int cuda_stream_expert_cache_copy_to_compact( + cuda_stream_expert_cache *cache, + uint32_t cache_slot, + uint32_t compact_slot, + char *compact_gate, + char *compact_up, + char *compact_down) { + const uint64_t gate_src = (uint64_t)cache_slot * cache->gate_expert_bytes; + const uint64_t down_src = (uint64_t)cache_slot * cache->down_expert_bytes; + const uint64_t gate_dst = (uint64_t)compact_slot * cache->gate_expert_bytes; + const uint64_t down_dst = (uint64_t)compact_slot * cache->down_expert_bytes; + return cuda_ok(cudaMemcpy(compact_gate + gate_dst, + cache->gate_ptr + gate_src, + (size_t)cache->gate_expert_bytes, + cudaMemcpyDeviceToDevice), + "streaming selected gate cache copy") && + cuda_ok(cudaMemcpy(compact_up + gate_dst, + cache->up_ptr + gate_src, + (size_t)cache->gate_expert_bytes, + cudaMemcpyDeviceToDevice), + "streaming selected up cache copy") && + cuda_ok(cudaMemcpy(compact_down + down_dst, + cache->down_ptr + down_src, + (size_t)cache->down_expert_bytes, + cudaMemcpyDeviceToDevice), + "streaming selected down cache copy"); +} + +static int cuda_stream_expert_cache_load_slot( + cuda_stream_expert_cache *cache, + const void *model_map, + uint64_t model_size, + uint32_t slot, + uint32_t layer, + uint32_t n_total_expert, + uint32_t expert, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + const uint64_t gate_src = + gate_offset + (uint64_t)expert * gate_expert_bytes; + const uint64_t up_src = + up_offset + (uint64_t)expert * gate_expert_bytes; + const uint64_t down_src = + down_offset + (uint64_t)expert * down_expert_bytes; + const uint64_t gate_dst = (uint64_t)slot * gate_expert_bytes; + const uint64_t down_dst = (uint64_t)slot * down_expert_bytes; + if (!cuda_model_copy_to_device_streamed(cache->gate_ptr + gate_dst, + model_map, + model_size, + gate_src, + gate_expert_bytes, + "cached moe_gate") || + !cuda_model_copy_to_device_streamed(cache->up_ptr + gate_dst, + model_map, + model_size, + up_src, + gate_expert_bytes, + "cached moe_up") || + !cuda_model_copy_to_device_streamed(cache->down_ptr + down_dst, + model_map, + model_size, + down_src, + down_expert_bytes, + "cached moe_down")) { + return 0; + } + cuda_stream_expert_cache_slot &entry = cache->slots[slot]; + entry.valid = 1; + entry.model_map = model_map; + entry.model_size = model_size; + entry.layer = layer; + entry.n_total_expert = n_total_expert; + entry.expert = expert; + entry.gate_offset = gate_offset; + entry.up_offset = up_offset; + entry.down_offset = down_offset; + entry.gate_expert_bytes = gate_expert_bytes; + entry.down_expert_bytes = down_expert_bytes; + entry.age = ++cache->tick; + return 1; +} + +static int cuda_stream_expert_cache_seed_one( + cuda_stream_expert_cache *cache, + const void *model_map, + uint64_t model_size, + uint32_t layer, + uint32_t n_total_expert, + uint32_t expert, + uint64_t gate_offset, + uint64_t up_offset, + uint64_t down_offset, + uint64_t gate_expert_bytes, + uint64_t down_expert_bytes) { + int cache_slot = cuda_stream_expert_cache_find(cache, + model_map, + model_size, + layer, + n_total_expert, + expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes); + if (cache_slot >= 0) { + cache->slots[(uint32_t)cache_slot].age = ++cache->tick; + return 1; + } + + const uint32_t load_slot = cuda_stream_expert_cache_lru_slot(cache); + const int append = !cache->slots[load_slot].valid; + if (!cuda_stream_expert_cache_load_slot(cache, + model_map, + model_size, + load_slot, + layer, + n_total_expert, + expert, + gate_offset, + up_offset, + down_offset, + gate_expert_bytes, + down_expert_bytes)) { + return 0; + } + if (append && cache->count < cache->capacity) cache->count++; + return 1; +} + static int cuda_stream_selected_cache_begin_load( const ds4_gpu_stream_expert_table *table, const int32_t *selected_ids, @@ -23252,33 +23885,116 @@ static int cuda_stream_selected_cache_begin_load( return 0; } + /* Consult the resident LRU expert cache (per-tier). On a hit, copy the + * expert's gate/up/down slab D2D from the LRU to the per-request compact + * cache (no SSD read). On a miss, stream the expert from SSD into a free + * (or least-recently-used) LRU slot, populate it, then copy D2D to the + * compact cache. Falls back to a direct SSD-to-compact load if the LRU + * could not be prepared or a cache populate step failed. */ + cuda_stream_expert_cache_note_size(table->gate_expert_bytes, + table->down_expert_bytes); + const uint32_t configured_cache_budget = + cuda_stream_expert_cache_configured_budget(); + const int use_global_cache = + cuda_stream_expert_cache_budget_visible_to_shared() && + configured_cache_budget != 0; + cuda_stream_expert_cache *expert_cache = use_global_cache ? + cuda_stream_expert_cache_prepare(table->gate_expert_bytes, + table->down_expert_bytes, + configured_cache_budget) : + NULL; + int expert_cache_disabled = expert_cache == NULL; + for (uint32_t i = 0; i < compact_ids.size(); i++) { const uint64_t expert = (uint32_t)compact_ids[i]; - const uint64_t gate_src = - table->gate_offset + expert * table->gate_expert_bytes; - const uint64_t up_src = - table->up_offset + expert * table->gate_expert_bytes; - const uint64_t down_src = - table->down_offset + expert * table->down_expert_bytes; const uint64_t gate_dst = (uint64_t)i * table->gate_expert_bytes; const uint64_t down_dst = (uint64_t)i * table->down_expert_bytes; - if (!cuda_model_copy_to_device_streamed( - sc->gate_ptr + gate_dst, - table->model_map, table->model_size, - gate_src, table->gate_expert_bytes, - "stream gate expert copy") || - !cuda_model_copy_to_device_streamed( - sc->up_ptr + gate_dst, - table->model_map, table->model_size, - up_src, table->gate_expert_bytes, - "stream up expert copy") || - !cuda_model_copy_to_device_streamed( - sc->down_ptr + down_dst, - table->model_map, table->model_size, - down_src, table->down_expert_bytes, - "stream down expert copy")) { - cuda_stream_selected_cache_invalidate(); - return 0; + int copied_from_global_cache = 0; + + if (!expert_cache_disabled) { + int cache_slot = + cuda_stream_expert_cache_find(expert_cache, + table->model_map, + table->model_size, + table->layer, + table->n_total_expert, + (uint32_t)expert, + table->gate_offset, + table->up_offset, + table->down_offset, + table->gate_expert_bytes, + table->down_expert_bytes); + if (cache_slot >= 0) { + expert_cache->slots[(uint32_t)cache_slot].age = + ++expert_cache->tick; + } else { + const uint32_t load_slot = + cuda_stream_expert_cache_lru_slot(expert_cache); + const int append = !expert_cache->slots[load_slot].valid; + if (cuda_stream_expert_cache_load_slot(expert_cache, + table->model_map, + table->model_size, + load_slot, + table->layer, + table->n_total_expert, + (uint32_t)expert, + table->gate_offset, + table->up_offset, + table->down_offset, + table->gate_expert_bytes, + table->down_expert_bytes)) { + if (append && expert_cache->count < expert_cache->capacity) { + expert_cache->count++; + } + cache_slot = (int)load_slot; + } else { + cuda_stream_expert_cache_invalidate(); + expert_cache_disabled = 1; + cache_slot = -1; + } + } + + if (cache_slot >= 0) { + copied_from_global_cache = + cuda_stream_expert_cache_copy_to_compact( + expert_cache, + (uint32_t)cache_slot, + i, + sc->gate_ptr, + sc->up_ptr, + sc->down_ptr); + if (!copied_from_global_cache) { + cuda_stream_expert_cache_invalidate(); + expert_cache_disabled = 1; + } + } + } + + if (!copied_from_global_cache) { + const uint64_t gate_src = + table->gate_offset + expert * table->gate_expert_bytes; + const uint64_t up_src = + table->up_offset + expert * table->gate_expert_bytes; + const uint64_t down_src = + table->down_offset + expert * table->down_expert_bytes; + if (!cuda_model_copy_to_device_streamed( + sc->gate_ptr + gate_dst, + table->model_map, table->model_size, + gate_src, table->gate_expert_bytes, + "stream gate expert copy") || + !cuda_model_copy_to_device_streamed( + sc->up_ptr + gate_dst, + table->model_map, table->model_size, + up_src, table->gate_expert_bytes, + "stream up expert copy") || + !cuda_model_copy_to_device_streamed( + sc->down_ptr + down_dst, + table->model_map, table->model_size, + down_src, table->down_expert_bytes, + "stream down expert copy")) { + cuda_stream_selected_cache_invalidate(); + return 0; + } } } if (!cuda_ok(cudaMemcpy(sc->slot_selected_ptr, @@ -27485,9 +28201,13 @@ extern "C" int ds4_gpu_stream_expert_cache_begin_selected_load( extern "C" uint32_t ds4_gpu_stream_expert_cache_budget_for_expert_size( uint64_t gate_expert_bytes, uint64_t down_expert_bytes) { - (void)gate_expert_bytes; - (void)down_expert_bytes; - return 0; + if (!cuda_stream_expert_cache_budget_visible_to_shared() || + cuda_stream_expert_cache_expert_bytes(gate_expert_bytes, + down_expert_bytes) == 0) { + return 0; + } + cuda_stream_expert_cache_note_size(gate_expert_bytes, down_expert_bytes); + return cuda_stream_expert_cache_configured_budget(); } extern "C" int ds4_gpu_tensor_copy_f32_to_f16(ds4_gpu_tensor *dst, uint64_t dst_offset, @@ -27598,20 +28318,26 @@ extern "C" void ds4_gpu_set_glm_model(bool enabled) { extern "C" void ds4_gpu_set_ssd_streaming(bool enabled) { g_ssd_streaming_mode = enabled ? 1 : 0; - /* Multi-tier: invalidate every tier's selected_cache, and when disabling - * free each tier's caches + staging pool on the device that owns them. - * cudaSetDevice before the release helpers so ssd_current() resolves to - * g_ssd[t]. Single-tier path (g_n_gpus == 0 at the time of the first - * call) is handled by the fallback loop iteration over n == 1, which - * runs on whatever device is current — matching the legacy behavior. */ + /* Multi-tier: reset runtime caps on EVERY tier and, when disabling, free + * each tier's caches on the device that owns them. cudaSetDevice before + * the release helpers so ssd_current() resolves to g_ssd[t]. The LRU + * itself is sized lazily on first use (cuda_stream_expert_cache_live_budget + * calls cudaMemGetInfo per-current-device), so each tier's cap comes out + * of that tier's free memory. */ const int n = g_n_gpus > 0 ? g_n_gpus : 1; int prev = -1; (void)cudaGetDevice(&prev); for (int t = 0; t < n; t++) { if (g_n_gpus > 0) (void)cudaSetDevice(g_gpu[t].device_id); + ds4_ssd_ctx *ssd = &g_ssd[t]; + ssd->runtime_cap = 0; + ssd->runtime_gate_bytes = 0; + ssd->runtime_down_bytes = 0; + ssd->memory_cap_notice = 0; cuda_stream_selected_cache_invalidate(); if (!g_ssd_streaming_mode) { cuda_stream_selected_cache_release(); + cuda_stream_expert_cache_release_all(); cuda_stream_selected_stage_release(); } } @@ -27619,7 +28345,25 @@ extern "C" void ds4_gpu_set_ssd_streaming(bool enabled) { } extern "C" void ds4_gpu_set_streaming_expert_cache_budget(uint32_t experts) { - (void)experts; + /* g_stream_expert_budget_override is GLOBAL — the user cap applies + * uniformly to every tier. */ + g_stream_expert_budget_override = experts; + /* Invalidate + release caches on every tier so each tier reallocates with + * the new budget on its own device. */ + const int n = g_n_gpus > 0 ? g_n_gpus : 1; + int prev = -1; + (void)cudaGetDevice(&prev); + for (int t = 0; t < n; t++) { + if (g_n_gpus > 0) (void)cudaSetDevice(g_gpu[t].device_id); + ds4_ssd_ctx *ssd = &g_ssd[t]; + ssd->runtime_cap = 0; + ssd->runtime_gate_bytes = 0; + ssd->runtime_down_bytes = 0; + ssd->memory_cap_notice = 0; + cuda_stream_selected_cache_invalidate(); + cuda_stream_expert_cache_release_all(); + } + if (prev >= 0) (void)cudaSetDevice(prev); } extern "C" void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) { @@ -27627,27 +28371,27 @@ extern "C" void ds4_gpu_set_streaming_expert_cache_expert_bytes(uint64_t bytes) } extern "C" uint32_t ds4_gpu_stream_expert_cache_configured_count(void) { - return 0; + if (!cuda_stream_expert_cache_budget_visible_to_shared()) return 0; + return cuda_stream_expert_cache_configured_budget(); } extern "C" uint32_t ds4_gpu_stream_expert_cache_current_count(void) { /* Status readout: report the active tier's resident count. Per-layer * dispatch has set the device; ssd_current() returns the right tier. */ - cuda_stream_selected_cache *sc = &ssd_current()->selected_cache; - return sc->valid ? sc->compact_count : 0; + return ssd_current()->expert_cache.count; } extern "C" void ds4_gpu_stream_expert_cache_reset_route_hotness(void) { } extern "C" void ds4_gpu_stream_expert_cache_release_resident(void) { - /* Release resident caches on every tier (each on its own device). */ + /* Release resident expert caches on every tier (each on its own device). */ const int n = g_n_gpus > 0 ? g_n_gpus : 1; int prev = -1; (void)cudaGetDevice(&prev); for (int t = 0; t < n; t++) { if (g_n_gpus > 0) (void)cudaSetDevice(g_gpu[t].device_id); - cuda_stream_selected_cache_release(); + cuda_stream_expert_cache_release_all(); } if (prev >= 0) (void)cudaSetDevice(prev); } @@ -27656,7 +28400,43 @@ extern "C" int ds4_gpu_stream_expert_cache_seed_selected( const ds4_gpu_stream_expert_table *table, const int32_t *selected_ids, uint32_t n_selected) { - (void)table; (void)selected_ids; (void)n_selected; + if (!g_ssd_streaming_mode) return 1; + if (!table || !selected_ids || n_selected == 0 || + n_selected > table->n_total_expert || + !cuda_stream_selected_ranges_valid(table)) { + return 0; + } + + cuda_stream_expert_cache *cache = + cuda_stream_expert_cache_prepare(table->gate_expert_bytes, + table->down_expert_bytes, + n_selected); + if (!cache) return 1; + for (uint32_t i = 0; i < n_selected; i++) { + if (selected_ids[i] < 0 || + (uint32_t)selected_ids[i] >= table->n_total_expert) { + fprintf(stderr, + "ds4: CUDA streaming seed selected expert id %d is outside 0..%u at layer %u\n", + selected_ids[i], + table->n_total_expert, + table->layer); + return 0; + } + if (!cuda_stream_expert_cache_seed_one(cache, + table->model_map, + table->model_size, + table->layer, + table->n_total_expert, + (uint32_t)selected_ids[i], + table->gate_offset, + table->up_offset, + table->down_offset, + table->gate_expert_bytes, + table->down_expert_bytes)) { + cuda_stream_expert_cache_invalidate(); + return 1; + } + } return 1; } From e26e2f06f54ac8c5b04be1accee0c50d3680d43e Mon Sep 17 00:00:00 2001 From: Anton Sokolchenko Date: Tue, 21 Jul 2026 14:26:55 +0000 Subject: [PATCH 4/4] fix(dspark): allocate DSpark scratch on exec_tier for SSD streaming 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 --- ds4.c | 87 ++++++++++++++++++++++++++++++++++++++--------------------- 1 file changed, 57 insertions(+), 30 deletions(-) diff --git a/ds4.c b/ds4.c index 0a8803f7f..d2aac8c2f 100644 --- a/ds4.c +++ b/ds4.c @@ -15881,28 +15881,51 @@ static bool metal_graph_configure_dspark_capture( return false; } + /* DSpark support weights, drafter kernels, and target-hidden capture all + * run on the configured executor tier (see ds4_session_prepare_dspark_draft, + * ds4.c:59635). The legacy ds4_gpu_tensor_alloc(bytes) always allocates on + * tier 0 (ds4_cuda.cu:2592), so under multi-tier SSD streaming — where the + * DSpark executor is tier 1 — every capture/drafter kernel that touched + * this scratch (hc_weighted_sum_tensor reading cur_hc on tier 1 against + * dspark_target_hidden / dspark_hc_mean_weights on tier 0) raised a sticky + * CUDA illegal-memory-access the moment the first target layer (40) was + * captured during decode. Allocate all DSpark scratch on exec_tier so the + * capture path stays on a single device. Single-tier (Metal/CPU) builds + * keep working because ds4_gpu_tensor_alloc_ptr_on(0, ...) collapses to + * the legacy allocator there (ds4.c:173). */ + const int exec_tier = + g->dspark_exec_tier >= 0 && g->dspark_exec_tier < DS4_MAX_GPUS + ? g->dspark_exec_tier : 0; + g->dspark_hc_mean_weights = - ds4_gpu_tensor_alloc((uint64_t)DS4_N_HC * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)DS4_N_HC * sizeof(float)); g->dspark_hc_mean_rows = - ds4_gpu_tensor_alloc((uint64_t)g->prefill_cap * - DS4_N_HC * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)g->prefill_cap * + DS4_N_HC * sizeof(float)); g->dspark_target_hidden = - ds4_gpu_tensor_alloc((uint64_t)dw->target_layer_count * - DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)dw->target_layer_count * + DS4_N_EMBD * sizeof(float)); g->dspark_target_hidden_batch = - ds4_gpu_tensor_alloc((uint64_t)dw->target_layer_count * - g->prefill_cap * - DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)dw->target_layer_count * + g->prefill_cap * + DS4_N_EMBD * sizeof(float)); if (dw->block_size != 0 && dw->block_size <= DS4_DSPARK_MAX_BLOCK_SIZE) { g->dspark_stage0_packed = - ds4_gpu_tensor_alloc(((uint64_t)dw->block_size + 1u) * - dw->target_layer_count * - DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + ((uint64_t)dw->block_size + 1u) * + dw->target_layer_count * + DS4_N_EMBD * sizeof(float)); } g->dspark_stage0_proj = - ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)DS4_N_EMBD * sizeof(float)); g->dspark_main_x = - ds4_gpu_tensor_alloc((uint64_t)DS4_N_EMBD * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)DS4_N_EMBD * sizeof(float)); if (!g->dspark_hc_mean_weights || !g->dspark_hc_mean_rows || !g->dspark_target_hidden || !g->dspark_target_hidden_batch || !g->dspark_stage0_proj || !g->dspark_main_x) { @@ -15911,20 +15934,26 @@ static bool metal_graph_configure_dspark_capture( if (dw->block_size != 0 && dw->block_size <= DS4_DSPARK_MAX_BLOCK_SIZE) { const uint64_t hc_dim = (uint64_t)DS4_N_HC * DS4_N_EMBD; g->dspark_draft_tokens = - ds4_gpu_tensor_alloc((uint64_t)dw->block_size * sizeof(int32_t)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)dw->block_size * sizeof(int32_t)); g->dspark_draft_hc = - ds4_gpu_tensor_alloc((uint64_t)dw->block_size * hc_dim * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)dw->block_size * hc_dim * sizeof(float)); g->dspark_target_hc = - ds4_gpu_tensor_alloc(hc_dim * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + hc_dim * sizeof(float)); g->dspark_stage_input_hc = - ds4_gpu_tensor_alloc((uint64_t)(dw->block_size + 1u) * - hc_dim * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)(dw->block_size + 1u) * + hc_dim * sizeof(float)); g->dspark_stage_output_hc = - ds4_gpu_tensor_alloc((uint64_t)dw->block_size * - hc_dim * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)dw->block_size * + hc_dim * sizeof(float)); g->dspark_position_ids = - ds4_gpu_tensor_alloc((uint64_t)(dw->block_size + 1u) * - sizeof(int32_t)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)(dw->block_size + 1u) * + sizeof(int32_t)); if (!g->dspark_draft_tokens || !g->dspark_draft_hc || !g->dspark_target_hc || !g->dspark_stage_input_hc || !g->dspark_stage_output_hc || !g->dspark_position_ids) { @@ -15933,8 +15962,9 @@ static bool metal_graph_configure_dspark_capture( if (dw->n_stages != 0 && g->raw_cap != 0) { for (uint32_t stage = 0; stage < dw->n_stages; stage++) { g->dspark_raw_cache[stage] = - ds4_gpu_tensor_alloc((uint64_t)g->raw_cap * - DS4_N_HEAD_DIM * sizeof(float)); + ds4_gpu_tensor_alloc_ptr_on(exec_tier, + (uint64_t)g->raw_cap * + DS4_N_HEAD_DIM * sizeof(float)); if (!g->dspark_raw_cache[stage]) return false; } g->dspark_cache_cap = g->raw_cap; @@ -55799,12 +55829,9 @@ static int ds4_engine_open_internal(ds4_engine **out, } if (opt->mtp_path && opt->mtp_path[0] && opt->distributed.role == DS4_DISTRIBUTED_NONE) { - if (e->ssd_streaming) { - fprintf(stderr, "ds4: --ssd-streaming is not compatible with --mtp yet\n"); - ds4_engine_close(e); - *out = NULL; - return 1; - } + /* SSD streaming + MTP/DSpark: previously refused upstream ("not compatible + * yet"). Allowed here — the support model loads resident while the main + * model streams; verify correctness end-to-end. */ model_open(&e->mtp_model, opt->mtp_path, graph_backend, true); ds4_dspark_summary dspark = {0}; e->support_kind =