CUDA: enable the CUB path on HIP via hipCUB - #26592
Conversation
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.
|
@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. |
|
So test this PR here with 4x gfx1200 9060xt and 1x gfx1151. No crash at 68727 PP context. |
|
@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 Thanks! |
|
@fairydreaming from AMD GitHub page directly |
|
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. |
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). |
|
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. |
|
HIP quality check CI found some missing CUB function call return value error checks, so I added them. Compilation worked fine with added |
Co-authored-by: Georgi Gerganov <ggerganov@gmail.com>
|
@fairydreaming thanks! I didn’t see any of the CIs run or get approved to run did you do a few manually in actions? |
|
@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. |
|
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? |
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 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?) |
|
Updated ROCm/TheRock#7625 |
|
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 |
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.
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.
|
Tested this PR on Strix Halo (Ryzen AI Max+ 395 / Radeon 8060S, gfx1151, ROCm 7.1.52801, Perf (llama-bench, UD-IQ4_XS 93.7 GB, tg64, r=3)
End-to-end at 24k-token prompts (llama-server, greedy): +38-53% depending on workload. Crash under HIP graph capture — reproducibleWith HIP graphs active, generation aborts after ~2k generated tokens: Mechanism (from timing, not yet from a debugger): the CUB temp-storage requirement for A possible fix: round the temp-storage request up to coarse buckets (e.g. next power of |
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. |
|
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: So on ROCm, hipCUB/rocPRIM's 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 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. |
@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. |
|
@IMbackK can you test this again, I think I found the issue, this should be resolved on ROCm 7.13 and 7.14 |
|
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. |
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
pp512prompt processing
tg32token generation
Qwen3.6-35B-A3B UD-Q4_K_XL
pp512prompt processing
tg32token generation
Requirements
@IMbackK @am17an @pwilkin