266 lines
9.1 KiB
Plaintext
266 lines
9.1 KiB
Plaintext
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/*
|
|
* topk_logits.cu — device-side exact top-k over the LM-head logits
|
|
*
|
|
* Target: sm_37 (K80, Kepler)
|
|
* Compile: nvcc -ptx -arch=sm_37 topk_logits.cu -o topk_logits.ptx
|
|
*
|
|
* Why: on a small-BAR K80 host the host link runs at ~0.4 GB/s, so copying
|
|
* the full 248k-vocab logits (1 MB) D2H costs ~4 ms per decoded token — the
|
|
* single largest non-compute slice of the token budget. Greedy decode only
|
|
* needs the argmax and the serving sampler (top_k=40) only needs the top 40
|
|
* (id, logit) pairs, so we reduce on the die and copy back k pairs instead.
|
|
*
|
|
* Exactness: selection is ITERATIVE ARGMAX — k rounds, each finding the
|
|
* global max of the remaining logits with the comparator
|
|
* (val_a > val_b) || (val_a == val_b && idx_a < idx_b)
|
|
* which reproduces, value-for-value and tie-for-tie, the host's serial
|
|
* for (j...) if (logits[j] > logits[best]) best = j;
|
|
* scan (lowest index wins ties) applied k times with previous winners
|
|
* removed. No float arithmetic is performed on the logits — only compares —
|
|
* so the returned values are the bit-exact LM-head outputs.
|
|
*
|
|
* Round structure (driven by fc_topk_logits_ctx in fc_cuda.c):
|
|
* 1. fc_topk_partial_argmax — grid of blocks, grid-stride scan, one
|
|
* (val, idx) candidate per block.
|
|
* 2. fc_topk_finish — single block reduces the per-block candidates,
|
|
* records the round winner in out_val/out_idx[slot], and (for k > 1)
|
|
* overwrites logits[winner] with -inf so the next round skips it.
|
|
* 3. fc_topk_restore — after the last round, writes the recorded winner
|
|
* values back so the logits buffer is left intact.
|
|
*
|
|
* PTX symbols emitted (for embedding into ../cuda/fc_cuda_ptx.h):
|
|
* FC_PTX_TOPK_LOGITS -> kernels
|
|
* fc_topk_partial_argmax, fc_topk_finish, fc_topk_restore
|
|
*
|
|
* Kepler notes: pure fp32, shared-memory tree reductions (no __shfl).
|
|
*/
|
|
|
|
extern "C" {
|
|
|
|
#define FC_TOPK_NEG_INF __int_as_float(0xff800000)
|
|
|
|
/* Comparator: does candidate (av, ai) beat incumbent (bv, bi)?
|
|
* Strictly greater value wins; equal value -> lower index wins. */
|
|
__device__ __forceinline__ int fc_topk_beats(float av, int ai,
|
|
float bv, int bi)
|
|
{
|
|
return (av > bv) || (av == bv && ai < bi);
|
|
}
|
|
|
|
/*
|
|
* Stage 1: per-block argmax candidates.
|
|
*
|
|
* x: [n] fp32 logits (read-only this stage)
|
|
* pval: [gridDim.x] fp32 — per-block best value (output)
|
|
* pidx: [gridDim.x] int32 — per-block best index (output)
|
|
*
|
|
* Launched with gridDim.x blocks (host caps at FC_TOPK_NPARTS), blockDim.x
|
|
* threads, grid-stride loop. Shared memory: blockDim.x * (4 + 4) bytes.
|
|
*/
|
|
__global__ void fc_topk_partial_argmax(const float *x, int n,
|
|
float *pval, int *pidx)
|
|
{
|
|
extern __shared__ float sh[];
|
|
float *s_val = sh;
|
|
int *s_idx = (int *)(sh + blockDim.x);
|
|
|
|
int tid = threadIdx.x;
|
|
int stride = gridDim.x * blockDim.x;
|
|
|
|
/* Thread-local scan in increasing index order: strict '>' keeps the
|
|
* first-seen (lowest) index on ties, same as the host loop. */
|
|
float best_v = FC_TOPK_NEG_INF;
|
|
int best_i = 0x7fffffff;
|
|
for (int i = blockIdx.x * blockDim.x + tid; i < n; i += stride) {
|
|
float v = x[i];
|
|
if (v > best_v) { best_v = v; best_i = i; }
|
|
}
|
|
|
|
s_val[tid] = best_v;
|
|
s_idx[tid] = best_i;
|
|
__syncthreads();
|
|
|
|
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
|
if (tid < s &&
|
|
fc_topk_beats(s_val[tid + s], s_idx[tid + s],
|
|
s_val[tid], s_idx[tid])) {
|
|
s_val[tid] = s_val[tid + s];
|
|
s_idx[tid] = s_idx[tid + s];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
if (tid == 0) {
|
|
pval[blockIdx.x] = s_val[0];
|
|
pidx[blockIdx.x] = s_idx[0];
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Stage 2: reduce the per-block candidates to the round winner.
|
|
*
|
|
* x: [n] fp32 logits — winner slot overwritten with -inf when clobber
|
|
* pval/pidx: stage-1 outputs, nparts entries
|
|
* out_val: [k] fp32 — winner value written at [slot]
|
|
* out_idx: [k] int32 — winner index written at [slot]
|
|
* clobber: nonzero -> x[winner] = -inf (skip it in later rounds)
|
|
*
|
|
* Launched with 1 block, blockDim.x >= nparts threads.
|
|
* Shared memory: blockDim.x * (4 + 4) bytes.
|
|
*/
|
|
__global__ void fc_topk_finish(float *x, const float *pval, const int *pidx,
|
|
int nparts, float *out_val, int *out_idx,
|
|
int slot, int clobber)
|
|
{
|
|
extern __shared__ float sh[];
|
|
float *s_val = sh;
|
|
int *s_idx = (int *)(sh + blockDim.x);
|
|
|
|
int tid = threadIdx.x;
|
|
if (tid < nparts) {
|
|
s_val[tid] = pval[tid];
|
|
s_idx[tid] = pidx[tid];
|
|
} else {
|
|
s_val[tid] = FC_TOPK_NEG_INF;
|
|
s_idx[tid] = 0x7fffffff;
|
|
}
|
|
__syncthreads();
|
|
|
|
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
|
if (tid < s &&
|
|
fc_topk_beats(s_val[tid + s], s_idx[tid + s],
|
|
s_val[tid], s_idx[tid])) {
|
|
s_val[tid] = s_val[tid + s];
|
|
s_idx[tid] = s_idx[tid + s];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
if (tid == 0) {
|
|
out_val[slot] = s_val[0];
|
|
out_idx[slot] = s_idx[0];
|
|
if (clobber)
|
|
x[s_idx[0]] = FC_TOPK_NEG_INF;
|
|
}
|
|
}
|
|
|
|
/*
|
|
* Stage 3: undo the -inf clobbers so the logits buffer stays intact.
|
|
* Launched with 1 block, k threads.
|
|
*/
|
|
__global__ void fc_topk_restore(float *x, const float *out_val,
|
|
const int *out_idx, int k)
|
|
{
|
|
int tid = threadIdx.x;
|
|
if (tid < k)
|
|
x[out_idx[tid]] = out_val[tid];
|
|
}
|
|
|
|
/* ============================================================
|
|
* Batched (multi-slot) top-k (L8): the SAME iterative-argmax rounds run
|
|
* for K slot rows of a [K][n] logits matrix IN PLACE — no per-slot D2D
|
|
* copy, one sync and one K*k-pair D2H for the whole batch instead of a
|
|
* serialized sync-per-user loop. Per row the scan order, comparator and
|
|
* reduction shapes are identical to the single-row kernels above, so each
|
|
* row's (id, val) pairs are bit-exact to fc_topk_logits run on that row.
|
|
*
|
|
* Layouts: x row s at x + s*n; stage-1 candidates at pval/pidx
|
|
* [s*nparts + part]; winners at out_val/out_idx[s*k + round].
|
|
* ============================================================ */
|
|
|
|
/* Stage 1: per-(block, slot) argmax candidates. Grid (nparts, nslots). */
|
|
__global__ void fc_topk_partial_argmax_slots(const float *x, int n,
|
|
float *pval, int *pidx)
|
|
{
|
|
extern __shared__ float sh[];
|
|
float *s_val = sh;
|
|
int *s_idx = (int *)(sh + blockDim.x);
|
|
|
|
const float *row = x + (size_t)blockIdx.y * n;
|
|
int tid = threadIdx.x;
|
|
int stride = gridDim.x * blockDim.x;
|
|
|
|
float best_v = FC_TOPK_NEG_INF;
|
|
int best_i = 0x7fffffff;
|
|
for (int i = blockIdx.x * blockDim.x + tid; i < n; i += stride) {
|
|
float v = row[i];
|
|
if (v > best_v) { best_v = v; best_i = i; }
|
|
}
|
|
|
|
s_val[tid] = best_v;
|
|
s_idx[tid] = best_i;
|
|
__syncthreads();
|
|
|
|
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
|
if (tid < s &&
|
|
fc_topk_beats(s_val[tid + s], s_idx[tid + s],
|
|
s_val[tid], s_idx[tid])) {
|
|
s_val[tid] = s_val[tid + s];
|
|
s_idx[tid] = s_idx[tid + s];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
if (tid == 0) {
|
|
pval[(size_t)blockIdx.y * gridDim.x + blockIdx.x] = s_val[0];
|
|
pidx[(size_t)blockIdx.y * gridDim.x + blockIdx.x] = s_idx[0];
|
|
}
|
|
}
|
|
|
|
/* Stage 2: per-slot round winner. Grid (1, nslots), blockDim >= nparts. */
|
|
__global__ void fc_topk_finish_slots(float *x, const float *pval,
|
|
const int *pidx, int n, int nparts,
|
|
float *out_val, int *out_idx, int k,
|
|
int slot, int clobber)
|
|
{
|
|
extern __shared__ float sh[];
|
|
float *s_val = sh;
|
|
int *s_idx = (int *)(sh + blockDim.x);
|
|
|
|
const int u = blockIdx.y;
|
|
const float *pv = pval + (size_t)u * nparts;
|
|
const int *pi = pidx + (size_t)u * nparts;
|
|
|
|
int tid = threadIdx.x;
|
|
if (tid < nparts) {
|
|
s_val[tid] = pv[tid];
|
|
s_idx[tid] = pi[tid];
|
|
} else {
|
|
s_val[tid] = FC_TOPK_NEG_INF;
|
|
s_idx[tid] = 0x7fffffff;
|
|
}
|
|
__syncthreads();
|
|
|
|
for (int s = blockDim.x / 2; s > 0; s >>= 1) {
|
|
if (tid < s &&
|
|
fc_topk_beats(s_val[tid + s], s_idx[tid + s],
|
|
s_val[tid], s_idx[tid])) {
|
|
s_val[tid] = s_val[tid + s];
|
|
s_idx[tid] = s_idx[tid + s];
|
|
}
|
|
__syncthreads();
|
|
}
|
|
|
|
if (tid == 0) {
|
|
out_val[(size_t)u * k + slot] = s_val[0];
|
|
out_idx[(size_t)u * k + slot] = s_idx[0];
|
|
if (clobber)
|
|
x[(size_t)u * n + s_idx[0]] = FC_TOPK_NEG_INF;
|
|
}
|
|
}
|
|
|
|
/* Stage 3: per-slot clobber restore. Grid (nslots), k threads. */
|
|
__global__ void fc_topk_restore_slots(float *x, const float *out_val,
|
|
const int *out_idx, int n, int k)
|
|
{
|
|
const int u = blockIdx.x;
|
|
int tid = threadIdx.x;
|
|
if (tid < k)
|
|
x[(size_t)u * n + out_idx[(size_t)u * k + tid]] =
|
|
out_val[(size_t)u * k + tid];
|
|
}
|
|
|
|
} /* extern "C" */
|