CharlesCNorton
logits-processor: sort-free block-parallel sampler (radix-select keep-threshold + Gumbel-max)
13163db | // logits-processor: the generation last-mile as loadable kernels. | |
| // | |
| // Two ops that every serving stack needs on the hot path but that the Kernel Hub | |
| // did not carry: | |
| // | |
| // apply_token_bitmask_inplace - structured/guided decoding. Given a packed | |
| // allow-mask per request (XGrammar convention: bit 1 = token allowed), set | |
| // every disallowed logit to -inf in place, so the following softmax cannot | |
| // select it. This is the kernel behind constrained JSON / tool-calling. | |
| // | |
| // sample - temperature, then the intersection of top-k / top-p (nucleus) / | |
| // min-p, then a multinomial draw, fused into one logits -> token op with a | |
| // counter-based RNG (deterministic in (seed, row, offset)). temperature <= 0 | |
| // is greedy. Because the three filters are all "keep the largest-probability | |
| // tokens above a threshold", their intersection is a single keep-threshold, so | |
| // no ordering of the vocabulary is needed. Greedy is a plain argmax; unfiltered | |
| // sampling is one Gumbel-max pass; the filtered case finds the keep-threshold | |
| // with a radix-select over the softmax values and draws by Gumbel-max over the | |
| // kept set. One block per row, fully parallel, no sort and no per-row serial | |
| // scan. | |
| namespace { | |
| constexpr int TPB = 256; | |
| constexpr int NWARP = TPB / 32; // per-warp privatized histograms | |
| constexpr int RADIX_BITS = 8; | |
| constexpr int RADIX_SIZE = 1 << RADIX_BITS; // 256 magnitude buckets per pass | |
| __device__ __forceinline__ float to_f(float x) { return x; } | |
| __device__ __forceinline__ float to_f(__half x) { return __half2float(x); } | |
| __device__ __forceinline__ float to_f(__nv_bfloat16 x) { return __bfloat162float(x); } | |
| __device__ __forceinline__ void set_ninf(float& x) { x = -INFINITY; } | |
| __device__ __forceinline__ void set_ninf(__half& x) { x = __float2half(-INFINITY); } | |
| __device__ __forceinline__ void set_ninf(__nv_bfloat16& x) { x = __float2bfloat16(-INFINITY); } | |
| __device__ __forceinline__ unsigned triple32(unsigned x) { | |
| x ^= x >> 17; x *= 0xed5ad4bbU; x ^= x >> 11; x *= 0xac4c1b51U; | |
| x ^= x >> 15; x *= 0x31848babU; x ^= x >> 14; return x; | |
| } | |
| // bit == 1 means allowed; disallowed logits become -inf. grid = (xtiles, B). | |
| template <typename T> | |
| __global__ void bitmask_kernel(T* logits, const int32_t* mask, int V, int mask_stride) { | |
| int row = blockIdx.y; | |
| const int32_t* m = mask + (long)row * mask_stride; | |
| T* lg = logits + (long)row * V; | |
| for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < V; i += gridDim.x * blockDim.x) { | |
| if (!((m[i >> 5] >> (i & 31)) & 1)) set_ninf(lg[i]); | |
| } | |
| } | |
| // e in [0, 1] -> a monotone uint key (positive IEEE floats order as their bits). | |
| __device__ __forceinline__ unsigned fkey(float e) { return __float_as_uint(e); } | |
| // Block-wide radix-select over the softmax values e[0..V). Returns the threshold | |
| // tau such that keeping {e_i >= tau} yields, from the largest e down, a cumulative | |
| // count (by_mass=false) or mass (by_mass=true) that first reaches `target`. Four | |
| // 8-bit passes resolve the full 32-bit key, so tau is the exact boundary value. | |
| // Called by every thread of the block; uses per-block shared histograms. | |
| __device__ float radix_threshold(const float* __restrict__ e, int V, bool by_mass, | |
| double target, unsigned* hist_c, double* hist_m, | |
| unsigned* prefix_sh, double* above_sh) { | |
| const int tid = threadIdx.x, nt = blockDim.x, warp = tid >> 5; | |
| unsigned prefix = 0; double above = 0.0; // committed high bits, mass/count above boundary | |
| for (int pass = 0; pass < 32 / RADIX_BITS; ++pass) { | |
| const int shift = 32 - RADIX_BITS * (pass + 1); | |
| const unsigned hi_mask = (shift + RADIX_BITS >= 32) ? 0u : (0xFFFFFFFFu << (shift + RADIX_BITS)); | |
| for (int b = tid; b < NWARP * RADIX_SIZE; b += nt) { hist_c[b] = 0u; hist_m[b] = 0.0; } | |
| __syncthreads(); | |
| // per-warp privatized histogram cuts shared-atomic contention (the first pass, | |
| // which scans all V, dominates; later passes touch only the boundary bucket). | |
| for (int i = tid; i < V; i += nt) { | |
| const unsigned key = fkey(e[i]); | |
| if ((key & hi_mask) == (prefix & hi_mask)) { | |
| const int b = warp * RADIX_SIZE + (int)((key >> shift) & (RADIX_SIZE - 1)); | |
| if (by_mass) atomicAdd(&hist_m[b], (double)e[i]); | |
| else atomicAdd(&hist_c[b], 1u); | |
| } | |
| } | |
| __syncthreads(); | |
| if (tid == 0) { | |
| double acc = above; int boundary = 0; | |
| for (int b = RADIX_SIZE - 1; b >= 0; --b) { | |
| double add = 0.0; | |
| if (by_mass) { for (int w = 0; w < NWARP; ++w) add += hist_m[w * RADIX_SIZE + b]; } | |
| else { unsigned cc = 0; for (int w = 0; w < NWARP; ++w) cc += hist_c[w * RADIX_SIZE + b]; add = (double)cc; } | |
| if (acc + add >= target) { boundary = b; break; } | |
| acc += add; | |
| } | |
| *prefix_sh = (prefix & hi_mask) | ((unsigned)boundary << shift); | |
| *above_sh = acc; // cumulative strictly above the boundary bucket | |
| } | |
| __syncthreads(); | |
| prefix = *prefix_sh; above = *above_sh; | |
| __syncthreads(); | |
| } | |
| return __uint_as_float(prefix); | |
| } | |
| // One block per row: max/argmax -> greedy early-out -> softmax sum -> keep-threshold | |
| // (radix-select) -> Gumbel-max draw over the kept set. | |
| template <typename T> | |
| __global__ void sample_kernel(const T* __restrict__ logits, const float* __restrict__ temp, | |
| const float* __restrict__ top_p, const int* __restrict__ top_k, | |
| const float* __restrict__ min_p, float* __restrict__ work, | |
| unsigned seed, unsigned offset, int64_t* __restrict__ out, int V) { | |
| const int row = blockIdx.x, tid = threadIdx.x, nt = blockDim.x; | |
| const T* lg = logits + (long)row * V; | |
| float* e = work + (long)row * V; | |
| __shared__ float red[TPB]; | |
| __shared__ int redi[TPB]; | |
| __shared__ unsigned hist_c[NWARP * RADIX_SIZE]; | |
| __shared__ double hist_m[NWARP * RADIX_SIZE]; | |
| __shared__ unsigned prefix_sh; __shared__ double above_sh; | |
| __shared__ float sh_scalar; | |
| // 1. max logit + argmax (tie-break: smaller index) | |
| float lmax = -INFINITY; int amax = 0; | |
| for (int i = tid; i < V; i += nt) { float v = to_f(lg[i]); if (v > lmax) { lmax = v; amax = i; } } | |
| red[tid] = lmax; redi[tid] = amax; __syncthreads(); | |
| for (int s = nt >> 1; s > 0; s >>= 1) { | |
| if (tid < s && (red[tid + s] > red[tid] || (red[tid + s] == red[tid] && redi[tid + s] < redi[tid]))) { | |
| red[tid] = red[tid + s]; redi[tid] = redi[tid + s]; | |
| } | |
| __syncthreads(); | |
| } | |
| const float gmax = red[0]; const int argmax = redi[0]; | |
| __syncthreads(); | |
| const float t = temp[row]; | |
| if (t <= 0.f) { if (tid == 0) out[row] = argmax; return; } // greedy | |
| const float scale = 1.f / t; | |
| // 2. softmax sum S; store e_i = exp((logit-gmax)/t) in work (e in (0,1], e_max=1) | |
| float s = 0.f; | |
| for (int i = tid; i < V; i += nt) { float ev = __expf((to_f(lg[i]) - gmax) * scale); e[i] = ev; s += ev; } | |
| red[tid] = s; __syncthreads(); | |
| for (int st = nt >> 1; st > 0; st >>= 1) { if (tid < st) red[tid] += red[tid + st]; __syncthreads(); } | |
| const float S = red[0]; | |
| __syncthreads(); | |
| const int tk = top_k[row]; const float tp = top_p[row]; const float mp = min_p[row]; | |
| const bool use_k = (tk > 0 && tk < V), use_p = (tp > 0.f && tp < 1.f), use_m = (mp > 0.f); | |
| // 3. keep-threshold tau in the e-domain (intersection = max of the per-filter thresholds) | |
| float tau = 0.f; | |
| if (use_m) tau = fmaxf(tau, mp); // p_i >= mp*p_max <=> e_i >= mp | |
| if (use_k) tau = fmaxf(tau, radix_threshold(e, V, false, (double)tk, hist_c, hist_m, &prefix_sh, &above_sh)); | |
| if (use_p) tau = fmaxf(tau, radix_threshold(e, V, true, (double)tp * (double)S, hist_c, hist_m, &prefix_sh, &above_sh)); | |
| if (tid == 0) sh_scalar = tau; __syncthreads(); tau = sh_scalar; __syncthreads(); | |
| // 4. Gumbel-max over {e_i >= tau}: token = argmax_i (log(e_i) + gumbel_i). The | |
| // max logit (e=1) always clears tau, so the kept set is never empty. | |
| float best = -INFINITY; int bi = argmax; | |
| for (int i = tid; i < V; i += nt) { | |
| if (e[i] >= tau) { | |
| unsigned h = triple32((unsigned)i * 0x9E3779B1u ^ seed ^ (offset * 0x85EBCA6Bu) ^ ((unsigned)row * 0xC2B2AE35u)); | |
| float u = fminf(fmaxf((h >> 8) * (1.0f / 16777216.0f), 1e-7f), 0.9999999f); | |
| float key = __logf(e[i]) + (-__logf(-__logf(u))); | |
| if (key > best || (key == best && i < bi)) { best = key; bi = i; } | |
| } | |
| } | |
| red[tid] = best; redi[tid] = bi; __syncthreads(); | |
| for (int st = nt >> 1; st > 0; st >>= 1) { | |
| if (tid < st && (red[tid + st] > red[tid] || (red[tid + st] == red[tid] && redi[tid + st] < redi[tid]))) { | |
| red[tid] = red[tid + st]; redi[tid] = redi[tid + st]; | |
| } | |
| __syncthreads(); | |
| } | |
| if (tid == 0) out[row] = redi[0]; | |
| } | |
| } // namespace | |
| void apply_token_bitmask_inplace(at::Tensor logits, at::Tensor bitmask) { | |
| TORCH_CHECK(logits.is_cuda() && logits.dim() == 2 && logits.is_contiguous(), | |
| "logits must be contiguous [B, V] CUDA"); | |
| TORCH_CHECK(bitmask.is_cuda() && bitmask.scalar_type() == at::kInt, | |
| "bitmask must be int32 CUDA"); | |
| const at::cuda::CUDAGuard guard(logits.device()); | |
| int B = logits.size(0), V = logits.size(1); | |
| auto bm = bitmask.contiguous(); | |
| int ms = bm.size(-1); | |
| TORCH_CHECK(ms >= (V + 31) / 32, "bitmask has too few words for V"); | |
| const int32_t* m = bm.data_ptr<int32_t>(); | |
| int gx = (V + TPB - 1) / TPB; if (gx > 512) gx = 512; if (gx < 1) gx = 1; | |
| dim3 grid(gx, B); | |
| cudaStream_t s = at::cuda::getCurrentCUDAStream(); | |
| auto st = logits.scalar_type(); | |
| if (st == at::kFloat) | |
| bitmask_kernel<float><<<grid, TPB, 0, s>>>(logits.data_ptr<float>(), m, V, ms); | |
| else if (st == at::kHalf) | |
| bitmask_kernel<__half><<<grid, TPB, 0, s>>>((__half*)logits.data_ptr(), m, V, ms); | |
| else if (st == at::kBFloat16) | |
| bitmask_kernel<__nv_bfloat16><<<grid, TPB, 0, s>>>((__nv_bfloat16*)logits.data_ptr(), m, V, ms); | |
| else | |
| TORCH_CHECK(false, "logits must be float/half/bf16"); | |
| } | |
| at::Tensor sample(at::Tensor logits, at::Tensor temperature, at::Tensor top_p, | |
| at::Tensor top_k, at::Tensor min_p, int64_t seed, int64_t offset) { | |
| TORCH_CHECK(logits.is_cuda() && logits.dim() == 2 && logits.is_contiguous(), | |
| "logits must be contiguous [B, V] CUDA"); | |
| const at::cuda::CUDAGuard guard(logits.device()); | |
| int B = logits.size(0), V = logits.size(1); | |
| cudaStream_t s = at::cuda::getCurrentCUDAStream(); | |
| auto fopt = logits.options().dtype(at::kFloat); | |
| auto tf = temperature.to(at::kFloat).contiguous(); | |
| auto pf = top_p.to(at::kFloat).contiguous(); | |
| auto kf = top_k.to(at::kInt).contiguous(); | |
| auto mf = min_p.to(at::kFloat).contiguous(); | |
| auto work = at::empty({B, V}, fopt); // one [B,V] scratch (the softmax values) | |
| auto out = at::empty({B}, logits.options().dtype(at::kLong)); | |
| auto st = logits.scalar_type(); | |
| if (st == at::kFloat) { LAUNCH(float) } | |
| else if (st == at::kHalf) { LAUNCH(__half) } | |
| else if (st == at::kBFloat16) { LAUNCH(__nv_bfloat16) } | |
| else TORCH_CHECK(false, "logits must be float/half/bf16"); | |
| return out; | |
| } | |