Pr/cuda moe cache

[giveen]

Overview

Do not merge until #287 is complete

• AI usage disclosure: Yes

[TheTom]

Re-reviewed against the merged base. Now that #287 is in, this is a 33-file, +9910/-323 diff instead of 209 files, which is far easier to read. Thanks for merging the base in and for dropping the dead k_idx restore, that resolves my earlier comment.

Local results, macOS/Metal

C++ build            : clean, 0 errors
metal shader compile : clean, 0 errors   (xcrun metal, not just cmake)
test-turbo-quant     : pass
test-chat            : pass
test-moe-cache-fit   : pass
test-arg-parser      : pass

--moe-cache shows up correctly in --help with all five modes. Nothing here regresses the non-CUDA path.

One ordering problem worth fixing before this lands

The !is_mla guard on auto-asymmetric is in #289, but the feature that needs it lands in #288. Traced across the stack:

Branch	!is_mla guard
merged base	absent
#288 (this PR)	absent
#289	present
#291	present

That matters because of your own Known Limitations entry: mixed K/V types do not work with the MoE cache, and the documented workaround is to pass identical types, specifically -ctk turbo4 -ctv turbo4 for DeepSeek-V4. But auto-asymmetric rewrites symmetric turbo K to q8_0 whenever the GQA ratio is 6 or higher, and DeepSeek-V4 is 64:1.

Confirmed on this exact build, Qwen3.6-27B, which sits right at the threshold:

$ llama-cli ... -ctk turbo4 -ctv turbo4
llama_kv_cache: auto-asymmetric: GQA ratio 6:1 (n_head=24, n_head_kv=4) —
  upgrading K from turbo4 to q8_0 to prevent quality degradation.

So between #288 landing and #289 landing, there is a window where --moe-cache exists and its documented workaround silently does not work on exactly the model the feature was benchmarked against. Moving that one-line guard from #289 into this PR closes the window, and it belongs here anyway since this is the PR that introduces the constraint.

Second thing, smaller but structural

test-moe-cache still does this on a non-CUDA build:

$ ./bin/test-moe-cache
SKIP: CUDA backend unavailable
$ echo $?
0

Exit 0 means ctest records a pass. On any non-CUDA runner this test is indistinguishable from a real one, while executing none of the 3260 lines it contains. That is what let the Metal and Vulkan backends ship uncompiled earlier in this stack.

Cheapest fix is to have it exit 77, which ctest understands as a genuine skip via SKIP_RETURN_CODE, so the summary says skipped rather than passed. Not blocking, but the current shape actively misleads.

What I still owe you

CUDA is still unverified on my side. Everything above is the non-CUDA path. I have not built this on a CUDA box, have not run test-moe-cache for real, and have not reproduced the soft-mode PP/TG table. Our bench has been occupied. I am not going to sign off on the CUDA half until I have actually run it, and I will report back here when I do.

Two questions on the numbers when I get there, so I can check the right things:

1. The soft-mode table shows PP improving 5.6% on DeepSeek V4 while dropping 6% on Laguna. The explanation given is that the dense-prioritised placement beats stock. Is that reproducible run to run, or within noise on a single -r 1 pass?
2. Since a cache hit routes an expert to the GPU and a miss leaves it on the CPU, two runs with the same seed and prompt can compute different experts on different devices depending on cache warmth. Is output expected to be bit-identical with the cache on versus off, or only close? Your testing guide says "byte-identical or very close", and those are quite different claims to test against.

[giveen]

@TheTom applying the fixes now.
The PP cost is a documented "regression" in exchange for a larger boost to TG. At high amount of TG speeds, the loss is negligible which results in a faster turn around time for the LLM. I'm seeing anywhere between (-)4-8% loss of PP in exchange for (+)20-30% TG, depending on the model.

[TheTom]

6d4a1db28 looks right. Both things I raised are addressed, and I checked the MLA guard properly rather than just eyeballing it.

Verified on macOS/Metal

build                : clean, 0 errors
metal shader compile : clean (xcrun metal, not just cmake)
test-turbo-quant     : pass
test-chat            : pass
test-moe-cache-fit   : pass
test-moe-cache       : exit 77  ->  ctest now reports SKIP, not PASS

The exit-77 change is exactly the right fix. A non-CUDA runner now says skipped instead of silently claiming a pass over 3260 lines it never ran.

On the MLA guard, which I did chase down

You used hparams.is_mla() alone:

if (k_is_turbo && !hparams.is_mla()) {

whereas #289 and llama-context.cpp:3736 both write hparams.is_mla() || arch == LLM_ARCH_DEEPSEEK4. I was ready to flag that as missing DeepSeek-V4, which would have been the one model that matters here. It is not missing, and I want to record why, because the reasoning is not obvious from the call site.

n_embd_head_k_mla_impl is assigned in exactly one place in the tree, dsv4_make_k_only() in llama-kv-cache-dsv4.cpp:831. That runs on the per-cache hparams copies at lines 1201 to 1217, before each llama_kv_cache is constructed with them. So by the time your guard executes inside the llama_kv_cache ctor, the member hparams is the DSv4-modified copy and is_mla() is already true.

The || arch == LLM_ARCH_DEEPSEEK4 in llama-context.cpp is needed there for the opposite reason: that check runs against model->hparams, the original, which never gets the fields set. Both spellings are correct in their own context. Nothing to change.

I did try an alternative fix that hoisted a shared is_mla above the block, and yours is better, since calling hparams.is_mla() at the site avoids introducing a declaration whose scope has to reach 90 lines up. Discarded mine.

On the tradeoff

Tom has said the PP-for-TG tradeoff is fine by him, decode speed is what he is after, so please treat my earlier question about whether the DeepSeek-V4 PP gain is reproducible as withdrawn. Not worth your time.

The one thing I would still like an answer on, because it changes what I test rather than whether I approve: with the cache on, is output expected to be bit-identical to cache-off, or only close? A hit computes an expert on the GPU and a miss on the CPU, so warmth alone could change results between runs. Your testing guide says "byte-identical or very close" and those need different oracles. If close is expected, I will compare distributions rather than diff tokens.

Where this leaves the PR

Non-CUDA is clean from my side and I have no blocking comments left. CUDA remains unverified by me, the bench has not been free. That is the only thing standing between this and my sign-off, and it is on me, not you.

[TheTom]

CUDA verification done at last, on a GB10 (Grace Blackwell, sm_121). Two findings, one of which I have pushed a fix for as 706d2a9a4.

1. The suite could not run at all on this hardware

test-moe-cache printed SKIP: CUDA backend unavailable on a machine with a working CUDA device that test-backend-ops happily uses. I instrumented find_cuda_device() to find out why:

[probe] dev_count=2
[probe] dev 0: reg='CUDA' type=2 desc='NVIDIA GB10'
[probe] dev 1: reg='CPU'  type=0 desc='CPU'

The lookup requires ggml_backend_dev_type(device) == GGML_BACKEND_DEVICE_TYPE_GPU. Type 2 is IGPU, not GPU:

GGML_BACKEND_DEVICE_TYPE_CPU   = 0
GGML_BACKEND_DEVICE_TYPE_GPU   = 1   // dedicated memory
GGML_BACKEND_DEVICE_TYPE_IGPU  = 2   // integrated, host memory
GGML_BACKEND_DEVICE_TYPE_ACCEL = 3

GB10 is unified memory, so IGPU is the correct classification. The consequence is that the entire MoE cache suite silently skips on Grace Blackwell, Jetson/Orin, and any integrated GPU, which is a meaningful slice of the hardware this feature is aimed at. test-backend-ops sidesteps it by excluding CPU rather than requiring GPU.

Fixed at both call sites in 706d2a9a4. With it, the suite runs:

37 OK
 2 SKIP   (cache-route-override: one CUDA device; cache-backend-unload: static backend)
 1 FAIL   (cache-shared-budget)
exit 1

Worth noting how this compounds with the exit-77 change from 6d4a1db28, which is otherwise the right fix: before, a skip reported as a pass; now it reports as a skip. Neither spelling would have told you the test was skipping on hardware that should have run it. The device filter was the actual leak.

2. cache-shared-budget fails, deterministically

Three consecutive runs, identical:

run1: cache-shared-budget: unexpected claim behavior
run1: cache-shared-budget: FAIL
run2: cache-shared-budget: unexpected claim behavior
run2: cache-shared-budget: FAIL
run3: cache-shared-budget: unexpected claim behavior
run3: cache-shared-budget: FAIL

Not flaky, so not a timing artifact. What the test does:

ggml_backend_dev_memory(cuda->device, &free_bytes, &total_bytes);
const std::string reserve = std::to_string(free_mib - 24);
set_env("GGML_CUDA_MOE_CACHE_RESERVE_MB", reserve.c_str());
...
const bool divided = count_field_at_least(shared_log, "granted=", 8) == 2 &&
    max_field_value(shared_log, "granted=") <= 16;

It reserves all but 24 MiB and expects two sessions to divide that remainder, each granted between 8 and 16 MiB. On this box free_mib is about 56,000, so the reserve is roughly 55,976 MiB.

I am not sure yet whether this is a product bug or a test assumption that does not hold on unified memory, and I do not want to guess. The reason I suspect the UMA path specifically is that the CUDA backend on this device computes availability through its own route:

ggml_backend_cuda_get_available_uma_memory: final available_memory_kb: 56951964

If the cache's budget logic derives from that while the test derives its reserve from ggml_backend_dev_memory, then reserving "free minus 24" would not leave 24 MiB and the granted values land outside the expected band. That is a hypothesis fitted to the evidence, not something I have traced.

What would settle it, in order of cost: print the actual granted= values on failure rather than only a boolean verdict, which would immediately show whether the split is wrong or the budget is; then check whether the cache's budget path and ggml_backend_dev_memory agree on a UMA device.

Happy to run either if you tell me which you would prefer. This box is available now.

What else passed

The 37 that pass include the substantive ones: the fused SwiGLU family, multi-token, invalidation, eviction, lifecycle, CPU overlap, shape liveness and inventory, admission policy, and the single-GPU auto path. test-moe-cache-fit and test-turbo-quant both pass, and the CUDA build is clean at -DCMAKE_CUDA_ARCHITECTURES=121.

So the cache machinery itself looks healthy on CUDA. The open item is the budget-sharing path on unified memory, plus whatever the granted values turn out to say.

[TheTom]

5c48fc288 fixes it. Verified on the GB10 that produced the failure.

CUDAToolkit: found 13.0.88, GGML_CUDA=1 compiled into the test
(cudaMemGetInfo present in the binary: 3 refs)

run 1:  cache-shared-budget: OK    38 OK, 0 FAIL, 2 SKIP    exit 0
run 2:  cache-shared-budget: OK    38 OK, 0 FAIL, 2 SKIP    exit 0

First fully green test-moe-cache run this hardware has ever produced. For the record, the progression was: suite skipped entirely (device-type filter), then 37 OK / 1 FAIL once it could run, now 38 OK / 0 FAIL. The two remaining skips are the legitimate ones, cache-route-override needing a second CUDA device and cache-backend-unload needing a dynamic backend.

And thank you for tracing the actual mechanism. I had only got as far as "the cache and the test seem to be reading different memory numbers on UMA" and explicitly flagged it as a hypothesis I had not confirmed. You confirmed it and fixed the right side: matching the cache's own query rather than loosening the expected band, which would have hidden it.

One thing to fix before this lands

cudaSetDevice() is called and never restored:

const cudaError_t err = cudaSetDevice(physical >= 0 ? (int) physical : 0);

and the call order in main is:

3244:    ok &= run_shared_budget(cuda, cpu, capture);
3245:    ok &= run_route_override(cuda_device, cuda, cpu, capture);

So run_shared_budget mutates the calling thread's current CUDA device and leaves it mutated, immediately before run_route_override, which is the one test that specifically exercises routing across two devices.

Neither of us can see this. On a single-GPU box cache-route-override prints SKIP (one CUDA device) and never runs, which is what happens on my GB10 and, I would guess, on your 5090. It would only surface on a two-card machine, which is the configuration the MoE cache was originally pitched for in the upstream docs.

Suggest saving and restoring:

int prev_device = 0;
cudaGetDevice(&prev_device);
// ... cudaSetDevice / cudaMemGetInfo ...
cudaSetDevice(prev_device);

Related and in the same two lines: when cuda_physical_device() returns -1 the code falls back to device 0, which on a multi-GPU box would report the wrong device's memory. Same blind spot.

Both are cheap, and I would rather they go in now than become a confusing multi-GPU-only failure later.

Where that leaves my review

Everything I can test is now green:

CUDA (GB10, sm_121)   build clean, test-moe-cache 38/0/2, test-moe-cache-fit pass, test-turbo-quant pass
macOS / Metal         build clean, shaders compile clean, test-chat / turbo-quant / moe-cache-fit pass

No blocking objections from me on the code. The remaining CI red on this PR is ubuntu-22-hip-quality-check, which is the pre-existing nodiscard bug in fattn-common.cuh that #293 fixes and that has nothing to do with this branch.

I have not benchmarked the soft-mode PP/TG table and am not planning to. Tom has said the tradeoff is what he wants, so re-measuring it would not change any decision.

[TheTom]

41a4eae09 verified on the GB10. Build clean, two consecutive runs:

run1 exit=0   38 OK, 0 FAIL, 2 SKIP
run2 exit=0   38 OK, 0 FAIL, 2 SKIP
  cache-shared-budget:  OK
  cache-route-override: SKIP (one CUDA device)

Both points covered, and you went further than I asked by restoring the device on the cudaMemGetInfo error path too, which I had not thought of. Using prev_device rather than 0 as the fallback is the right call for the same reason.

Being straight about the limits of this particular check: the save/restore only matters on a two-GPU box, and this one has a single device, so cache-route-override still skips. What I have shown is that the change does not break the single-GPU path, not that it fixes the multi-GPU one. The reasoning for the latter is sound and the code is plainly correct by inspection, but neither of us has hardware that exercises it. Worth someone with two cards running the suite before that path is trusted.

Sign-off from me

No blocking objections to #288. Everything I have hardware for is green:

CUDA, GB10 sm_121	build clean; test-moe-cache 38/0/2; test-moe-cache-fit pass; test-turbo-quant pass
macOS, Metal	build clean; shaders compile clean via xcrun metal; test-chat, test-turbo-quant, test-moe-cache-fit pass

For anyone reading this thread later, the three test-infrastructure defects found along the way are worth remembering as a set, because each one hid the next:

1. test-moe-cache returned 0 when it skipped, so ctest recorded a pass. Fixed to 77.
2. It required GGML_BACKEND_DEVICE_TYPE_GPU, so it silently skipped on every unified-memory and integrated device including GB10 and Jetson. Fixed to accept IGPU.
3. It derived its VRAM reserve from ggml_backend_dev_memory while the cache itself queries cudaMemGetInfo, which disagree on UMA. Fixed to match the cache.

Only after all three did a single line of the 3260-line suite actually execute on this hardware, and then it passed. The feature was fine the whole time; the harness could not see it.

Remaining CI red on this PR is ubuntu-22-hip-quality-check, the pre-existing nodiscard bug in fattn-common.cuh that #293 fixes and that is unrelated to this branch.

[TheTom]

Heads-up and an apology: I rewrote this branch's commit metadata and force-pushed, which I should have asked about first.

What changed: only the author/committer email on my own three commits, which were wrongly attributed to a work address. Your two commits are untouched and still authored by you, and every tree is byte-identical, so there is no content difference at all. I used --force-with-lease pinned to the previous SHA, so a concurrent push from you would have been rejected rather than overwritten.

What it means for you: if you have this branch checked out locally, the SHAs moved:

419c17937  ->  22a370e89
git fetch origin && git reset --hard origin/pr/cuda-moe-cache

Sorry for the churn. I will not force-push here again.

The one upside is that it triggered a completely fresh CI run, which re-runs windows (x64-openblas). That job failed once on test-thread-safety with a heap-corruption exit, and the evidence said flake rather than a real defect: the test is untouched by this PR, contains no reference to the cache, the cache is CUDA-only so it cannot even execute on a CPU-only openblas build, the sibling windows (x64-cpu-static) passes the same suite, and #289 with the same cache code passes openblas. A clean second run would settle it, and I could not re-run the single job because the stuck self-hosted queue keeps the workflow from ever completing.

Current state on 22a370e89: no failures yet, everything else either green or still running.

[giveen]

@TheTom No worries, I dont need "credit" or anything :) I just like working on the project.

[TheTom]

Appreciated. To be clear though, the apology was for force-pushing your branch without asking, not about credit. The authorship line was me saying I had not clobbered your work, and it came out sounding like the other thing.

Either way it will not happen again. CI is running clean on the new head so far, only the pre-existing HIP nodiscard red.