CharlesCNorton
spec-decode-ops: grid-stride cast_kernel rows so M may exceed the 65535 gridDim.y cap
335d7ff | // Fused sampling and speculative-decoding verification. | |
| // | |
| // fused_sample / fused_filter implement the transformers logits pipeline | |
| // (repetition penalty -> temperature -> top-k -> top-p -> min-p -> sample) | |
| // with one pass structure: an optional segmented descending sort, then a | |
| // single kernel that resolves all filter cutoffs in one walk (each filter is | |
| // a prefix of the sorted order) and samples by the Gumbel-argmax identity, | |
| // which requires no normalization and no multinomial. Sort-free fast paths | |
| // cover greedy decoding and unfiltered sampling. | |
| // | |
| // spec_verify implements canonical speculative-decoding rejection sampling | |
| // (Leviathan et al.; Chen et al.): accept draft token x_i with probability | |
| // min(1, p_t(x_i) / p_d(x_i)); on first rejection sample from the residual | |
| // distribution max(p_t - p_d, 0) renormalized; if all k drafts are accepted, | |
| // sample a bonus token from the target's position-k distribution. The output | |
| // token stream is distributed exactly as the target model's distribution. | |
| // | |
| // Randomness is counter-based Philox (seed, offset). For a fixed seed, | |
| // offset, tensor shape, and launch geometry (constant in this file), results | |
| // are bitwise reproducible on a given architecture. | |
| namespace { | |
| constexpr int kThreads = 256; | |
| constexpr float kNegInf = -FLT_MAX; | |
| template <typename T> | |
| __device__ __forceinline__ float ldf(const T* p, int64_t i); | |
| template <> | |
| __device__ __forceinline__ float ldf<float>(const float* p, int64_t i) { return p[i]; } | |
| template <> | |
| __device__ __forceinline__ float ldf<__nv_bfloat16>(const __nv_bfloat16* p, int64_t i) { | |
| return __bfloat162float(p[i]); | |
| } | |
| template <> | |
| __device__ __forceinline__ float ldf<__half>(const __half* p, int64_t i) { | |
| return __half2float(p[i]); | |
| } | |
| __device__ __forceinline__ float gumbel_from_uniform(float u) { | |
| // curand_uniform returns u in (0, 1]; 1 - u lies in [0, 1), so the result | |
| // is bounded above and degrades to -inf (never selected) at the endpoint. | |
| return -__logf(-__logf(1.0f - u)); | |
| } | |
| __device__ float block_sum(float v) { | |
| __shared__ float smem[kThreads]; | |
| smem[threadIdx.x] = v; | |
| __syncthreads(); | |
| for (int s = kThreads / 2; s > 0; s >>= 1) { | |
| if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x + s]; | |
| __syncthreads(); | |
| } | |
| float r = smem[0]; | |
| __syncthreads(); | |
| return r; | |
| } | |
| __device__ void block_argmax(float v, int64_t idx, float* out_v, int64_t* out_i) { | |
| __shared__ float sv[kThreads]; | |
| __shared__ int64_t si[kThreads]; | |
| sv[threadIdx.x] = v; | |
| si[threadIdx.x] = idx; | |
| __syncthreads(); | |
| for (int s = kThreads / 2; s > 0; s >>= 1) { | |
| if (threadIdx.x < s) { | |
| if (sv[threadIdx.x + s] > sv[threadIdx.x] || | |
| (sv[threadIdx.x + s] == sv[threadIdx.x] && si[threadIdx.x + s] < si[threadIdx.x])) { | |
| sv[threadIdx.x] = sv[threadIdx.x + s]; | |
| si[threadIdx.x] = si[threadIdx.x + s]; | |
| } | |
| } | |
| __syncthreads(); | |
| } | |
| *out_v = sv[0]; | |
| *out_i = si[0]; | |
| __syncthreads(); | |
| } | |
| // --------------------------------------------------------------------------- | |
| // Pre-pass: cast to f32 workspace, apply repetition penalty. | |
| // prev_tokens: [M, P] int64, -1 padded. HF semantics: l > 0 ? l / r : l * r. | |
| // --------------------------------------------------------------------------- | |
| template <typename T> | |
| __global__ void cast_kernel(const T* __restrict__ logits, int64_t M, int64_t V, | |
| float* __restrict__ out) { | |
| // grid-stride over rows so M may exceed the 65535 gridDim.y hardware cap. | |
| for (int64_t row = blockIdx.y; row < M; row += gridDim.y) { | |
| for (int64_t v = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; v < V; | |
| v += (int64_t)gridDim.x * blockDim.x) { | |
| out[row * V + v] = ldf(logits, row * V + v); | |
| } | |
| } | |
| } | |
| __global__ void iota_kernel(int* __restrict__ out, int64_t M, int64_t V) { | |
| for (int64_t i = blockIdx.x * (int64_t)blockDim.x + threadIdx.x; i < M * V; | |
| i += (int64_t)gridDim.x * blockDim.x) { | |
| out[i] = (int)(i % V); | |
| } | |
| } | |
| // separate launch so every base value is written before the scatter reads it | |
| __global__ void penalty_kernel(float* __restrict__ out, int64_t V, | |
| const int64_t* __restrict__ prev, int64_t P, float rep) { | |
| const int64_t row = blockIdx.x; | |
| for (int64_t j = threadIdx.x; j < P; j += blockDim.x) { | |
| const int64_t t = prev[row * P + j]; | |
| if (t >= 0 && t < V) { | |
| float l = out[row * V + t]; | |
| out[row * V + t] = l > 0.f ? l / rep : l * rep; | |
| } | |
| } | |
| } | |
| // --------------------------------------------------------------------------- | |
| // Post-sort pass: resolve cutoffs on the sorted row, then Gumbel-argmax over | |
| // the kept prefix. One block per row. | |
| // | |
| // Kept prefix length L = min(top_k cutoff, top_p cutoff, min_p cutoff), where | |
| // the top_p cutoff uses probabilities normalized over the top_k prefix | |
| // (matching the warper order temperature -> top_k -> top_p -> min_p). | |
| // --------------------------------------------------------------------------- | |
| __global__ void sorted_sample_kernel(const float* __restrict__ sorted_logits, | |
| const int* __restrict__ sorted_idx, int64_t M, int64_t V, | |
| float inv_temp, int64_t top_k, float top_p, float min_p, | |
| uint64_t seed, uint64_t offset, bool write_mask, | |
| float* __restrict__ masked_out, | |
| int64_t* __restrict__ tokens) { | |
| const int64_t row = blockIdx.x; | |
| const float* sl = sorted_logits + row * V; | |
| const int* si = sorted_idx + row * V; | |
| const float lmax = sl[0] * inv_temp; | |
| // top-k with tie inclusion: transformers removes only values strictly | |
| // below the k-th, so equal-valued tokens past position k stay kept. | |
| int64_t kk = V; | |
| if (top_k > 0 && top_k < V) { | |
| const float kth = sl[top_k - 1]; | |
| __shared__ int64_t s_kk; | |
| if (threadIdx.x == 0) s_kk = V; | |
| __syncthreads(); | |
| const int64_t start = (top_k / kThreads) * kThreads; | |
| for (int64_t base = start; base < V; base += blockDim.x) { | |
| const int64_t i = base + threadIdx.x; | |
| if (i >= top_k && i < V && sl[i] < kth) { | |
| atomicMin((unsigned long long*)&s_kk, (unsigned long long)i); | |
| } | |
| __syncthreads(); | |
| if (s_kk < base + (int64_t)kThreads) break; | |
| } | |
| kk = s_kk; | |
| __syncthreads(); | |
| } | |
| // Z over the top_k prefix (post-temperature). | |
| float part = 0.f; | |
| for (int64_t i = threadIdx.x; i < kk; i += blockDim.x) { | |
| part += __expf(sl[i] * inv_temp - lmax); | |
| } | |
| const float Z = block_sum(part); | |
| // top_p cutoff: smallest prefix with cumulative probability >= top_p | |
| // (the crossing token is kept). Sequential chunk walk with running sum. | |
| __shared__ int64_t s_L; | |
| if (threadIdx.x == 0) s_L = kk; | |
| __syncthreads(); | |
| if (top_p < 1.0f) { | |
| float cum = 0.f; | |
| for (int64_t base = 0; base < kk; base += blockDim.x) { | |
| const int64_t i = base + threadIdx.x; | |
| float p = 0.f; | |
| if (i < kk) p = __expf(sl[i] * inv_temp - lmax) / Z; | |
| // inclusive scan within the chunk | |
| __shared__ float sc[kThreads]; | |
| sc[threadIdx.x] = p; | |
| __syncthreads(); | |
| for (int s = 1; s < kThreads; s <<= 1) { | |
| float add = (threadIdx.x >= s) ? sc[threadIdx.x - s] : 0.f; | |
| __syncthreads(); | |
| sc[threadIdx.x] += add; | |
| __syncthreads(); | |
| } | |
| const float inc = cum + sc[threadIdx.x]; | |
| // first position whose inclusive cumulative >= top_p ends the prefix | |
| if (i < kk && inc >= top_p) { | |
| const float prev_cum = inc - p; | |
| if (prev_cum < top_p) atomicMin((unsigned long long*)&s_L, (unsigned long long)(i + 1)); | |
| } | |
| cum += sc[kThreads - 1]; | |
| __syncthreads(); | |
| if (cum >= top_p) break; | |
| } | |
| } | |
| __syncthreads(); | |
| int64_t L = s_L; | |
| // min_p cutoff in logit space: keep l_i/T >= lmax + ln(min_p). | |
| if (min_p > 0.f) { | |
| const float thresh = lmax + __logf(min_p); | |
| // binary-search-free: walk chunks to find first violation | |
| __shared__ int64_t s_Lm; | |
| if (threadIdx.x == 0) s_Lm = L; | |
| __syncthreads(); | |
| for (int64_t base = 0; base < L; base += blockDim.x) { | |
| const int64_t i = base + threadIdx.x; | |
| if (i < L && sl[i] * inv_temp < thresh) { | |
| atomicMin((unsigned long long*)&s_Lm, (unsigned long long)i); | |
| } | |
| __syncthreads(); | |
| if (s_Lm <= base + blockDim.x) break; | |
| } | |
| L = s_Lm > 0 ? s_Lm : 1; // min_tokens_to_keep = 1 | |
| } | |
| __syncthreads(); | |
| if (write_mask) { | |
| // masked entries carry true -inf, matching the transformers chain's | |
| // filter_value; -FLT_MAX would evaluate identically under softmax but | |
| // fails isinf-based consumers. | |
| float* mo = masked_out + row * V; | |
| for (int64_t i = threadIdx.x; i < V; i += blockDim.x) { | |
| mo[si[i]] = (i < L) ? sl[i] * inv_temp : -INFINITY; | |
| } | |
| return; | |
| } | |
| // Gumbel-argmax over the kept prefix. | |
| curandStatePhilox4_32_10_t st; | |
| curand_init(seed, (uint64_t)row * kThreads + threadIdx.x, offset, &st); | |
| float best = kNegInf; | |
| int64_t besti = 0; | |
| for (int64_t i = threadIdx.x; i < L; i += blockDim.x) { | |
| const float g = gumbel_from_uniform(curand_uniform(&st)); | |
| const float v = sl[i] * inv_temp + g; | |
| if (v > best) { best = v; besti = i; } | |
| } | |
| float bv; int64_t bi; | |
| block_argmax(best, besti, &bv, &bi); | |
| if (threadIdx.x == 0) tokens[row] = (int64_t)si[bi]; | |
| } | |
| // --------------------------------------------------------------------------- | |
| // Sort-free fast paths: greedy argmax, and unfiltered Gumbel sampling. | |
| // --------------------------------------------------------------------------- | |
| template <typename T> | |
| __global__ void greedy_kernel(const T* __restrict__ logits, int64_t M, int64_t V, | |
| int64_t* __restrict__ tokens) { | |
| const int64_t row = blockIdx.x; | |
| float best = kNegInf; | |
| int64_t besti = 0; | |
| for (int64_t v = threadIdx.x; v < V; v += blockDim.x) { | |
| const float l = ldf(logits, row * V + v); | |
| if (l > best || (l == best && v < besti)) { best = l; besti = v; } | |
| } | |
| float bv; int64_t bi; | |
| block_argmax(best, besti, &bv, &bi); | |
| if (threadIdx.x == 0) tokens[row] = bi; | |
| } | |
| template <typename T> | |
| __global__ void gumbel_kernel(const T* __restrict__ logits, int64_t M, int64_t V, float inv_temp, | |
| uint64_t seed, uint64_t offset, int64_t* __restrict__ tokens) { | |
| const int64_t row = blockIdx.x; | |
| curandStatePhilox4_32_10_t st; | |
| curand_init(seed, (uint64_t)row * kThreads + threadIdx.x, offset, &st); | |
| float best = kNegInf; | |
| int64_t besti = 0; | |
| for (int64_t v = threadIdx.x; v < V; v += blockDim.x) { | |
| const float g = gumbel_from_uniform(curand_uniform(&st)); | |
| const float val = ldf(logits, row * V + v) * inv_temp + g; | |
| if (val > best) { best = val; besti = v; } | |
| } | |
| float bv; int64_t bi; | |
| block_argmax(best, besti, &bv, &bi); | |
| if (threadIdx.x == 0) tokens[row] = bi; | |
| } | |
| // --------------------------------------------------------------------------- | |
| // Speculative verification. | |
| // Phase 1: log-sum-exp of every target and draft row (grid parallel). | |
| // Phase 2: one block per sequence runs the sequential accept loop; on first | |
| // rejection samples the residual max(p_t - p_d, 0) by inverse-CDF walk; on | |
| // full acceptance samples the bonus position by Gumbel-argmax. | |
| // --------------------------------------------------------------------------- | |
| template <typename T> | |
| __global__ void lse_kernel(const T* __restrict__ logits, int64_t rows, int64_t V, float inv_temp, | |
| float* __restrict__ lse, float* __restrict__ rowmax) { | |
| const int64_t row = blockIdx.x; | |
| float mx = kNegInf; | |
| for (int64_t v = threadIdx.x; v < V; v += blockDim.x) { | |
| mx = fmaxf(mx, ldf(logits, row * V + v) * inv_temp); | |
| } | |
| __shared__ float smx[kThreads]; | |
| smx[threadIdx.x] = mx; | |
| __syncthreads(); | |
| for (int s = kThreads / 2; s > 0; s >>= 1) { | |
| if (threadIdx.x < s) smx[threadIdx.x] = fmaxf(smx[threadIdx.x], smx[threadIdx.x + s]); | |
| __syncthreads(); | |
| } | |
| mx = smx[0]; | |
| __syncthreads(); | |
| float part = 0.f; | |
| for (int64_t v = threadIdx.x; v < V; v += blockDim.x) { | |
| part += __expf(ldf(logits, row * V + v) * inv_temp - mx); | |
| } | |
| const float Z = block_sum(part); | |
| if (threadIdx.x == 0) { | |
| lse[row] = __logf(Z) + mx; | |
| rowmax[row] = mx; | |
| } | |
| } | |
| template <typename T> | |
| __global__ void verify_kernel(const T* __restrict__ tgt, const T* __restrict__ drf, | |
| const int64_t* __restrict__ draft_tokens, int64_t B, int64_t k, | |
| int64_t V, float inv_temp, const float* __restrict__ lse_t, | |
| const float* __restrict__ lse_d, uint64_t seed, uint64_t offset, | |
| int64_t* __restrict__ accept_len, | |
| int64_t* __restrict__ out_tokens) { | |
| const int64_t b = blockIdx.x; | |
| curandStatePhilox4_32_10_t st; | |
| curand_init(seed, (uint64_t)b * kThreads + threadIdx.x, offset, &st); | |
| // thread 0 draws the accept/reject uniforms so consumption is position-only | |
| __shared__ int64_t s_reject_at; | |
| __shared__ float s_u; | |
| if (threadIdx.x == 0) s_reject_at = k; | |
| __syncthreads(); | |
| for (int64_t i = 0; i < k; i++) { | |
| if (threadIdx.x == 0 && s_reject_at == k) { | |
| const int64_t x = draft_tokens[b * k + i]; | |
| const float lt = ldf(tgt, (b * (k + 1) + i) * V + x) * inv_temp - lse_t[b * (k + 1) + i]; | |
| const float ld = ldf(drf, (b * k + i) * V + x) * inv_temp - lse_d[b * k + i]; | |
| const float ratio = __expf(lt - ld); | |
| const float u = curand_uniform(&st); | |
| if (u >= ratio) { | |
| s_reject_at = i; | |
| s_u = curand_uniform(&st); | |
| } else { | |
| out_tokens[b * (k + 1) + i] = x; | |
| } | |
| } | |
| __syncthreads(); | |
| if (s_reject_at < k) break; | |
| } | |
| const int64_t rej = s_reject_at; | |
| if (rej == k) { | |
| // all accepted: bonus token from target position k by Gumbel-argmax | |
| float best = kNegInf; | |
| int64_t besti = 0; | |
| const int64_t roff = (b * (k + 1) + k) * V; | |
| for (int64_t v = threadIdx.x; v < V; v += blockDim.x) { | |
| const float g = gumbel_from_uniform(curand_uniform(&st)); | |
| const float val = ldf(tgt, roff + v) * inv_temp + g; | |
| if (val > best) { best = val; besti = v; } | |
| } | |
| float bv; int64_t bi; | |
| block_argmax(best, besti, &bv, &bi); | |
| if (threadIdx.x == 0) { | |
| out_tokens[b * (k + 1) + k] = bi; | |
| accept_len[b] = k; | |
| } | |
| return; | |
| } | |
| // rejection at position rej: sample from max(p_t - p_d, 0) / R | |
| const int64_t toff = (b * (k + 1) + rej) * V; | |
| const int64_t doff = (b * k + rej) * V; | |
| const float lt_lse = lse_t[b * (k + 1) + rej]; | |
| const float ld_lse = lse_d[b * k + rej]; | |
| float part = 0.f; | |
| for (int64_t v = threadIdx.x; v < V; v += blockDim.x) { | |
| const float pt = __expf(ldf(tgt, toff + v) * inv_temp - lt_lse); | |
| const float pd = __expf(ldf(drf, doff + v) * inv_temp - ld_lse); | |
| part += fmaxf(pt - pd, 0.f); | |
| } | |
| const float R = block_sum(part); | |
| __shared__ int64_t s_pick; | |
| if (threadIdx.x == 0) s_pick = -1; | |
| __syncthreads(); | |
| if (R > 0.f) { | |
| // inverse-CDF: sequential chunk walk over residual mass. The qualifying | |
| // intervals partition the mass, so at most one thread writes. | |
| const float target_mass = s_u * R; | |
| __shared__ float s_cum; | |
| if (threadIdx.x == 0) s_cum = 0.f; | |
| __syncthreads(); | |
| for (int64_t base = 0; base < V && s_pick < 0; base += blockDim.x) { | |
| const int64_t v = base + threadIdx.x; | |
| float r = 0.f; | |
| if (v < V) { | |
| const float pt = __expf(ldf(tgt, toff + v) * inv_temp - lt_lse); | |
| const float pd = __expf(ldf(drf, doff + v) * inv_temp - ld_lse); | |
| r = fmaxf(pt - pd, 0.f); | |
| } | |
| __shared__ float sc[kThreads]; | |
| sc[threadIdx.x] = r; | |
| __syncthreads(); | |
| for (int s = 1; s < kThreads; s <<= 1) { | |
| float add = (threadIdx.x >= s) ? sc[threadIdx.x - s] : 0.f; | |
| __syncthreads(); | |
| sc[threadIdx.x] += add; | |
| __syncthreads(); | |
| } | |
| const float inc = s_cum + sc[threadIdx.x]; | |
| if (v < V && r > 0.f && inc >= target_mass && (inc - r) < target_mass) { | |
| s_pick = v; | |
| } | |
| if (threadIdx.x == 0) s_cum += sc[kThreads - 1]; | |
| __syncthreads(); | |
| } | |
| } | |
| __syncthreads(); | |
| if (s_pick < 0) { | |
| // numerically empty residual (p_t == p_d): fall back to target sampling | |
| float best = kNegInf; | |
| int64_t besti = 0; | |
| for (int64_t v = threadIdx.x; v < V; v += blockDim.x) { | |
| const float g = gumbel_from_uniform(curand_uniform(&st)); | |
| const float val = ldf(tgt, toff + v) * inv_temp + g; | |
| if (val > best) { best = val; besti = v; } | |
| } | |
| float bv; int64_t bi; | |
| block_argmax(best, besti, &bv, &bi); | |
| if (threadIdx.x == 0) s_pick = bi; | |
| __syncthreads(); | |
| } | |
| if (threadIdx.x == 0) { | |
| out_tokens[b * (k + 1) + rej] = s_pick; | |
| accept_len[b] = rej; | |
| } | |
| } | |
| // greedy verification: accept while target argmax equals draft token | |
| template <typename T> | |
| __global__ void verify_greedy_kernel(const T* __restrict__ tgt, | |
| const int64_t* __restrict__ draft_tokens, int64_t B, | |
| int64_t k, int64_t V, int64_t* __restrict__ accept_len, | |
| int64_t* __restrict__ out_tokens) { | |
| const int64_t b = blockIdx.x; | |
| __shared__ int64_t s_am; | |
| int64_t alen = -1; | |
| for (int64_t i = 0; i <= k; i++) { | |
| if (alen >= 0) break; | |
| const int64_t roff = (b * (k + 1) + i) * V; | |
| float best = kNegInf; | |
| int64_t besti = 0; | |
| for (int64_t v = threadIdx.x; v < V; v += blockDim.x) { | |
| const float l = ldf(tgt, roff + v); | |
| if (l > best || (l == best && v < besti)) { best = l; besti = v; } | |
| } | |
| float bv; int64_t bi; | |
| block_argmax(best, besti, &bv, &bi); | |
| if (threadIdx.x == 0) s_am = bi; | |
| __syncthreads(); | |
| if (i == k || s_am != draft_tokens[b * k + i]) { | |
| if (threadIdx.x == 0) { | |
| out_tokens[b * (k + 1) + i] = s_am; | |
| accept_len[b] = i; | |
| } | |
| alen = i; | |
| } else { | |
| if (threadIdx.x == 0) out_tokens[b * (k + 1) + i] = s_am; | |
| } | |
| __syncthreads(); | |
| } | |
| } | |
| int64_t grid_x_for(int64_t n) { | |
| int64_t g = (n + kThreads - 1) / kThreads; | |
| return std::min<int64_t>(g, 1024); | |
| } | |
| template <typename T> | |
| const T* tptr(const torch::Tensor& t); | |
| template <> | |
| const float* tptr<float>(const torch::Tensor& t) { return t.const_data_ptr<float>(); } | |
| template <> | |
| const __nv_bfloat16* tptr<__nv_bfloat16>(const torch::Tensor& t) { | |
| return reinterpret_cast<const __nv_bfloat16*>(t.const_data_ptr<at::BFloat16>()); | |
| } | |
| template <> | |
| const __half* tptr<__half>(const torch::Tensor& t) { | |
| return reinterpret_cast<const __half*>(t.const_data_ptr<at::Half>()); | |
| } | |
| bool needs_sort(int64_t top_k, double top_p, double min_p, int64_t V) { | |
| return (top_k > 0 && top_k < V) || top_p < 1.0 || min_p > 0.0; | |
| } | |
| template <typename T> | |
| void run_sample(const torch::Tensor& logits, const c10::optional<torch::Tensor>& prev, | |
| double temperature, int64_t top_k, double top_p, double min_p, double rep, | |
| int64_t seed, int64_t offset, bool filter_only, torch::Tensor& out_tokens_or_mask, | |
| cudaStream_t stream) { | |
| const int64_t M = logits.size(0), V = logits.size(1); | |
| const float inv_temp = temperature > 0.0 ? (float)(1.0 / temperature) : 1.0f; | |
| auto opts_f = logits.options().dtype(torch::kFloat32); | |
| if (!filter_only && temperature == 0.0) { | |
| greedy_kernel<T><<<M, kThreads, 0, stream>>>(tptr<T>(logits), M, V, | |
| out_tokens_or_mask.data_ptr<int64_t>()); | |
| return; | |
| } | |
| const bool sortless = !needs_sort(top_k, top_p, min_p, V) && (!prev.has_value() || rep == 1.0); | |
| if (!filter_only && sortless) { | |
| gumbel_kernel<T><<<M, kThreads, 0, stream>>>(tptr<T>(logits), M, V, inv_temp, | |
| (uint64_t)seed, (uint64_t)offset, | |
| out_tokens_or_mask.data_ptr<int64_t>()); | |
| return; | |
| } | |
| // pre-pass into f32 workspace | |
| TORCH_CHECK(M * V < INT32_MAX, "M * V must fit in int32 for the segmented sort"); | |
| auto work = torch::empty({M, V}, opts_f); | |
| { | |
| dim3 grid((unsigned)grid_x_for(V), (unsigned)std::min<int64_t>(M, 65535)); | |
| cast_kernel<T><<<grid, kThreads, 0, stream>>>(tptr<T>(logits), M, V, work.data_ptr<float>()); | |
| if (prev.has_value() && rep != 1.0) { | |
| penalty_kernel<<<M, kThreads, 0, stream>>>(work.data_ptr<float>(), V, | |
| prev->const_data_ptr<int64_t>(), prev->size(1), | |
| (float)rep); | |
| } | |
| } | |
| // segmented descending sort with indices | |
| auto sorted = torch::empty({M, V}, opts_f); | |
| auto idx_in = torch::empty({M, V}, logits.options().dtype(torch::kInt32)); | |
| iota_kernel<<<grid_x_for(M * V), kThreads, 0, stream>>>(idx_in.data_ptr<int>(), M, V); | |
| auto idx_out = torch::empty({M, V}, logits.options().dtype(torch::kInt32)); | |
| auto offs = torch::arange(0, (M + 1) * V, V, logits.options().dtype(torch::kInt32)).contiguous(); | |
| size_t temp_bytes = 0; | |
| cub::DeviceSegmentedRadixSort::SortPairsDescending( | |
| nullptr, temp_bytes, work.const_data_ptr<float>(), sorted.data_ptr<float>(), | |
| idx_in.const_data_ptr<int>(), idx_out.data_ptr<int>(), M * V, M, | |
| offs.const_data_ptr<int>(), offs.const_data_ptr<int>() + 1, 0, 32, stream); | |
| auto temp = torch::empty({(int64_t)temp_bytes}, logits.options().dtype(torch::kUInt8)); | |
| cub::DeviceSegmentedRadixSort::SortPairsDescending( | |
| temp.data_ptr(), temp_bytes, work.const_data_ptr<float>(), sorted.data_ptr<float>(), | |
| idx_in.const_data_ptr<int>(), idx_out.data_ptr<int>(), M * V, M, | |
| offs.const_data_ptr<int>(), offs.const_data_ptr<int>() + 1, 0, 32, stream); | |
| sorted_sample_kernel<<<M, kThreads, 0, stream>>>( | |
| sorted.const_data_ptr<float>(), idx_out.const_data_ptr<int>(), M, V, inv_temp, | |
| top_k, (float)top_p, (float)min_p, (uint64_t)seed, (uint64_t)offset, filter_only, | |
| filter_only ? out_tokens_or_mask.data_ptr<float>() : nullptr, | |
| filter_only ? nullptr : out_tokens_or_mask.data_ptr<int64_t>()); | |
| } | |
| } // namespace | |
| void fused_sample(torch::Tensor& tokens, torch::Tensor const& logits, | |
| c10::optional<torch::Tensor> const& prev_tokens, double temperature, | |
| int64_t top_k, double top_p, double min_p, double repetition_penalty, | |
| int64_t seed, int64_t offset) { | |
| TORCH_CHECK(logits.is_cuda() && logits.dim() == 2 && logits.is_contiguous(), | |
| "logits must be contiguous [M, V] CUDA"); | |
| TORCH_CHECK(tokens.is_cuda() && tokens.dtype() == torch::kInt64 && | |
| tokens.numel() == logits.size(0), | |
| "tokens must be int64 [M]"); | |
| if (prev_tokens.has_value()) { | |
| TORCH_CHECK(prev_tokens->is_cuda() && prev_tokens->dtype() == torch::kInt64 && | |
| prev_tokens->dim() == 2 && prev_tokens->size(0) == logits.size(0) && | |
| prev_tokens->is_contiguous(), | |
| "prev_tokens must be contiguous int64 [M, P]"); | |
| } | |
| const at::cuda::CUDAGuard guard(logits.device()); | |
| cudaStream_t stream = at::cuda::getCurrentCUDAStream(); | |
| if (logits.dtype() == torch::kFloat32) { | |
| run_sample<float>(logits, prev_tokens, temperature, top_k, top_p, min_p, repetition_penalty, | |
| seed, offset, false, tokens, stream); | |
| } else if (logits.dtype() == torch::kBFloat16) { | |
| run_sample<__nv_bfloat16>(logits, prev_tokens, temperature, top_k, top_p, min_p, | |
| repetition_penalty, seed, offset, false, tokens, stream); | |
| } else if (logits.dtype() == torch::kHalf) { | |
| run_sample<__half>(logits, prev_tokens, temperature, top_k, top_p, min_p, | |
| repetition_penalty, seed, offset, false, tokens, stream); | |
| } else { | |
| TORCH_CHECK(false, "logits must be f32, bf16, or f16"); | |
| } | |
| } | |
| void fused_filter(torch::Tensor& out_logits, torch::Tensor const& logits, | |
| c10::optional<torch::Tensor> const& prev_tokens, double temperature, | |
| int64_t top_k, double top_p, double min_p, double repetition_penalty) { | |
| TORCH_CHECK(logits.is_cuda() && logits.dim() == 2 && logits.is_contiguous(), | |
| "logits must be contiguous [M, V] CUDA"); | |
| TORCH_CHECK(out_logits.is_cuda() && out_logits.dtype() == torch::kFloat32 && | |
| out_logits.sizes() == logits.sizes() && out_logits.is_contiguous(), | |
| "out_logits must be contiguous f32 [M, V]"); | |
| TORCH_CHECK(temperature > 0.0, "fused_filter requires temperature > 0"); | |
| const at::cuda::CUDAGuard guard(logits.device()); | |
| cudaStream_t stream = at::cuda::getCurrentCUDAStream(); | |
| if (logits.dtype() == torch::kFloat32) { | |
| run_sample<float>(logits, prev_tokens, temperature, top_k, top_p, min_p, repetition_penalty, | |
| 0, 0, true, out_logits, stream); | |
| } else if (logits.dtype() == torch::kBFloat16) { | |
| run_sample<__nv_bfloat16>(logits, prev_tokens, temperature, top_k, top_p, min_p, | |
| repetition_penalty, 0, 0, true, out_logits, stream); | |
| } else if (logits.dtype() == torch::kHalf) { | |
| run_sample<__half>(logits, prev_tokens, temperature, top_k, top_p, min_p, | |
| repetition_penalty, 0, 0, true, out_logits, stream); | |
| } else { | |
| TORCH_CHECK(false, "logits must be f32, bf16, or f16"); | |
| } | |
| } | |
| void spec_verify(torch::Tensor& accept_len, torch::Tensor& out_tokens, | |
| torch::Tensor const& target_logits, torch::Tensor const& draft_logits, | |
| torch::Tensor const& draft_tokens, double temperature, int64_t seed, | |
| int64_t offset) { | |
| TORCH_CHECK(target_logits.is_cuda() && target_logits.dim() == 3 && target_logits.is_contiguous(), | |
| "target_logits must be contiguous [B, k+1, V] CUDA"); | |
| TORCH_CHECK(draft_logits.is_cuda() && draft_logits.dim() == 3 && draft_logits.is_contiguous(), | |
| "draft_logits must be contiguous [B, k, V] CUDA"); | |
| const int64_t B = target_logits.size(0), V = target_logits.size(2); | |
| const int64_t k = draft_logits.size(1); | |
| TORCH_CHECK(target_logits.size(1) == k + 1, "target_logits must cover k+1 positions"); | |
| TORCH_CHECK(draft_logits.size(2) == V, "vocab mismatch"); | |
| TORCH_CHECK(draft_tokens.is_cuda() && draft_tokens.dtype() == torch::kInt64 && | |
| draft_tokens.sizes() == torch::IntArrayRef({B, k}), | |
| "draft_tokens must be int64 [B, k]"); | |
| TORCH_CHECK(accept_len.dtype() == torch::kInt64 && accept_len.numel() == B, | |
| "accept_len must be int64 [B]"); | |
| TORCH_CHECK(out_tokens.dtype() == torch::kInt64 && | |
| out_tokens.sizes() == torch::IntArrayRef({B, k + 1}), | |
| "out_tokens must be int64 [B, k+1]"); | |
| TORCH_CHECK(target_logits.dtype() == draft_logits.dtype(), "dtype mismatch"); | |
| const at::cuda::CUDAGuard guard(target_logits.device()); | |
| cudaStream_t stream = at::cuda::getCurrentCUDAStream(); | |
| if (temperature == 0.0) { | |
| if (target_logits.dtype() == torch::kFloat32) { | |
| verify_greedy_kernel<float><<<B, kThreads, 0, stream>>>( | |
| tptr<float>(target_logits), draft_tokens.const_data_ptr<int64_t>(), B, k, V, | |
| accept_len.data_ptr<int64_t>(), out_tokens.data_ptr<int64_t>()); | |
| } else if (target_logits.dtype() == torch::kBFloat16) { | |
| verify_greedy_kernel<__nv_bfloat16><<<B, kThreads, 0, stream>>>( | |
| tptr<__nv_bfloat16>(target_logits), draft_tokens.const_data_ptr<int64_t>(), B, k, V, | |
| accept_len.data_ptr<int64_t>(), out_tokens.data_ptr<int64_t>()); | |
| } else { | |
| verify_greedy_kernel<__half><<<B, kThreads, 0, stream>>>( | |
| tptr<__half>(target_logits), draft_tokens.const_data_ptr<int64_t>(), B, k, V, | |
| accept_len.data_ptr<int64_t>(), out_tokens.data_ptr<int64_t>()); | |
| } | |
| return; | |
| } | |
| const float inv_temp = (float)(1.0 / temperature); | |
| auto opts_f = target_logits.options().dtype(torch::kFloat32); | |
| auto lse_t = torch::empty({B * (k + 1)}, opts_f); | |
| auto lse_d = torch::empty({B * k}, opts_f); | |
| auto mx_t = torch::empty({B * (k + 1)}, opts_f); | |
| auto mx_d = torch::empty({B * k}, opts_f); | |
| if (target_logits.dtype() == torch::kFloat32) { | |
| DISPATCH_LSE(float) | |
| } else if (target_logits.dtype() == torch::kBFloat16) { | |
| DISPATCH_LSE(__nv_bfloat16) | |
| } else { | |
| DISPATCH_LSE(__half) | |
| } | |
| } | |