CUDA/HIP: partial-top-k + hipCUB large-N for argsort/top-k on ROCm#504
Draft
geometric[bot] wants to merge 5 commits into
Draft
CUDA/HIP: partial-top-k + hipCUB large-N for argsort/top-k on ROCm#504geometric[bot] wants to merge 5 commits into
geometric[bot] wants to merge 5 commits into
Conversation
The GPU top-K + log-prob extraction (geometric_draft_topk_cuda.cu) was compiled only on CUDA, so DFLASH_GPU_DRAFT_TOPK was a no-op on ROCm and AMD paid a per-step vocab x n_tokens D2H + CPU heap extract every speculation step. The kernel body is plain arithmetic (no tensor cores, no CUDA-only intrinsics), so the same .cu compiles unchanged for HIP. Compile geometric_draft_topk_cuda.cu directly with LANGUAGE HIP through the hip_compat <cuda_runtime.h> shim (the shared-.cu pattern already used by deepseek4_hc_cuda.cu), instead of adding a separate .hip.cu wrapper translation unit. Rename the backend guard macro DFLASH27B_HAVE_DRAFT_TOPK_CUDA -> DFLASH27B_HAVE_DRAFT_TOPK so the consumers (qwen35_dflash_target, test_dflash) are backend-neutral. test_draft_topk_cuda is now built on the HIP backend as well (CUDA spellings mapped by the same shim) and validated against the CPU reference on gfx1151 (Radeon 8060S, wave32, ROCm 6.4.4). Also qualify the README GPU-flag defaults to the server harness.
The TOP_K path full-sorts every row via bitonic argsort regardless of k, so its latency is independent of k. Make it partial: - argsort.cu: cache keys in shared memory and bound the bitonic sort by the padded k, so work scales with k rather than ncols. - argsort.cu: GGML_ARGSORT_SHFL_XOR — for bitonic inner stages whose stride is below the wavefront width, both compare-exchange partners live in the same wave, so the swap is done register-to-register via __shfl_xor with no shared memory round-trip or __syncthreads() (wave32 on RDNA3.5/RDNA4). - top-k.cu: dedicated k=1 argmax kernel that skips the sort entirely. On HIP the CUB fast path is unavailable (GGML_CUDA_USE_CUB undefined), so the bitonic kernel is the hot path on AMD. test-backend-ops -o TOP_K passes on gfx1201 and gfx1151. Microbench A/B (ne=[1000,16]): k=1 15.59->2.59us (6.0x), k=10 ->4.58us (3.4x), k=400 ->6.56us (2.4x); geomean ~3.5-3.8x.
The GPU argsort/top-k path uses a single-block bitonic sort (block_dims =
ncols_pad), so it caps at 1024 threads/block. On CUDA, ncols > 1024 is handled
by the CUB device-sort path; on HIP that path was compiled out (GGML_CUDA_USE_CUB
is CUDA-only), so ggml_backend_cuda_supports_op reported TOP_K/ARGSORT as
unsupported for ncols > 1024 and they fell back to the CPU reference — 275 of the
test-backend-ops cases (full-vocab sampling, large sorts) on every ROCm build.
Enable the existing device-sort path on HIP via hipCUB (rocPRIM), which provides
the same DeviceRadixSort / DeviceSegmentedSort / DeviceSegmentedRadixSort API:
* common.cuh: define GGML_CUDA_USE_HIPCUB on HIP.
* argsort.cu/.cuh: include <hipcub/hipcub.hpp> and compile argsort_f32_i32_cuda_cub
on HIP too; hipCUB has no CCCL strided iterator, so the init_offsets segment
path is used (already the CUB fallback).
* top-k.cu: on HIP, route ncols > 1024 through the device sort + slice-first-k,
mirroring the CUB branch; the fast partial-bitonic top-k still handles the
common ncols <= 1024 decode/verify case unchanged.
* ggml-cuda.cu: supports_op returns true for TOP_K/ARGSORT when either CUB or
hipCUB is available.
test-backend-ops -o TOP_K and -o ARGSORT: 0 not-supported / 0 fail on gfx1201
(R9700) and gfx1151 (Strix Halo) — the 275 previously-unsupported ncols>1024
cases (up to ne=[262144]) now run on-GPU and are bit-for-bit vs the CPU
reference. Requires the rocm hipcub + rocprim headers (ship with ROCm).
…ds_swizzle (~1.19x)
Two schedule-level improvements to the ncols<=1024 partial-bitonic top-k, both
correctness-preserving (comparator/merge networks byte-identical → bit-exact
TOP_K set) and validated with interleaved A/B rebuilds on gfx1201 and gfx1151:
* Template the smallk / 2wave / shared bitonic kernels on kpad (and warp width)
with #pragma unroll, dispatched via switch(kpad) over the power-of-two ladder,
so the fixed-length shuffle/merge chains unroll fully and drop loop overhead.
Fully general — nothing keyed on the benchmark's ncols or k.
* Intra-wave xor shuffles (j < 32) issue a single ds_swizzle_b32 (bitmask mode,
(j<<10)|0x1f) instead of __shfl_xor, which on RDNA lowers to ds_bpermute_b32
and builds a per-lane address VGPR first. These kernels are latency-bound on
the dependent shuffle chains, so removing the address-compute shortens the
critical path. j==32 (wave64 tail) falls back to __shfl_xor.
test-backend-ops -o TOP_K: 0 fail on gfx1201 (R9700) and gfx1151 (Strix Halo).
Perf (ne=[1000,16], geomean over k=1/10/40/400 vs the prior partial-bitonic):
gfx1201 ~1.18x, gfx1151 ~1.19x; the k>1 cases gain 1.21-1.26x (k=1 argmax was
already at its floor and is unchanged).
…~1.06-1.10x)
Follow-up tuning of the partial-bitonic top-k, correctness-preserving (comparator/
merge networks byte-identical -> bit-exact TOP_K set), validated with interleaved
A/B rebuilds on gfx1201 and gfx1151.
The intra-wave xor shuffles were issuing through the LDS permute crossbar
(ds_swizzle_b32 / ds_bpermute_b32). Move the ones expressible as pure VALU
cross-lane gathers off the crossbar onto the ALU pipe, shortening the dependent
shuffle chains these latency-bound kernels sit on:
* xor ^1 / ^2 (intra-quad) -> DPP mov_dpp quad_perm (0xB1 / 0x4E).
* xor ^16 (widest intra-wave stride, entry stage of every 32-lane chain) ->
v_permlanex16_b32 in its identity-cross form.
* ^4 / ^8 stay on ds_swizzle_b32; the wave64 j==32 tail falls back to __shfl_xor.
* plus a middle-cross-wave-stride group-merge fuse in the shared/2wave Phase-A/B.
All cross-lane encodings verified bit-identical to __shfl_xor on gfx1201 and
gfx1151. test-backend-ops -o TOP_K: 0 fail on both archs. Perf (ne=[1000,16],
geomean over k=1/10/40/400 vs the prior partial-bitonic): gfx1201 ~1.06x,
gfx1151 ~1.10x (k=1 and k=400 gain most: gfx1151 1.13x / 1.15x).
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
2 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.
Base is
base/hip-draft-topk-488(upstream #488's commit, @cheese-cakee's "compile DFlash GPU draft top-K for HIP"), so the diff is only our four commits.Optimizes the ggml
TOP_K/ARGSORTCUDA-HIP kernels for RDNA (CUB unavailable on HIP → the bitonic path is the hot path on AMD; this is where DSpark's indexer/draft/verify top-K lands).Impact — kernel microbench (
test-backend-ops perf -o TOP_K, ne=[1000,16])End-to-end — DeepSeek-V4-Flash + DSpark decode (gfx1151, spec-decode on)
Measured on the real workload (DS4 main + DSpark drafter,
dflash_server, decode of 200 tokens):+12% end-to-end decode. The win comes from the verify step (95.8 → 83.8 ms/step): the DSpark verify runs the indexer
top_k=512across all 43 layers (batched over the draft chain), so faster TOP_K/ARGSORT compounds. accept_rate unchanged (0.31) — top-K speed doesn't affect correctness/acceptance.Changes (4 commits, all correctness-preserving; comparator/merge networks byte-identical → bit-exact TOP_K set)
argsort.cu/top-k.cu: shared-mem key cache + partial bitonic bounded by padded k (work scales with k, not ncols);GGML_ARGSORT_SHFL_XORregister-to-register sub-wavefront compare-exchange; dedicated k=1 argmax;smallk/2wavespecializations.common.cuh/argsort.*/ggml-cuda.cu: enable the device-sort path on HIP via hipCUB (rocPRIM). Previouslysupports_opgated TOP_K/ARGSORT to ne[0]≤1024 on every ROCm build (275 test-backend-ops cases fell back to CPU); now they run on-GPU (cases up to ne=[262144] verified bit-exact vs CPU). Partial-bitonic still handles ncols≤1024 (~5× faster than hipCUB there).#pragma unroll+switch(kpad));ds_swizzle_b32for intra-wave xor shuffles.quad_perm(xor ^1/^2) +permlanex16(^16) instead of the LDS-crossbards_swizzle/ds_bpermute, shortening the dependent shuffle chains (small extra gain, ~+0.5% e2e).Validation
test-backend-ops test -o TOP_K/-o ARGSORT:2/2 backends passed, 0 fail, on gfx1201 (R9700) and gfx1151 (Strix Halo).Requires the ROCm
hipcub+rocprimheaders (ship with ROCm).