ROCm: add radix TOP_K for long rows - #27466
Conversation
github-actions
Bot
added
ggml
|
Hi @jadenmach2, 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. |
ggml-gh-bot
Bot
added
the
draft
github-actions
Bot
removed
the
draft
743d1a4 to
ff7cd32
Compare
IMbackK
left a comment
IMbackK
left a comment
There was a problem hiding this comment.
While we would like to have just one code path here, i think this is a good idea take as a stop gap until the hipCUB hipGraph interaction is fixed in rocm
| if (ncols > 1024) { | ||
| top_k_radix_cuda(pool, src0_d, dst_d, ncols, nrows, k, stream); | ||
| } else { | ||
| #endif |
There was a problem hiding this comment.
missing comment
// defined(GGML_USE_HIP)
| return true; | ||
| #else | ||
| return op->src[0]->ne[0] <= 1024; | ||
| #endif |
There was a problem hiding this comment.
missing comment
// defined(GGML_USE_HIP) || defined(GGML_CUDA_USE_CUB)
| cudaMemcpyDeviceToDevice, stream)); | ||
| #if defined(GGML_USE_HIP) | ||
| } | ||
| #endif |
There was a problem hiding this comment.
missing comment
// defined(GGML_USE_HIP)
|
I tested the performance of this on NVIDIA RTX PRO 6000 Max-Q and it's pretty good, beats CUB argsort-based and DeviceTopK-based implementations for almost all tested shapes (and it's close enough for the few slower ones). Maybe we could simply use it as the default TOP_K implementation for now? Reddish is this PR slower, blueish is this PR faster. PR 27466 vs argsort (k = 1024)
PR 27466 vs argsort (k = 64)
PR 27466 vs DeviceTopK (k = 64)
CC @ORippler |
|
Id like to also test it on RDNA/CDNA vs hipCUB before doing so, but yeah sure. |
@IMbackK Sure, it's just a vibe-coded thing that accepts two CSV files with ncols,nrows,time columns (I use time per run gathered with sed from the test-backend-ops perf output |
There was a problem hiding this comment.
Maybe we could simply use it as the default TOP_K implementation for now?
Id like to also test it on RDNA/CDNA vs hipCUB before doing so, but yeah sure.
I'd also like to see this tested
- for smaller k (typically used for backend-sampling like 20/40)
- on more SKUs (can test on what I have available)
- on more OSs (Windows may take longer to launch the 5 kernels here as opposed to Linux)
if we want to make it the default path
| top_k_radix_state * states = states_alloc.get(); | ||
| int * histograms = histograms_alloc.get(); | ||
|
|
||
| top_k_radix_init<<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows, k); |
There was a problem hiding this comment.
we should really have a ceildiv function in the cuda backend 😄
|
|
||
| top_k_radix_init<<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows, k); | ||
|
|
||
| const dim3 histogram_grid(blocks_per_row, nrows); |
There was a problem hiding this comment.
This is going to hit the 16bit limit of griddim.y/z (on CUDA at least). Either guard dispatch, chunk, or fold into griddim.x:
There was a problem hiding this comment.
Thank you for the review, I will push the changes
|
I want to add a testpoint with Q3.8-27B, 2xR9700, any advice on the settings I should use to trigger the paths? Or just a simple baseline with long context from 60k to 120k np 1..3 with MTP from 0..4? Current Settings Baseline
|
RDNA3.5 (gfx1151 / Strix Halo APU) datapoint — +1New SKU: Radeon 8060S iGPU (gfx1151, RDNA3.5), ROCm 7.2.1, on current master. End-to-end — Qwen3.8-Flash-Next UD-IQ4_XS (its QSA indexer runs
Long-context decode collapse gone, and the radix kernels capture into HIP graphs fine. vs hipCUB (re: @IMbackK) — I also tried a hipCUB Happy to run more shapes if useful. Nice work! |
|
2X R9700, Powercapped at 230W, Settings
Results Detailed
Image
|
|
you can defiantly find shapes where this performs poorly against the cub path (1 row), so from a hip perspective we may want to keep it around until amd fixes the graph capture problem, or roll another own solution DetailsGFX908:
GFX1100:
GFX1201
|
IMbackK
left a comment
IMbackK
left a comment
There was a problem hiding this comment.
Its good as is for the purposes of supporting ncols > 1024 on hip.
|
Tested this PR on Strix Halo (Ryzen AI Max+ 395 / Radeon 8060S, gfx1151, ROCm 7.1), applied onto master Correctness: HIP graph capture: safe. Three consecutive 2500-token speculative-decoding generations with graphs enabled, zero incidents. This is notable because the hipCUB alternative (#26592) aborts capture deterministically — Single-row decode (llama-bench tg64, UD-IQ4_XS 93.7 GB): parity with the hipCUB path, both a big win over master's CPU fallback:
Multi-row wide case — tuning opportunity: with speculative decoding, TOP_K arrives as ~65 rows × 24k cols during draft verification. There the hipCUB segmented sort is ~27% faster end-to-end on our workload (21.3 vs 16.1 tok/s on a file-rewrite task at 24k-token context). Possibly the Net: for graph compatibility this PR is currently the only working GPU TOP_K on ROCm, and single-row performance matches CUB. Would love to see it land. |
IMbackK
added
the
merge ready
|
We probably want @ggml-org/ggml-cuda to sign off on this one. |
|
Running GLM 5.3, on 2x gfx1201 with offloading, with larger prompts breaks reproducible greedy decoding. (GLM has topk 2048). |
|
@ORippler think this is good enough? |
|
By the way is Edit: when |
|
Ties are allowed to be in different order: Lines 2415 to 2420 in 075e1a2 |
ORippler
left a comment
ORippler
left a comment
There was a problem hiding this comment.
We probably want @ggml-org/ggml-cuda to sign off on this one.
@ORippler think this is good enough?
Left some perf comments for NVGPUs. However, given that this is currently AMD only, they do not have to be considered necessarily.
I'd recommend expanding the test-suite a bit more potentially (test ties, test > 2**16 ncols etc.), or at least run these in a one-off fashion locally.
| const int row = blockIdx.x / blocks_per_row; | ||
| const int row_block = blockIdx.x % blocks_per_row; |
There was a problem hiding this comment.
if block_per_row fits in uint32_t range, using fastdiv can be beneficial (at least on NVGPUs)
There was a problem hiding this comment.
its the same for amdgcn isa
| const int row = blockIdx.x / blocks_per_row; | ||
| const int row_block = blockIdx.x % blocks_per_row; |
There was a problem hiding this comment.
same comment as above w.r.t fastdiv
There was a problem hiding this comment.
I missed that, thanks for the review @ORippler
| const dim3 row_grid(blocks_per_row * nrows); | ||
| for (int shift = 32 - RADIX_BITS; shift >= 0; shift -= RADIX_BITS) { | ||
| top_k_radix_histogram<BLOCK_SIZE, RADIX_BITS> | ||
| <<<row_grid, BLOCK_SIZE, 0, stream>>>( | ||
| src, states, histograms, ncols, blocks_per_row, shift); | ||
| top_k_radix_select<BLOCK_SIZE, RADIX_BITS> | ||
| <<<nrows, BLOCK_SIZE, 0, stream>>>(histograms, states, blocks_per_row, shift); | ||
| } |
There was a problem hiding this comment.
This should be expressible via persistent kernels + cooperative groups -> whole grid syncs on NVGPUs
|
Aight, let's take working over perfect and merge this since the lack of TOPK is really hurting ROCm right now. Can add extra tests in a followup. |
ROCm has no such header. vendors/hip.h already includes <hipcub/hipcub.hpp> and aliases cub to hipcub, which is why this branch had carried the include removed rather than guarded. Resolving the ggml-org#27466 conflict by splicing the fork's HIP block into upstream's file took upstream's header along with it and reinstated the include, breaking the gfx1151 build at top-k.cu:5. Guard it instead of deleting it, so a CUDA build keeps upstream's include and the next merge cannot quietly undo this again. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
This kernel seems to cause a segfault when run with HIP graphs enabled, a simple |
|
After this PR Qwen 3.8 Next degenerates into only outputting "/////////" forever on ROCm backend with larger context windows. |
Cant repo, could you describe your setup to repo that? |
cant repo that either at 100k context could you provide more information? |
I'm using Unsloth's Q4_K_XL https://huggingface.co/unsloth/Qwen3.8-Flash-Next-GGUF/tree/main/UD-Q4_K_XL with Machine is Strix Halo, I unfortunately can't paste exact prompts/context for repro |
|
rocm version, and compile time options please. |
Right, sorry; |
* ROCm: add radix TOP_K for long rows
shouldn't this be -DGGML_HIP_UMA ? instead of UMAD |
13 participants
Overview
This adds a ROCm TOP_K path for rows larger than 1024 elements. The current non-CUB path uses the bitonic implementation for small rows, but larger rows are reported as unsupported.
The new path uses exact 8-bit radix selection. Small rows continue to use the existing bitonic path, and the implementation does not depend on hipCUB.
The main motivation is DeepSeek-V4 long-context decoding, where the lightning indexer produces a TOP_K row over the KV history.
Performance
I tested this with DeepSeek-V4 Flash GGUF (155 GB) on three MI250X devices using ROCm 7.2. The numbers below are median generation throughput over 10 repetitions:
Testing
test-backend-ops test -o TOP_K -b ROCm0Requirements