diff --git a/server/CMakeLists.txt b/server/CMakeLists.txt index 5c258bdc1..ab325bed8 100644 --- a/server/CMakeLists.txt +++ b/server/CMakeLists.txt @@ -404,7 +404,11 @@ if(DFLASH27B_GPU_BACKEND STREQUAL "hip") elseif(DFLASH27B_GPU_BACKEND STREQUAL "cuda") target_sources(dflash_common PRIVATE src/flashprefill_select.cpp - src/flashprefill.cpp) + src/flashprefill.cpp + src/common/draft_topk_cuda.cu) + # 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) # 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. @@ -591,6 +595,14 @@ if(DFLASH27B_TESTS) target_include_directories(test_flashprefill_kernels PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) target_link_libraries(test_flashprefill_kernels PRIVATE dflash_common CUDA::cudart) endif() + # GPU draft top-K kernel vs CPU reference (extract_draft_topk). CUDA only: + # draft_topk_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_draft_topk_cuda.cpp") + add_executable(test_draft_topk_cuda test/test_draft_topk_cuda.cpp) + target_include_directories(test_draft_topk_cuda PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/src) + target_link_libraries(test_draft_topk_cuda PRIVATE dflash_common CUDA::cudart) + add_test(NAME draft_topk_cuda COMMAND test_draft_topk_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/scripts/bench_llm.py b/server/scripts/bench_llm.py index 4722a7f3f..b8b31c91d 100644 --- a/server/scripts/bench_llm.py +++ b/server/scripts/bench_llm.py @@ -35,7 +35,7 @@ DRAFT = None TEST_DFLASH = os.environ.get("DFLASH_BIN", str(ROOT / "build" / f"test_dflash{BIN_SUFFIX}")) TEST_GENERATE = os.environ.get("DFLASH_BIN_AR", str(ROOT / "build" / f"test_generate{BIN_SUFFIX}")) -TOKENIZER = os.environ.get("DFLASH_TOKENIZER", "Qwen/Qwen3.5-27B") +TOKENIZER = os.environ.get("DFLASH_TOKENIZER", "Qwen/Qwen3.6-27B") TMPDIR = Path(tempfile.gettempdir()) / "dflash_bench" TMPDIR.mkdir(parents=True, exist_ok=True) @@ -53,8 +53,8 @@ def _gsm_gold(x): BENCHES = [ - ("HumanEval", "openai_humaneval", None, "test", lambda x: x["prompt"], None, N_GEN), - ("GSM8K", "gsm8k", "main", "test", lambda x: f"Question: {x['question']}\nAnswer: ", _gsm_gold, 1024), + ("HumanEval", "openai/openai_humaneval", None, "test", lambda x: x["prompt"], None, N_GEN), + ("GSM8K", "openai/gsm8k", "main", "test", lambda x: f"Question: {x['question']}\nAnswer: ", _gsm_gold, 1024), ("Math500", "HuggingFaceH4/MATH-500", None, "test", lambda x: f"Problem: {x['problem']}\nSolution: Put your final answer in \\boxed{{}}.\n", lambda x: x["answer"], 2048), ] diff --git a/server/src/common/draft_topk_cuda.cu b/server/src/common/draft_topk_cuda.cu new file mode 100644 index 000000000..1ca62b568 --- /dev/null +++ b/server/src/common/draft_topk_cuda.cu @@ -0,0 +1,411 @@ +// GPU top-K + log-prob extraction for DDTree draft distributions. +// See draft_topk_cuda.h for the contract. Mirrors extract_draft_topk (ddtree.cpp). + +#include "draft_topk_cuda.h" + +#include +#include +#include +#include +#include + +namespace dflash::common { + +namespace { + +constexpr int kMaxK = 8; // ddtree_K is 8 in practice; K>kMaxK → CPU fallback +constexpr int kBlock = 256; // threads per block (power of two for the reduction) +constexpr int kMaxSplit = 128; // max vocab splits per position (combine-block cap) + +// The workload is tiny in rows (n_positions ~ 15) but huge in vocab (~152k), +// so it is purely DRAM-bandwidth bound: ~9 MB to read per call. We +// split each position's vocab scan across `split` blocks (a 2D grid of +// n_positions × split) so the whole device stays busy, then a cheap second +// kernel merges the `split` partials per position into the final top-K. + +// Merge two sorted-descending top-K lists (av/ai, bv/bi) into dst (dv/di), +// keeping the K largest; ties resolve toward the lower id. dst may alias av/ai +// safely only via the temporary below, so callers pass a distinct dst when the +// inputs come from shared memory. +template +__device__ __forceinline__ void merge_topk(const float * av, const int * ai, + const float * bv, const int * bi, + float * dv, int * di) { + int ia = 0, ib = 0; +#pragma unroll + for (int k = 0; k < K; k++) { + bool takeA; + if (ib >= K) takeA = true; + else if (ia >= K) takeA = false; + else if (av[ia] > bv[ib]) takeA = true; + else if (av[ia] < bv[ib]) takeA = false; + else takeA = (ai[ia] <= bi[ib]); // tie → lower id + if (takeA) { dv[k] = av[ia]; di[k] = ai[ia]; ia++; } + else { dv[k] = bv[ib]; di[k] = bi[ib]; ib++; } + } +} + +// Fold one logit (raw value `LRAW`, vocab id `ID`) into a thread's running +// online-logsumexp (lmax,lsum) and its register-resident sorted-descending +// top-K (topv,topi). K is a compile-time constant so the unrolled bubble uses +// only fixed indices — the arrays stay in registers instead of spilling to +// local memory the way a data-dependent insertion index would. The new value +// enters at slot K-1 and bubbles up only past strictly-smaller entries, so on +// ties the earlier (lower-id, since ids ascend within a thread) entry wins — +// matching the CPU heap's first-wins behaviour. +#define DFLASH_TOPK_CONSUME(LRAW, ID) \ + do { \ + const float _l = (LRAW) * inv_t; \ + if (_l > lmax) { lsum = lsum * __expf(lmax - _l) + 1.0f; lmax = _l; } \ + else { lsum += __expf(_l - lmax); } \ + if (_l > topv[K - 1]) { \ + topv[K - 1] = _l; topi[K - 1] = (ID); \ + _Pragma("unroll") \ + for (int _p = K - 1; _p > 0; --_p) { \ + if (topv[_p] > topv[_p - 1]) { \ + const float _tv = topv[_p]; \ + topv[_p] = topv[_p - 1]; topv[_p - 1] = _tv; \ + const int _ti = topi[_p]; \ + topi[_p] = topi[_p - 1]; topi[_p - 1] = _ti; \ + } \ + } \ + } \ + } while (0) + +// ---- pass 1: per-(position, split) partial logsumexp + local top-K --------- +// Grid: (n_positions, split). Block s of row `row` scans its contiguous vocab +// chunk and writes one partial (max, sum, sorted top-K) to part_* at index +// row*split + s. Contiguous chunks keep the strided per-warp reads coalesced. +// VEC selects 16-byte float4 loads (one coalesced transaction per 4 logits) — +// only used when every row's base pointer is 16-byte aligned (vocab % 4 == 0 +// and an aligned tensor); otherwise the scalar path is used. +template +__global__ void draft_topk_partial(const float * __restrict__ logits, + int vocab, float inv_t, int split, + float * __restrict__ part_max, + float * __restrict__ part_sum, + float * __restrict__ part_v, + int32_t * __restrict__ part_i) { + const int row = blockIdx.x; + const int s = blockIdx.y; + const int tid = threadIdx.x; + const float * __restrict__ li = logits + (size_t)row * vocab; + + float lmax = -FLT_MAX; + float lsum = 0.0f; + float topv[K]; + int topi[K]; +#pragma unroll + for (int k = 0; k < K; k++) { topv[k] = -FLT_MAX; topi[k] = -1; } + + if (VEC) { + // Partition the float4s of the row into `split` contiguous chunks. + const int vocab4 = vocab >> 2; // # of full float4s + const int chunk4 = (vocab4 + split - 1) / split; + const int b4 = s * chunk4; + const int e4 = min(b4 + chunk4, vocab4); + const float4 * __restrict__ li4 = reinterpret_cast(li); + for (int j4 = b4 + tid; j4 < e4; j4 += kBlock) { + const float4 f = li4[j4]; + const int base = j4 << 2; + DFLASH_TOPK_CONSUME(f.x, base + 0); + DFLASH_TOPK_CONSUME(f.y, base + 1); + DFLASH_TOPK_CONSUME(f.z, base + 2); + DFLASH_TOPK_CONSUME(f.w, base + 3); + } + // Tail elements past the last full float4 (only when vocab % 4 != 0); + // the last split owns them so no id is scanned twice. + if (s == split - 1) { + for (int j = (vocab4 << 2) + tid; j < vocab; j += kBlock) + DFLASH_TOPK_CONSUME(li[j], j); + } + } else { + const int chunk = (vocab + split - 1) / split; + const int begin = s * chunk; + const int end = min(begin + chunk, vocab); + for (int j = begin + tid; j < end; j += kBlock) + DFLASH_TOPK_CONSUME(li[j], j); + } + + // ---- block reduction over kBlock threads ------------------------------ + __shared__ float s_max[kBlock]; + __shared__ float s_sum[kBlock]; + __shared__ float s_topv[kBlock * K]; + __shared__ int32_t s_topi[kBlock * K]; + + s_max[tid] = lmax; + s_sum[tid] = lsum; +#pragma unroll + for (int k = 0; k < K; k++) { + s_topv[tid * K + k] = topv[k]; + s_topi[tid * K + k] = topi[k]; + } + __syncthreads(); + + for (int stride = kBlock / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + const float am = s_max[tid], as = s_sum[tid]; + const float bm = s_max[tid + stride], bs = s_sum[tid + stride]; + const float m = fmaxf(am, bm); + s_sum[tid] = as * __expf(am - m) + bs * __expf(bm - m); + s_max[tid] = m; + + float dv[K]; int di[K]; + merge_topk(&s_topv[tid * K], &s_topi[tid * K], + &s_topv[(tid + stride) * K], &s_topi[(tid + stride) * K], + dv, di); +#pragma unroll + for (int k = 0; k < K; k++) { + s_topv[tid * K + k] = dv[k]; + s_topi[tid * K + k] = di[k]; + } + } + __syncthreads(); + } + + if (tid == 0) { + const int idx = row * split + s; + part_max[idx] = s_max[0]; + part_sum[idx] = s_sum[0]; +#pragma unroll + for (int k = 0; k < K; k++) { + part_v[(size_t)idx * K + k] = s_topv[k]; + part_i[(size_t)idx * K + k] = s_topi[k]; + } + } +} + +#undef DFLASH_TOPK_CONSUME + +// ---- pass 2: merge the `split` partials per position into the final top-K -- +// Grid: n_positions blocks of blockDim = pow2_ceil(split) threads. Thread t +// loads partial t (or an identity when t >= split), then a shared-memory tree +// reduction merges all partials. log_prob[k] = top_v[k] - log_z. +template +__global__ void draft_topk_combine(const float * __restrict__ part_max, + const float * __restrict__ part_sum, + const float * __restrict__ part_v, + const int32_t * __restrict__ part_i, + int split, + float * __restrict__ out_lp, + int32_t * __restrict__ out_ids) { + const int row = blockIdx.x; + const int tid = threadIdx.x; // blockDim is pow2_ceil(split) + + __shared__ float s_max[kMaxSplit]; + __shared__ float s_sum[kMaxSplit]; + __shared__ float s_topv[kMaxSplit * K]; + __shared__ int32_t s_topi[kMaxSplit * K]; + + if (tid < split) { + const int idx = row * split + tid; + s_max[tid] = part_max[idx]; + s_sum[tid] = part_sum[idx]; +#pragma unroll + for (int k = 0; k < K; k++) { + s_topv[tid * K + k] = part_v[(size_t)idx * K + k]; + s_topi[tid * K + k] = part_i[(size_t)idx * K + k]; + } + } else { + s_max[tid] = -FLT_MAX; + s_sum[tid] = 0.0f; +#pragma unroll + for (int k = 0; k < K; k++) { + s_topv[tid * K + k] = -FLT_MAX; + s_topi[tid * K + k] = -1; + } + } + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + const float am = s_max[tid], as = s_sum[tid]; + const float bm = s_max[tid + stride], bs = s_sum[tid + stride]; + const float m = fmaxf(am, bm); + s_sum[tid] = as * __expf(am - m) + bs * __expf(bm - m); + s_max[tid] = m; + + float dv[K]; int di[K]; + merge_topk(&s_topv[tid * K], &s_topi[tid * K], + &s_topv[(tid + stride) * K], &s_topi[(tid + stride) * K], + dv, di); +#pragma unroll + for (int k = 0; k < K; k++) { + s_topv[tid * K + k] = dv[k]; + s_topi[tid * K + k] = di[k]; + } + } + __syncthreads(); + } + + if (tid == 0) { + const float log_z = s_max[0] + logf(s_sum[0]); +#pragma unroll + for (int k = 0; k < K; k++) { + out_lp[(size_t)row * K + k] = s_topv[k] - log_z; + out_ids[(size_t)row * K + k] = s_topi[k]; + } + } +} + +// Per-device scratch for the [n_positions × K] outputs plus the +// [n_positions × split × K] pass-1 partials, grown as needed. The decode loop +// is single-threaded, so a plain static cache is safe and avoids a +// cudaMalloc/cudaFree on every step. +struct Scratch { + int device = -1; + size_t cap = 0; // output elements (n_positions*K) + size_t part_cap = 0; // partial-list elements (n_positions*split*kMaxK) + float * d_lp = nullptr; + int32_t * d_ids = nullptr; + float * d_pmax = nullptr; // [n_positions*split] + float * d_psum = nullptr; // [n_positions*split] + float * d_pv = nullptr; // [n_positions*split*kMaxK] + int32_t * d_pi = nullptr; // [n_positions*split*kMaxK] +}; +Scratch g_scratch; + +void free_scratch() { + if (g_scratch.d_lp) cudaFree(g_scratch.d_lp); + if (g_scratch.d_ids) cudaFree(g_scratch.d_ids); + if (g_scratch.d_pmax) cudaFree(g_scratch.d_pmax); + if (g_scratch.d_psum) cudaFree(g_scratch.d_psum); + if (g_scratch.d_pv) cudaFree(g_scratch.d_pv); + if (g_scratch.d_pi) cudaFree(g_scratch.d_pi); + g_scratch = Scratch{}; +} + +// Allocate output + partial buffers. n = n_positions*K outputs; +// n_parts = n_positions*split partials (each carrying K entries). +bool ensure_scratch(int device, size_t n, size_t n_parts) { + const size_t n_part_lists = n_parts * (size_t)kMaxK; // upper bound on K + if (g_scratch.device == device && g_scratch.cap >= n && + g_scratch.part_cap >= n_part_lists) + return true; + free_scratch(); + if (cudaMalloc(&g_scratch.d_lp, n * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_ids, n * sizeof(int32_t)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_pmax, n_parts * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_psum, n_parts * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_pv, n_part_lists * sizeof(float)) != cudaSuccess) goto fail; + if (cudaMalloc(&g_scratch.d_pi, n_part_lists * sizeof(int32_t)) != cudaSuccess) goto fail; + g_scratch.device = device; + g_scratch.cap = n; + g_scratch.part_cap = n_part_lists; + return true; +fail: + free_scratch(); + return false; +} + +inline int pow2_ceil(int x) { + int p = 1; + while (p < x) p <<= 1; + return p; +} + +// Choose how many blocks to split each position's vocab scan across. The work +// is bandwidth bound, so we want enough total blocks (n_positions × split) to +// saturate the device while keeping each chunk large enough that the strided +// reads stay coalesced and per-block overhead stays amortized. +int pick_split(int vocab, int n_positions) { + if (const char * v = std::getenv("DFLASH_TOPK_SPLIT")) { + int s = std::atoi(v); + if (s >= 1 && s <= kMaxSplit) return s; + } + // Aim for ~240 total blocks (≈3 waves on a ~80-SM device, measured sweet + // spot for vocab~150k), but cap the chunk floor at ~2k elements so a block + // never scans too little to be worth launching. + int by_blocks = (240 + n_positions - 1) / n_positions; + int by_chunk = vocab / 2048; + int split = by_blocks < by_chunk ? by_blocks : by_chunk; + if (split < 1) split = 1; + if (split > kMaxSplit) split = kMaxSplit; + return split; +} + +} // namespace + +bool extract_draft_topk_cuda(const void * d_logits, + int n_positions, int vocab, int K, + float * out_log_probs, + int32_t * out_token_ids, + float temperature) { + if (!d_logits || n_positions <= 0 || vocab <= 0 || K <= 0 || K > kMaxK) return false; + + cudaPointerAttributes attr{}; + if (cudaPointerGetAttributes(&attr, d_logits) != cudaSuccess) { + cudaGetLastError(); // clear the error so we don't poison the next CUDA call + return false; + } + if (attr.type != cudaMemoryTypeDevice) return false; + + int prev = 0; + cudaGetDevice(&prev); + const int dev = attr.device; + if (dev != prev) cudaSetDevice(dev); + + static const bool kProfile = std::getenv("DFLASH_TOPK_PROFILE") != nullptr; + bool ok = false; + const int split = pick_split(vocab, n_positions); + const size_t n = (size_t)n_positions * K; + const size_t n_parts = (size_t)n_positions * split; + if (ensure_scratch(dev, n, n_parts)) { + const float inv_t = 1.0f / fmaxf(1e-3f, temperature); + cudaEvent_t e_k0, e_k1, e_c1; + if (kProfile) { cudaEventCreate(&e_k0); cudaEventCreate(&e_k1); cudaEventCreate(&e_c1); cudaEventRecord(e_k0); } + + const dim3 grid1(n_positions, split); + const int comb_block = pow2_ceil(split); + const float * lp_in = static_cast(d_logits); + // float4 loads are safe only when every row base is 16-byte aligned: + // the tensor base aligned and a vocab stride that is a multiple of 4. + const bool use_vec = (vocab % 4 == 0) && + (reinterpret_cast(lp_in) % 16 == 0); + // K (and the vectorization flag) are compile-time template parameters + // so the per-thread/per-partial top-K stays register-resident; dispatch + // the runtime K to its instantiation. K>kMaxK is already rejected above. +#define DFLASH_TOPK_LAUNCH(KV, VEC) \ + draft_topk_partial<<>>( \ + lp_in, vocab, inv_t, split, \ + g_scratch.d_pmax, g_scratch.d_psum, g_scratch.d_pv, g_scratch.d_pi); \ + draft_topk_combine<<>>( \ + g_scratch.d_pmax, g_scratch.d_psum, g_scratch.d_pv, g_scratch.d_pi, \ + split, g_scratch.d_lp, g_scratch.d_ids); +#define DFLASH_TOPK_CASE(KV) \ + case KV: \ + if (use_vec) { DFLASH_TOPK_LAUNCH(KV, true) } \ + else { DFLASH_TOPK_LAUNCH(KV, false) } \ + break; + switch (K) { + DFLASH_TOPK_CASE(1) DFLASH_TOPK_CASE(2) DFLASH_TOPK_CASE(3) DFLASH_TOPK_CASE(4) + DFLASH_TOPK_CASE(5) DFLASH_TOPK_CASE(6) DFLASH_TOPK_CASE(7) DFLASH_TOPK_CASE(8) + default: break; + } +#undef DFLASH_TOPK_CASE +#undef DFLASH_TOPK_LAUNCH + + if (kProfile) cudaEventRecord(e_k1); + if (cudaGetLastError() == cudaSuccess && cudaDeviceSynchronize() == cudaSuccess) { + const cudaError_t e1 = cudaMemcpy(out_log_probs, g_scratch.d_lp, + n * sizeof(float), cudaMemcpyDeviceToHost); + const cudaError_t e2 = cudaMemcpy(out_token_ids, g_scratch.d_ids, + n * sizeof(int32_t), cudaMemcpyDeviceToHost); + ok = (e1 == cudaSuccess && e2 == cudaSuccess); + } + if (kProfile) { + cudaEventRecord(e_c1); cudaEventSynchronize(e_c1); + float k_ms = 0, c_ms = 0; + cudaEventElapsedTime(&k_ms, e_k0, e_k1); + cudaEventElapsedTime(&c_ms, e_k1, e_c1); + std::fprintf(stderr, "[topk] kernels=%.3f ms sync+copy=%.3f ms (n_pos=%d vocab=%d split=%d)\n", + k_ms, c_ms, n_positions, vocab, split); + cudaEventDestroy(e_k0); cudaEventDestroy(e_k1); cudaEventDestroy(e_c1); + } + } + if (!ok) cudaGetLastError(); + if (dev != prev) cudaSetDevice(prev); + return ok; +} + +} // namespace dflash::common diff --git a/server/src/common/draft_topk_cuda.h b/server/src/common/draft_topk_cuda.h new file mode 100644 index 000000000..36ec0c8fe --- /dev/null +++ b/server/src/common/draft_topk_cuda.h @@ -0,0 +1,34 @@ +// GPU top-K + log-prob extraction for DDTree draft distributions. +// +// Drop-in device-side replacement for extract_draft_topk (ddtree.cpp): instead +// of D2H-ing the full [vocab × n_positions] draft logits and running an OpenMP +// heap top-K + online-logsumexp on the CPU, this runs the whole thing on the +// GPU directly on the logits tensor's device buffer and copies back only the +// [n_positions × K] results. +// +// Semantics match extract_draft_topk exactly: +// scaled = logit * (1 / max(1e-3, temperature)) +// log_z = logsumexp_j(scaled_j) (per position) +// out_log_probs = top-K scaled logits (desc), minus log_z +// out_token_ids = matching vocab ids; ties broken toward the lower id +// +// Returns false (caller must fall back to the CPU path) when CUDA is +// unavailable, the pointer is not device memory, K is out of range, or any +// CUDA call fails. Only compiled into CUDA builds; see CMakeLists.txt. + +#pragma once + +#include + +namespace dflash::common { + +// d_logits: device pointer to row-major [n_positions][vocab] f32 logits (the +// position stride is `vocab` floats — pass an offset pointer to skip +// leading positions). out_* are HOST buffers of size n_positions*K. +bool extract_draft_topk_cuda(const void * d_logits, + int n_positions, int vocab, int K, + float * out_log_probs, + int32_t * out_token_ids, + float temperature); + +} // namespace dflash::common diff --git a/server/src/qwen35/qwen35_dflash_target.cpp b/server/src/qwen35/qwen35_dflash_target.cpp index 88be8c979..65e181928 100644 --- a/server/src/qwen35/qwen35_dflash_target.cpp +++ b/server/src/qwen35/qwen35_dflash_target.cpp @@ -5,6 +5,7 @@ #include "graph_builders.h" #include "step_graph.h" #include "attn_masks.h" +#include "common/draft_topk_cuda.h" // gpu_runtime_compat.h maps the raw cudaStream_t / cudaMemcpy* symbols used // below (rollback_to / rollback_to_tree) onto their HIP equivalents. Without // it the file only compiles on CUDA via a transitive ; HIP @@ -278,13 +279,38 @@ bool Qwen35DFlashTarget::verify_tree( return false; } - // Posterior = per-node target argmax. Tree-shaped graphs can return -1 from - // the GPU argmax shortcut, so compute argmax on CPU from full logits. + // Posterior = per-node target argmax. + // + // The verify graph already computes a batched per-node GPU argmax + // (sg_.argmax_tokens, built by build_target_step_tree). When the caller does + // not need the full logits (greedy decode, logits_out == nullptr) we read + // those N_actual int32s directly and skip the vocab×N_actual D2H + CPU + // argmax entirely — eliminates the verify-logits transfer hotspot. + // + // Historically the GPU argmax shortcut has returned -1 for tree-shaped + // verify graphs on some builds; guard against that by validating every row + // and falling back to the CPU path for the step if any index is bad. + // Escape hatch: DFLASH_GPU_VERIFY_ARGMAX=0 forces the legacy CPU path. + static const bool kGpuVerifyArgmax = []() { + const char * v = std::getenv("DFLASH_GPU_VERIFY_ARGMAX"); + return v == nullptr || v[0] != '0'; + }(); const int vocab = (int)sg_.logits->ne[0]; + posterior_out.resize(N_actual); + + if (kGpuVerifyArgmax && !logits_out && sg_.argmax_tokens) { + ggml_backend_tensor_get(sg_.argmax_tokens, posterior_out.data(), 0, + sizeof(int32_t) * N_actual); + bool ok = true; + for (int i = 0; i < N_actual; i++) { + if (posterior_out[i] < 0 || posterior_out[i] >= vocab) { ok = false; break; } + } + if (ok) return true; // fast path; otherwise fall through to CPU argmax + } + std::vector logits((size_t)vocab * N_actual); ggml_backend_tensor_get(sg_.logits, logits.data(), 0, sizeof(float) * (size_t)vocab * N_actual); - posterior_out.resize(N_actual); for (int i = 0; i < N_actual; i++) { const float * row = logits.data() + (size_t)i * vocab; int am = 0; float best = row[0]; @@ -631,12 +657,28 @@ bool Qwen35DFlashTarget::project_hidden_to_topk( if (st != GGML_STATUS_SUCCESS) return false; const int vocab = (int)proj_sg_.logits->ne[0]; + top_log_probs.assign((size_t)n_tokens * K, 0.0f); + top_token_ids.assign((size_t)n_tokens * K, 0); + +#ifdef DFLASH27B_HAVE_DRAFT_TOPK_CUDA + // GPU path: top-K + logsumexp directly on the logits device buffer, skipping + // the vocab×n_tokens D2H and the CPU heap extract. Falls back to the CPU path + // on any failure. Escape hatch: DFLASH_GPU_DRAFT_TOPK=0. + static const bool kGpuDraftTopk = []() { + const char * v = std::getenv("DFLASH_GPU_DRAFT_TOPK"); + return v == nullptr || v[0] != '0'; + }(); + if (kGpuDraftTopk && + extract_draft_topk_cuda(proj_sg_.logits->data, n_tokens, vocab, K, + top_log_probs.data(), top_token_ids.data(), + temperature)) { + return true; + } +#endif + std::vector logits((size_t)vocab * n_tokens); ggml_backend_tensor_get(proj_sg_.logits, logits.data(), 0, sizeof(float) * (size_t)vocab * n_tokens); - - top_log_probs.assign((size_t)n_tokens * K, 0.0f); - top_token_ids.assign((size_t)n_tokens * K, 0); extract_draft_topk(logits.data(), n_tokens, vocab, K, top_log_probs.data(), top_token_ids.data(), temperature); return true; diff --git a/server/test/test_dflash.cpp b/server/test/test_dflash.cpp index ffb91df41..b524f5b81 100644 --- a/server/test/test_dflash.cpp +++ b/server/test/test_dflash.cpp @@ -270,6 +270,7 @@ using dflash::common::free_qwen35_layer_split_shards; #include "qwen35_layer_split_dflash_target.h" #include "common/dflash_spec_decode.h" #include "common/gguf_mmap.h" +#include "common/draft_topk_cuda.h" using dflash::common::is_eos_tok; // ─── Layer-split daemon — extracted to src/qwen35/layer_split_daemon.{h,cpp} ─ @@ -3271,16 +3272,35 @@ int main(int argc, char ** argv) { } } else { // DDTree K>1: need real log-probs for best-first tree scoring. - // Transfer full logits for positions 1..q_len-1. - if (!draft_hidden_bridge) { - ggml_backend_tensor_get(draft_sg.logits, draft_logits_buf.data(), 0, - sizeof(float) * vocab * q_len); + bool topk_done = false; +#ifdef DFLASH27B_HAVE_DRAFT_TOPK_CUDA + // GPU path: top-K + logsumexp on the draft logits device buffer + // (positions 1..q_len-1), no full-vocab D2H. Escape: DFLASH_GPU_DRAFT_TOPK=0. + static const bool kGpuDraftTopk = [](){ + const char * v = std::getenv("DFLASH_GPU_DRAFT_TOPK"); + return v == nullptr || v[0] != '0'; + }(); + if (kGpuDraftTopk && !draft_hidden_bridge) { + topk_done = dflash::common::extract_draft_topk_cuda( + (const float *)draft_sg.logits->data + (size_t)vocab, + L, vocab, ddtree_K, + ddtree_top_log_probs.data(), + ddtree_top_token_ids.data(), + ddtree_temp); + } +#endif + if (!topk_done) { + // Transfer full logits for positions 1..q_len-1, extract on CPU. + if (!draft_hidden_bridge) { + ggml_backend_tensor_get(draft_sg.logits, draft_logits_buf.data(), 0, + sizeof(float) * vocab * q_len); + } + extract_draft_topk(draft_logits_buf.data() + (size_t)vocab, + L, vocab, ddtree_K, + ddtree_top_log_probs.data(), + ddtree_top_token_ids.data(), + ddtree_temp); } - extract_draft_topk(draft_logits_buf.data() + (size_t)vocab, - L, vocab, ddtree_K, - ddtree_top_log_probs.data(), - ddtree_top_token_ids.data(), - ddtree_temp); } } auto T_draft_logits = sync_us(); @@ -3420,16 +3440,40 @@ int main(int argc, char ** argv) { T_verify_compute = sync_us(); tt_verify_compute += std::chrono::duration(T_verify_compute - T_verify_set).count(); - // DDTree test mode reads full logits and computes posterior on - // CPU. The GPU argmax shortcut has returned -1 for tree-shaped - // verify graphs on some builds, which makes the harness stop - // after the root even though logits are valid. This is test-only; - // server decode paths are unaffected. + // DDTree posterior: per-node argmax over the verify logits. + // default : full vocab×N D2H + CPU argmax (legacy). + // GPU_VERIFY_ARGMAX=1: read the in-graph batched GPU argmax + // (tree_verify_argmax) — N int32s, no bulk D2H. + // GPU_VERIFY_ARGMAX=2: run BOTH and report per-step mismatches + // (validates the historical "-1 / tie" concern). + static const int kGpuVerifyArgmax = [](){ + const char * v = std::getenv("DFLASH_GPU_VERIFY_ARGMAX"); + return v ? std::atoi(v) : 0; + }(); std::vector posterior(N_actual); - ggml_backend_tensor_get(sg.logits, verify_logits_buf.data(), 0, - sizeof(float) * (size_t)vocab * N_actual); - for (int i = 0; i < N_actual; i++) { - posterior[i] = argmax_f32(verify_logits_buf.data() + (size_t)i * vocab, vocab); + bool logits_resident = false; // verify_logits_buf populated this step? + if (kGpuVerifyArgmax == 1 && sg.argmax_tokens) { + ggml_backend_tensor_get(sg.argmax_tokens, posterior.data(), 0, + sizeof(int32_t) * N_actual); + } else { + ggml_backend_tensor_get(sg.logits, verify_logits_buf.data(), 0, + sizeof(float) * (size_t)vocab * N_actual); + logits_resident = true; + for (int i = 0; i < N_actual; i++) { + posterior[i] = argmax_f32(verify_logits_buf.data() + (size_t)i * vocab, vocab); + } + if (kGpuVerifyArgmax == 2 && sg.argmax_tokens) { + std::vector gpu_post(N_actual); + ggml_backend_tensor_get(sg.argmax_tokens, gpu_post.data(), 0, + sizeof(int32_t) * N_actual); + int mism = 0, first = -1; + for (int i = 0; i < N_actual; i++) + if (gpu_post[i] != posterior[i]) { mism++; if (first < 0) first = i; } + if (mism) + std::fprintf(stderr, "[verify_argmax cmp] step=%d N=%d mismatches=%d " + "first@%d gpu=%d cpu=%d\n", n_draft_steps, N_actual, mism, + first, gpu_post[first], posterior[first]); + } } auto T_verify_logits_ddtree = sync_us(); tt_verify_logits += std::chrono::duration( @@ -3440,10 +3484,18 @@ int main(int argc, char ** argv) { int bonus_node_idx = 0; std::vector accepted = follow_verified_tree(tree, posterior.data(), next_token, &bonus_node_idx); if (g_sampler.temp > 0.0f) { + // Sampling needs the bonus node's full logit row. If we took the + // GPU-argmax path we skipped the bulk D2H, so fetch just that row. std::vector bonus_logits(vocab); - std::memcpy(bonus_logits.data(), - verify_logits_buf.data() + (size_t)bonus_node_idx * vocab, - (size_t)vocab * sizeof(float)); + if (logits_resident) { + std::memcpy(bonus_logits.data(), + verify_logits_buf.data() + (size_t)bonus_node_idx * vocab, + (size_t)vocab * sizeof(float)); + } else { + ggml_backend_tensor_get(sg.logits, bonus_logits.data(), + (size_t)bonus_node_idx * vocab * sizeof(float), + (size_t)vocab * sizeof(float)); + } next_token = sample_logits(bonus_logits.data(), vocab, g_sampler, out_all, g_sampler_rng); } const int accept_depth = (int)accepted.size(); // includes root diff --git a/server/test/test_draft_topk_cuda.cpp b/server/test/test_draft_topk_cuda.cpp new file mode 100644 index 000000000..06e76228a --- /dev/null +++ b/server/test/test_draft_topk_cuda.cpp @@ -0,0 +1,177 @@ +// Correctness test for 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 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) +// - log-probs within a small bf16/float tolerance +// +// Exact float ties (where two vocab entries share a logit) are vanishingly +// unlikely with random normal logits, but if one does occur the two paths may +// order the tied ids differently; we treat an id swap as OK when the matching +// log-probs are equal within tolerance. +// +// Build: registered in server/CMakeLists.txt under DFLASH27B_TESTS (CUDA only). +// Run: ./test_draft_topk_cuda (exit 0 = pass, non-zero = fail) + +#include "../src/common/draft_topk_cuda.h" +#include "../src/common/ddtree.h" + +#include + +#include +#include +#include +#include +#include + +using dflash::common::extract_draft_topk; +using dflash::common::extract_draft_topk_cuda; + +namespace { + +// Tolerance on log-prob magnitude. CPU uses double-free float exp/log; the GPU +// kernel does the same in f32 with a different reduction order, so small drift +// is expected — 2e-3 is comfortably above observed error and well below the gap +// between distinct top-K logits. +constexpr float kLogProbTol = 2e-3f; + +struct Case { + int n; + int vocab; + int K; + float temp; +}; + +bool run_case(const Case & c, unsigned seed) { + const size_t n_logits = (size_t)c.n * c.vocab; + const size_t n_out = (size_t)c.n * c.K; + + std::vector h_logits(n_logits); + std::mt19937 rng(seed); + std::normal_distribution dist(0.f, 4.f); + for (auto & x : h_logits) x = dist(rng); + + // CPU reference (the production path). + std::vector cpu_lp(n_out); + std::vector cpu_ids(n_out); + extract_draft_topk(h_logits.data(), c.n, c.vocab, c.K, + cpu_lp.data(), cpu_ids.data(), c.temp); + + // GPU kernel: logits must live in device memory. + float * d_logits = nullptr; + cudaError_t err = cudaMalloc(&d_logits, n_logits * sizeof(float)); + if (err != cudaSuccess) { + printf(" cudaMalloc failed: %s\n", cudaGetErrorString(err)); + return false; + } + cudaMemcpy(d_logits, h_logits.data(), n_logits * sizeof(float), + cudaMemcpyHostToDevice); + + std::vector gpu_lp(n_out); + std::vector gpu_ids(n_out); + bool ok = extract_draft_topk_cuda(d_logits, c.n, c.vocab, c.K, + gpu_lp.data(), gpu_ids.data(), c.temp); + cudaFree(d_logits); + + if (!ok) { + printf(" FAIL: extract_draft_topk_cuda returned false\n"); + return false; + } + + int id_mismatch = 0; + int tie_swap = 0; + float max_lp_err = 0.f; + for (int r = 0; r < c.n; r++) { + for (int k = 0; k < c.K; k++) { + const size_t i = (size_t)r * c.K + k; + const float lp_err = std::fabs(gpu_lp[i] - cpu_lp[i]); + max_lp_err = std::fmax(max_lp_err, lp_err); + + if (gpu_ids[i] != cpu_ids[i]) { + // Accept as a tie reorder only if the log-prob at this rank is + // identical within tolerance (both paths picked equal-logit + // entries, just in a different order). + if (lp_err <= kLogProbTol) { + tie_swap++; + } else { + id_mismatch++; + if (id_mismatch <= 8) { + printf(" id mismatch pos=%d rank=%d gpu=%d(lp=%.5f) " + "cpu=%d(lp=%.5f)\n", + r, k, gpu_ids[i], gpu_lp[i], + cpu_ids[i], cpu_lp[i]); + } + } + } + } + } + + const bool pass = (id_mismatch == 0) && (max_lp_err <= kLogProbTol); + printf(" [%s] n=%d vocab=%d K=%d temp=%.2f id_mismatch=%d tie_swap=%d " + "max_lp_err=%.3e\n", + pass ? "PASS" : "FAIL", c.n, c.vocab, c.K, c.temp, + id_mismatch, tie_swap, max_lp_err); + return pass; +} + +} // namespace + +int main() { + int dev_count = 0; + if (cudaGetDeviceCount(&dev_count) != cudaSuccess || dev_count == 0) { + printf("SKIP: no CUDA device available\n"); + return 0; + } + + // The kernel supports K up to kMaxK (=8 in 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. + {15, 151936, 8, 1.0f}, + {1, 151936, 8, 1.0f}, + {15, 151936, 8, 0.7f}, // temperature scaling + {15, 151936, 8, 2.0f}, + // Small/edge shapes to stress the kernel's split-K / tail handling. + {7, 1024, 8, 1.0f}, + {32, 4096, 8, 1.0f}, + {3, 257, 8, 1.0f}, // vocab barely above K, non-power-of-two + {1, 151936, 1, 1.0f}, // K=1 (argmax + log_z) + {15, 151936, 4, 1.0f}, + }; + + int failures = 0; + int idx = 0; + for (const Case & c : cases) { + if (!run_case(c, /*seed=*/1234u + idx)) failures++; + idx++; + } + + // Fallback contract: K beyond the kernel's supported range must return false + // (not silently produce wrong output) so the caller can use the CPU path. + { + const int n = 4, vocab = 4096, big_K = 64; + std::vector h(n * vocab, 0.f); + float * d = nullptr; + if (cudaMalloc(&d, h.size() * sizeof(float)) == cudaSuccess) { + cudaMemcpy(d, h.data(), h.size() * sizeof(float), cudaMemcpyHostToDevice); + std::vector lp(n * big_K); + std::vector ids(n * big_K); + bool ret = extract_draft_topk_cuda(d, n, vocab, big_K, + lp.data(), ids.data(), 1.0f); + cudaFree(d); + const bool pass = !ret; // expect false + printf(" [%s] fallback contract: K=%d (>kMaxK) returned %s\n", + pass ? "PASS" : "FAIL", big_K, ret ? "true" : "false"); + if (!pass) failures++; + idx++; + } + } + + if (failures) { + printf("\nFAILED: %d/%d cases\n", failures, idx); + return 1; + } + printf("\nALL PASS: %d/%d cases\n", idx, idx); + return 0; +}