Skip to content

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

Open
Geramy wants to merge 2 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 2 commits into
ggml-org:masterfrom
Geramy:Geramy/argsort_top_k_cub_on_hip

Conversation

@Geramy

@Geramy Geramy commented Aug 4, 2026 •
edited
Loading

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

@Geramy
CUDA: enable the CUB path on HIP via hipCUB
e1b2fba 
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.
@github-actions github-actions Bot added ggml changes relating to the ggml tensor library for machine learning CUDA Related to the CUDA backend labels Aug 4, 2026
This was referenced Aug 4, 2026
Closed
Open
fairydreaming
fairydreaming reviewed Aug 4, 2026
View reviewed changes
Comment thread ggml/src/ggml-cuda/cub-compat.cuh Outdated
@fairydreaming

fairydreaming commented Aug 4, 2026 •
edited
Loading

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 •
edited
Loading

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 •
edited
Loading

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 •
edited
Loading

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

fairydreaming commented Aug 4, 2026

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

fairydreaming commented Aug 4, 2026

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.

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

Reviewers

@IMbackK IMbackK Awaiting requested review from IMbackK IMbackK is a code owner
@fairydreaming fairydreaming Awaiting requested review from fairydreaming

At least 2 approving reviews are required to merge this pull request.

Assignees

No one assigned

Labels

CUDA Related to the CUDA backend ggml changes relating to the ggml tensor library for machine learning

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@Geramy @fairydreaming @Neresco @remeh