Skip to content

CUDA: enable the CUB path on HIP via hipCUB - #26592

Open
Geramy wants to merge 11 commits into
ggml-org:masterfrom
Geramy:Geramy/argsort_top_k_cub_on_hip
Open

CUDA: enable the CUB path on HIP via hipCUB#26592
Geramy wants to merge 11 commits into
ggml-org:masterfrom
Geramy:Geramy/argsort_top_k_cub_on_hip

Conversation

@Geramy

@Geramy Geramy commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Overview

Since the CUB-based ARGSORT/TOP_K paths were never running on a HIP based device, I have enabled them and added a few missing hip graph functions. Originally those ops fell back to the shared-memory bitonic sort. That sort uses one thread per element, so it is limited to rows of 1024 elements; anything wider was reported as unsupported and ran on the CPU.

Additional information

I have added stream-capture functions missing from vendors/hip.h which is required in CUB argsort.

DeepSeek-V4-Flash UD-IQ1_S — 284B MoE, 78 GB

testdepthmain (t/s)PR (t/s)Δ (t/s)change
pp512
prompt processing
0162.28 ± 4.83163.17 ± 3.21+0.89+0.55% ± 3.59
2048156.02 ± 1.28156.67 ± 0.85+0.65+0.41% ± 0.99
4096139.08 ± 0.48146.72 ± 1.66+7.64+5.49% ± 1.25
6192132.30 ± 0.88140.59 ± 0.99+8.29+6.26% ± 1.03
tg32
token generation
016.05 ± 0.1516.11 ± 0.11+0.06+0.40% ± 1.16
204814.24 ± 0.1314.34 ± 0.01+0.11+0.74% ± 0.95
409614.64 ± 0.2515.55 ± 0.12+0.91+6.22% ± 2.00
619214.68 ± 0.0815.65 ± 0.15+0.96+6.56% ± 1.20

Qwen3.6-35B-A3B UD-Q4_K_XL

testdepthmain (t/s)PR (t/s)Δ (t/s)change
pp512
prompt processing
02064.78 ± 967.802061.34 ± 948.84-3.44-0.17% ± 65.58
20481890.68 ± 868.161884.03 ± 856.98-6.65-0.35% ± 64.41
40961766.93 ± 781.341748.30 ± 774.43-18.63-1.05% ± 61.93
61921644.69 ± 716.901636.84 ± 703.89-7.84-0.48% ± 60.94
tg32
token generation
054.66 ± 4.5654.57 ± 3.53-0.09-0.17% ± 10.54
204855.32 ± 4.6955.28 ± 4.62-0.04-0.08% ± 11.90
409655.09 ± 4.8955.03 ± 4.82-0.06-0.11% ± 12.46
619254.69 ± 4.9454.61 ± 4.91-0.08-0.15% ± 12.72

Requirements

  1. AMD GPU and 2* AMD GPUs to test tensor splitting.
  2. 110GB+ VRAM
  3. Deepseek v4 flash

@IMbackK @am17an @pwilkin

The CUB-based ARGSORT/TOP_K (and sum/mean/cumsum) paths were compiled
out for HIP, so those ops fell back to the shared-memory bitonic sort.
That sort uses one thread per element, so it is limited to rows of 1024
elements; anything wider was reported as unsupported and ran on the CPU.
For deepseek4 that meant every lightning-indexer TOP_K above 1024
context went to the host, adding a graph split per layer.

hipCUB provides the device-wide sorts these paths need - the rocPRIM
backend has DeviceRadixSort, DeviceSegmentedRadixSort,
DeviceSegmentedSort, DeviceReduce and DeviceScan - so include it through
a small cub-compat shim that aliases the cub namespace, and define
GGML_CUDA_USE_CUB for HIP as well.

hipCUB has no DeviceTopK and does not define CCCL_*, so
CUB_TOP_K_AVAILABLE and STRIDED_ITERATOR_AVAILABLE stay disabled there:
TOP_K keeps using the argsort + copy path and argsort keeps using the
init_offsets kernel instead of a strided iterator.

vendors/hip.h gains the stream-capture defines that the CUB argsort path
needs now that it is compiled for HIP.
Comment thread ggml/src/ggml-cuda/cub-compat.cuh Outdated
@fairydreaming

fairydreaming commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

@ggml-org/ci Looks like CI hip workflows need to be updated to install hipcub dev package for this. Not sure about HIP Windows release, but I guess that's something to worry about later.

@Neresco

Neresco commented Aug 4, 2026

Copy link
Copy Markdown

So test this PR here with 4x gfx1200 9060xt and 1x gfx1151.
Slight increase in VRAM usage over the Prompt Processing (PP).
~400MB over 65k Context but i have not closely observed.
Two times my Window Manager crashed before i recognize the reason.
Oddly not the GPU's where the Monitors are plugged in had too much VRAM.

No crash at 68727 PP context.
It let it run further now.

@Geramy

Geramy commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

@fairydreaming isn't hipcub included in rocm?

@fairydreaming

fairydreaming commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

@fairydreaming isn't hipcub included in rocm?

@Geramy No idea, my adventure with AMD GPUs ended during Radeon R9 290 times. By the way, could you run some test-backend-ops test_top_k test cases with large shapes on AMD GPU with this PR to make sure there are no issues for huge tensors? For example these two:

test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {262144, 8192, 1, 1}, 1024));
test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {1048576, 512, 1, 1}, 2048));

Thanks!

@Geramy

Geramy commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

@fairydreaming from AMD GitHub page directly hipCUB is included with the ROCm Core SDK on Linux and Windows.
Your tests also passed, I'm moving this from draft to review.

TOP_K(type=f32,ne=[262144,8192,1,1],k=1024,ties=0): OK
TOP_K(type=f32,ne=[1048576,512,1,1],k=2048,ties=0): OK
11/11 tests passed
3/3 backends passed
OK

@Geramy
Geramy marked this pull request as ready for review August 4, 2026 19:18
@Geramy
Geramy requested review from a team and IMbackK as code owners August 4, 2026 19:18
@Neresco

Neresco commented Aug 4, 2026

Copy link
Copy Markdown

Cline is working now including reasoning and cache with up to 172k tokens for me over rpc.

Edit: by 193k tokens ca. 600mb VRAM increase per GPU i can observe.
So avoid too tight fits for AMD GPU's.

@Geramy
Geramy requested a review from fairydreaming August 4, 2026 19:30
@fairydreaming

Copy link
Copy Markdown
Contributor

@fairydreaming from AMD GitHub page directly `hipCUB is included with the ROCm Core SDK on Linux and Windows.

Well for some reason hipCUB header files were missing in the CI container during compilation, I see that it installs Ubuntu rocblas-dev and hipblas-dev packages, so maybe they also need to install libhipcub-dev or something.

@Geramy

Geramy commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

@fairydreaming from AMD GitHub page directly `hipCUB is included with the ROCm Core SDK on Linux and Windows.

Well for some reason hipCUB header files were missing in the CI container during compilation, I see that it installs Ubuntu rocblas-dev and hipblas-dev packages, so maybe they also need to install libhipcub-dev or something.

Yeah perhaps it depends on the ROCm version we are on too. Who will or can take care of that?

@fairydreaming

Copy link
Copy Markdown
Contributor

@fairydreaming from AMD GitHub page directly `hipCUB is included with the ROCm Core SDK on Linux and Windows.

Well for some reason hipCUB header files were missing in the CI container during compilation, I see that it installs Ubuntu rocblas-dev and hipblas-dev packages, so maybe they also need to install libhipcub-dev or something.

Yeah perhaps it depends on the ROCm version we are on too. Who will or can take care of that?

@Geramy I don't know, if no one fixes it I'll see what I can do tomorrow (going to sleep soon).

@remeh

remeh commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

Confirmed to work on a Strix Halo: with this patch the tg speed stays correct even after >4-5k context (instead of having a big drop to 5 tok/s). 🎉

However, I'm wondering if the prefill speed doesn't get worse faster than with #26493, but that can be totally wrong, don't read too much into this.

@fairydreaming
fairydreaming requested a review from a team as a code owner August 5, 2026 06:38
Comment thread ggml/src/ggml-cuda/vendors/hip.h Outdated
@github-actions github-actions Bot added the devops improvements to build systems and github actions label Aug 5, 2026
@fairydreaming

Copy link
Copy Markdown
Contributor

HIP quality check CI found some missing CUB function call return value error checks, so I added them. Compilation worked fine with added hipcub-dev package.

Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
@Geramy

Geramy commented Aug 5, 2026

Copy link
Copy Markdown
Contributor Author

@fairydreaming thanks! I didn’t see any of the CIs run or get approved to run did you do a few manually in actions?

@fairydreaming

Copy link
Copy Markdown
Contributor

@Geramy Yeah, in this PR I have to approve workflows manually to run and they reset after each new commit. Probably you can see them in Actions history (ran about 4 hours ago).

Now let's wait for someone from ggml-cuda to review and approve this.

@IMbackK

IMbackK commented Aug 13, 2026

Copy link
Copy Markdown
Contributor

the problem is def an interaction between hipGraph and cub i can reproduce the problem with a separate toy kernel

This pr also works fine with hipgraph disabled

@Geramy

Geramy commented Aug 13, 2026

Copy link
Copy Markdown
Contributor Author

the problem is def an interaction between hipGraph and cub i can reproduce the problem with a separate toy kernel

This pr also works fine with hipgraph disabled

Right which is a plus it works without hipgraph but I don't think we can just disable hipgraph on CDNA and/or RDNA3, that would create a regression in my opinion, what are your thoughts on the next steps, wait for amd?

@IMbackK

IMbackK commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

Right which is a plus it works without hipgraph but I don't think we can just disable hipgraph on CDNA and/or RDNA3, that would create a regression in my opinion, what are your thoughts on the next steps, wait for amd?

pretty much, or we can look taking #26493 or #27466 (haven't looked at the implementation in the later at all) as a stop gap

@fairydreaming

fairydreaming commented Aug 25, 2026

Copy link
Copy Markdown
Contributor

@IMbackK Can you link the corresponding ROCm bug here so that we know when it's fixed? (Edit: found it, but perhaps it should be reported in https://github.com/ROCm/TheRock now?)

@Geramy

Geramy commented Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

Updated ROCm/TheRock#7625

@Geramy

Geramy commented Aug 25, 2026

Copy link
Copy Markdown
Contributor Author

pretty much, or we can look taking #26493 or #27466 (haven't looked at the implementation in the later at all) as a stop gap

@IMbackK I'll start taking a look at these implementations.

@IMbackK
IMbackK marked this pull request as draft August 26, 2026 08:37
@fairydreaming

Copy link
Copy Markdown
Contributor

pretty much, or we can look taking #26493 or #27466 (haven't looked at the implementation in the later at all) as a stop gap

@IMbackK I'll start taking a look at these implementations.

@Geramy There's another one in #27342 (also there's a talk whether to include it in the PR or not)

@IMbackK

IMbackK commented Aug 27, 2026

Copy link
Copy Markdown
Contributor

having only the rocprim path for large shapes is not really an option since we want to support older versions of rocm too, i think #27466 is the best stop-gap

Aristo94 added a commit to Aristo94/EngramHalo.cpp that referenced this pull request Aug 27, 2026
ggml_top_k/argsort had no HIP path for ne0 > 1024, so sparse-attention
indexers (qwen4exp QSA: 12 calls per decoded token) fell back to the CPU,
collapsing long-context decode. Adds a wide selection kernel tuned for
wave32/RDNA 3.5. Earlier hipCUB-based attempts: ggml-org#26592, ggml-org#26388.
Aristo94 added a commit to Aristo94/EngramHalo.cpp that referenced this pull request Aug 28, 2026
ggml_top_k/argsort had no HIP path for ne0 > 1024, so sparse-attention
indexers (qwen4exp QSA: 12 calls per decoded token) fell back to the CPU,
collapsing long-context decode. Adds a wide selection kernel tuned for
wave32/RDNA 3.5. Earlier hipCUB-based attempts: ggml-org#26592, ggml-org#26388.
@drluoto

drluoto commented Aug 28, 2026

Copy link
Copy Markdown

Tested this PR on Strix Halo (Ryzen AI Max+ 395 / Radeon 8060S, gfx1151, ROCm 7.1.52801,
hipCUB 7.1.0 from Ubuntu packages), applied onto master ca3d5a3, because it turns out
to matter enormously for Qwen3.8-Flash-Next (qwen4exp): the QSA indexer runs ggml_top_k
over the full context in 12 layers per token, and without CUB the ne[0] <= 1024 limit
sends every one of those to the CPU. The decode-vs-depth curve has a knee at exactly
d1024.

Perf (llama-bench, UD-IQ4_XS 93.7 GB, tg64, r=3)

depth master this PR gain
256 23.12 23.07
1024 20.99 22.71 +8%
4096 18.44 21.83 +18%
16384 14.57 17.62 +21%

End-to-end at 24k-token prompts (llama-server, greedy): +38-53% depending on workload.
Correctness: test-backend-ops test -b ROCm0 passes for TOP_K, ARGSORT, SSM_SCAN,
CUMSUM, MEAN and SUM.

Crash under HIP graph capture — reproducible

With HIP graphs active, generation aborts after ~2k generated tokens:

ggml/src/ggml-cuda/ggml-cuda.cu:107: ROCm error
ROCm error: operation not permitted when stream is capturing

Mechanism (from timing, not yet from a debugger): the CUB temp-storage requirement for
the sort grows with n_kv; when the pool needs a fresh hipMalloc while a decode graph
is being captured, HIP rejects it. It surfaces mid-generation exactly when n_kv crosses
a pool-growth boundary. GGML_CUDA_DISABLE_GRAPHS=1 is a full workaround — stable for
hours under load here.

A possible fix: round the temp-storage request up to coarse buckets (e.g. next power of
two) so pool growth becomes rare, and/or pre-size from ne[0] padded to the context
limit on first use, so capture never observes an allocation. Happy to test patches on
gfx1151 — the repro is deterministic here.

@Geramy

Geramy commented Aug 28, 2026

Copy link
Copy Markdown
Contributor Author

Tested this PR on Strix Halo (Ryzen AI Max+ 395 / Radeon 8060S, gfx1151, ROCm 7.1.52801, hipCUB 7.1.0 from Ubuntu packages), applied onto master ca3d5a3, because it turns out to matter enormously for Qwen3.8-Flash-Next (qwen4exp): the QSA indexer runs ggml_top_k over the full context in 12 layers per token, and without CUB the ne[0] <= 1024 limit sends every one of those to the CPU. The decode-vs-depth curve has a knee at exactly d1024.

Perf (llama-bench, UD-IQ4_XS 93.7 GB, tg64, r=3)

depth master this PR gain
256 23.12 23.07 —
1024 20.99 22.71 +8%
4096 18.44 21.83 +18%
16384 14.57 17.62 +21%
End-to-end at 24k-token prompts (llama-server, greedy): +38-53% depending on workload. Correctness: test-backend-ops test -b ROCm0 passes for TOP_K, ARGSORT, SSM_SCAN, CUMSUM, MEAN and SUM.

Crash under HIP graph capture — reproducible

With HIP graphs active, generation aborts after ~2k generated tokens:

ggml/src/ggml-cuda/ggml-cuda.cu:107: ROCm error
ROCm error: operation not permitted when stream is capturing

Mechanism (from timing, not yet from a debugger): the CUB temp-storage requirement for the sort grows with n_kv; when the pool needs a fresh hipMalloc while a decode graph is being captured, HIP rejects it. It surfaces mid-generation exactly when n_kv crosses a pool-growth boundary. GGML_CUDA_DISABLE_GRAPHS=1 is a full workaround — stable for hours under load here.

A possible fix: round the temp-storage request up to coarse buckets (e.g. next power of two) so pool growth becomes rare, and/or pre-size from ne[0] padded to the context limit on first use, so capture never observes an allocation. Happy to test patches on gfx1151 — the repro is deterministic here.

I think the next step is for me to try to find the issue in TheRock see what I can do there, I don't believe growing the pool is the proper stable solution, its a bit of a workaround. Let me see how far I get in ROCm first and I'll report back Sunday ish.

@drluoto

drluoto commented Aug 28, 2026

Copy link
Copy Markdown

Follow-up: pinned down the graph-capture crash. It is not pool allocation growth — I tested that theory by rounding every pool request in the CUB paths up to power-of-two buckets (argsort.cu ×4, top-k.cu ×1) and the crash reproduces unchanged. The failing call is the sort itself:

ROCm error: operation not permitted when stream is capturing
  in function argsort_f32_i32_cuda_cub at ggml/src/ggml-cuda/argsort.cu:157
  DeviceSegmentedRadixSort::SortPairsDescending(d_temp_storage, ...)

So on ROCm, hipCUB/rocPRIM's DeviceSegmentedRadixSort is itself not capture-safe — the existing is_capturing switch (added for CCCL's DeviceSegmentedSort limitation) routes to the very call that HIP rejects. The CCCL comment about DeviceSegmentedRadixSort supporting capture holds for NVIDIA's implementation but evidently not for rocPRIM's (ROCm 7.1, gfx1151; deterministic after ~2k generated tokens once a captured decode graph contains a batched TOP_K > 1024 wide).

Suggestion: on HIP, either gate the CUB paths out of captured streams entirely (fall back to disabling graph capture for graphs containing wide TOP_K/ARGSORT), or document GGML_CUDA_DISABLE_GRAPHS=1 as required. With graphs disabled this PR is stable for hours under load here and the wide-TOP_K win stands (+21% decode at 16k depth for qwen4exp).

For the qwen4exp use case specifically, #27466's native radix kernel may compose better with HIP graphs since it has no library-internal stream operations — testing that now and will report there.

@Geramy

Geramy commented Aug 28, 2026

Copy link
Copy Markdown
Contributor Author

Follow-up: pinned down the graph-capture crash. It is not pool allocation growth — I tested that theory by rounding every pool request in the CUB paths up to power-of-two buckets (argsort.cu ×4, top-k.cu ×1) and the crash reproduces unchanged. The failing call is the sort itself:

ROCm error: operation not permitted when stream is capturing
  in function argsort_f32_i32_cuda_cub at ggml/src/ggml-cuda/argsort.cu:157
  DeviceSegmentedRadixSort::SortPairsDescending(d_temp_storage, ...)

So on ROCm, hipCUB/rocPRIM's DeviceSegmentedRadixSort is itself not capture-safe — the existing is_capturing switch (added for CCCL's DeviceSegmentedSort limitation) routes to the very call that HIP rejects. The CCCL comment about DeviceSegmentedRadixSort supporting capture holds for NVIDIA's implementation but evidently not for rocPRIM's (ROCm 7.1, gfx1151; deterministic after ~2k generated tokens once a captured decode graph contains a batched TOP_K > 1024 wide).

Suggestion: on HIP, either gate the CUB paths out of captured streams entirely (fall back to disabling graph capture for graphs containing wide TOP_K/ARGSORT), or document GGML_CUDA_DISABLE_GRAPHS=1 as required. With graphs disabled this PR is stable for hours under load here and the wide-TOP_K win stands (+21% decode at 16k depth for qwen4exp).

For the qwen4exp use case specifically, #27466's native radix kernel may compose better with HIP graphs since it has no library-internal stream operations — testing that now and will report there.

┌────────────────────┬─────────────────┬────────────────────────┬──────────────┐
│      rocPRIM       │ ROCPRIM_VERSION │ is_graph_capture guard │ capture-safe │
├────────────────────┼─────────────────┼────────────────────────┼──────────────┤
│ 4.2.0 (ROCm 7.2.x) │ 400200          │ NO                     │ ❌ crashes   │
├────────────────────┼─────────────────┼────────────────────────┼──────────────┤
│ 4.4.0 (ROCm 7.13)  │ 400400          │ YES                    │ ✅           │
└────────────────────┴─────────────────┴────────────────────────┴──────────────┘

@drluoto your crash is related to the rocPRIM version aka ROCm, this works properly on 7.13 and 7.14, which I believe llama.cpp has been upgraded to recently or is about to be upgraded to.

@Geramy

Geramy commented Aug 28, 2026

Copy link
Copy Markdown
Contributor Author

@IMbackK can you test this again, I think I found the issue, this should be resolved on ROCm 7.13 and 7.14

@Geramy

Geramy commented Sep 1, 2026

Copy link
Copy Markdown
Contributor Author

this doesn't exist on 7.13 or 7.14 I've tested again and amd tested on the MI325X which also does not have any problem running it, my new guard I added prevents it compiling against 7.2.4 routing completely against the bad binary causing the segfault.

@IMbackK IMbackK left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I will retest for the failure on rocm 10.0.0.
Regardless of this we need to add a find_package for hipCUB in ggml-hip/CMakeLists.txt when the cub option is set, as hipCUB need not be available in system installations of rocm.

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

Labels

CUDA Related to the CUDA backend devops improvements to build systems and github actions ggml changes relating to the ggml tensor library for machine learning testing Everything test related

Projects

None yet

Development

Successfully merging this pull request may close these issues.