ggml-cuda: enable CUB via hipCUB on HIP for device-wide primitives - #26388
ggml-cuda: enable CUB via hipCUB on HIP for device-wide primitives#26388singulared wants to merge 1 commit into
Conversation
The device-wide CUB primitives (DeviceRadixSort / DeviceSegmentedRadixSort / DeviceSegmentedSort / DeviceReduce / DeviceScan) were compiled only for CUDA (GGML_CUDA_USE_CUB gated behind !defined(GGML_USE_HIP)). On HIP this forced argsort / top_k over rows > 1024 onto the CPU backend, and sum / mean / cumsum / ssm-scan onto their fallback paths. hipCUB exposes a CUB-compatible interface over rocPRIM, so the same code compiles on HIP with a namespace alias. This enables the GPU path via __has_include(<hipcub/hipcub.hpp>) (opt-out with GGML_HIP_NO_HIPCUB). Motivation: DeepSeek sparse-attention (DSA / lightning indexer, DeepSeek-V4) runs a top_k every decode step; on HIP that fell back to the CPU with cost growing with context depth. Measured on gfx1151 (Strix Halo), deep-context decode improves and the per-token CPU round-trip is removed. Requires rocPRIM/hipCUB >= 7.2: 6.4.3 fails on RDNA (gfx1151, wave32) with a 'WarpSize > 32 without DPP broadcasts' static assert because it defaults __AMDGCN_WAVEFRONT_SIZE to 64 under LLVM that no longer predefines it. Validated: test-backend-ops -o ARGSORT/TOP_K/SUM/MEAN/CUMSUM pass on ROCm0 (shapes to 1M columns, exercising the CUB path). Stream-capture guard in argsort mirrored via hipStreamIsCapturing mappings in vendors/hip.h.
github-actions
Bot
added
ggml
|
Hi @singulared, thanks for your contribution! Per our contribution guidelines, the automated PR checker found the following issue(s) that need your attention:
Please note that maintainers reserve the right to make final decisions on PRs. If you believe there is a mistake, please comment below. |
|
Oops, my mistake. I use AI generated PR body. Will recreate it latter. |
|
I tested this PR to solve #26399 (missing topk on ROCm). On latest ROCm nightly it does not compile without modification because llama.cpp header files define a __trap macro and ROCm hipcub header files define a __trap function, so there is a conflict. There is probably a better way to solve it, but I got it to compile by guarding the include like this: Once it compiled, this PR also fixed the missing topk problem for ROCm for me on RDNA 3.5/Strix Halo. |
|
Thanks for your review. I will take a look at this today, slightly later. |
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.
2 participants
Overview
Enable the CUB code paths on HIP by using hipCUB (the CUB-compatible API that ships with rocPRIM), so AMD builds stop falling back to the slow/CPU paths for device-wide primitives.
Today
GGML_CUDA_USE_CUBis only defined for CUDA (CUDART_VERSION >= 11070), so on HIP:argsort/top_kwithncols > 1024exceed the single-block bitonic path and fall back to the CPU backend. For DeepSeek-Sparse-Attention models (DeepSeek-V4, DeepSeek-3.2, GLM-DSA, MiniMax-M3) the per-token indexertop_kruns over the whole context, so this fallback is hit on every token, with cost growing as context grows.sum,mean,cumsumandssm_scanuse their inefficient fallbacks.ggml/src/ggml-cuda/sum.cualready notes this:Implementation
common.cuhenables the existing CUB paths when hipCUB is available:hipcub-devkeep building exactly as before.-DGGML_HIP_NO_HIPCUB.namespace cub = hipcub;and pulls in the specific symbols it uses, rather than dumping the namespace globally.vendors/hip.hgains thecudaStreamIsCapturing/cudaStreamCaptureStatusmappings the CUB paths need.8 files changed, 53 insertions(+), 7 deletions(-).
Correctness
test-backend-ops testpasses on the ROCm backend for every affected op:ARGSORT,TOP_K,SUM,MEAN,CUMSUM,SSM_SCAN.Performance
Hardware: Radeon 8060S iGPU (gfx1151, Strix Halo), 128 GB unified memory, ROCm 7.2.4, Debian.
Both arms built from the same commit, differing only by this patch.
llama-bench -p 0 -n 64 -d <depth> -fa 1 -mmp 0, ROCm backend.DeepSeek-V4-Flash (284B-A13B, IQ2_XXS, KV
q4_0,-b/-ub 2048), tg t/s:Depth 0 is the natural control: with no context there is no indexer
top_kwork and the two builds are identical. The gain appears as soon as context exists and stays ≥25% out to 128K — consistent with removing a fallback whose cost scales with depth.Models that do not use these ops are unaffected, as expected:
topk-moepath, ≤256 cols)Those two are within run-to-run noise (arms were run sequentially, so the second arm sees slightly warmer hardware).
Requirements