Skip to content

perf(hip): enable the fused GPU sampler on the HIP backend#10

Open
DeanoC wants to merge 1 commit into
mainfrom
feat/hip-gpu-sampler
Open

perf(hip): enable the fused GPU sampler on the HIP backend#10
DeanoC wants to merge 1 commit into
mainfrom
feat/hip-gpu-sampler

Conversation

@DeanoC

@DeanoC DeanoC commented Jul 7, 2026

Copy link
Copy Markdown
Collaborator

What

Enables the fused GPU sampler — the sample_logits chain (rep/freq/pres
penalties → softmax(temp) → greedy/multinomial draw) added in Luce-Org#478 — on the
HIP backend, so it runs on-device on AMD (gfx1100 / gfx1151 / gfx1201)
instead of falling back to the CPU chain.

geometric_sampler_cuda.cu, its header, the sampler.cpp / qwen35_backend.cpp
dispatch, and test_gpu_sampler_cuda.cpp already live in the tree (Luce-Org#478) but
were wired for the CUDA backend only. This PR adds the HIP wiring — nothing
in the kernel or the dispatch changes.

The whole diff is 2 files

  • server/CMakeLists.txt — compile geometric_sampler_cuda.cu as
    LANGUAGE HIP under the hip backend (option(DFLASH_GPU_SAMPLER), ON) and
    define DFLASH27B_HAVE_GPU_SAMPLER so sampler.cpp / qwen35_backend.cpp
    take the GPU dispatch branch; build test_gpu_sampler_cuda on hip too (the
    CUDA-spelled test compiles as HIP via hip_compat/, same as
    test_flashprefill_kernels).
  • server/hip_compat/cuda_runtime.h — add cudaPointerAttributes /
    cudaPointerGetAttributes / cudaMemoryTypeDevice (the sampler uses them to
    run where the device logits live), and map CUDA's 32-bit-mask __shfl_xor_sync
    onto HIP's mask-less __shfl_xor. HIP's _sync form static_asserts against
    the implicit 32→64-bit mask promotion; the all-lanes-active warp reductions
    here don't need the mask, and the kernel's 32-lane assumption holds on RDNA
    (gfx11xx/gfx12xx are wave32) — same idiom as rms_norm_hip.cu.

No source changes to the kernel or dispatch — the AMD adaptation is entirely the
shim + build, so future edits to the sampler need no HIP-side follow-up.

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

Correctnesstest_gpu_sampler_cuda: 22/22 pass. Greedy and
greedy+penalties match the CPU argmax bit-for-bit; temperature and top_p-assist
draw distributions match the analytic softmax / nucleus (L1 < 0.03).

Per-call microbench (DFLASH_SAMPLER_BENCH=1, vocab = 151936, 1000 iters),
GPU sampling from the device logits tensor — the production decode path
(logits_on_device=true):

config CPU GPU (devptr) speedup
greedy (temp=0) 183.4 µs 54.1 µs 3.39×
temp=0.8 (full vocab) 218.8 µs 92.1 µs 2.38×
temp=0.8, rep_pen=1.2 303.7 µs 181.4 µs 1.67×

The GPU +H2D path measures ~1.0× (copying 152k logits H→D per token eats the
win) — which is exactly why the decode loop keeps logits on device and samples
in place.

Real-world end-to-end impact

The per-call speedup (3.4× on R9700) is not the e2e story. Running the same
A/B as the original CUDA PR (Luce-Org#478)
— its exact repro command, GPU sampler off
vs on (Qwen3.6-27B Q4_K_M + draft, budget 22, 10 HumanEval prompts, n_gen 256,
mean of 2 runs) — on both AMD GPUs in the box:

DFLASH_SAMP=0.8,1.0,0,1.1,42 DFLASH_GPU_SAMPLE=<0|1> python -m server.scripts.bench_llm --bench HumanEval

GPU AR DFlash (SAMPLE=0) DFlash (SAMPLE=1) e2e Δ per-call (devptr)
RX 7900 XTX (gfx1100) 25.4 59.35 60.21 +1.5% (+2.1%/step) 2.65× / 1.37×
Radeon AI PRO R9700 (gfx1201) 21.8 52.19 53.19 +1.9% (+1.4%/step) 3.4× / 2.4×

AR is unchanged (it ignores DFLASH_SAMP) and AL barely moves (≤0.7%) on both, so
the raw and per-step deltas agree. 22/22 sampler + 10/10 draft-topk correctness
pass on both arches.

The e2e gain is set by the decode denominator, not the kernel's per-call
speed.
The 7900 XTX has a lower per-call sampler speedup than the R9700
(2.65× vs 3.4×) yet the same 2% e2e — because DFlash decode is
matmul/bandwidth-bound and the sampler is a tiny O(vocab) slice of per-token
time. For contrast, Luce-Org#478 measured **
+33%** on an RTX 3090 (DFlash 69.33 →
92.30): its far faster decode shrinks that denominator, so the same fixed
sampler + per-token D→H copy the kernel removes becomes a large fraction (and its
synchronous copy stalls a fast draft/target pipeline). Net on AMD: a small but
consistent ~2% e2e, free (bit-exact greedy, distribution-correct sampling), and
never a regression.

Greedy decode is unaffected either way (argmax selects the token and never calls
sample_logits).

Behavioral note (identical to the CUDA backend)

Under temperature sampling the GPU draw is a valid sample from the same
distribution
as the CPU chain (the test checks the draw distribution against
the analytic softmax, L1 < 0.03) but is not token-identical to it for a
given seed: the GPU inverse-CDF walks vocab ids in grid-stride order vs the
CPU's contiguous order, so the same uniform can land on a different token. This
is by design (see the kModeSample comment in geometric_sampler_cuda.cu) and
is exactly how the kernel behaves on CUDA — this PR only enables the existing
kernel on HIP. Greedy decode is unaffected (argmax is deterministic). Minor doc
nit for a separate follow-up: sampler.cpp's "reproducible whether the GPU
sampler is on or off" comment refers to the RNG stream advancing by one draw
regardless of path, not to token-identical output vs the CPU chain.

Runtime controls

The fused GPU sampler (sample_logits chain: rep/freq/pres penalties ->
softmax(temp) -> greedy/multinomial draw) added in Luce-Org#478 — its kernel
(geometric_sampler_cuda.cu), header, the sampler.cpp / qwen35_backend.cpp
dispatch, and test_gpu_sampler_cuda — was wired for the CUDA backend only. This
enables it on HIP so AMD (gfx1100/gfx1151/gfx1201) samples on-device instead of
falling back to the CPU chain. No kernel or dispatch changes; the whole diff is
the HIP build wiring + shim:

  * CMakeLists: compile geometric_sampler_cuda.cu as LANGUAGE HIP under the hip
    backend (option DFLASH_GPU_SAMPLER, ON) and define DFLASH27B_HAVE_GPU_SAMPLER
    so the backends take the GPU dispatch branch; build test_gpu_sampler_cuda on
    hip too (the CUDA-spelled test compiles as HIP via hip_compat/, like
    test_flashprefill_kernels).
  * hip_compat/cuda_runtime.h: add cudaPointerAttributes /
    cudaPointerGetAttributes / cudaMemoryTypeDevice, and map CUDA's 32-bit-mask
    __shfl_xor_sync onto HIP's mask-less __shfl_xor (HIP's _sync form
    static_asserts against the implicit 32->64-bit mask promotion; the
    all-lanes-active wave32 reductions here don't need the mask, matching
    rms_norm_hip.cu). The kernel's 32-lane warp assumption holds on RDNA wave32.

Validated on the AMD Radeon AI PRO R9700 (gfx1201, RDNA4, ROCm 7.1.1):
  * test_gpu_sampler_cuda: 22/22 pass (greedy/penalties match CPU argmax
    bit-for-bit; temp + top_p-assist draws match analytic softmax/nucleus).
  * per-call microbench (vocab=151936, 1000 iters), GPU sampling from the device
    logits tensor (the production decode path, logits_on_device=true):
      greedy (temp=0)        54.1 us vs CPU 183.4 us  = 3.39x
      temp=0.8 (full vocab)  92.1 us vs CPU 218.8 us  = 2.38x
      temp=0.8 rep_pen=1.2  181.4 us vs CPU 303.7 us  = 1.67x
    (GPU+H2D is ~1.0x — copying 152k logits H2D per token eats the win, which is
    why the decode loop samples in place.)
  * dflash_server e2e (Qwen3.6-27B Q4_K_M + draft, temp=0.8): DFLASH_GPU_SAMPLE=1
    vs =0 both coherent, no regression (21.9 vs 21.7 tok/s), no HIP errors.
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