perf(cuda): optimize sampling based decode

README.md

@@ -283,6 +283,39 @@ The draft-token top-K extraction and the per-step verify argmax used to run on t
 
 To reproduce the benchmark: baseline `DFLASH_GPU_DRAFT_TOPK=0 DFLASH_GPU_VERIFY_ARGMAX=0`, optimized `DFLASH_GPU_DRAFT_TOPK=1 DFLASH_GPU_VERIFY_ARGMAX=1`, both via `python server/scripts/bench_llm.py --bench HumanEval`.
 
+| Env | Default | Effect |
+|---|---|---|
+| `DFLASH_TOPK_SPLIT=N` | auto-tuned | Override the split-K factor (blocks per draft position) for the GPU top-K kernel; auto-tune aims for ~240 total blocks across the device. Useful to re-sweep on a GPU with a different SM count. |
+| `DFLASH_TOPK_PROFILE=1` | off | Print per-launch CUDA event timing (partial pass + combine pass) for the GPU top-K kernel to stderr. |
+
+**GPU sampler (DFlash)**
+
+The CPU `sample_logits` chain (repetition/frequency/presence penalty → softmax(temp) → top_p nucleus → multinomial draw) requires a full vocab-wide D2H logits copy every token. `geometric_sampler_cuda.cu` ports penalty application, the softmax reductions, and the draw onto the GPU, reading logits straight off the device tensor in the qwen35 decode loop (skipping that D2H). It's **on by default** at runtime on CUDA builds (opt out with `DFLASH_GPU_SAMPLE=0`).
+
+Coverage is config-dependent, chosen by measurement rather than what's merely *possible* on the GPU:
+- **Greedy and plain temperature/penalty sampling** (no `top_k`/`top_p` truncation) run entirely on the GPU — one kernel launch does penalties, softmax, and the multinomial draw, then a 4-byte D2H copy for the result.
+- **Pure `top_p` nucleus sampling** (no `top_k`) is GPU-*assisted*: the GPU computes penalties+softmax and hands back the normalized probability vector, and the CPU does the nucleus search (`std::nth_element`-based binary search — O(vocab), not the O(vocab log vocab) a full sort would cost) on that already-computed vector.
+- **`top_k` (with or without `top_p`)** always stays on the CPU. Its cost is already cheap — `partial_sort` scales with `k`, not vocab — so a GPU round trip (kernel launch + D2H copy) measured as a net *regression*, not just a non-win.
+
+Per-call sampler-only latency at the Qwen3 vocab (151,936), measured on an RTX 3090 (GPU column reflects `DFLASH_GPU_SAMPLE=1`, CPU column reflects `=0`):
+
+| Config | CPU | GPU | Speedup |
+|---|---|---|---|
+| greedy (temp=0) | 746 µs | 139 µs | ~5.4× |
+| temp=0.8 (no truncation) | 1092 µs | 235 µs | ~4.6× |
+| temp=0.8 + top_p=0.9 | 4915 µs | 3504 µs | ~1.4× (GPU-assisted) |
+| temp=0.8 + top_k=40 | 283 µs | 273 µs | ~1.0× (top_k stays CPU-only either way) |
+
+| Env / flag | Default | Effect |
+|---|---|---|
+| `DFLASH_GPU_SAMPLE=0` | on | Opt out of the GPU `sample_logits` path at runtime (on by default on CUDA builds). Falls back to the CPU chain per call when the config is unsupported or on any CUDA error. |
+| `DFLASH_GPU_SAMPLER` (CMake option) | `ON` | Build-time switch; compiles `src/common/geometric_sampler_cuda.cu` into `dflash_common`. Configure with `-DDFLASH_GPU_SAMPLER=OFF` to drop the kernel entirely. |
+| `--samp=temp,top_p,top_k,rep_pen,seed[,freq,pres]` **(for `test_dflash`)** | greedy | Exercise the sampler chain (and its GPU port, gated by `DFLASH_GPU_SAMPLE`) in the positional (non-daemon) harness instead of greedy decode. Same field order as the daemon's ` samp=` request-line tail. |
+| `DFLASH_SAMP=temp,top_p,top_k,rep_pen,seed[,freq,pres]` **(for `bench_llm.py`)** | off (greedy) | Forward the same sampler tail to every DFlash bench call instead of greedy decode. AR (`test_generate`) is greedy-only and ignores this. |
+| `DFLASH_N_SAMPLE=N` **(for `bench_llm.py`)** | `10` | Overrides how many prompts are drawn per benchmark dataset. |
+
+End-to-end repro: `DFLASH_SAMP=0.8,1.0,0,1.1,42 python server/scripts/bench_llm.py --bench HumanEval` (GPU sampler on by default) vs the same command with `DFLASH_GPU_SAMPLE=0` (CPU-only).
+
 **Prefill compression (PFlash)**
 
 | Flag / env | Default | Effect |

server/CMakeLists.txt

@@ -415,6 +415,16 @@ elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda")
     # PUBLIC so consumers (e.g. the test_dflash executable) also see the macro
     # and take the GPU draft top-K path instead of the CPU fallback.
     target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_DRAFT_TOPK_CUDA=1)
+    # GPU port of the sample_logits chain. Compiled in by default; the path is
+    # then opted into at runtime via the DFLASH_GPU_SAMPLE env var. Turn the
+    # whole thing off at configure time with -DDFLASH_GPU_SAMPLER=OFF.
+    option(DFLASH_GPU_SAMPLER "Build the CUDA sample_logits path" ON)
+    if(DFLASH_GPU_SAMPLER)
+        target_sources(dflash_common PRIVATE src/common/geometric_sampler_cuda.cu)
+        # PUBLIC so the backends and test executables that call sample_logits()
+        # also compile their GPU dispatch branch.
+        target_compile_definitions(dflash_common PUBLIC DFLASH27B_HAVE_GPU_SAMPLER=1)
+    endif()
     # Multi-arch: scan all resolved arches and compile every applicable
     # flashprefill kernel variant.  This lets a single binary run on mixed
     # GPUs (e.g. Volta sm_70 + Pascal sm_61) without "no kernel image" errors.
@@ -624,6 +634,14 @@ if(DFLASH27B_TESTS)
         target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common CUDA::cudart)
         add_test(NAME draft_topk_cuda COMMAND test_draft_topk_cuda)
     endif()
+    # GPU port of the sample_logits chain vs the CPU reference. CUDA only:
+    # geometric_sampler_cuda.cu is compiled into dflash_common solely on the cuda backend.
+    if(DFLASH27B_GPU_BACKEND STREQUAL "cuda" AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_gpu_sampler_cuda.cpp")
+        add_executable(test_gpu_sampler_cuda test/test_gpu_sampler_cuda.cpp)
+        target_include_directories(test_gpu_sampler_cuda PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src)
+        target_link_libraries(test_gpu_sampler_cuda PRIVATE dflash_common CUDA::cudart)
+        add_test(NAME gpu_sampler_cuda COMMAND test_gpu_sampler_cuda)
+    endif()
     if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/test/test_kv_quant.cpp")
         add_executable(test_kv_quant test/test_kv_quant.cpp)
         target_include_directories(test_kv_quant PRIVATE ${DFLASH27B_SRC_INCLUDE_DIRS})

server/scripts/bench_llm.py

@@ -41,7 +41,12 @@
 
 N_GEN = 256
 BUDGET = 22  # default; overridden by --budget CLI arg
-N_SAMPLE = 10
+N_SAMPLE = int(os.environ.get("DFLASH_N_SAMPLE", "10"))
+# Optional sampler tail for the DFlash run, "temp,top_p,top_k,rep_pen,seed[,freq,pres]".
+# When set, exercises the sample_logits chain (and its GPU port, on by default,
+# opt out with DFLASH_GPU_SAMPLE=0) instead of the default greedy path. AR
+# (test_generate) is greedy-only and ignores this.
+SAMP = os.environ.get("DFLASH_SAMP", "").strip()
 
 def _gsm_gold(x):
     """Extract numeric answer after #### from GSM8K answer field."""
@@ -137,22 +142,21 @@ def _auto_max_ctx(n_prompt, n_gen: int = N_GEN):
 def run_df(path: Path, n_prompt, n_gen: int = N_GEN):
     max_ctx = _auto_max_ctx(n_prompt, n_gen)
     out_bin = TMPDIR / f"df_out.bin"
-    r = _run_checked(
-        [
-            TEST_DFLASH,
-            TARGET,
-            DRAFT,
-            str(path),
-            str(n_gen),
-            str(out_bin),
-            "--fast-rollback",
-            "--ddtree",
-            f"--ddtree-budget={BUDGET}",
-            f"--max-ctx={max_ctx}",
-        ],
-        timeout=300,
-        label="test_dflash",
-    )
+    cmd = [
+        TEST_DFLASH,
+        TARGET,
+        DRAFT,
+        str(path),
+        str(n_gen),
+        str(out_bin),
+        "--fast-rollback",
+        "--ddtree",
+        f"--ddtree-budget={BUDGET}",
+        f"--max-ctx={max_ctx}",
+    ]
+    if SAMP:
+        cmd.append(f"--samp={SAMP}")
+    r = _run_checked(cmd, timeout=300, label="test_dflash")
     tps = re.search(r"(\d+(?:\.\d+)?)\s+tok/s", r.stdout)
     al = re.search(r"avg commit/step=(\d+(?:\.\d+)?)", r.stdout)
     if not (tps and al):