cuda : fix CUB argsort corruption caused by in-place keys

[TheArchitectit]

Overview

argsort_f32_i32_cuda_cub() calls the one-shot CUB radix-sort API with d_keys_in == d_keys_out. That API is not in-place: it ping-pongs double buffers internally, and aliasing its ends lets it overwrite its own input mid-pass, corrupting the permutation on sm_50 with CCCL 2.x. The fix allocates a distinct keys-out scratch buffer at all six call sites.

Problem

argsort_f32_i32_cuda_cub() (ggml/src/ggml-cuda/argsort.cu) calls the one-shot DeviceRadixSort::SortPairs / SortPairsDescending API with d_keys_in == d_keys_out: the same scratch buffer (temp_keys) is passed as both key input and key output at every CUB call site (size query + actual sort, ASC + DESC, nrows == 1 + segmented variants).

The one-shot API is not an in-place sort: internally it runs a double-buffer ping-pong, and with the two ends aliased it partially overwrites its own input mid-pass (observed with CCCL 2.x on Maxwell). The result is a corrupted permutation: a mixture of valid indices and stale values from the keys buffer.

How it manifests

Two properties make this rare in practice:

• normal sampling argsorts on the CPU, so the CUDA CUB path is only exercised by ops that sort on the GPU (e.g. GGML_OP_TOP_K in backend sampling);
• the CUB path only runs when ncols > 1024 (smaller inputs take the bitonic path).

Observed failure: backend top_k over a 248,320-column vocab (DFlash speculative decoding on Qwen3.8-27B) on Maxwell (sm_50, CUDA 12.5, CCCL 2.x). The corrupted indices were fed to get_rows, which gathered with a garbage row id -> illegal memory access -> CUDA error: an illegal memory access (Xid 31, sticky context error). Non-deterministic across runs (depends on key values and scheduling).

Evidence from debugging the repro:

• compute-sanitizer memcheck pinned the fault to k_get_rows_float reading an unmapped address ~2.1 GB below the nearest allocation;
• bounds-check instrumentation on the get_rows ids input showed the top_k output containing float bit patterns instead of indices, and an A/B test (aliased vs distinct keys-out) localized the corruption to the argsort.

Fix

Allocate a distinct keys-out scratch buffer (temp_keys_out) from the same transient pool and pass it as d_keys_out at all 6 CUB call sites. Peak memory grows by at most one chunk (~64 MB) of transient scratch during the sort.

Testing

• Repro that crashed within 3 tokens (top_k ncols = 248320, temp 0.8, 4-way GPU split) now completes full generation; with instrumentation re-enabled, all top_k outputs validated as in-bounds indices (69/69 calls in a 40-token run).
• 3/3 stress runs (3 seeds x 200 tokens each) complete with zero CUDA errors and zero MMU faults.

Requirements

• I have read and agree with the contributing guidelines
• AI usage disclosure: YES — an AI assistant helped with sanitizer/log analysis and polishing this description; the diagnosis, fix, and every test result are the author's own and were verified hands-on against the repro.

[ggml-gh-bot]

Hi @TheArchitectit, thanks for your contribution!

Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:

• PR Template not respected: Please respect the template when creating a new pull request. Make sure to fill out all required sections.

• Multiple open PRs from a new contributor: We limit new contributors (those without a previously merged PR) to 1 open PR at a time. You currently have 2 open PRs.

---

Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below.

[fairydreaming]

@ORippler Can d_keys_in == d_keys_out cause problems on old CCCL release (2.5?) like reported in this PR?

[fairydreaming]

@TheArchitectit Tried to reproduce this in test-backend-ops with CUDA 12.5.1 but no luck:

  TOP_K(type=f32,ne=[248320,1,1,1],k=3,ties=0): OK
  TOP_K(type=f32,ne=[248320,1,1,1],k=3,ties=1): OK

compute-sanitizer shows no errors.

Did you use some custom CCCL installation?

[TheArchitectit]

@fairydreaming No custom CCCL — stock CUDA 12.5.20240515 from /opt/cuda-12.5 (CUB_VERSION 200400, i.e. CCCL 2.4.0), driver 580.178.04. Hardware: 4x Tesla M10 (sm_50) + 1x GTX 1080, model split 4-way.

On non-reproducibility with test-backend-ops: the corruption is non-deterministic and value-dependent — the aliased ping-pong only clobbers input ranges that overlap live output writes mid-pass, so synthetic TOP_K on a single row often passes. In our repro it crashed within 3 generated tokens with real logit distributions (vocab 248,320, temp 0.8, backend top_k on the speculative path), pinned via compute-sanitizer to k_get_rows_float reading ~2.1 GB below the nearest allocation, and localized to the argsort via an A/B on aliased vs distinct keys-out. Note the single-row ne=[248320,1,1,1] case exercises only one of the six patched call sites. Happy to provide the repro command + binary, or a screen-share if that's easier.

Re: the new-contributor PR limit flagged above — plan: I will keep my other open PR, #28097, in draft until this one is reviewed/approved. It is a companion to #27836 and blocked on it anyway, so it loses nothing by waiting. If you would prefer it closed and re-filed later instead, say the word.

[fairydreaming]

@TheArchitectit I see, if you have an example prompt and command that allows to reproduce this please share them. I mean if there is such a failure mode then it should appear long ago, so I'd like to reproduce this locally and investigate more.

[TheArchitectit]

@fairydreaming Here's the exact repro. The failure lives on the DFlash speculative-decoding path, which is why a synthetic single-row TOP_K in test-backend-ops won't hit it — three things have to line up: (1) ncols > 1024 so the CUB path (not bitonic) runs, (2) real logit distributions rather than synthetic keys, and (3) backend top_k on the speculative path (normal sampling argsorts on the CPU, so it never touches this code).

Command (llama.cpp-xhtoken fork, but the argsort code path is upstream-identical):

llama-server \
  -m Qwen3.8-27B-UD-Q4_K_M.gguf \
  --spec-draft-model Qwen3.8-27B-DFlash-bootstrap-Q8_0.gguf \
  --spec-type draft-dflash --spec-draft-n-max 3 \
  --device CUDA1,CUDA2,CUDA3,CUDA4 -ngl 99 -sm layer -ts 1,1,1,1 \
  -c 65536 -ctk q4_0 -ctv q4_0 -fa on --parallel 1 \
  --temp 0.7 --top-p 0.8 --top-k 20

Hardware: 4× Tesla M10 (sm_50) + 1× GTX 1080, CUDA 12.5 (CCCL 2.4.0), driver 580.178.04.

Crash signature (non-deterministic, within 1–3 generated tokens):

E CUDA error: an illegal memory access was encountered
E   current device: 3, in function ggml_backend_cuda_synchronize at .../ggml-cuda.cu:2530

The device varies run-to-run (3, then 4 on the next attempt) — consistent with a corrupted permutation fed to get_rows, landing on whichever die holds that KV slice.

Why test-backend-ops passes: ne=[248320,1,1,1] is a single row, which exercises only one of the six patched call sites, and with synthetic keys the ping-pong aliasing doesn't clobber live output. The corruption is value- and schedule-dependent, so a fixed synthetic input is exactly the case that won't reproduce.

Minimal repro without the full 27B model: the same top_k (ncols 248320) driven through the speculative path is what triggers it. I can post the two GGUF files and the exact prompt I used, or screen-share, whichever is easier.

[fairydreaming]

@TheArchitectit I tried the current mainline llama.cpp with the following command:

./bin/llama-server -m ~/ggufs/Qwen3.8-27B-IQ3_XXS.gguf --spec-draft-model ~/ggufs/Qwen3.8-27B-DFlash2-BF16.gguf --spec-type draft-dflash --spec-draft-n-max 3 -ngl 99 -sm layer -c 65536 -ctk q4_0 -ctv q4_0 -fa on --parallel 1 --temp 0.7 --top-p 0.8 --top-k 20

worked without any issues. Please reproduce the issue in mainline llama.cpp.