perf(hip): enable the fused GPU sampler on the HIP backend#10
Open
DeanoC wants to merge 1 commit into
Open
Conversation
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.
This was referenced Jul 7, 2026
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 fused GPU sampler — the
sample_logitschain (rep/freq/prespenalties →
softmax(temp)→ greedy/multinomial draw) added in Luce-Org#478 — on theHIP 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, thesampler.cpp/qwen35_backend.cppdispatch, and
test_gpu_sampler_cuda.cppalready live in the tree (Luce-Org#478) butwere 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— compilegeometric_sampler_cuda.cuasLANGUAGE HIPunder the hip backend (option(DFLASH_GPU_SAMPLER), ON) anddefine
DFLASH27B_HAVE_GPU_SAMPLERsosampler.cpp/qwen35_backend.cpptake the GPU dispatch branch; build
test_gpu_sampler_cudaon hip too (theCUDA-spelled test compiles as HIP via
hip_compat/, same astest_flashprefill_kernels).server/hip_compat/cuda_runtime.h— addcudaPointerAttributes/cudaPointerGetAttributes/cudaMemoryTypeDevice(the sampler uses them torun where the device logits live), and map CUDA's 32-bit-mask
__shfl_xor_synconto HIP's mask-less
__shfl_xor. HIP's_syncformstatic_asserts againstthe 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)
Correctness —
test_gpu_sampler_cuda: 22/22 pass. Greedy andgreedy+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):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 HumanEvalAR is unchanged (it ignores
DFLASH_SAMP) and AL barely moves (≤0.7%) on both, sothe 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+33%** on an RTX 3090 (DFlash 69.33 →matmul/bandwidth-bound and the sampler is a tiny O(vocab) slice of per-token
time. For contrast, Luce-Org#478 measured **
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
kModeSamplecomment ingeometric_sampler_cuda.cu) andis 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 GPUsampler 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
DFLASH_GPU_SAMPLE=0— disable at runtime (CPU fallback).-DDFLASH_GPU_SAMPLER=OFF— drop it at configure time.top_k>0andtop_p∈(0,1)deliberately fall back to the CPU chain (unchanged from perf(cuda): port sample_logits chain to a fused GPU sampler kernel Luce-Org/lucebox#478).