Repository navigation
CUDA: improve top-k algorithm selection - #28713
Conversation
Replaces CUB's per-row DeviceTopKKernel with a grid-over-rows radix select, gated on GGML_CUDA_TOPK_RADIX_MIN_ROWS. On qwen4exp at 34,816 tokens this cuts top-k from 1,671,253 launches / 5,761.8 ms to 2,329 / 941.8 ms.
github-actions
Bot
added
ggml
|
Please see discussion over at #28547 for a principled analysis that can be used to determine a threshold. Might make sense to add a data-point for DGX Spark here |
Acknowledged. I will check the PR #28547. Performance with various -lzm options on Qwen 3.8 Flash Next over DGX Spark
Six runs per arm over two rounds. Cold cache before every arm. Prompts from a 2.09M-word prose corpus so no two requests share n-grams — filler from a small vocabulary saturates the reachable n-gram set after one request and inflates the demand-paged arm. On/off order flipped between rounds. The on arm is I/O-bound, demand-paging a 26.8 GiB embedding table one fault at a time, so removing launches buys nothing there — hence its runs interleaving. An independent 2×2 at pp8192 agrees: 675.2 → 800.8 resident, 182.4 → 193.6 faulting. Decode is unchanged in every arm (~17.9–22.3 tok/s): ggml_top_k is not on this model's decode path at batch 1.
Summary:
|
| for (int i = 0; i < nrows; i++) { | ||
| top_k_cub(pool, src0_d + i * ncols, dst_d + i * k, ncols, k, stream); | ||
| if (nrows >= top_k_radix_min_rows() && ncols > 1024) { | ||
| top_k_radix_cuda(pool, src0_d, dst_d, ncols, nrows, k, stream); |
There was a problem hiding this comment.
Looks like this allocates the memory for all rows at once.
llama.cpp/ggml/src/ggml-cuda/top-k.cu
Line 191 in a011ccd
|
Seems like these are targeting the same file, can you verify: Does the compile flags resolve it or is there a need for the other additions here? |
|
Feel free to reopen if the PR would change the heuristic of #28547 |
|
#29278 is closed as [fairydreaming] will not continue working on that. @praneshgo , can you reopen this PR and take over. |
Yes, I plan to work on it this week. |
|
Reopening this PR and marking it as draft as #29278 is closed. |
Replace the nrows/ncols special case with the decision boundary from ggml-org#28547 (as implemented in ggml-org#29278): bitonic for short rows, radix select for several long rows, and DeviceTopK or CUB argsort for a single long row. The thresholds stay overridable at build time. Two refinements on top of that boundary: - bitonic stays in use for rows up to a padded 1024 while the rows fit in one wave of blocks (nrows <= number of SMs); radix select pays a fixed cost of about a dozen launches that only amortizes over more rows - with DeviceTopK available, it handles up to two rows Radix select now processes rows in chunks so its scratch memory stays bounded, and the bitonic path keeps its chunking. HIP and MUSA keep their previous thresholds. Add perf cases around the bitonic/radix crossover to test-backend-ops.
praneshgo
changed the title
|
@gaugarg-nv @ORippler can you please review this PR? |
What causes perf regression for row=1, CCCL 3.4.3? |
gaugarg-nv
left a comment
gaugarg-nv
left a comment
There was a problem hiding this comment.
Looks like this check is no longer needed-
llama.cpp/ggml/src/ggml-cuda/ggml-cuda.cu
Lines 5693 to 5695 in b9acf13
The chunking helpers are now shared across paths in the latest commit. |
ORippler
left a comment
ORippler
left a comment
There was a problem hiding this comment.
Thanks for taking this up! Please add a test to test-backend-ops to exercise the chunked paths for both radix-sort & bitonic sort in addition to the comments made below.
…TOP_K_NROWS_THRESHOLD
The tests are added now in the latest code. |
| return ggml_is_contiguous_rows(op->src[0]); | ||
| case GGML_OP_TOP_K: | ||
| #if defined(GGML_USE_HIP) || defined(GGML_CUDA_USE_CUB) | ||
| return true; |
There was a problem hiding this comment.
I think we should double check the limits here in a follow-up PR (for CCCL 2.8, sorting based algos tend to limit to n_cols <= INT_MAX, radix-sort implemented in the CUDA backend also supports only n_cols <= INT_MAX for now)
There was a problem hiding this comment.
The latest code now limits supports_op to ne[0] <= INT_MAX for TOP_K in all builds, and for argsort in CUB builds. argsort in builds without CUB was already capped by its shared-memory check.
Upstream brings CUDA top-k/argsort/mmq fixes (PR ggml-org#28713, ggml-org#29953, ggml-org#30147, ggml-org#29453), a DFlash output-head fix (ggml-org#30111) and a UI fix (ggml-org#29668). Conflicts and resolution: - top-k.cu / argsort.cu: upstream reworked the top-k selection and the same radix path the fork added. Take upstream's shape-based dispatch and int64_t fixes, keep the fork workspace cap in the shared ggml_cuda_chunk_nrows (GGML_CUDA_ARGSORT_CHUNK_MB) and the fork benchmark-only GGML_CUDA_TOPK_IMPL selector used by bench-topk.py. - mmq.cu: upstream PR ggml-org#29953 fixes the same mul_mat_id padding over-read the fork patched in 1386ec2, but sizes the padding from the chosen J tile in both branches. Take upstream's version. - test-backend-ops.cpp: keep the fork GGML_TOPK_BENCH shapes and add upstream's chunk-spanning cases. Assisted-by: pi (deepseek-flash)
…-10-10) Conflicts: - ggml-backend.cpp: upstream moved input copies to ggml_backend_sched_copy_input with host weights last (ggml-org#29943); the GGML_SCHED_BATCH_INPUTS fast path now runs in the first loop before falling back to it. - ggml-cuda/fattn-common.cuh: launch_fattn takes upstream async_kv_preload, then our parallel_blocks_fixed (GQA-6 tile launcher updated). - ggml-cuda/fattn-mma-f16.cuh: upstream swizzle refactor (ggml-org#29612); its helpers take ncols2 like the rest of the fork's config getters. - ggml-cuda/gated_delta_net.cu: keep the fork's kernel and launch geometry (its columns_per_block template argument is not upstream's cols_per_warp, ggml-org#30087). - ggml-cuda/mmq.cu: upstream per-tile src1 padding (ggml-org#29953); ggml_cuda_mul_mat_q_src1_nbytes and the MMQ input cache use the worst-case padding so producers and projections with other tiles can share the buffer. - ggml-cuda/mmq-vec-dot.cuh: fork vec_dot kernels pass GGML_PREC_Q8 to the config helpers (ggml-org#30168). - ggml-cuda/mmvf.cu/.cuh: keep both includes and the HIP declarations, upstream warp_size argument. - ggml-cuda/rope.cu: upstream grid-stride rms_norm+mul+rope (ggml-org#28175) with the fork's M-RoPE theta (nchannels instead of gridDim.y) and rope_store_cast. - ggml-cuda/top-k.cu: upstream shape-based selection (ggml-org#28713), RDNA4 tournament kept for few long rows. - src/llama-batch.*: mixed token/embd batches (ggml-org#29622) copy embd rows per entry; other batches keep pointing at the ext rows. - src/llama-context.cpp: expert copy callback also on the second MTP scheduler. - src/llama-graph.cpp: build_rs keeps the GET_ROWS of the ubatch states (GDN fusions, deferred state) and gathers the extra states from the tail like upstream's custom getter path (ggml-org#29856).
Make the existing radix selector available in hipCUB builds and dispatch it only above measured width/batch cutoffs. Gather boundary ties in input order so the selected set matches stable hipCUB sorting. Preserve the existing single-row and other-device paths and the original non-CUB radix behavior. All 569 operator cases pass, including extreme scores and cutoff boundaries. The five target-model workloads retain all 40972800 logits and token IDs; WikiText2 perplexity remains 2.0259. Native PP2048 at depth40000 measures 358.986 +/-17.688 ->392.294 +/-0.692 tok/s in one five-sample pair; zero-depth prefill is unchanged within variation. The deep control is noisy, although every candidate sample exceeds every control sample. Assisted-by: Codex (cherry picked from commit ec05a66) [rebase: top-k.cu change dropped - upstream ggml-org#28713 already dispatches radix top-k on HIP for nrows > 1, ncols > 1024; only the tests are kept]
cuda : top-k selection by shape for temp-11259 (port of upstream ggml-org#28713)
* CUDA: radix top-k for large row counts Replaces CUB's per-row DeviceTopKKernel with a grid-over-rows radix select, gated on GGML_CUDA_TOPK_RADIX_MIN_ROWS. On qwen4exp at 34,816 tokens this cuts top-k from 1,671,253 launches / 5,761.8 ms to 2,329 / 941.8 ms. * CUDA: select the TOP_K implementation by shape Replace the nrows/ncols special case with the decision boundary from ggml-org#28547 (as implemented in ggml-org#29278): bitonic for short rows, radix select for several long rows, and DeviceTopK or CUB argsort for a single long row. The thresholds stay overridable at build time. Two refinements on top of that boundary: - bitonic stays in use for rows up to a padded 1024 while the rows fit in one wave of blocks (nrows <= number of SMs); radix select pays a fixed cost of about a dozen launches that only amortizes over more rows - with DeviceTopK available, it handles up to two rows Radix select now processes rows in chunks so its scratch memory stays bounded, and the bitonic path keeps its chunking. HIP and MUSA keep their previous thresholds. Add perf cases around the bitonic/radix crossover to test-backend-ops. * CUDA: make top-k comments less verbose * CUDA: remove the TOP_K width limit from supports_op * CUDA: use DeviceTopK for single-row TOP_K if available * CUDA: avoid ncols overflow in the TOP_K bitonic check * CUDA: share the row chunking helper between argsort and top-k * CUDA: do the TOP_K radix blocks_per_row math in int64_t * CUDA: rename GGML_CUDA_TOP_K_NROWS_THRESHOLD_DEVICETOPK to GGML_CUDA_TOP_K_NROWS_THRESHOLD * CUDA: share one sort helper between the bitonic and CUB TOP_K paths * CUDA: update the TOP_K TODO, threshold and chunking comments * tests: add TOP_K cases that span several row chunks * CUDA: use int64_t col in the TOP_K radix loops, fix threshold comment * CUDA: limit TOP_K and ARGSORT support to ne[0] <= INT_MAX --------- Co-authored-by: praneshgo <227579474+praneshgo@users.noreply.github.com> Co-authored-by: Pranesh Gonegandla <pgonegandla@nvidia.com>
Overview
This PR supersedes #29278.
It starts from #28547's boundary, as in #29278, with the same three build-time thresholds, and makes a few refinements on top of them.
Refinement 1 (CUDA only): bitonic for a padded width up to 1024 while the rows fit in one wave of SMs, set by a new fourth threshold,
GGML_CUDA_TOP_K_NCOLS_THRESHOLD_BITONIC_FEW_ROWS(1024; 0 on HIP/MUSA). This fixes #29278's own CUDA 12.8 losses at 1024 columns × 2–128 rows (55–78%).Refinement 2: DeviceTopK for up to 2 rows when it's available.
GGML_CUDA_TOP_K_NROWS_THRESHOLDnow defaults to 2 with DeviceTopK and stays 1 without it. This fixes the 2-row losses in #29278's 13.2 chart (52–80%).Refinement 3: in CUDA builds with DeviceTopK (CCCL ≥ 3.4.3), a single row always uses DeviceTopK, the same as master. Rows ≥2 and builds without DeviceTopK keep the rules above. This fixes the single-row losses at ≤512 columns (down to 45% on the RTX 5090, 62% on GB10 and 81% on GB300).
Testing
Correctness was checked with
test-backend-ops -o TOP_Kon the patched build:Performance
Method: the same harness as #29278. 1027 shapes, graphs on, runs capped at 1024, master vs patched, two passes each.
DGX Spark
Default CCCL
CCCL v3.4.3
RTX 5090
Default CCCL
CCCL v3.4.3
GB300
Native CCCL 3.4.3
Patched vs master per cell (master time / patched time); geomean, with the lowest cell in brackets.
Adding some E2E perf numbers below
E2E perf numbers
GB300
DGX Station GB300 — Qwen3.8-Flash-Next UD-IQ4_XS, CUDA 13.4.2 (bundled CCCL 3.4.3)
llama-bench -fa 1 -ngl 99 -r 3(pp:-b 2048 -ub 2048); batched:llama-batched-bench -c 65792 -b 2048 -ub 512 -kvu -npp 1024 -ntg 64 -npl 3,4GB10
DGX Spark GB10 — Qwen3.8-Flash-Next UD-IQ4_XS, CUDA 13.0, CCCL 3.4.3
llama-bench -fa 1 -ngl 99 -b 2048 -ub 2048 -r 3; batched:llama-batched-bench -c 65792 -b 2048 -ub 512 -kvu -npp 1024 -ntg 64 -npl 3,4Requirements