Skip to content

cuda: enable streaming auto cache (implement recommended_working_set_size) - #488

Open
riccardo-galbani wants to merge 1 commit into
antirez:mainfrom
riccardo-galbani:fix-cuda-auto-cache-guard
Open

cuda: enable streaming auto cache (implement recommended_working_set_size)#488
riccardo-galbani wants to merge 1 commit into
antirez:mainfrom
riccardo-galbani:fix-cuda-auto-cache-guard

Conversation

@riccardo-galbani

Copy link
Copy Markdown

ds4_backend_supports_streaming_auto_cache() only allowed the SSD
streaming auto cache planner to run under DS4_BACKEND_METAL, or
under DS4_BACKEND_CUDA when built with DS4_ROCM_BUILD — which a
plain make cuda-generic (nvcc) build never defines. As a result,
CUDA users always had to pass --ssd-streaming-cache-experts
explicitly, and ds4_gpu_recommended_working_set_size() in
ds4_cuda.cu was an unimplemented stub returning 0.

This implements the CUDA working set size using cudaMemGetInfo's
total device memory (the closest analogue to Metal's
recommendedMaxWorkingSetSize), and extends the guard so CUDA can use
the same auto cache planner Metal already has.

Tested on a CUDA GPU with 8GB VRAM with DeepSeek V4 Flash (the
project's reference model):

ds4: SSD streaming auto cache budget
ds4: cuda recommends 7.62 GiB working set
ds4: using 80% total for model + cached experts: 6.10 GiB
ds4: non-routed weights: 8.20 GiB
ds4: routed expert size: 6.75 MiB
ds4: cached expert count: 1 (0.01 GiB)

@jhohertz

jhohertz commented Jul 3, 2026

Copy link
Copy Markdown

This worked for me on a 7800 XT but I also needed some stuff from #461 like these flags: DS4_ROCM_STREAM_FREE_RESERVE_GIB=1 DS4_CUDA_Q8_F16_CACHE_MB=0.

It's not exactly fast, ~33 t/s PP, 5-6 on TG out of the gate. I just wanted to see if I could. :)

@riccardo-galbani

Copy link
Copy Markdown
Author

Tested on NVIDIA GB10 (DGX Spark, Grace Blackwell, 128GB unified
memory, sm_121, CUDA 13.0, driver 580.159.03) — the unified-memory
open question raised earlier is now resolved with real data.

cudaMemGetInfo reports the full system memory correctly on this
hardware-coherent unified-memory system:

ds4: SSD streaming auto cache budget
ds4: cuda recommends 121.69 GiB working set
ds4: using 80% total for model + cached experts: 97.35 GiB
ds4: non-routed weights: 8.20 GiB
ds4: routed expert size: 6.75 MiB
ds4: cached expert count: 11008 (72.56 GiB)

The initial plan (11008 experts) is then correctly refined at runtime
by the existing capping logic based on actually available memory:

ds4: CUDA streaming expert cache capped from 11008 to 4747 experts
(available 47.29 GiB, reserve 16.00 GiB, 6.75 MiB/expert)

With the auto-sized cache, generation reaches 8.4 t/s once the expert
cache is warm (q2 model, short prompts, no manual tuning flags at
all). The same binary auto-sizes correctly on an 8GB discrete laptop
GPU and on this 128GB unified-memory system.

So the PR is validated on both memory architectures: discrete VRAM
(RTX 4070 Laptop 8GB) and hardware-coherent unified memory (GB10).

@iSevenDays

Copy link
Copy Markdown

This PR won't work for multi-gpu systems.
Here is my variant for multi-gpu:

extern "C" uint64_t ds4_gpu_recommended_working_set_size(void) {
    /* GLM graph memory guard: on this backend the model weights are
     * distributed across all devices by the multi-tier placement, so the
     * relevant budget is the aggregate VRAM. */
    int n = 0;
    if (cudaGetDeviceCount(&n) != cudaSuccess || n <= 0) return 0;
    size_t free_b = 0, total_b = 0;
    if (cudaMemGetInfo(&free_b, &total_b) != cudaSuccess) return 0;
    return (uint64_t)total_b * (uint64_t)n;        // ← aggregate across all GPUs
}

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants