perf(hip): enable the split-K GPU draft top-K kernel on the HIP backend#11
Open
DeanoC wants to merge 1 commit into
Open
perf(hip): enable the split-K GPU draft top-K kernel on the HIP backend#11DeanoC wants to merge 1 commit into
DeanoC wants to merge 1 commit into
Conversation
The split-K draft top-K + logsumexp kernel (geometric_draft_topk_cuda.cu, the spec-decode draft path) added alongside the sampler in Luce-Org#478 was wired for the CUDA backend only, so on AMD the draft top-K fell back to the CPU reference (extract_draft_topk). This enables the GPU kernel on HIP so AMD (gfx1100/gfx1151/gfx1201) uses it too. * CMakeLists: compile geometric_draft_topk_cuda.cu as LANGUAGE HIP under the hip backend and define DFLASH27B_HAVE_DRAFT_TOPK_CUDA so callers take the GPU path; build test_draft_topk_cuda on hip too (CUDA-spelled source compiles as HIP via hip_compat/). Stacked on the sampler HIP PR, which adds the shared hip_compat pointer-attribute shims this kernel also uses. * test_draft_topk_cuda: add a DFLASH_TOPK_BENCH-gated CPU-vs-GPU microbench (mirrors test_gpu_sampler_cuda's DFLASH_SAMPLER_BENCH) so the win is measurable on any backend. The kernel needs no source or shim changes: it has NO warp-level intrinsics (a plain kBlock-wide __syncthreads reduction, so no wavefront-size assumption), and its only runtime calls (cudaMalloc / cudaMemcpy / cudaPointer*) resolve through the existing shim. Validated on the AMD Radeon AI PRO R9700 (gfx1201, RDNA4, ROCm 7.1.1): * test_draft_topk_cuda: 10/10 pass — top-K ids bit-for-bit vs the CPU reference (id_mismatch=0) across Qwen vocab (151936) and n/K/temp shapes, plus the K>kMaxK CPU-fallback contract. max log-prob err <= 3.2e-4. * CPU-vs-GPU microbench (DFLASH_TOPK_BENCH=1, K=8, 500 iters, logits resident on device = the decode steady state; GPU time = launch+sync+D2H of the n*K result): n=15 vocab=151936 CPU 695.2 us -> GPU 164.9 us = 4.22x n=1 vocab=151936 CPU 594.4 us -> GPU 153.6 us = 3.87x
6f21200 to
b5f3f83
Compare
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
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.
What
Enables the split-K GPU draft top-K + logsumexp kernel
(
geometric_draft_topk_cuda.cu, the spec-decode draft path added alongside thesampler in Luce-Org#478) on the HIP backend. It was wired for the CUDA backend only,
so on AMD the draft top-K fell back to the CPU reference (
extract_draft_topk);this lets AMD (gfx1100 / gfx1151 / gfx1201) use the GPU kernel.
Changes (2 files)
server/CMakeLists.txt— compilegeometric_draft_topk_cuda.cuasLANGUAGE HIPunder the hip backend + defineDFLASH27B_HAVE_DRAFT_TOPK_CUDAso callers take the GPU path; build
test_draft_topk_cudaon hip too (theCUDA-spelled test compiles as HIP via
hip_compat/, liketest_gpu_sampler_cuda).server/test/test_draft_topk_cuda.cpp— add aDFLASH_TOPK_BENCH-gatedCPU-vs-GPU microbench (mirrors
test_gpu_sampler_cuda'sDFLASH_SAMPLER_BENCH)so the win is measurable on any backend.
No kernel or shim changes. The kernel has no warp-level intrinsics — a
plain
kBlock-wide__syncthreadsreduction, so no wavefront-size assumption —and its only runtime calls (
cudaMalloc/cudaMemcpy/cudaPointer*) resolvethrough the existing shim.
Validation — AMD Radeon AI PRO R9700 (gfx1201, RDNA4, ROCm 7.1.1)
Correctness —
test_draft_topk_cuda: 10/10 pass. Top-K ids match the CPUreference bit-for-bit (
id_mismatch=0) across Qwen vocab (151936) and a range ofn/K/temp, plus theK>kMaxKCPU-fallback contract. Max log-prob error≤ 3.2e-4 (float exp/log rounding).
CPU-vs-GPU microbench (
DFLASH_TOPK_BENCH=1, K=8, 500 iters). Logits stayresident on the device across iters — the decode steady state, where the draft
logits are produced on the GPU. GPU time is the full per-step call
(launch + sync + D2H of the small n×K result):
extract_draft_topk)geometric_extract_draft_topk_cuda)(This is a full-call CPU-vs-GPU comparison, distinct from upstream's kernel-only
392 µs → 36 µs = 10.8× figure, which measured the split-K kernel vs the old
single-block kernel on an RTX 3090.)
Real-world end-to-end impact
The per-call microbench is one thing; the honest question is the decode-level
gain. A/B on the standard
bench_he.pyworkload (Qwen3.6-27B Q4_K_M + draft,greedy spec-decode, ddtree budget 22, 10 HumanEval prompts, n_gen 256, 2 runs
each), toggling only
DFLASH_GPU_DRAFT_TOPK:Run-to-run noise <0.5%. Acceptance length is bit-identical (5.82 in every
run) — the GPU kernel is a pure, exact speedup with no output change. The
+2.7% e2e (vs 4.2× per-call) is expected: DFlash decode is dominated by the 27B
target forward pass (Q4_K matmuls + KV attention, matmul/bandwidth-bound), and
the draft top-K is an O(vocab) sliver of per-token time — so a large per-call
multiplier lands as a few % end-to-end. It's free (bit-exact) and never regresses.