ROCm: add radix TOP_K for long rows - #27466
Open
jadenmach2 wants to merge 2 commits into
Open
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
jadenmach2
force-pushed
the
rocm-radix-topk
branch
from
August 21, 2026 02:10
743d1a4 to
ff7cd32
Compare
jadenmach2
marked this pull request as ready for review
August 21, 2026 02:11
IMbackK
requested changes
Aug 26, 2026
IMbackK
left a comment
IMbackK
left a comment
Contributor
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 |
Contributor
There was a problem hiding this comment.
missing comment
// defined(GGML_USE_HIP)
| return true; | ||
| #else | ||
| return op->src[0]->ne[0] <= 1024; | ||
| #endif |
Contributor
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 |
Contributor
There was a problem hiding this comment.
missing comment
// defined(GGML_USE_HIP)
Contributor
|
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 |
Contributor
|
Id like to also test it on RDNA/CDNA vs hipCUB before doing so, but yeah sure. |
Contributor
@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 |
ORippler
reviewed
Aug 27, 2026
Collaborator
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); |
Collaborator
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); |
Collaborator
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:
Contributor
Author
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
|
Aristo94
added a commit
to Aristo94/EngramHalo.cpp
that referenced
this pull request
Aug 28, 2026
Contributor
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! |
Aristo94
added a commit
to Aristo94/EngramHalo.cpp
that referenced
this pull request
Aug 28, 2026
|
2X R9700, Powercapped at 230W, Settings
Results Detailed
Image
|
Contributor
|
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
approved these changes
Aug 28, 2026
IMbackK
left a comment
IMbackK
left a comment
Contributor
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
Member
|
We probably want @ggml-org/ggml-cuda to sign off on this one. |
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
8 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.
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