(WIP) CUDA/HIP: GPU partial top-k for ROCm (fix full-sort/CPU fallback) - #27860
Closed
zihaomu wants to merge 2 commits into
Closed
(WIP) CUDA/HIP: GPU partial top-k for ROCm (fix full-sort/CPU fallback)#27860zihaomu wants to merge 2 commits into
zihaomu wants to merge 2 commits into
Conversation
ggml_cuda_op_top_k had a fast partial-top-k path only under CUB_TOP_K_AVAILABLE
(NVIDIA CCCL >= 3.2 DeviceTopK). On ROCm/HIP that is never defined, so:
- supports_op reported TOP_K unsupported for ne[0] > 1024, pushing it to the CPU
backend (a per-token GPU<->CPU roundtrip + CPU sort), and
- where it did run, it full-sorted every row.
This collapses Qwen3-Next/qwen4exp QSA long-context decode on Strix Halo, while a
same-bandwidth NVIDIA GB10 (DeviceTopK) stays flat.
- top-k.cu: add a HIP branch using rocPRIM partial_sort_copy on packed (score,index)
keys (nth_element + sort of the first k). k == ncols falls back to bitonic
(rocPRIM requires middle < size).
- ggml-cuda.cu: enable TOP_K in supports_op on HIP when k < ncols (or ncols <= 1024),
and disable graph capture for graphs whose TOP_K uses the partial path (it syncs
the stream), mirroring the MUL_MAT_ID guard.
Validation on gfx1151 / ROCm 7.2.1:
- test-backend-ops -o TOP_K: 517/517 OK (incl. ncols up to 16384).
- qwen4exp UD-IQ4_XS single-stream, graphs ON: S_TG at N_KV 640/2176/8320 =
21.5/20.7/19.3 t/s (was 21.5/8.8/8.1) -> long-context collapse eliminated.
github-actions
Bot
added
ggml
|
are you aware of #27466 ? |
Replace the rocPRIM partial_sort_copy top-k path with hipCUB radix sort-pairs. rocPRIM partial_sort is not stream-capture-safe, which forced disabling HIP graphs for the decode path. hipCUB DeviceRadixSort (single row) and DeviceSegmentedRadixSort (multi-row) are capture-safe, so single-row decode top-k now runs inside CUDA graphs. - top-k.cu: sort (score, index) pairs descending with radix and keep the first k indices of each row; drop the packed-u64 key and the per-row loop. CUB sorts float keys directly, so no manual bit-packing is needed. - ggml-cuda.cu: supports_op returns true for TOP_K on ROCm for any ncols/nrows. - ggml-cuda.cu: instantiating a HIP graph that contains a large segmented sort overflows the ROCm runtime's graph builder, so disable CUDA graphs only for the multi-row (nrows>1) top-k case; single-row decode keeps graphs. test-backend-ops -o TOP_K: 517/517 pass; perf -o TOP_K runs clean under graphs.
zihaomu
marked this pull request as draft
August 28, 2026 09:46
Contributor
|
closed in favor of #27466 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
ggml_cuda_op_top_k had a fast partial-top-k path only under CUB_TOP_K_AVAILABLE (NVIDIA CCCL >= 3.2 DeviceTopK). On ROCm/HIP that is never defined, so:
supports_op reported TOP_K unsupported for ne[0] > 1024, pushing it to the CPU backend (a per-token GPU<->CPU roundtrip + CPU sort), and
where it did run, it full-sorted every row. This collapses Qwen3-Next/qwen4exp QSA long-context decode on Strix Halo, while a same-bandwidth NVIDIA GB10 (DeviceTopK) stays flat.
top-k.cu: add a HIP branch using rocPRIM partial_sort_copy on packed (score,index) keys (nth_element + sort of the first k). k == ncols falls back to bitonic (rocPRIM requires middle < size).
ggml-cuda.cu: enable TOP_K in supports_op on HIP when k < ncols (or ncols <= 1024), and disable graph capture for graphs whose TOP_K uses the partial path (it syncs the stream), mirroring the MUL_MAT_ID guard.
Validation on gfx1151 / ROCm 7.2.1:
The following is speed testing result of
Qwen3.8-Flash-Next-GGUFon AMD Strix Halo 395 Chip:cc @zhangnju
Overview
Additional information
Requirements