perf(cuda): optimize sampling based decode

[geometric]

Summary

• Ports the CPU sample_logits chain (repetition/frequency/presence penalty, softmax(temp), top_p nucleus, multinomial draw) to CUDA so the Qwen3.5 AR decode loop can sample straight off the device logits tensor
  instead of paying a full vocab-wide (151,936-float) D2H copy every token.

• Penalty application, the softmax reductions, and the draw are now one fused kernel (mode-selected: greedy / sample / emit-probs) using warp-shuffle block reductions, replacing what used to be a separate penalty kernel plus a shared-memory tree reduction.

• top_k truncation stays on the CPU but we use partial_sort instead of a full sort, and pure top_p is
  GPU-assisted: the GPU computes penalties+softmax and hands back the probability vector for the CPU to truncate.

• CPU-side top_p also changed independently of the GPU work: nucleus truncation no longer does a full O(vocab log vocab) sort. It now usesnucleus_cutoff, an std::nth_element-based recursive bisection that
  finds the cutoff index in O(vocab) total work regardless of where the nucleus lands.

• The GPU path is on by default for CUDA builds; opt out with DFLASH_GPU_SAMPLE=0.

Impact

Kernel Level

DFLASH_SAMPLER_BENCH=1 ./test_gpu_sampler_cuda

[microbench] vocab=151936 iters=1000 (per call; >1.0x = GPU faster)
  greedy (temp=0)        CPU   133.87 us | GPU+H2D   132.36 us (1.01x) | GPU devptr    47.34 us (2.83x)
  temp=0.8 (full vocab)  CPU   233.28 us | GPU+H2D   232.35 us (1.00x) | GPU devptr   146.70 us (1.59x)
  temp=0.8 rep_pen=1.2   CPU   245.39 us | GPU+H2D   245.75 us (1.00x) | GPU devptr   167.37 us (1.47x)

End-to-End

Baseline

Compute on CPU

Run via: DFLASH_SAMP=0.8,1.0,0,1.1,42 DFLASH_GPU_SAMPLE=0 python -m server.scripts.bench_llm --bench HumanEval

Note: In the above command temp=0.8, top-p=1.0, top-k=0, rep_pen=1.1,r andom_seed=42

[bench] === SUMMARY ===
Task                AR    DFlash      AL   Speedup     Score
HumanEval        34.45     69.33    5.96     2.01x    

New Kernel

Run via: DFLASH_SAMP=0.8,1.0,0,1.1,42 DFLASH_GPU_SAMPLE=1 python -m server.scripts.bench_llm --bench HumanEval
Results: An approx ~32% increase in tok/s with similar Acceptance Lengths.

[bench] === SUMMARY ===
Task                AR    DFlash      AL   Speedup     Score
HumanEval        34.45     92.30    6.06     2.68x          

Implementation

• server/src/common/sampler.cpp — added nucleus_cutoff (nth_element-based O(vocab) bisection, replaces a full sort for top_p) and draw_from_weights (deduplicates the final weighted CDF draw); wired the two GPU dispatch points (full-GPU for greedy/temp, GPU-assisted for pure top_p) into sample_logits.

• server/src/common/geometric_sampler_cuda.cu/h — new/rewritten: single fused geometric_sample_kernel (mode-selected greedy / sample / emit-probs) doing penalty application, softmax reductions, and the multinomial draw in one launch, with warp-shuffle block reductions and a per-device pick_block_size.

• server/src/qwen35/qwen35_backend.cpp — AR decode now calls geometric_sample_logits_cuda directly on the device logits tensor when the sampler config is GPU-supported, skipping the vocab-wide D2H copy the CPU chain otherwise needs.

• server/CMakeLists.txt — added the DFLASH_GPU_SAMPLER build option (default ON) that compiles geometric_sampler_cuda.cu into dflash_common, and registered the test_gpu_sampler_cuda ctest target.

• server/test/test_gpu_sampler_cuda.cpp — new correctness test: GPU vs CPU agreement for greedy, greedy+penalties, temperature-sample distribution, and top_k/top_p CPU-fallback signaling.

• server/test/test_dflash.cpp — added --samp=temp,top_p,top_k,rep_pen,seed[,freq,pres] to the positional (non-daemon) harness so benchmarks can exercise the sampler chain.

• server/scripts/bench_llm.py — added DFLASH_SAMP (forwards the sampler tail to test_dflash --samp=) and DFLASH_N_SAMPLE (overrides prompts-per-dataset) env vars.

• README.md — documented GPU sampler coverage, runtime/build flags, and the benchmark table below.

Runtime Flags / Configuration

Default-on paths:

• DFLASH_GPU_SAMPLE — on by default on CUDA builds; handles greedy and plain temperature/penalty sampling entirely on GPU, and assists pure top_p.
  Disable path:
• DFLASH_GPU_SAMPLE=0 — opt out at runtime; every call falls back to the CPU chain.
• -DDFLASH_GPU_SAMPLER=OFF (CMake option, default ON) — drop geometric_sampler_cuda.cu from the build entirely.
  Debug/profiling-only flags:
• --samp=temp,top_p,top_k,rep_pen,seed[,freq,pres] (test_dflash positional harness) — exercise the sampler chain instead of greedy decode.
• DFLASH_SAMP=temp,top_p,top_k,rep_pen,seed[,freq,pres] / DFLASH_N_SAMPLE=N (bench_llm.py) — forward the same sampler tail to every DFlash bench call, and override the per-dataset prompt count.
  top_k (with or without top_p) is intentionally never routed to the GPU
  — its CPU partial_sort cost scales with k, not vocab, and a GPU round
  trip (kernel launch + D2H copy) measured as a net regression, not just a
  non-win. This is a deliberate, measurement-driven exclusion, not a gap.

Notes

top_p support directly on the GPU kernel (rather than GPU-assisted) is
deliberately out of scope for this PR;

[pramodith]

@claude review this PR

[b-albar]

Isn't this redundant ? Look like it is already computed above (in lz). Can't it be just be stored in sh[t] above before the reduction and avoid computing it again here ?

[pramodith]

Yes I actually created a new branch called geometric_ai/sampling_kernels and pushed this change over there before submitting the PR to lucebox

[b-albar]

What about an approximation like __expf ? Don't know if it is that faster and if an approximation is acceptable.

[pramodith]

I remember asking claude to try out the mufn fast-path exp and I think that it either said that expf already mapped to the fast path or that the results with both were similar.

[b-albar]

This is a lot of loops like this one, would vectorializing those help like in a float4 ? Don't know if its memory limited though

[pramodith]

Looking into this, I'm actually not sure how loop unrolling and vectorization happen in raw cuda.

[pramodith]

Thanks for the pointer! Coalescing the reads and writes to 128 bits i.e. 4 floats instead of just one gave us a nice boost in speed.

[DeanoC]

Superseded — this work landed upstream as Luce-Org#478 (commit 029efb7) and is now in main (the fork was fast-forwarded to upstream). Every substantive file here is byte-identical to upstream/main except geometric_sampler_cuda.cu, where upstream is a strict superset (it has the additional float4 coalesced read/write optimization layered on top of this branch's version). Closing as redundant.

Note: upstream wires both GPU kernels from this PR (the fused sampler and the split-K geometric_draft_topk_cuda) for the CUDA backend only. The fused sampler is now enabled on the HIP/R9700 path in #10; the draft top-K kernel is still CUDA-only.