// 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.
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <cstdint>
#include <math.h>

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();
#define LAUNCH(CT) sample_kernel<CT><<<B, TPB, 0, s>>>( \
    reinterpret_cast<const CT*>(logits.const_data_ptr()), tf.data_ptr<float>(), pf.data_ptr<float>(), \
    kf.data_ptr<int>(), mf.data_ptr<float>(), work.data_ptr<float>(), \
    (unsigned)seed, (unsigned)offset, out.data_ptr<int64_t>(), V);
  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");
#undef LAUNCH
  return out;
}