diff --git a/README.md b/README.md index e7aa0f8e2..136d166b1 100644 --- a/README.md +++ b/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 | diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 1b894c0af..803c156a2 100644 --- a/server/CMakeLists.txt +++ b/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}) diff --git a/server/deps/llama.cpp b/server/deps/llama.cpp index 0aaf1dee4..9cd9e1edb 160000 --- a/server/deps/llama.cpp +++ b/server/deps/llama.cpp @@ -1 +1 @@ -Subproject commit 0aaf1dee404c93d69abfee1ce1f03981ea5bda1e +Subproject commit 9cd9e1edb6d244ac82b9dc925d67c3bf9c9f3480 diff --git a/server/scripts/bench_llm.py b/server/scripts/bench_llm.py index b8b31c91d..8348e06a2 100644 --- a/server/scripts/bench_llm.py +++ b/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): diff --git a/server/src/common/geometric_sampler_cuda.cu b/server/src/common/geometric_sampler_cuda.cu new file mode 100644 index 000000000..8bbd67972 --- /dev/null +++ b/server/src/common/geometric_sampler_cuda.cu @@ -0,0 +1,510 @@ +// CUDA port of the sample_logits chain. See geometric_sampler_cuda.h for the contract and +// the CPU/GPU split rationale, and sampler.cpp for the reference CPU chain this +// mirrors: rep_penalty -> freq/pres_penalty -> softmax(temp) -> draw. +// top_p (nucleus) is not implemented here — cfg.top_p in (0,1) falls back to +// the CPU chain (see geometric_sampler_cuda.h). +// +// The per-call workload is one logit row (vocab ~150k). That is small enough +// that a single thread block handles the whole row: it keeps every reduction +// and the inverse-CDF scan in shared memory with no cross-block +// synchronization, which is far simpler and — for one row per token — fast +// enough. (The bandwidth-bound multi-block split used by +// geometric_draft_topk_cuda.cu pays off only for many rows at once.) +// +// A few structural choices below are lifted from ggml-cuda's softmax kernel +// (ggml/src/ggml-cuda/softmax.cu): +// * penalty application, greedy argmax, temp-sample draw, and (when +// top_k/top_p force CPU-side truncation) raw probability emission are all +// one kernel (geometric_sample_kernel, mode-selected) instead of separate +// penalty/softmax kernels — pass 0 (penalties) and pass 1 (max/argmax) +// are identical work regardless of what happens after, so duplicating +// them across kernels bought nothing; +// * penalty application is fused into pass 0 of that kernel instead of a +// separate launch — the penalty set is tiny (<= rep_window ids), so +// folding it in costs a few extra loop iterations for a handful of +// threads, not a new pass over the vocab, and it removes a whole kernel +// launch + sync boundary per token; +// * block-wide max/sum/argmax reductions use warp-shuffle (__shfl_xor_sync) +// first, falling back to a small (<=32-slot) shared-memory reduction only +// across warps, mirroring ggml's warp_reduce_*/block reduction idiom. +// This cuts block-wide __syncthreads() calls per reduction from +// log2(nthreads) (10, for a 1024-thread block) down to a handful, and +// drops the per-thread argmax-index shared array entirely (moved to a +// fixed 32-slot static array), shrinking the dynamic shared-memory +// footprint below. + +#include "geometric_sampler_cuda.h" + +#include +#include +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +constexpr int kWarpSize = 32; + +// Warp-level reductions via register shuffles — no shared memory, no +// __syncthreads(). Mirrors ggml-cuda's warp_reduce_sum/warp_reduce_max +// (ggml/src/ggml-cuda/common.cuh); we only need sum and argmax variants here +// (the kernel below always tracks argmax alongside max — see its pass 1). +__device__ __forceinline__ float warp_reduce_sum(float x) { +#pragma unroll + for (int off = kWarpSize / 2; off > 0; off >>= 1) + x += __shfl_xor_sync(0xffffffffu, x, off, kWarpSize); + return x; +} +__device__ __forceinline__ void warp_reduce_argmax(float & val, int & idx) { +#pragma unroll + for (int off = kWarpSize / 2; off > 0; off >>= 1) { + const float ov = __shfl_xor_sync(0xffffffffu, val, off, kWarpSize); + const int oi = __shfl_xor_sync(0xffffffffu, idx, off, kWarpSize); + if (ov > val || (ov == val && oi < idx)) { val = ov; idx = oi; } + } +} + +// Block-wide reductions: each warp reduces internally via the shuffles above +// (no sync), lane 0 of each warp parks its partial in a small scratch array +// (one slot per warp, so <=32 slots for any block <=1024 threads), then the +// first warp finishes the reduction over those partials and broadcasts the +// result back through the same scratch slot. `*_scratch` must have >=32 +// (or, for argmax, 32 float + 32 int) entries; callers pass a fixed-size +// __shared__ array so this needs no dynamic-shared-memory sizing. The +// trailing sync (after every thread has read the broadcast value) guards +// against a subsequent block_reduce_* call reusing the same scratch slots +// before all threads are done reading this one. +__device__ __forceinline__ float block_reduce_sum(float x, float * scratch) { + const int lane = threadIdx.x & (kWarpSize - 1); + const int wid = threadIdx.x / kWarpSize; + const int nwarps = (blockDim.x + kWarpSize - 1) / kWarpSize; + x = warp_reduce_sum(x); + if (lane == 0) scratch[wid] = x; + __syncthreads(); + x = (threadIdx.x < nwarps) ? scratch[threadIdx.x] : 0.0f; + if (wid == 0) x = warp_reduce_sum(x); + if (threadIdx.x == 0) scratch[0] = x; + __syncthreads(); + const float result = scratch[0]; + __syncthreads(); + return result; +} +__device__ __forceinline__ void block_reduce_argmax(float & val, int & idx, + float * val_scratch, int * idx_scratch) { + const int lane = threadIdx.x & (kWarpSize - 1); + const int wid = threadIdx.x / kWarpSize; + const int nwarps = (blockDim.x + kWarpSize - 1) / kWarpSize; + warp_reduce_argmax(val, idx); + if (lane == 0) { val_scratch[wid] = val; idx_scratch[wid] = idx; } + __syncthreads(); + if (threadIdx.x < nwarps) { val = val_scratch[threadIdx.x]; idx = idx_scratch[threadIdx.x]; } + else { val = -FLT_MAX; idx = INT_MAX; } + if (wid == 0) warp_reduce_argmax(val, idx); + if (threadIdx.x == 0) { val_scratch[0] = val; idx_scratch[0] = idx; } + __syncthreads(); + val = val_scratch[0]; + idx = idx_scratch[0]; + __syncthreads(); +} + +// geometric_sample_kernel's per-thread shared-memory footprint: two float +// arrays (sh, s_off), each sized [blockDim.x], used only by the inverse-CDF +// draw (each thread's chunk mass, then its exclusive prefix offset — a true +// per-thread array, not reducible to a scalar). The max/argmax/sum +// reductions above no longer need a per-thread shared array (they use the +// fixed 32-slot statics), so this is 2 floats/thread instead of the previous +// 2 floats + 1 int32. float, not double: the reductions sum ~vocab +// already-float terms (softmax probabilities in (0,1]), and float32 error is +// negligible for a stochastic sampling draw — while avoiding double-precision +// arithmetic, which on consumer GPUs (e.g. RTX 3090) runs at a small fraction +// of float32 throughput. +constexpr size_t kSampleShmemPerThread = 2 * sizeof(float); + +// Per-device sample-kernel block size (threads per row), cached after the +// first call — same single-threaded-decode-loop assumption as Scratch below. +// Queried rather than hardcoded because maxThreadsPerBlock and +// sharedMemPerBlock vary across GPUs (e.g. older compute-capability parts cap +// at 512 threads and/or less shared memory than the 1024-thread/8-byte-per- +// thread shape that fits an RTX 3090). +struct BlockCfg { + int device = -1; + int block = 0; +}; +BlockCfg g_block_cfg; + +int pick_block_size(int device) { + if (g_block_cfg.device == device) return g_block_cfg.block; + int max_threads = 1024; + size_t smem_cap = 48 * 1024; // conservative fallback if the query fails + cudaDeviceProp prop{}; + if (cudaGetDeviceProperties(&prop, device) == cudaSuccess) { + max_threads = prop.maxThreadsPerBlock; + smem_cap = prop.sharedMemPerBlock; + } else { + cudaGetLastError(); + } + int block = 1; + while (block * 2 <= max_threads) block *= 2; // largest power of two <= max_threads + // Halve further if the reductions' shared-memory arrays wouldn't fit + // (stays a power of two so the block-reduction loop below still works). + while (block > 32 && (size_t)block * kSampleShmemPerThread > smem_cap) block /= 2; + g_block_cfg = {device, block}; + return block; +} + +// Applies the (already-uploaded) sparse penalty set to `work[]` in place — +// pass 0 of geometric_sample_kernel below (shared by all three modes). +// `m` is tiny (<= cfg.rep_window token ids), so a grid-stride loop over it +// costs a handful of extra iterations for a handful of threads, not a new +// pass over the vocab; folding it in here removes what used to be a separate +// geometric_apply_penalties launch + sync boundary before every sample/probs +// call. `rep_pen` is the raw multiplicative penalty (or 1.0 to disable); +// `add[i]` folds freq_pen*count + pres_pen for pen_ids[i]. Order matches the +// CPU chain: multiplicative repetition penalty first, then the additive +// frequency/presence subtraction. Callers must __syncthreads() after this +// (only if m>0) before any thread reads work[] for the reductions below, so +// every thread's penalty writes are visible. +__device__ __forceinline__ void apply_penalties_inplace(float * __restrict__ work, + const int32_t * __restrict__ pen_ids, + const float * __restrict__ pen_add, + int m, float rep_pen, int rep_active) { + for (int i = threadIdx.x; i < m; i += blockDim.x) { + const int id = pen_ids[i]; + float v = work[id]; + if (rep_active) v = (v > 0.0f) ? v / rep_pen : v * rep_pen; + v -= pen_add[i]; + work[id] = v; + } +} + +// geometric_sample_kernel's three mutually-exclusive endings: +// kModeGreedy - write the argmax token id, nothing else (temp<=0). +// kModeSample - draw one token via inverse-CDF and write its id. +// kModeEmitProbs - write the full normalized probability vector (used when +// top_k/top_p forces CPU-side truncation, so the CPU +// doesn't need to redo penalties+softmax itself). +// All three share pass 0 (penalty fusion) and pass 1 (max/argmax) below; +// kModeSample and kModeEmitProbs also share pass 2 (the Z reduction). Only +// kModeGreedy and kModeEmitProbs skip the inverse-CDF draw's per-thread `sh`/ +// `s_off` array, so only kModeSample launches need to size dynamic shared +// memory for it (see the call sites) — kModeGreedy/kModeEmitProbs never touch +// `geometric_smem` and can launch with 0 dynamic shared memory. Note +// out_probs is written ONLY in kModeEmitProbs: the greedy/sample paths never +// materialize the full probability vector in global memory at all, since the +// draw happens entirely from registers/shared memory inside the kernel. +constexpr int kModeGreedy = 0; +constexpr int kModeSample = 1; +constexpr int kModeEmitProbs = 2; + +// Dynamic shared memory for geometric_sample_kernel's kModeSample path, +// sized at launch to blockDim.x * kSampleShmemPerThread (see +// pick_block_size): two float arrays (sh, s_off), used only by the +// inverse-CDF draw (the block-wide reductions use the fixed 32-slot statics +// declared inside the kernel). +extern __shared__ unsigned char geometric_smem[]; + +// Single-block sampler/softmax kernel. `work[]` holds the pre-penalty logits +// (penalties are applied in place as pass 0, in-kernel — see +// apply_penalties_inplace). Behaviour selected by `mode` (see above); ties in +// the greedy argmax go to the lowest token id, matching the CPU manual-argmax +// and DFLASH_GPU_ARGMAX behaviour. +__global__ void geometric_sample_kernel(float * __restrict__ work, int vocab, + float inv_t, int mode, + double r_uniform, + int32_t * __restrict__ out_token, + float * __restrict__ out_probs, + const int32_t * __restrict__ pen_ids, + const float * __restrict__ pen_add, + int pen_m, float rep_pen, int rep_active) { + const int t = threadIdx.x; + const int nthreads = blockDim.x; + + apply_penalties_inplace(work, pen_ids, pen_add, pen_m, rep_pen, rep_active); + if (pen_m > 0) __syncthreads(); + + const int chunk = (vocab + nthreads - 1) / nthreads; + const int begin = t * chunk; + const int end = min(begin + chunk, vocab); + + float * sh = reinterpret_cast(geometric_smem); + float * s_off = sh + nthreads; + __shared__ float s_val_scratch[kWarpSize]; + __shared__ int s_idx_scratch[kWarpSize]; + + // ---- pass 1: max logit + argmax (lowest id on ties), warp-shuffle first, + // falling back to the <=32-slot statics only across warps (see + // block_reduce_argmax) — far fewer __syncthreads() than a full + // shared-memory tree reduction over all nthreads slots. Computed in every + // mode (kModeEmitProbs doesn't need the index, but the extra shuffle over + // an int alongside the float is negligible next to the O(vocab) scan + // above it, and keeping one shared pass 1 avoids a second reduction + // helper just to drop the index). + float lmax = -FLT_MAX; + int largmax = vocab; // sentinel so a real id always wins + for (int i = begin; i < end; i++) { + const float v = work[i]; + if (v > lmax || (v == lmax && i < largmax)) { lmax = v; largmax = i; } + } + block_reduce_argmax(lmax, largmax, s_val_scratch, s_idx_scratch); + const float xmax = lmax * inv_t; + const int argmax = largmax; + + if (mode == kModeGreedy) { + if (t == 0) *out_token = argmax; + return; + } + + // ---- pass 2: softmax denominator Z = sum exp(x_i - xmax) -------------- + // float throughout, not double: expf() (float) is ~19x faster than exp() + // (double) on consumer GPUs (e.g. RTX 3090), where FP64 throughput is a + // small fraction of FP32. The sums are ~vocab terms in (0,1], and float32 + // error is negligible for a stochastic sampling draw. + float lz = 0.0f; + for (int i = begin; i < end; i++) + lz += expf(work[i] * inv_t - xmax); + const float Z = block_reduce_sum(lz, s_val_scratch); + + if (mode == kModeEmitProbs) { + for (int i = begin; i < end; i++) + out_probs[i] = expf(work[i] * inv_t - xmax) / Z; + return; + } + + // ---- kModeSample: multinomial inverse-CDF draw over the full distribution + // Each thread owns a contiguous id chunk, so ascending-id order is exactly + // thread 0's chunk, then thread 1's, ... A serial exclusive scan over the + // nthreads per-thread masses gives each thread its CDF offset; the one whose + // [offset, offset+mass) straddles the target re-scans its chunk for the id. + // (This per-thread array can't be collapsed into the warp-shuffle + // reductions above — thread 0 needs every individual chunk mass, not just + // their sum, to compute the prefix offsets.) + // Each thread's chunk mass is exactly the partial sum it already accumulated + // for Z in pass 2 — block_reduce_sum takes `lz` by value, so it is still + // intact here. Reuse it instead of a second O(vocab) expf pass over the row. + const float pm = lz; + sh[t] = pm; + __syncthreads(); + + // Thread 0: seed the safety default (used only if fp rounding leaves no + // straddling thread) and compute the exclusive prefix offsets. + if (t == 0) { + *out_token = argmax; + float acc = 0.0f; + for (int k = 0; k < nthreads; k++) { s_off[k] = acc; acc += sh[k]; } + } + __syncthreads(); + + // r_uniform is a host-drawn double; keep this comparison in double (a + // single scalar op, not a reduction, so it doesn't cost the FP64 penalty). + const double targetv = r_uniform * (double)Z; + if (targetv >= (double)s_off[t] && targetv < (double)s_off[t] + pm) { + float acc = s_off[t]; + for (int i = begin; i < end; i++) { + acc += expf(work[i] * inv_t - xmax); + if (targetv < acc) { *out_token = i; break; } + } + } +} + +// Per-device persistent scratch. The decode loop is single-threaded, so a plain +// static cache avoids a cudaMalloc/cudaFree per token (mirrors geometric_draft_topk_cuda). +struct Scratch { + int device = -1; + int vocab_cap = 0; + int pen_cap = 0; + float * d_work = nullptr; // [vocab] mutable logit copy + float * d_probs = nullptr; // [vocab] softmax output (geometric_compute_probs_cuda) + int32_t * d_pen_id = nullptr; // [pen_cap] + float * d_pen_add = nullptr; // [pen_cap] + int32_t * d_out = nullptr; // [1] +}; +Scratch g_scratch; + +void free_scratch() { + if (g_scratch.d_work) cudaFree(g_scratch.d_work); + if (g_scratch.d_probs) cudaFree(g_scratch.d_probs); + if (g_scratch.d_pen_id) cudaFree(g_scratch.d_pen_id); + if (g_scratch.d_pen_add) cudaFree(g_scratch.d_pen_add); + if (g_scratch.d_out) cudaFree(g_scratch.d_out); + g_scratch = Scratch{}; +} + +bool ensure_scratch(int device, int vocab, int pen) { + const bool ok = g_scratch.device == device && + g_scratch.vocab_cap >= vocab && + g_scratch.pen_cap >= pen; + if (ok) return true; + free_scratch(); + if (cudaMalloc(&g_scratch.d_out, sizeof(int32_t)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_work, (size_t)vocab * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_probs, (size_t)vocab * sizeof(float)) != cudaSuccess) goto fail; + if (pen > 0) { + if (cudaMalloc(&g_scratch.d_pen_id, (size_t)pen * sizeof(int32_t)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_pen_add, (size_t)pen * sizeof(float)) != cudaSuccess) goto fail; + } + g_scratch.device = device; + g_scratch.vocab_cap = vocab; + g_scratch.pen_cap = pen; + return true; +fail: + free_scratch(); + return false; +} + +// Uploads `logits` (host or device) into g_scratch.d_work and the sparse +// penalty set into g_scratch.d_pen_id/d_pen_add — the shared prefix of +// geometric_sample_logits_cuda and geometric_compute_probs_cuda. Penalty +// *application* itself is no longer done here: it's fused into pass 0 of +// geometric_sample_kernel (see apply_penalties_inplace), so this function +// only stages data, no kernel launch. Writes the penalty count to `*out_m`. +// Assumes `dev` is already the current device. Returns false on any CUDA +// allocation/copy error. +bool stage_and_penalize(int dev, const float * logits, int vocab, const SamplerCfg & cfg, + const std::vector & history, bool logits_on_device, int * out_m) { + std::vector pen_id; + std::vector pen_add; + const bool rep_active = cfg.rep_pen > 1.0f; + const bool add_active = (cfg.freq_pen != 0.0f || cfg.pres_pen != 0.0f); + if ((rep_active || add_active) && !history.empty()) { + const int win = std::min((int)history.size(), cfg.rep_window); + const int from = (int)history.size() - win; + std::unordered_map counts; + for (int i = from; i < (int)history.size(); i++) counts[history[i]]++; + pen_id.reserve(counts.size()); + pen_add.reserve(counts.size()); + for (const auto & kv : counts) { + if (kv.first < 0 || kv.first >= vocab) continue; + pen_id.push_back(kv.first); + pen_add.push_back(add_active ? (cfg.freq_pen * kv.second + cfg.pres_pen) : 0.0f); + } + } + const int m = (int)pen_id.size(); + + if (!ensure_scratch(dev, vocab, m)) return false; + if (cudaMemcpy(g_scratch.d_work, logits, (size_t)vocab * sizeof(float), + logits_on_device ? cudaMemcpyDeviceToDevice + : cudaMemcpyHostToDevice) != cudaSuccess) { + return false; + } + if (m > 0) { + if (cudaMemcpy(g_scratch.d_pen_id, pen_id.data(), (size_t)m * sizeof(int32_t), + cudaMemcpyHostToDevice) != cudaSuccess) return false; + if (cudaMemcpy(g_scratch.d_pen_add, pen_add.data(), (size_t)m * sizeof(float), + cudaMemcpyHostToDevice) != cudaSuccess) return false; + } + *out_m = m; + return true; +} + +} // namespace + +bool gpu_sampler_enabled() { + static const bool on = []() { + const char * v = std::getenv("DFLASH_GPU_SAMPLE"); + if (v == nullptr || v[0] == '\0') return true; // on by default + return v[0] != '0'; // "0" (or "0...") opts out + }(); + return on; +} + +int geometric_sample_logits_cuda(const float * logits, + int vocab, + const SamplerCfg & cfg, + const std::vector & history, + double r_uniform, + bool logits_on_device) { + if (!logits || vocab <= 0) return -1; + // top_k stays on the CPU (a single-row partial_sort beats a per-token GPU + // select); signal fallback. + if (cfg.top_k > 0 && cfg.top_k < vocab) return -1; + // top_p (nucleus) is not implemented on this kernel; signal fallback (see + // geometric_sampler_cuda.h for why). + if (cfg.top_p > 0.0f && cfg.top_p < 1.0f) return -1; + + // Pick the device. For a device pointer, derive it from the allocation so we + // run where the logits live; otherwise use the current device. + int dev = 0; + if (logits_on_device) { + cudaPointerAttributes attr{}; + if (cudaPointerGetAttributes(&attr, logits) != cudaSuccess) { cudaGetLastError(); return -1; } + if (attr.type != cudaMemoryTypeDevice) return -1; + dev = attr.device; + } else { + cudaGetDevice(&dev); + } + int prev = 0; + cudaGetDevice(&prev); + if (dev != prev) cudaSetDevice(dev); + + int result = -1; + int m = 0; + if (stage_and_penalize(dev, logits, vocab, cfg, history, logits_on_device, &m)) { + const int mode = (cfg.temp > 0.0f) ? kModeSample : kModeGreedy; + const float inv_t = 1.0f / fmaxf(1e-3f, cfg.temp); + const int block = pick_block_size(dev); + // Only kModeSample's inverse-CDF draw touches geometric_smem; skip + // sizing it for the (cheaper, no-shared-mem) greedy launch. + const size_t shmem_bytes = (mode == kModeSample) ? (size_t)block * kSampleShmemPerThread : 0; + geometric_sample_kernel<<<1, block, shmem_bytes>>>(g_scratch.d_work, vocab, inv_t, mode, + r_uniform, g_scratch.d_out, /*out_probs=*/nullptr, + g_scratch.d_pen_id, g_scratch.d_pen_add, m, + cfg.rep_pen, cfg.rep_pen > 1.0f ? 1 : 0); + int32_t tok = -1; + if (cudaGetLastError() == cudaSuccess && + cudaMemcpy(&tok, g_scratch.d_out, sizeof(int32_t), + cudaMemcpyDeviceToHost) == cudaSuccess) { + result = tok; + } + } + if (result < 0) cudaGetLastError(); // clear so we don't poison the next call + if (dev != prev) cudaSetDevice(prev); + return result; +} + +bool geometric_compute_probs_cuda(const float * logits, + int vocab, + const SamplerCfg & cfg, + const std::vector & history, + float * out_probs, + bool logits_on_device) { + if (!logits || vocab <= 0 || !out_probs || cfg.temp <= 0.0f) return false; + + int dev = 0; + if (logits_on_device) { + cudaPointerAttributes attr{}; + if (cudaPointerGetAttributes(&attr, logits) != cudaSuccess) { cudaGetLastError(); return false; } + if (attr.type != cudaMemoryTypeDevice) return false; + dev = attr.device; + } else { + cudaGetDevice(&dev); + } + int prev = 0; + cudaGetDevice(&prev); + if (dev != prev) cudaSetDevice(dev); + + bool ok = false; + int m = 0; + if (stage_and_penalize(dev, logits, vocab, cfg, history, logits_on_device, &m)) { + const float inv_t = 1.0f / fmaxf(1e-3f, cfg.temp); + const int block = pick_block_size(dev); + // kModeEmitProbs never touches geometric_smem (no inverse-CDF draw), + // so no dynamic shared memory needed. + geometric_sample_kernel<<<1, block>>>(g_scratch.d_work, vocab, inv_t, kModeEmitProbs, + /*r_uniform=*/0.0, /*out_token=*/nullptr, g_scratch.d_probs, + g_scratch.d_pen_id, g_scratch.d_pen_add, m, + cfg.rep_pen, cfg.rep_pen > 1.0f ? 1 : 0); + if (cudaGetLastError() == cudaSuccess && + cudaMemcpy(out_probs, g_scratch.d_probs, (size_t)vocab * sizeof(float), + cudaMemcpyDeviceToHost) == cudaSuccess) { + ok = true; + } + } + if (!ok) cudaGetLastError(); // clear so we don't poison the next call + if (dev != prev) cudaSetDevice(prev); + return ok; +} + +} // namespace dflash::common diff --git a/server/src/common/geometric_sampler_cuda.h b/server/src/common/geometric_sampler_cuda.h new file mode 100644 index 000000000..9ae4b8d42 --- /dev/null +++ b/server/src/common/geometric_sampler_cuda.h @@ -0,0 +1,92 @@ +// CUDA port of the shared sample_logits chain (see sampler.cpp / sampler.h). +// +// Rationale for what is / isn't on the GPU +// ---------------------------------------- +// The model already produces logits on the GPU. The CPU sampler forces a full +// ~vocab-wide D2H copy of those logits every token (the existing greedy path +// already dodges this with DFLASH_GPU_ARGMAX). The vocab-wide work — penalty +// application, the softmax max/sum-exp reductions, and the multinomial +// inverse-CDF draw — is data-parallel and is what geometric_sample_logits_cuda +// moves to the GPU, entirely, for greedy and plain temperature/penalty +// sampling. A few things deliberately stay on (or partly on) the CPU, decided +// by measurement rather than what's merely possible on the GPU: +// * building the repetition/frequency penalty index from the (<=rep_window) +// token history — a few hundred elements, dwarfed by a kernel launch; +// * the single scalar RNG draw (a std::mt19937_64 uniform) — kept on the host +// so the random stream stays reproducible and identical to the CPU path; +// * top_k selection (with or without top_p) — a CPU partial_sort scales with +// k, not vocab, and is already cheap; a GPU round trip (kernel launch + +// D2H copy) measured as a net *regression* here, not just a non-win, so +// cfg.top_k>0 always returns -1 from geometric_sample_logits_cuda and the +// caller (sample_logits) never routes it through geometric_compute_probs_cuda +// either; +// * top_p (nucleus) without top_k — geometric_sample_logits_cuda still +// returns -1 for this (no on-GPU nucleus search), but geometric_compute_probs_cuda +// lets the CPU skip straight to its O(vocab) std::nth_element-based nucleus +// search instead of also recomputing penalties+softmax; measured as a real +// (if modest, ~1.4x) end-to-end win, since that search's cost dominates +// either way. + +#pragma once + +#include "sampler.h" + +#include +#include + +namespace dflash::common { + +// GPU sample_logits. Returns the chosen token id, or -1 when the caller should +// fall back to the CPU sample_logits (unsupported config such as cfg.top_k>0 +// or cfg.top_p in (0,1), or any CUDA error). +// +// logits : pointer to vocab contiguous floats. +// logits_on_device : true -> `logits` is a device pointer (e.g. a ggml CUDA +// tensor's ->data); the D2H copy is skipped and we +// read straight from device memory (the real win). +// false -> `logits` is host memory and is uploaded H2D. +// r_uniform : a pre-drawn uniform in [0,1) from the caller's RNG; +// ignored for greedy (cfg.temp <= 0). +int geometric_sample_logits_cuda(const float * logits, + int vocab, + const SamplerCfg & cfg, + const std::vector & history, + double r_uniform, + bool logits_on_device); + +// Computes the post-penalty, post-softmax(temp) probability vector on the GPU +// and writes it to `out_probs` (vocab floats, host memory) — the same +// vocab-wide, data-parallel prefix geometric_sample_logits_cuda computes +// internally, exposed here for callers that need to do their own top_k/top_p +// truncation afterward (unsupported inside the kernel above). Lets the CPU +// skip re-deriving penalties+softmax when cfg.top_k>0 or cfg.top_p in (0,1) +// forces truncation to happen on the host. +// +// Returns false (out_probs left untouched) if cfg.temp <= 0 (softmax is +// undefined for greedy — callers should use plain argmax instead) or on any +// CUDA error, so the caller falls back to computing everything on the CPU. +bool geometric_compute_probs_cuda(const float * logits, + int vocab, + const SamplerCfg & cfg, + const std::vector & history, + float * out_probs, + bool logits_on_device); + +// True unless the env var DFLASH_GPU_SAMPLE is explicitly set to "0" — the GPU +// path is enabled by default. Cached after the first call. Lets call sites +// gate the GPU path at runtime. +bool gpu_sampler_enabled(); + +// Whether geometric_sample_logits_cuda can (and should) handle this config +// entirely on the GPU. false for top_k>0 or top_p in (0,1) — not because +// geometric_sample_logits_cuda can't be asked to try (it always returns -1 +// for those anyway), but so callers like sample_logits know to reach for +// geometric_compute_probs_cuda instead for the pure-top_p case. See the +// rationale comment at the top of this file for the per-case reasoning. +inline bool gpu_sampler_supports(const SamplerCfg & cfg) { + if (cfg.top_k > 0) return false; // top_k: not implemented on GPU + if (cfg.top_p > 0.0f && cfg.top_p < 1.0f) return false; // top_p: not implemented on GPU + return true; +} + +} // namespace dflash::common diff --git a/server/src/common/sampler.cpp b/server/src/common/sampler.cpp index 3927916a3..b00c44ce1 100644 --- a/server/src/common/sampler.cpp +++ b/server/src/common/sampler.cpp @@ -9,13 +9,152 @@ #include #include +#ifdef DFLASH27B_HAVE_GPU_SAMPLER +#include "geometric_sampler_cuda.h" +#endif + namespace dflash::common { +namespace { + +// Binary search (quickselect via std::nth_element) for the smallest index +// `cut` such that the descending-by-value prefix cand[0,cut) has cumulative +// mass >= target, where "mass" of an element is given by `mass_of`. At each +// level, partitioning [lo,hi) at its midpoint puts exactly the top (mid-lo) +// elements of that range into [lo,mid) (in some order) — whichever half still +// contains the boundary is recursed into and the other is discarded. Mutates +// `cand` in place; only the final small base-case range ends up sorted, since +// order doesn't matter for the caller's draw, only which elements make the +// cut. Each level's cost is proportional to its (shrinking) range, so total +// work is O(cand.size()), not O(cand.size() log cand.size()) like a full sort, +// regardless of where the cutoff lands. +template +size_t nucleus_cutoff(std::vector> & cand, double target, MassFn mass_of) { + constexpr size_t kBaseCase = 64; + size_t lo = 0, hi = cand.size(); + while (hi - lo > kBaseCase) { + const size_t mid = lo + (hi - lo) / 2; + std::nth_element(cand.begin() + lo, cand.begin() + mid, cand.begin() + hi, + [](auto & a, auto & b){ return a.first > b.first; }); + double mass = 0.0; + for (size_t i = lo; i < mid; i++) mass += mass_of(cand[i]); + if (mass >= target) { + hi = mid; // cutoff lies within [lo, mid) + } else { + target -= mass; // [lo, mid) fully included; keep searching [mid, hi) + lo = mid; + } + } + // Base case: small enough range left, sort it and walk the exact cumulative + // cutoff (cand[0, lo) is already fully confirmed included from above). + std::sort(cand.begin() + lo, cand.begin() + hi, + [](auto & a, auto & b){ return a.first > b.first; }); + size_t cut = hi; + double cum = 0.0; + for (size_t i = lo; i < hi; i++) { + cum += mass_of(cand[i]); + if (cum >= target) { cut = i + 1; break; } + } + return cut; +} + +// Draws a token from `cand`, whose .first fields are proportional +// probabilities (need not already sum to 1 or be sorted). `r_uniform` is a +// pre-drawn uniform in [0,1) supplied by the caller (drawn once per +// sample_logits call) so every path — GPU, GPU-assisted top_p, or CPU — +// consumes the same single RNG value. +int draw_from_weights(const std::vector> & cand, double r_uniform) { + double Z = 0.0; + for (auto & c : cand) Z += c.first; + const double r = r_uniform * Z; + double acc = 0.0; + for (auto & c : cand) { + acc += c.first; + if (r <= acc) return c.second; + } + return cand.back().second; +} + +#ifdef DFLASH27B_HAVE_GPU_SAMPLER +// Given probabilities the GPU already computed (penalties + softmax(temp) +// applied, summing to ~1) for a pure top_p (no top_k) config, find the +// nucleus and draw. Skips all exp()/Z bookkeeping the raw-logit path needs, +// since the input is already normalized. +// +// Only worth calling for top_p without top_k: top_k's CPU cost is already +// cheap (partial_sort scales with k, not vocab — measured ~270-300us at +// vocab=151936), so a GPU round trip (kernel + D2H copy, ~500-800us) makes it +// slower, not faster. top_p's CPU cost without top_k is dominated by +// nucleus_cutoff's O(vocab) std::nth_element passes regardless of who +// computed the input probabilities, so skipping the CPU-side exp() pass here +// is a net win (measured ~1.4x faster end-to-end at vocab=151936). +int sample_from_gpu_probs(std::vector & probs, double top_p, double r_uniform) { + std::vector> cand(probs.size()); + for (size_t i = 0; i < probs.size(); i++) cand[i] = {probs[i], (int)i}; + + double Z = 0.0; + for (auto & c : cand) Z += c.first; + const double target = top_p * Z; + const size_t cut = nucleus_cutoff(cand, target, [](auto & c){ return (double)c.first; }); + cand.resize(cut); + + return draw_from_weights(cand, r_uniform); +} +#endif + +} // namespace + int sample_logits(const float * logits_in, int vocab, const SamplerCfg & cfg, const std::vector & history, std::mt19937_64 & rng) { + // Draw the single uniform up front, exactly once, and only when we will + // actually sample (temp>0; greedy returns before any draw). It is then + // threaded through every path below — the full-GPU draw, the GPU-assisted + // top_p draw, and the CPU draw all consume this same value — so the RNG + // stream advances identically no matter which path resolves the token + // (including a GPU→CPU fallback), keeping decode reproducible whether the + // GPU sampler is on or off. + double r_uniform = 0.0; + if (cfg.temp > 0.0f) { + std::uniform_real_distribution u(0.0, 1.0); + r_uniform = u(rng); + } + +#ifdef DFLASH27B_HAVE_GPU_SAMPLER + // GPU path (on by default; set DFLASH_GPU_SAMPLE=0 to disable). top_k>0, + // top_p in (0,1) (both unsupported on the GPU, see geometric_sampler_cuda.h), + // and any CUDA error return -1 and fall through to the CPU chain below. + if (gpu_sampler_enabled() && gpu_sampler_supports(cfg)) { + const int g = geometric_sample_logits_cuda(logits_in, vocab, cfg, history, r_uniform, + /*logits_on_device=*/false); + if (g >= 0) return g; + } +#endif + const bool need_top_k = cfg.top_k > 0 && cfg.top_k < vocab; + const bool need_top_p = cfg.top_p > 0.0f && cfg.top_p < 1.0f; + +#ifdef DFLASH27B_HAVE_GPU_SAMPLER + // Reaching here means the GPU either can't fully handle this config + // (top_k/top_p above) or is disabled. For pure top_p (no top_k), the GPU + // can still compute the shared, vocab-wide penalty+softmax prefix and + // hand back normalized probabilities, letting the CPU skip straight to + // nucleus_cutoff — worth it because that search's O(vocab) cost dominates + // regardless of who computed its input. top_k (with or without top_p) is + // deliberately excluded: its CPU cost is already cheap (partial_sort + // scales with k, not vocab), so the GPU round trip would make it slower, + // not faster (measured regression, not just "no win"). + if (cfg.temp > 0.0f && need_top_p && !need_top_k && gpu_sampler_enabled()) { + std::vector gpu_probs(vocab); + if (geometric_compute_probs_cuda(logits_in, vocab, cfg, history, + gpu_probs.data(), /*logits_on_device=*/false)) { + return sample_from_gpu_probs(gpu_probs, cfg.top_p, r_uniform); + } + // else: fall through to the full CPU chain below. + } +#endif + std::vector> cand(vocab); for (int i = 0; i < vocab; i++) cand[i] = {logits_in[i], i}; @@ -48,22 +187,45 @@ int sample_logits(const float * logits_in, } } - if (cfg.top_k > 0 && cfg.top_k < vocab) { + // temp=0 → deterministic argmax (after penalties have been applied above). + // Independent of top_k/top_p (the single highest-logit token is always the + // answer), so this skips sorting/truncation entirely: an O(vocab) max scan + // beats the O(vocab log vocab) sort below. Ties go to the lowest token id + // (max_element returns the first of equal maxima), matching the GPU kernel. + if (cfg.temp <= 0.0f) { + return std::max_element(cand.begin(), cand.end(), + [](auto & a, auto & b){ return a.first < b.first; })->second; + } + + // Only sort/truncate when top_k or top_p actually need it. Softmax and the + // inverse-CDF draw below are order-independent, so the common case (plain + // temperature sampling, no truncation) skips sorting entirely. + if (need_top_k) { std::partial_sort(cand.begin(), cand.begin() + cfg.top_k, cand.end(), [](auto & a, auto & b){ return a.first > b.first; }); cand.resize(cfg.top_k); - } else { - std::sort(cand.begin(), cand.end(), - [](auto & a, auto & b){ return a.first > b.first; }); } - // temp=0 → deterministic argmax (after penalties have been applied above). - if (cfg.temp <= 0.0f) { - return cand.front().second; + const float inv_t = 1.0f / std::max(1e-3f, cfg.temp); + const float maxv_logit = need_top_k + ? cand.front().first + : std::max_element(cand.begin(), cand.end(), + [](auto & a, auto & b){ return a.first < b.first; })->first; + const float maxv = maxv_logit * inv_t; + + if (need_top_p && !need_top_k) { + // Nucleus cutoff over the full (untruncated) vocab. Z is the true + // full-vocab softmax denominator (one O(vocab) exp() pass, needed + // regardless to know the absolute mass threshold). + double Z = 0.0; + for (auto & c : cand) Z += std::exp((double)c.first * inv_t - maxv); + const double target = (double)cfg.top_p * Z; + const size_t cut = nucleus_cutoff(cand, target, [&](auto & c) { + return std::exp((double)c.first * inv_t - maxv); + }); + cand.resize(cut); } - const float inv_t = 1.0f / std::max(1e-3f, cfg.temp); - const float maxv = cand.front().first * inv_t; double Z = 0.0; std::vector probs(cand.size()); for (size_t i = 0; i < cand.size(); i++) { @@ -72,7 +234,9 @@ int sample_logits(const float * logits_in, } for (auto & p : probs) p = (float)(p / Z); - if (cfg.top_p > 0.0f && cfg.top_p < 1.0f) { + // top_k+top_p combined: cut the already top_k-truncated (and thus + // already-sorted) subset further to the top_p nucleus within it. + if (need_top_p && need_top_k) { double cum = 0.0; size_t cut = probs.size(); for (size_t i = 0; i < probs.size(); i++) { @@ -80,19 +244,14 @@ int sample_logits(const float * logits_in, if (cum >= cfg.top_p) { cut = i + 1; break; } } probs.resize(cut); cand.resize(cut); - double zz = 0.0; - for (auto p : probs) zz += p; - for (auto & p : probs) p = (float)(p / zz); } - std::uniform_real_distribution u(0.0, 1.0); - const double r = u(rng); - double acc = 0.0; - for (size_t i = 0; i < probs.size(); i++) { - acc += probs[i]; - if (r <= acc) return cand[i].second; - } - return cand.back().second; + // Draw from the final candidate set. Same CDF walk as draw_from_weights + // above (it renormalizes internally, which is a no-op cost here since + // probs already sums to ~1) — reuse it instead of re-deriving the same + // loop a second time. + for (size_t i = 0; i < cand.size(); i++) cand[i].first = probs[i]; + return draw_from_weights(cand, r_uniform); } bool parse_sampler_token(std::string & line, SamplerCfg & out) { diff --git a/server/src/qwen35/qwen35_backend.cpp b/server/src/qwen35/qwen35_backend.cpp index 4c9a727da..d37c1b758 100644 --- a/server/src/qwen35/qwen35_backend.cpp +++ b/server/src/qwen35/qwen35_backend.cpp @@ -8,6 +8,10 @@ #include "peer_access.h" #include "attn_masks.h" #include "common/sampler.h" +#ifdef DFLASH27B_HAVE_GPU_SAMPLER +#include "common/geometric_sampler_cuda.h" +#include +#endif #include "common/io_utils.h" #include "common/restore_delta.h" #include "qwen3/qwen3_drafter.h" @@ -1587,10 +1591,44 @@ bool Qwen35Backend::do_ar_decode(int committed, int n_gen, }(); int32_t next_tok; if (sampler_.needs_logit_processing()) { - ggml_backend_tensor_get(sg_.logits, logits_buf.data(), 0, - sizeof(float) * vocab); - next_tok = sample_logits(logits_buf.data(), vocab, sampler_, - out_tokens, sampler_rng_); +#ifdef DFLASH27B_HAVE_GPU_SAMPLER + // GPU sample straight from the device logits tensor, skipping the + // full ~vocab-wide D2H copy (the same payoff DFLASH_GPU_ARGMAX gets + // for greedy). Falls back to the CPU chain on -1. + int g_tok = -1; + if (gpu_sampler_enabled() && gpu_sampler_supports(sampler_) && + sg_.logits && sg_.logits->data) { + // Draw the uniform from a copy of the RNG so the real stream is + // not advanced yet; we only commit that single draw if the GPU + // path succeeds (below). On fallback the stream is untouched and + // sample_logits() draws the identical value it would with the + // GPU disabled, so the token stream is reproducible either way. + double r = 0.0; + if (sampler_.temp > 0.0f) { + std::mt19937_64 rng_peek = sampler_rng_; + std::uniform_real_distribution u(0.0, 1.0); + r = u(rng_peek); + } + g_tok = geometric_sample_logits_cuda( + static_cast(sg_.logits->data), vocab, sampler_, + out_tokens, r, /*logits_on_device=*/true); + } + if (g_tok >= 0) { + // Commit the single uniform the GPU draw consumed, so the RNG + // advances by exactly one draw — matching the CPU/GPU-off path. + if (sampler_.temp > 0.0f) { + std::uniform_real_distribution u(0.0, 1.0); + (void)u(sampler_rng_); + } + next_tok = g_tok; + } else +#endif + { + ggml_backend_tensor_get(sg_.logits, logits_buf.data(), 0, + sizeof(float) * vocab); + next_tok = sample_logits(logits_buf.data(), vocab, sampler_, + out_tokens, sampler_rng_); + } } else if (kGpuArgmaxAR && sg_.argmax_tokens) { int32_t tok_i = 0; ggml_backend_tensor_get(sg_.argmax_tokens, &tok_i, 0, sizeof(int32_t)); diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index 261dc9fc5..7f66db35e 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -901,6 +901,15 @@ int main(int argc, char ** argv) { else if (std::strncmp(argv[i], "--max-ctx=", 10) == 0) { g_max_ctx_override = std::atoi(argv[i] + 10); } + // Sampler params for the positional (non-daemon) path, so benchmarks can + // exercise the sample_logits chain (and its GPU port). Same field order + // as the daemon ` samp=` tail: temp,top_p,top_k,rep_pen,seed[,freq,pres]. + else if (std::strncmp(argv[i], "--samp=", 7) == 0) { + std::string fake = std::string(" samp=") + (argv[i] + 7); + if (parse_sampler_token(fake, g_sampler) && g_sampler.seed != 0) { + g_sampler_rng.seed(g_sampler.seed); + } + } // KV cache type flags (mirror llama-cli -ctk / -ctv). // Set the env var before resolve_kv_types() reads it inside create_target_cache. else if (std::strcmp(argv[i], "--cache-type-k") == 0 || std::strcmp(argv[i], "-ctk") == 0) { diff --git a/server/test/test_draft_topk_cuda.cpp b/server/test/test_draft_topk_cuda.cpp index 851bd17b0..ba847ccc8 100644 --- a/server/test/test_draft_topk_cuda.cpp +++ b/server/test/test_draft_topk_cuda.cpp @@ -1,6 +1,6 @@ // Correctness test for geometric_extract_draft_topk_cuda (GPU) vs extract_draft_topk (CPU). // -// The GPU kernel in src/common/draft_topk_cuda.cu is a drop-in replacement for +// The GPU kernel in src/common/geometric_draft_topk_cuda.cu is a drop-in replacement for // the CPU top-K + online-logsumexp path in ddtree.cpp. This test feeds the same // random logits to both and asserts the GPU results match the CPU reference: // - token ids identical (rank by rank, per position) @@ -124,7 +124,7 @@ int main() { return 0; } - // The kernel supports K up to kMaxK (=8 in draft_topk_cuda.cu); larger K is + // The kernel supports K up to kMaxK (=8 in geometric_draft_topk_cuda.cu); larger K is // handled by a documented CPU fallback (returns false), checked separately. const Case cases[] = { // Realistic decode shape: Qwen3.5 vocab, small position batch. diff --git a/server/test/test_gpu_sampler_cuda.cpp b/server/test/test_gpu_sampler_cuda.cpp new file mode 100644 index 000000000..41aef3ae4 --- /dev/null +++ b/server/test/test_gpu_sampler_cuda.cpp @@ -0,0 +1,368 @@ +// Correctness tests for geometric_sample_logits_cuda (src/common/geometric_sampler_cuda.cu) +// vs the CPU sample_logits chain (src/common/sampler.h/.cpp). CUDA only — the GPU +// sampler is compiled into dflash_common solely on the cuda backend (DFLASH_GPU_SAMPLER, +// default ON). All tests self-skip at runtime when no CUDA device is present. +// +// Build: registered in server/CMakeLists.txt under DFLASH27B_TESTS (CUDA only). +// Run: ./test_gpu_sampler_cuda (exit 0 = pass, non-zero = fail) + +#include "common/sampler.h" +#include "common/geometric_sampler_cuda.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace dflash::common; + +// ─── Test framework (ds4 style) ──────────────────────────────────────── + +static int test_failures = 0; +static int test_count = 0; + +#define TEST_ASSERT_MSG(expr, msg) do { \ + test_count++; \ + if (!(expr)) { \ + test_failures++; \ + std::fprintf(stderr, " FAIL: %s:%d: %s — %s\n", __FILE__, __LINE__, #expr, msg); \ + } \ +} while (0) + +#define RUN_TEST(fn) do { \ + std::fprintf(stderr, " %s ...", #fn); \ + int before = test_failures; \ + fn(); \ + if (test_failures == before) std::fprintf(stderr, " ok\n"); \ + else std::fprintf(stderr, "\n"); \ +} while (0) + +// ═══════════════════════════════════════════════════════════════════════ +// GPU sampler (geometric_sampler_cuda.cu) — compared against the CPU chain. +// All tests self-skip when the build has no GPU sampler or no CUDA device. +// ═══════════════════════════════════════════════════════════════════════ + +static bool gpu_sampler_test_available() { + int n = 0; + if (cudaGetDeviceCount(&n) != cudaSuccess) { cudaGetLastError(); return false; } + return n > 0; +} + +// Deterministic logits: a smooth but non-monotonic spread over the vocab so the +// argmax is unambiguous and the softmax has spread-out mass. +static std::vector gpu_test_logits(int vocab, uint64_t seed) { + std::mt19937_64 rng(seed); + std::uniform_real_distribution u(-6.0f, 6.0f); + std::vector v(vocab); + for (auto & x : v) x = u(rng); + return v; +} + +// Analytic softmax over `logits` at temperature `temp`. +static std::vector cpu_softmax(const std::vector & logits, float temp) { + const double inv_t = 1.0 / std::max(1e-3f, temp); + double mx = -1e300; + for (float l : logits) mx = std::max(mx, (double)l * inv_t); + std::vector p(logits.size()); + double z = 0.0; + for (size_t i = 0; i < logits.size(); i++) { p[i] = std::exp((double)logits[i] * inv_t - mx); z += p[i]; } + for (auto & x : p) x /= z; + return p; +} + +// Analytic nucleus (top_p) distribution: the full softmax, restricted to the +// smallest descending-probability prefix whose cumulative mass reaches `top_p`, +// then renormalized within that nucleus (mass 0 elsewhere). This is the exact +// distribution the CPU nucleus_cutoff + draw is supposed to reproduce. +static std::vector cpu_top_p_dist(const std::vector & logits, float temp, double top_p) { + auto p = cpu_softmax(logits, temp); + std::vector idx(p.size()); + for (size_t i = 0; i < idx.size(); i++) idx[i] = (int)i; + std::sort(idx.begin(), idx.end(), [&](int a, int b){ return p[a] > p[b]; }); + std::vector q(p.size(), 0.0); + double cum = 0.0, znuc = 0.0; + std::vector keep; + for (int i : idx) { keep.push_back(i); cum += p[i]; if (cum >= top_p) break; } + for (int i : keep) znuc += p[i]; + for (int i : keep) q[i] = p[i] / znuc; + return q; +} + +// Greedy (temp=0): GPU must pick exactly the same token as the CPU chain, with +// and without penalties active. +static void test_gpu_sampler_greedy_matches_cpu() { + if (!gpu_sampler_test_available()) { std::fprintf(stderr, " (skip: no CUDA) "); return; } + const int vocab = 257; // not a multiple of the block size on purpose + for (uint64_t seed = 1; seed <= 6; seed++) { + auto logits = gpu_test_logits(vocab, seed); + SamplerCfg cfg; // temp=0 → greedy + std::vector history; + std::mt19937_64 rng(seed); + const int cpu_tok = sample_logits(logits.data(), vocab, cfg, history, rng); + const int gpu_tok = geometric_sample_logits_cuda(logits.data(), vocab, cfg, history, + 0.0, /*on_device=*/false); + TEST_ASSERT_MSG(gpu_tok == cpu_tok, "greedy GPU token must equal CPU argmax"); + } +} + +static void test_gpu_sampler_greedy_penalties_match_cpu() { + if (!gpu_sampler_test_available()) { std::fprintf(stderr, " (skip: no CUDA) "); return; } + const int vocab = 200; + auto logits = gpu_test_logits(vocab, 99); + // A history that makes penalties bite on what would otherwise be the argmax. + std::vector history = {3, 3, 7, 12, 7, 50, 50, 50, 100}; + struct { float rep, freq, pres; } cases[] = { + {1.5f, 0.0f, 0.0f}, + {1.0f, 0.8f, 0.0f}, + {1.0f, 0.0f, 1.2f}, + {1.3f, 0.5f, 0.7f}, + }; + for (auto c : cases) { + SamplerCfg cfg; // temp=0 → greedy after penalties + cfg.rep_pen = c.rep; cfg.freq_pen = c.freq; cfg.pres_pen = c.pres; + std::mt19937_64 rng(7); + const int cpu_tok = sample_logits(logits.data(), vocab, cfg, history, rng); + const int gpu_tok = geometric_sample_logits_cuda(logits.data(), vocab, cfg, history, + 0.0, /*on_device=*/false); + TEST_ASSERT_MSG(gpu_tok == cpu_tok, "greedy+penalties GPU token must equal CPU"); + } +} + +// top_k>0 is intentionally CPU-only: the GPU entry must signal fallback (-1). +static void test_gpu_sampler_top_k_falls_back() { + if (!gpu_sampler_test_available()) { std::fprintf(stderr, " (skip: no CUDA) "); return; } + const int vocab = 128; + auto logits = gpu_test_logits(vocab, 5); + SamplerCfg cfg; cfg.temp = 1.0f; cfg.top_k = 10; + std::vector history; + const int gpu_tok = geometric_sample_logits_cuda(logits.data(), vocab, cfg, history, + 0.3, /*on_device=*/false); + TEST_ASSERT_MSG(gpu_tok == -1, "top_k>0 must return -1 (CPU fallback)"); +} + +// L1 distance between an empirical histogram of GPU draws and the analytic +// distribution. Pulls the uniform from the same RNG family the CPU path uses. +static double gpu_empirical_l1(const std::vector & logits, const SamplerCfg & cfg, + const std::vector & analytic, int n_draws) { + std::mt19937_64 rng(1234); + std::uniform_real_distribution u(0.0, 1.0); + std::vector hist(logits.size(), 0); + std::vector history; + int valid = 0; + for (int i = 0; i < n_draws; i++) { + const int tok = geometric_sample_logits_cuda(logits.data(), (int)logits.size(), cfg, + history, u(rng), /*on_device=*/false); + if (tok >= 0 && tok < (int)logits.size()) { hist[tok]++; valid++; } + } + if (valid == 0) return 1e9; + double l1 = 0.0; + for (size_t k = 0; k < logits.size(); k++) + l1 += std::fabs((double)hist[k] / valid - analytic[k]); + return l1; +} + +// Temperature sampling (no truncation): the GPU draw distribution must match the +// analytic softmax — i.e. the same distribution the CPU multinomial samples. +static void test_gpu_sampler_temperature_distribution() { + if (!gpu_sampler_test_available()) { std::fprintf(stderr, " (skip: no CUDA) "); return; } + const int vocab = 24; + auto logits = gpu_test_logits(vocab, 42); + SamplerCfg cfg; cfg.temp = 1.0f; // full vocab + auto analytic = cpu_softmax(logits, cfg.temp); + const double l1 = gpu_empirical_l1(logits, cfg, analytic, 120000); + TEST_ASSERT_MSG(l1 < 0.03, "GPU temp-sample dist must match analytic softmax (L1<0.03)"); +} + +// top_p in (0,1) is intentionally unimplemented on the GPU (removed; see +// geometric_sampler_cuda.h): the GPU entry must signal fallback (-1), same +// contract as top_k>0. +static void test_gpu_sampler_top_p_falls_back() { + if (!gpu_sampler_test_available()) { std::fprintf(stderr, " (skip: no CUDA) "); return; } + const int vocab = 128; + auto logits = gpu_test_logits(vocab, 5); + SamplerCfg cfg; cfg.temp = 1.0f; cfg.top_p = 0.8f; + std::vector history; + const int gpu_tok = geometric_sample_logits_cuda(logits.data(), vocab, cfg, history, + 0.3, /*on_device=*/false); + TEST_ASSERT_MSG(gpu_tok == -1, "top_p in (0,1) must return -1 (CPU fallback)"); +} + +// geometric_compute_probs_cuda is the shared prefix behind the GPU-assisted +// top_p path: it must return the same penalty+softmax(temp) probability vector +// the CPU chain computes. Checks a plain-temp config and one with penalties. +static void test_gpu_compute_probs_matches_softmax() { + if (!gpu_sampler_test_available()) { std::fprintf(stderr, " (skip: no CUDA) "); return; } + const int vocab = 300; // not a multiple of the block size on purpose + auto logits = gpu_test_logits(vocab, 21); + + struct { float temp, rep, freq, pres; std::vector hist; } cases[] = { + {0.8f, 1.0f, 0.0f, 0.0f, {}}, + {1.0f, 1.4f, 0.3f, 0.5f, {3, 3, 7, 12, 7, 50, 50, 100, 100, 100}}, + }; + for (auto & c : cases) { + SamplerCfg cfg; cfg.temp = c.temp; cfg.rep_pen = c.rep; + cfg.freq_pen = c.freq; cfg.pres_pen = c.pres; + std::vector probs(vocab); + const bool ok = geometric_compute_probs_cuda(logits.data(), vocab, cfg, c.hist, + probs.data(), /*on_device=*/false); + TEST_ASSERT_MSG(ok, "geometric_compute_probs_cuda must succeed for temp>0"); + if (!ok) continue; + // Reference: apply the same penalties on the CPU, then softmax(temp). + std::vector penalized = logits; + if (!c.hist.empty()) { + std::unordered_map counts; + for (int id : c.hist) counts[id]++; + for (auto & kv : counts) { + if (kv.first < 0 || kv.first >= vocab) continue; + float & v = penalized[kv.first]; + if (cfg.rep_pen > 1.0f) v = (v > 0.0f) ? v / cfg.rep_pen : v * cfg.rep_pen; + v -= cfg.freq_pen * kv.second + cfg.pres_pen; + } + } + auto analytic = cpu_softmax(penalized, cfg.temp); + double l1 = 0.0, sum = 0.0; + for (int i = 0; i < vocab; i++) { l1 += std::fabs((double)probs[i] - analytic[i]); sum += probs[i]; } + TEST_ASSERT_MSG(l1 < 1e-3, "GPU probs must match CPU penalty+softmax (L1<1e-3)"); + TEST_ASSERT_MSG(std::fabs(sum - 1.0) < 1e-3, "GPU probs must sum to ~1"); + } +} + +// GPU-assisted top_p: sample_logits routes pure top_p (no top_k) through +// geometric_compute_probs_cuda + the CPU nucleus_cutoff/draw. The empirical +// draw distribution must match the analytic nucleus distribution — this is the +// only end-to-end check of nucleus_cutoff and sample_from_gpu_probs. +static void test_gpu_sampler_top_p_distribution() { + if (!gpu_sampler_test_available()) { std::fprintf(stderr, " (skip: no CUDA) "); return; } + const int vocab = 48; + auto logits = gpu_test_logits(vocab, 7); + SamplerCfg cfg; cfg.temp = 1.0f; cfg.top_p = 0.9f; + std::vector history; + auto analytic = cpu_top_p_dist(logits, cfg.temp, cfg.top_p); + + std::mt19937_64 rng(2024); + std::vector hist(vocab, 0); + const int n = 120000; + int valid = 0; + for (int i = 0; i < n; i++) { + const int tok = sample_logits(logits.data(), vocab, cfg, history, rng); + if (tok >= 0 && tok < vocab) { hist[tok]++; valid++; } + } + TEST_ASSERT_MSG(valid > 0, "top_p sampling must produce valid draws"); + double l1 = 0.0; + for (int k = 0; valid > 0 && k < vocab; k++) + l1 += std::fabs((double)hist[k] / valid - analytic[k]); + TEST_ASSERT_MSG(l1 < 0.03, "GPU-assisted top_p dist must match analytic nucleus (L1<0.03)"); +} + +// Direct CPU-vs-GPU agreement on the same draws is not bit-exact (different +// inverse-CDF orderings), but both must agree on the *argmax* of their empirical +// histograms — the most-likely token — under temperature sampling. +static void test_gpu_sampler_modal_token_matches_cpu() { + if (!gpu_sampler_test_available()) { std::fprintf(stderr, " (skip: no CUDA) "); return; } + const int vocab = 32; + auto logits = gpu_test_logits(vocab, 11); + SamplerCfg cfg; cfg.temp = 0.9f; + std::vector history; + + auto modal = [&](bool gpu) { + std::mt19937_64 rng(555); + std::uniform_real_distribution u(0.0, 1.0); + std::vector hist(vocab, 0); + for (int i = 0; i < 60000; i++) { + const int tok = gpu + ? geometric_sample_logits_cuda(logits.data(), vocab, cfg, history, u(rng), false) + : sample_logits(logits.data(), vocab, cfg, history, rng); + if (tok >= 0 && tok < vocab) hist[tok]++; + } + return (int)(std::max_element(hist.begin(), hist.end()) - hist.begin()); + }; + TEST_ASSERT_MSG(modal(true) == modal(false), "GPU and CPU agree on the modal token"); +} + +// Per-call latency microbench (gated by env DFLASH_SAMPLER_BENCH=1). Isolates +// the three regimes that explain the end-to-end numbers: the CPU chain, the GPU +// path fed host logits (pays a full-vocab H2D every call), and the GPU path fed +// a device pointer (the integrated path that skips the copy). +static void gpu_sampler_microbench() { + if (!gpu_sampler_test_available()) { std::fprintf(stderr, "[microbench] no CUDA device\n"); return; } + const int vocab = 151936; // Qwen3 vocab + const int iters = 1000; + auto logits = gpu_test_logits(vocab, 1); + std::vector history; + + auto now = []{ return std::chrono::steady_clock::now(); }; + auto us = [](auto a, auto b){ return std::chrono::duration(b - a).count(); }; + + // Persistent device copy for the device-pointer regime (no per-call H2D). + float * d_logits = nullptr; + bool have_dev = cudaMalloc(&d_logits, (size_t)vocab * sizeof(float)) == cudaSuccess && + cudaMemcpy(d_logits, logits.data(), (size_t)vocab * sizeof(float), + cudaMemcpyHostToDevice) == cudaSuccess; + + volatile long sink = 0; + auto bench = [&](const char * label, const SamplerCfg & cfg) { + std::mt19937_64 rng(1); + std::uniform_real_distribution u(0.0, 1.0); + for (int i = 0; i < 30; i++) sink += sample_logits(logits.data(), vocab, cfg, history, rng); + for (int i = 0; i < 30; i++) sink += geometric_sample_logits_cuda(logits.data(), vocab, cfg, history, u(rng), false); + auto t0 = now(); + for (int i = 0; i < iters; i++) sink += sample_logits(logits.data(), vocab, cfg, history, rng); + auto t1 = now(); + for (int i = 0; i < iters; i++) sink += geometric_sample_logits_cuda(logits.data(), vocab, cfg, history, u(rng), false); + auto t2 = now(); + double dev_us = -1.0; + if (have_dev) { + for (int i = 0; i < 30; i++) sink += geometric_sample_logits_cuda(d_logits, vocab, cfg, history, u(rng), true); + auto t3 = now(); + for (int i = 0; i < iters; i++) sink += geometric_sample_logits_cuda(d_logits, vocab, cfg, history, u(rng), true); + auto t4 = now(); + dev_us = us(t3, t4) / iters; + } + const double cpu_us = us(t0, t1) / iters; + const double gpuh_us = us(t1, t2) / iters; + std::fprintf(stderr, + " %-22s CPU %8.2f us | GPU+H2D %8.2f us (%.2fx) | GPU devptr %8.2f us (%.2fx)\n", + label, cpu_us, gpuh_us, cpu_us / gpuh_us, + dev_us, dev_us > 0 ? cpu_us / dev_us : 0.0); + }; + + std::fprintf(stderr, "[microbench] vocab=%d iters=%d (per call; >1.0x = GPU faster)\n", vocab, iters); + { SamplerCfg c; bench("greedy (temp=0)", c); } + { SamplerCfg c; c.temp = 0.8f; bench("temp=0.8 (full vocab)", c); } + // top_p is unimplemented on the GPU (always falls back to CPU), so it's not benched here. + { SamplerCfg c; c.temp = 0.8f; c.rep_pen = 1.2f; + std::vector h; for (int i = 0; i < 200; i++) h.push_back(i * 7 % vocab); + history = h; bench("temp=0.8 rep_pen=1.2", c); history.clear(); } + + if (have_dev) cudaFree(d_logits); + (void)sink; +} + +// ─── main ──────────────────────────────────────────────────────────── + +int main() { + if (const char * b = std::getenv("DFLASH_SAMPLER_BENCH"); b && b[0] == '1') { + gpu_sampler_microbench(); + return 0; + } + + std::fprintf(stderr, "=== test_gpu_sampler_cuda ===\n"); + + RUN_TEST(test_gpu_sampler_greedy_matches_cpu); + RUN_TEST(test_gpu_sampler_greedy_penalties_match_cpu); + RUN_TEST(test_gpu_sampler_top_k_falls_back); + RUN_TEST(test_gpu_sampler_temperature_distribution); + RUN_TEST(test_gpu_sampler_top_p_falls_back); + RUN_TEST(test_gpu_compute_probs_matches_softmax); + RUN_TEST(test_gpu_sampler_top_p_distribution); + RUN_TEST(test_gpu_sampler_modal_token_matches_cpu); + + std::fprintf(stderr, "\n%d tests, %d failures\n", test_count, test_failures); + return (test_failures == 0) ? 0 : 1; +}