Skip to content

perf(hip): enable the split-K GPU draft top-K kernel on the HIP backend#11

Open
DeanoC wants to merge 1 commit into
feat/hip-gpu-samplerfrom
feat/hip-draft-topk
Open

perf(hip): enable the split-K GPU draft top-K kernel on the HIP backend#11
DeanoC wants to merge 1 commit into
feat/hip-gpu-samplerfrom
feat/hip-draft-topk

Conversation

@DeanoC

@DeanoC DeanoC commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

What

Enables the split-K GPU draft top-K + logsumexp kernel
(geometric_draft_topk_cuda.cu, the spec-decode draft path added alongside the
sampler 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.

Stacked on #10 (the fused-sampler HIP PR). #10 adds the shared
hip_compat pointer-attribute shims (cudaPointerAttributes /
cudaPointerGetAttributes / cudaMemoryTypeDevice) that this kernel also
uses. Merge #10 first; this PR's base will then retarget to main.

Changes (2 files)

  • server/CMakeLists.txt — compile geometric_draft_topk_cuda.cu as
    LANGUAGE HIP under the hip backend + define DFLASH27B_HAVE_DRAFT_TOPK_CUDA
    so callers take the GPU path; build test_draft_topk_cuda on hip too (the
    CUDA-spelled test compiles as HIP via hip_compat/, like
    test_gpu_sampler_cuda).
  • server/test/test_draft_topk_cuda.cpp — 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.

No kernel or shim changes. The kernel 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.

Validation — AMD Radeon AI PRO R9700 (gfx1201, RDNA4, ROCm 7.1.1)

Correctnesstest_draft_topk_cuda: 10/10 pass. Top-K ids match the CPU
reference bit-for-bit (id_mismatch=0) across Qwen vocab (151936) and a range of
n/K/temp, plus the K>kMaxK CPU-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 stay
resident 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):

shape CPU (extract_draft_topk) GPU (geometric_extract_draft_topk_cuda) speedup
n=15, vocab=151936 695.2 µs 164.9 µs 4.22×
n=1, vocab=151936 594.4 µs 153.6 µs 3.87×

(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.py workload (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:

draft top-K decode tok/s AL Δ
CPU (baseline) 50.24 5.82
GPU 51.59 5.82 +2.7%

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.

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
@DeanoC DeanoC force-pushed the feat/hip-draft-topk branch from 6f21200 to b5f3f83 Compare July 7, 2026 16:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant