// 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.

#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include <cub/cub.cuh>
#include <curand_kernel.h>
#include <cuda_bf16.h>
#include <cfloat>
#include <cstdint>

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);

#define DISPATCH_LSE(T)                                                                       \
  lse_kernel<T><<<B * (k + 1), kThreads, 0, stream>>>(tptr<T>(target_logits), B * (k + 1), V, \
                                                      inv_temp, lse_t.data_ptr<float>(),      \
                                                      mx_t.data_ptr<float>());                \
  lse_kernel<T><<<B * k, kThreads, 0, stream>>>(tptr<T>(draft_logits), B * k, V, inv_temp,    \
                                                lse_d.data_ptr<float>(),                      \
                                                mx_d.data_ptr<float>());                      \
  verify_kernel<T><<<B, kThreads, 0, stream>>>(                                               \
      tptr<T>(target_logits), tptr<T>(draft_logits), draft_tokens.const_data_ptr<int64_t>(),  \
      B, k, V, inv_temp, lse_t.const_data_ptr<float>(), lse_d.const_data_ptr<float>(),        \
      (uint64_t)seed, (uint64_t)offset, accept_len.data_ptr<int64_t>(),                       \
      out_tokens.data_ptr<int64_t>());

  if (target_logits.dtype() == torch::kFloat32) {
    DISPATCH_LSE(float)
  } else if (target_logits.dtype() == torch::kBFloat16) {
    DISPATCH_LSE(__nv_bfloat16)
  } else {
    DISPATCH_LSE(__half)
  }
#undef DISPATCH_LSE
}