TopK Bitonic Kimi K3 (256k)

"""Custom top-k kernel for B200 (SM100).

Algorithm: threshold-filter selection (no library ordering/select calls).
  - Pack (value, index) -> sortable uint64 key: key = (ord(-v) << 32) | idx;
    ascending key order == descending value order (ties broken by index).
  - Pass A: each thread maxes its strided share of its chunk and atomically
    bumps a 4096-bin histogram (bins = ord32 >> 20) of thread maxima.
  - tau: the largest bin floor whose top-cumulative count reaches k -- a
    provable lower bound on the true k-th value (those k thread maxima are k
    distinct row elements >= tau). Costs one atomic per thread, no sorting.
  - Pass B (L2-hot re-read): warp-ballot compaction of elements >= tau into a
    capped smem stage (stage aliases the dead histogram). Ties at tau and the
    cap are value-safe; pathological overflow re-reads the chunk exactly.
  - Phase C: bitonic top-128 only over ~k candidates (solo warp if <= 128).
  - Rows split into C chunks when batch is small: chunk top-k partials land in
    scratch; the last block per row (atomic ticket) merges partials in the
    same launch. k==1 uses a plain min-key reduction. k>128 falls back to
    multi-round extraction over the exact bitonic kernel.
"""
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline

CUDA_SRC = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/EmptyTensor.h>
#include <cuda_runtime.h>
#include <math_constants.h>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <vector>

#define PAD_KEY 0xFFFFFFFFFFFFFFFFULL
#define DEVFN __device__ __forceinline__
#define NBINS 4096

DEVFN uint32_t ord_float(float f) {
  uint32_t u = __float_as_uint(f);
  return u ^ ((uint32_t)((int32_t)u >> 31) | 0x80000000u);
}
DEVFN float unord_float(uint32_t u) {
  uint32_t b = (u & 0x80000000u) ? (u ^ 0x80000000u) : ~u;
  return __uint_as_float(b);
}
DEVFN float key_value(uint64_t key) {  // decode original value
  return -unord_float((uint32_t)(key >> 32));
}
DEVFN uint64_t make_key(float v, uint32_t idx) {
  return ((uint64_t)ord_float(-v) << 32) | (uint64_t)idx;
}
DEVFN uint64_t umin64(uint64_t a, uint64_t b) { return a < b ? a : b; }
DEVFN uint64_t umax64(uint64_t a, uint64_t b) { return a > b ? a : b; }

// Bitonic network: sort 128 keys ascending across a warp (4 per lane),
// element index = lane*4 + r.
DEVFN void sort128(uint64_t k[4], int lane) {
#pragma unroll
  for (int size = 2; size <= 128; size <<= 1) {
#pragma unroll
    for (int s = size >> 1; s >= 1; s >>= 1) {
      if (s <= 2) {
#pragma unroll
        for (int r = 0; r < 4; ++r) {
          if ((r & s) == 0) {
            int idx = (lane << 2) | r;
            bool asc = (idx & size) == 0;
            uint64_t a = k[r], b = k[r ^ s];
            uint64_t mn = umin64(a, b), mx = umax64(a, b);
            k[r] = asc ? mn : mx;
            k[r ^ s] = asc ? mx : mn;
          }
        }
      } else {
        int src = s >> 2;
#pragma unroll
        for (int r = 0; r < 4; ++r) {
          uint64_t o = __shfl_xor_sync(0xffffffffu, k[r], src);
          int idx = (lane << 2) | r;
          bool asc = (idx & size) == 0;
          bool low = (idx & s) == 0;
          uint64_t mn = umin64(k[r], o), mx = umax64(k[r], o);
          k[r] = (asc == low) ? mn : mx;
        }
      }
    }
  }
}


// Bitonic network: sort 64 keys ascending across a warp (2 per lane),
// element index = lane*2 + r.
DEVFN void sort64p(uint64_t k[2], int lane) {
#pragma unroll
  for (int size = 2; size <= 64; size <<= 1) {
#pragma unroll
    for (int s = size >> 1; s >= 1; s >>= 1) {
      if (s == 1) {
        int idx = (lane << 1);
        bool asc = (idx & size) == 0;
        uint64_t a = k[0], b = k[1];
        k[0] = asc ? umin64(a, b) : umax64(a, b);
        k[1] = asc ? umax64(a, b) : umin64(a, b);
      } else {
        int src = s >> 1;
#pragma unroll
        for (int r = 0; r < 2; ++r) {
          uint64_t o = __shfl_xor_sync(0xffffffffu, k[r], src);
          int idx = (lane << 1) | r;
          bool asc = (idx & size) == 0;
          bool low = (idx & s) == 0;
          k[r] = (asc == low) ? umin64(k[r], o) : umax64(k[r], o);
        }
      }
    }
  }
}

// Halve sorted tile into running top-128 list and re-sort (ascending).
DEVFN void merge128(uint64_t run[4], const uint64_t tile[4], int lane) {
#pragma unroll
  for (int r = 0; r < 4; ++r) {
    uint64_t partner = __shfl_xor_sync(0xffffffffu, tile[3 - r], 0x1F);
    run[r] = umin64(run[r], partner);
  }
#pragma unroll
  for (int s = 64; s >= 1; s >>= 1) {
    if (s <= 2) {
#pragma unroll
      for (int r = 0; r < 4; ++r) {
        if ((r & s) == 0) {
          uint64_t a = run[r], b = run[r ^ s];
          run[r] = umin64(a, b);
          run[r ^ s] = umax64(a, b);
        }
      }
    } else {
      int src = s >> 2;
#pragma unroll
      for (int r = 0; r < 4; ++r) {
        uint64_t o = __shfl_xor_sync(0xffffffffu, run[r], src);
        int idx = (lane << 2) | r;
        bool low = (idx & s) == 0;
        run[r] = low ? umin64(run[r], o) : umax64(run[r], o);
      }
    }
  }
}

// Merge two ascending-key sorted 64-lists into their ascending top-64:
// c[i] = min(A[i], B[63-i]) then bitonic completion (uniform ascending).
DEVFN void topk_merge64(uint64_t a[2], const uint64_t b[2], int lane) {
  uint64_t bf0 = __shfl_xor_sync(0xffffffffu, b[1], 0x1F);
  uint64_t bf1 = __shfl_xor_sync(0xffffffffu, b[0], 0x1F);
  a[0] = umin64(a[0], bf0);
  a[1] = umin64(a[1], bf1);
#pragma unroll
  for (int s = 32; s >= 2; s >>= 1) {
    int src = s >> 1;
    uint64_t o0 = __shfl_xor_sync(0xffffffffu, a[0], src);
    uint64_t o1 = __shfl_xor_sync(0xffffffffu, a[1], src);
    bool low = (((lane << 1) & s) == 0);
    a[0] = low ? umin64(a[0], o0) : umax64(a[0], o0);
    a[1] = low ? umin64(a[1], o1) : umax64(a[1], o1);
  }
  uint64_t mn = umin64(a[0], a[1]), mx = umax64(a[0], a[1]);
  a[0] = mn;
  a[1] = mx;
}

// Merge the W lists in smem down to smem[0] (ascending top-128).
template <int W>
DEVFN void tree_merge(uint64_t (*smem)[128], int warp, int lane) {
  uint64_t run[4];
#pragma unroll
  for (int r = 0; r < 4; ++r) run[r] = smem[warp][lane * 4 + r];
  for (int active = W >> 1; active >= 1; active >>= 1) {
    if (warp < active) {
      uint64_t other[4];
#pragma unroll
      for (int r = 0; r < 4; ++r) other[r] = smem[warp + active][lane * 4 + r];
      merge128(run, other, lane);
#pragma unroll
      for (int r = 0; r < 4; ++r) smem[warp][lane * 4 + r] = run[r];
    }
    __syncthreads();
  }
}

// Resolve: lowest bin floor whose top-cumulative count reaches k.
// bins: 4096 counters over (ord32 >> 20), block-wide; result written to *out.
template <int W>
DEVFN void resolve_tord(const uint32_t* __restrict__ bins, int k, int warp, int lane,
                        uint32_t* slice_sum, uint32_t* slice_above, uint32_t* out) {
  constexpr int SPW = NBINS / W;
  constexpr int BPL = SPW / 32;
  {
    uint32_t s = 0;
#pragma unroll
    for (int i = 0; i < BPL; ++i) s += bins[warp * SPW + lane * BPL + i];
#pragma unroll
    for (int d = 16; d > 0; d >>= 1) s += __shfl_xor_sync(0xffffffffu, s, d);
    if (lane == 0) slice_sum[warp] = s;
  }
  __syncthreads();
  if (warp == 0) {
    uint32_t v = (lane < W) ? slice_sum[lane] : 0u;
#pragma unroll
    for (int d = 1; d < W; d <<= 1) {
      uint32_t y = __shfl_down_sync(0xffffffffu, v, d);
      if (lane + d < W) v += y;
    }
    uint32_t excl = __shfl_down_sync(0xffffffffu, v, 1);
    if (lane < W - 1) slice_above[lane] = excl;
    else if (lane == W - 1) slice_above[lane] = 0u;
  }
  __syncthreads();
  const uint32_t above = slice_above[warp];
  if (above < (uint32_t)k && above + slice_sum[warp] >= (uint32_t)k) {
    // owner warp: lane l covers bins [warp*SPW + (31-l)*BPL, +BPL), lane 0 = top
    uint32_t ls = 0;
#pragma unroll
    for (int i = 0; i < BPL; ++i) ls += bins[warp * SPW + (31 - lane) * BPL + i];
#pragma unroll
    for (int d = 1; d < 32; d <<= 1) {
      uint32_t y = __shfl_up_sync(0xffffffffu, ls, d);
      if (lane >= d) ls += y;
    }
    uint32_t before_sh = __shfl_up_sync(0xffffffffu, ls, 1);
    uint32_t rb = above + ((lane == 0) ? 0u : before_sh);
    int lo = (31 - lane) * BPL;
    int target = -1;
#pragma unroll
    for (int j = BPL - 1; j >= 0; --j) {
      uint32_t c = bins[warp * SPW + lo + j];
      uint32_t nb = rb + c;
      if (target < 0 && nb >= (uint32_t)k) target = lo + j;
      rb = nb;
    }
    unsigned m = __ballot_sync(0xffffffffu, target >= 0);
    int src = __ffs(m) - 1;
    if (lane == src) *out = ((uint32_t)(warp * SPW + target)) << 20;
  }
}

// ---------------------------------------------------------------------------
// Threshold-filter top-k kernel. grid = (C chunks, R rows), block = W*32.
// ---------------------------------------------------------------------------
template <int W, bool VEC>
__global__ void __launch_bounds__(W * 32) topk_thresh_kernel(
    const float* __restrict__ x, int n, int k, int C, int chunk_len,
    float* __restrict__ out_v, long long* __restrict__ out_i,
    unsigned long long* __restrict__ scratch,   // R*C*k keys (C>1 only)
    unsigned long long* __restrict__ counters) {
  constexpr int T = W * 32;
  constexpr int CAP = 2048;
  constexpr int SPW = NBINS / W;    // bins per warp slice
  constexpr int BPL = SPW / 32;     // bins per lane within slice
  const int row = blockIdx.y;
  const int chunk = blockIdx.x;
  const int tid = threadIdx.x;
  const int lane = tid & 31;
  const int warp = tid >> 5;

  __shared__ uint64_t smem[W][128];
  __shared__ uint64_t stage_raw[CAP];
  uint32_t* bins = reinterpret_cast<uint32_t*>(stage_raw);  // aliases stage
  __shared__ uint32_t slice_sum[W];
  __shared__ uint32_t slice_above[W];
  __shared__ uint32_t tord_s;
  __shared__ int count_s;
  __shared__ int overflow_s;
  __shared__ bool last_s;

  const float* xrow = x + (size_t)row * n;
  const int e0 = chunk * chunk_len;
  const int e1 = min(n, e0 + chunk_len);
  const int v0 = e0 >> 2;
  const int v1 = (e1 + 3) >> 2;   // vec4 range [v0, v1)
  const int iters = (v1 - v0 + T - 1) / T;

  // zero histogram (+ misc) -- bins aliases stage[0..2048)
#pragma unroll
  for (int i = 0; i < NBINS / T; ++i) bins[tid + i * T] = 0u;
  if (tid == 0) { count_s = 0; overflow_s = 0; tord_s = 0u; }
  __syncthreads();

  // ---------------- Phase A: per-thread maxima + histogram ----------------
  {
    float tm = -CUDART_INF_F;
    int v = v0 + tid;
    if (VEC) {
      int it2 = iters >> 1;
      for (int i = 0; i < it2; ++i, v += 2 * T) {
        bool va = (v < v1), vb = (v + T < v1);
        bool fa = va && ((v + 1) << 2) <= n;
        bool fb = vb && ((v + T + 1) << 2) <= n;
        float4 a = fa ? *(const float4*)(xrow + (v << 2)) : float4{-CUDART_INF_F, -CUDART_INF_F, -CUDART_INF_F, -CUDART_INF_F};
        float4 b = fb ? *(const float4*)(xrow + ((v + T) << 2)) : float4{-CUDART_INF_F, -CUDART_INF_F, -CUDART_INF_F, -CUDART_INF_F};
        tm = fmaxf(tm, fmaxf(fmaxf(a.x, a.y), fmaxf(a.z, a.w)));
        tm = fmaxf(tm, fmaxf(fmaxf(b.x, b.y), fmaxf(b.z, b.w)));
      }
      if (iters & 1) {
        if (v < v1 && ((v + 1) << 2) <= n) {
          float4 f = *(const float4*)(xrow + (v << 2));
          tm = fmaxf(tm, fmaxf(f.x, fmaxf(f.y, fmaxf(f.z, f.w))));
        }
        v += T;
      }
    } else {
      for (int i = 0; i < iters; ++i, v += T) {
        if (v >= v1) break;
#pragma unroll
        for (int r = 0; r < 4; ++r) {
          int e = (v << 2) + r;
          if (e < e1) tm = fmaxf(tm, xrow[e]);
        }
      }
    }
    atomicAdd(&bins[ord_float(tm) >> 20], 1u);
  }
  __syncthreads();

  // ---------------- tau resolve: lowest bin floor with cumtop >= k ---------
  resolve_tord<W>(bins, k, warp, lane, slice_sum, slice_above, &tord_s);
  __syncthreads();
  const uint32_t tord = tord_s;
  const float tau_ = unord_float(tord);
  const float tau = (tau_ == tau_) ? tau_ : -CUDART_INF_F;

  // ---------------- Phase B: collect keys with v >= tau ----------------
  int itersB = iters;
  if (VEC) itersB = (iters + 1) >> 1;
  for (int i = 0; i < itersB; ++i) {
    int v = v0 + tid + (VEC ? i * 2 * T : i * T);
    float vals[4];
    bool anyq = false;
    bool valid = false;
    if (VEC) {
      valid = v < v1 && ((v + 1) << 2) <= n;
      int v2 = v + T;
      bool valid2 = v2 < v1 && ((v2 + 1) << 2) <= n;
      float vals2[4];
      bool anyq2 = false;
      if (valid) {
        float4 f = __ldcs((const float4*)(xrow + (v << 2)));
        vals[0] = f.x; vals[1] = f.y; vals[2] = f.z; vals[3] = f.w;
        anyq = fmaxf(f.x, fmaxf(f.y, fmaxf(f.z, f.w))) >= tau;
      }
      if (valid2) {
        float4 f = __ldcs((const float4*)(xrow + (v2 << 2)));
        vals2[0] = f.x; vals2[1] = f.y; vals2[2] = f.z; vals2[3] = f.w;
        anyq2 = fmaxf(f.x, fmaxf(f.y, fmaxf(f.z, f.w))) >= tau;
      }
      unsigned vec2q = __ballot_sync(0xffffffffu, (valid && anyq) || (valid2 && anyq2));
      if (vec2q) {
        // process the two tiles with the same per-element compaction
#pragma unroll
        for (int half = 0; half < 2; ++half) {
          int vv = half ? v2 : v;
          bool valh = half ? valid2 : valid;
          bool anyh = half ? anyq2 : anyq;
          const float* vv4 = half ? vals2 : vals;
          if (__ballot_sync(0xffffffffu, valh && anyh)) {
#pragma unroll
            for (int r = 0; r < 4; ++r) {
              bool q = valh && (vv4[r] >= tau);
              unsigned mask = __ballot_sync(0xffffffffu, q);
              if (mask) {
                int base = 0;
                if (lane == 0) base = atomicAdd(&count_s, __popc(mask));
                base = __shfl_sync(0xffffffffu, base, 0);
                if (base + __popc(mask) <= CAP) {
                  if (q) {
                    int pos = base + __popc(mask & ((1u << lane) - 1u));
                    stage_raw[pos] = make_key(vv4[r], (uint32_t)((vv << 2) + r));
                  }
                } else {
                  overflow_s = 1;
                }
              }
            }
          }
        }
      }
      continue;
    }
    valid = v < v1;
    if (valid) {
      float mx = -CUDART_INF_F;
#pragma unroll
      for (int r = 0; r < 4; ++r) {
        int e = (v << 2) + r;
        vals[r] = (e < e1) ? xrow[e] : -CUDART_INF_F;
        mx = fmaxf(mx, vals[r]);
      }
      anyq = mx >= tau;
    }
    unsigned vecq = __ballot_sync(0xffffffffu, valid && anyq);
    if (vecq) {
#pragma unroll
      for (int r = 0; r < 4; ++r) {
        bool q = valid && (((v << 2) + r) < e1) && (vals[r] >= tau);
        unsigned mask = __ballot_sync(0xffffffffu, q);
        if (mask) {
          int base = 0;
          if (lane == 0) base = atomicAdd(&count_s, __popc(mask));
          base = __shfl_sync(0xffffffffu, base, 0);
          if (base + __popc(mask) <= CAP) {
            if (q) {
              int pos = base + __popc(mask & ((1u << lane) - 1u));
              stage_raw[pos] = make_key(vals[r], (uint32_t)((v << 2) + r));
            }
          } else {
            overflow_s = 1;
          }
        }
      }
    }
  }
  __syncthreads();
  const int count = count_s;
  const bool overflow = overflow_s != 0;

  // ---------------- Phase C: chunk top-k -------------------------------
  if (!overflow && count <= 64) {
    if (warp == 0) {
      uint64_t keys[2];
#pragma unroll
      for (int r = 0; r < 2; ++r) {
        int e = lane * 2 + r;
        keys[r] = (e < count) ? stage_raw[e] : PAD_KEY;
      }
      sort64p(keys, lane);
#pragma unroll
      for (int r = 0; r < 2; ++r) smem[0][lane * 2 + r] = keys[r];
    }
    __syncthreads();
  } else if (!overflow && count <= 128 && k <= 64) {
    // two parallel sort64 halves + sorted top-64 merge
    uint64_t keys[2];
#pragma unroll
    for (int r = 0; r < 2; ++r) {
      int e = warp * 64 + lane * 2 + r;
      keys[r] = (warp <= 1 && e < count) ? stage_raw[e] : PAD_KEY;
    }
    if (warp <= 1) {
      sort64p(keys, lane);
#pragma unroll
      for (int r = 0; r < 2; ++r) smem[warp][lane * 2 + r] = keys[r];
    }
    __syncthreads();
    if (warp == 0) {
      uint64_t b2[2];
      b2[0] = smem[1][lane * 2 + 0];
      b2[1] = smem[1][lane * 2 + 1];
      topk_merge64(keys, b2, lane);
#pragma unroll
      for (int r = 0; r < 2; ++r) smem[0][lane * 2 + r] = keys[r];
    }
    __syncthreads();
  } else if (!overflow && count <= 128) {
    if (warp == 0) {
      uint64_t keys[4];
#pragma unroll
      for (int r = 0; r < 4; ++r) {
        int e = lane * 4 + r;
        keys[r] = (e < count) ? stage_raw[e] : PAD_KEY;
      }
      sort128(keys, lane);
#pragma unroll
      for (int r = 0; r < 4; ++r) smem[0][lane * 4 + r] = keys[r];
    }
    __syncthreads();
  } else {
    uint64_t run[4];
#pragma unroll
    for (int r = 0; r < 4; ++r) run[r] = PAD_KEY;
    if (!overflow) {
      int ttiles = (count + 127) >> 7;
      for (int t = warp; t < ttiles; t += W) {
        uint64_t keys[4];
        int base = t << 7;
#pragma unroll
        for (int r = 0; r < 4; ++r) {
          int e = base + lane * 4 + r;
          keys[r] = (e < count) ? stage_raw[e] : PAD_KEY;
        }
        sort128(keys, lane);
        merge128(run, keys, lane);
      }
    } else {
      int ttiles = (e1 - e0 + 127) >> 7;
      for (int t = warp; t < ttiles; t += W) {
        uint64_t keys[4];
        int base = e0 + (t << 7);
#pragma unroll
        for (int r = 0; r < 4; ++r) {
          int e = base + lane * 4 + r;
          keys[r] = PAD_KEY;
          if (e < e1) keys[r] = make_key(xrow[e], (uint32_t)e);
        }
        sort128(keys, lane);
        merge128(run, keys, lane);
      }
    }
#pragma unroll
    for (int r = 0; r < 4; ++r) smem[warp][lane * 4 + r] = run[r];
    __syncthreads();
    tree_merge<W>(smem, warp, lane);
  }

  float* vrow = out_v + (size_t)row * k;
  long long* irow = out_i + (size_t)row * k;

  if (C == 1) {
    if (tid < (unsigned)k) {
      uint64_t key = smem[0][tid];
      vrow[tid] = key_value(key);
      irow[tid] = (long long)(uint32_t)key;
    }
    return;
  }

  if (tid < (unsigned)k) scratch[(size_t)row * C * k + chunk * k + tid] = smem[0][tid];
  __threadfence();
  if (tid == 0) {
    unsigned long long old = atomicAdd(&counters[row], 1ULL);
    last_s = (old == (unsigned long long)(C - 1));
    if (last_s) counters[row] = 0ULL;
  }
  __syncthreads();
  if (!last_s) return;

  const unsigned long long* prow = scratch + (size_t)row * C * k;
  const int m = C * k;
  if (m <= 128) {
    if (warp == 0) {
      uint64_t keys[4];
#pragma unroll
      for (int r = 0; r < 4; ++r) {
        int e = lane * 4 + r;
        keys[r] = (e < m) ? prow[e] : PAD_KEY;
      }
      sort128(keys, lane);
#pragma unroll
      for (int r = 0; r < 4; ++r) smem[0][lane * 4 + r] = keys[r];
    }
    __syncthreads();
  } else if (k == 64 && (C & (C - 1)) == 0 && C <= 32 && m <= CAP) {
    // fast path: partial lists pre-sorted -> tournament of sorted 64-merges
    for (int e = tid; e < m; e += T) stage_raw[e] = prow[e];
    __syncthreads();
    for (int active = C >> 1; active >= 1; active >>= 1) {
      if (warp < active) {
        uint64_t a[2], b[2];
        const uint64_t* A = (const uint64_t*)stage_raw + (size_t)warp * 64 + lane * 2;
        const uint64_t* B = (const uint64_t*)stage_raw + (size_t)(warp + active) * 64 + lane * 2;
        a[0] = A[0]; a[1] = A[1]; b[0] = B[0]; b[1] = B[1];
        topk_merge64(a, b, lane);
        uint64_t* D = (uint64_t*)stage_raw + (size_t)warp * 64 + lane * 2;
        D[0] = a[0]; D[1] = a[1];
      }
      __syncthreads();
    }
    if (tid < (unsigned)k) smem[0][tid] = stage_raw[tid];
    __syncthreads();
  } else {
    const int ttiles = (m + 127) >> 7;
    uint64_t run[4];
#pragma unroll
    for (int r = 0; r < 4; ++r) run[r] = PAD_KEY;
    for (int t = warp; t < ttiles; t += W) {
      uint64_t keys[4];
      int base = t << 7;
#pragma unroll
      for (int r = 0; r < 4; ++r) {
        int e = base + lane * 4 + r;
        keys[r] = (e < m) ? prow[e] : PAD_KEY;
      }
      sort128(keys, lane);
      merge128(run, keys, lane);
    }
#pragma unroll
    for (int r = 0; r < 4; ++r) smem[warp][lane * 4 + r] = run[r];
    __syncthreads();
    tree_merge<W>(smem, warp, lane);
  }
  if (tid < (unsigned)k) {
    uint64_t key = smem[0][tid];
    vrow[tid] = key_value(key);
    irow[tid] = (long long)(uint32_t)key;
  }
}

// ---------------------------------------------------------------------------
// Exact bitonic-stream kernel (used for k>128 multi-round fallback).
// ---------------------------------------------------------------------------
template <int W, bool VEC>
__global__ void __launch_bounds__(W * 32) topk_kernel(
    const float* __restrict__ x, int n, int kw, int C, int out_stride,
    float* __restrict__ out_v, long long* __restrict__ out_i,
    unsigned long long* __restrict__ scratch,
    unsigned long long* __restrict__ counters,
    const unsigned long long* __restrict__ exclude,
    unsigned long long* __restrict__ next_exclude) {
  const int row = blockIdx.y;
  const int chunk = blockIdx.x;
  const int tid = threadIdx.x;
  const int lane = tid & 31;
  const int warp = tid >> 5;

  __shared__ uint64_t smem[W][128];
  __shared__ bool last_s;

  const int tiles_total = (n + 127) >> 7;
  const float* xrow = x + (size_t)row * n;

  uint64_t run[4];
#pragma unroll
  for (int r = 0; r < 4; ++r) run[r] = PAD_KEY;
  for (int j = warp; ; j += W) {
    int t = chunk + j * C;
    if (t >= tiles_total) break;
    int base = t << 7;
    uint64_t keys[4];
    if (VEC && base + 128 <= n) {
      float4 f = *(const float4*)(xrow + base + lane * 4);
      uint32_t e = base + lane * 4;
      keys[0] = make_key(f.x, e + 0);
      keys[1] = make_key(f.y, e + 1);
      keys[2] = make_key(f.z, e + 2);
      keys[3] = make_key(f.w, e + 3);
    } else {
#pragma unroll
      for (int r = 0; r < 4; ++r) {
        int e = base + lane * 4 + r;
        keys[r] = PAD_KEY;
        if (e < n) keys[r] = make_key(xrow[e], (uint32_t)e);
      }
    }
    if (exclude != nullptr) {
      unsigned long long ex = exclude[row];
#pragma unroll
      for (int r = 0; r < 4; ++r)
        if (keys[r] <= ex) keys[r] = PAD_KEY;
    }
    sort128(keys, lane);
    merge128(run, keys, lane);
  }
#pragma unroll
  for (int r = 0; r < 4; ++r) smem[warp][lane * 4 + r] = run[r];
  __syncthreads();
  tree_merge<W>(smem, warp, lane);

  float* vrow = out_v + (size_t)row * out_stride;
  long long* irow = out_i + (size_t)row * out_stride;

  if (C == 1) {
    if (tid < (unsigned)kw) {
      uint64_t key = smem[0][tid];
      vrow[tid] = key_value(key);
      irow[tid] = (long long)(uint32_t)key;
    }
    if (next_exclude != nullptr && tid == 0) next_exclude[row] = smem[0][kw - 1];
    return;
  }
  if (tid < (unsigned)kw) scratch[(size_t)row * C * kw + chunk * kw + tid] = smem[0][tid];
  __threadfence();
  if (tid == 0) {
    unsigned long long old = atomicAdd(&counters[row], 1ULL);
    last_s = (old == (unsigned long long)(C - 1));
    if (last_s) counters[row] = 0ULL;
  }
  __syncthreads();
  if (!last_s) return;

  const unsigned long long* prow = scratch + (size_t)row * C * kw;
  int m = C * kw;
  int ttiles = (m + 127) >> 7;
#pragma unroll
  for (int r = 0; r < 4; ++r) run[r] = PAD_KEY;
  for (int t = warp; t < ttiles; t += W) {
    uint64_t keys[4];
    int base = t << 7;
#pragma unroll
    for (int r = 0; r < 4; ++r) {
      int e = base + lane * 4 + r;
      keys[r] = (e < m) ? prow[e] : PAD_KEY;
    }
    sort128(keys, lane);
    merge128(run, keys, lane);
  }
#pragma unroll
  for (int r = 0; r < 4; ++r) smem[warp][lane * 4 + r] = run[r];
  __syncthreads();
  tree_merge<W>(smem, warp, lane);
  if (tid < (unsigned)kw) {
    uint64_t key = smem[0][tid];
    vrow[tid] = key_value(key);
    irow[tid] = (long long)(uint32_t)key;
  }
  if (next_exclude != nullptr && tid == 0) next_exclude[row] = smem[0][kw - 1];
}

// argmax (k==1) kernel: plain min-key reduction
template <int W, bool VEC>
__global__ void __launch_bounds__(W * 32) argmax_kernel(
    const float* __restrict__ x, int n, int C, int chunk_len,
    float* __restrict__ out_v, long long* __restrict__ out_i,
    unsigned long long* __restrict__ scratch,   // R*C (C>1)
    unsigned long long* __restrict__ counters) {
  const int row = blockIdx.y;
  const int chunk = blockIdx.x;
  const int tid = threadIdx.x;
  const int lane = tid & 31;
  const int warp = tid >> 5;
  __shared__ uint64_t smemk[W];
  __shared__ bool last_s;
  const float* xrow = x + (size_t)row * n;
  const int e0 = chunk * chunk_len;
  const int e1 = min(n, e0 + chunk_len);
  const int v0 = e0 >> 2;
  const int v1 = (e1 + 3) >> 2;
  const int iters = (v1 - v0 + (W * 32) - 1) / (W * 32);

  uint64_t best = PAD_KEY;
  for (int i = 0; i < iters; ++i) {
    int v = v0 + tid + i * (W * 32);
    if (v >= v1) break;
    if (VEC && ((v + 1) << 2) <= n) {
      float4 f = __ldcs((const float4*)(xrow + (v << 2)));
      uint32_t e = (uint32_t)(v << 2);
      best = umin64(best, make_key(f.x, e));
      best = umin64(best, make_key(f.y, e + 1));
      best = umin64(best, make_key(f.z, e + 2));
      best = umin64(best, make_key(f.w, e + 3));
    } else {
#pragma unroll
      for (int r = 0; r < 4; ++r) {
        int e = (v << 2) + r;
        if (e < e1) best = umin64(best, make_key(xrow[e], (uint32_t)e));
      }
    }
  }
#pragma unroll
  for (int d = 16; d > 0; d >>= 1)
    best = umin64(best, __shfl_xor_sync(0xffffffffu, best, d));
  if (lane == 0) smemk[warp] = best;
  __syncthreads();
  if (warp == 0) {
    uint64_t b = (lane < W) ? smemk[lane] : PAD_KEY;
#pragma unroll
    for (int d = 16; d > 0; d >>= 1)
      b = umin64(b, __shfl_xor_sync(0xffffffffu, b, d));
    if (lane == 0) {
      if (C == 1) {
        out_v[row] = key_value(b);
        out_i[row] = (long long)(uint32_t)b;
        return;
      }
      scratch[(size_t)row * C + chunk] = b;
    }
  }
  __syncthreads();
  if (C == 1) return;
  __threadfence();
  if (tid == 0) {
    unsigned long long old = atomicAdd(&counters[row], 1ULL);
    last_s = (old == (unsigned long long)(C - 1));
    if (last_s) counters[row] = 0ULL;
  }
  __syncthreads();
  if (!last_s) return;
  uint64_t b = PAD_KEY;
  const unsigned long long* prow = scratch + (size_t)row * C;
  for (int c = tid; c < C; c += W * 32) b = umin64(b, prow[c]);
#pragma unroll
  for (int d = 16; d > 0; d >>= 1) b = umin64(b, __shfl_xor_sync(0xffffffffu, b, d));
  if (lane == 0) smemk[warp] = b;
  __syncthreads();
  if (warp == 0) {
    uint64_t bb = (lane < W) ? smemk[lane] : PAD_KEY;
#pragma unroll
    for (int d = 16; d > 0; d >>= 1) bb = umin64(bb, __shfl_xor_sync(0xffffffffu, bb, d));
    if (lane == 0) {
      out_v[row] = key_value(bb);
      out_i[row] = (long long)(uint32_t)bb;
    }
  }
}

// ---------------------------------------------------------------------------
static torch::Tensor g_scratch;
static torch::Tensor g_counters;
static torch::Tensor g_roundbuf;

static inline int cdiv(int a, int b) { return (a + b - 1) / b; }

std::vector<at::Tensor> topk_forward(const at::Tensor& x, int64_t k_) {
  TORCH_CHECK(x.is_cuda() && x.dtype() == at::kFloat, "expect fp32 cuda tensor");
  TORCH_CHECK(x.dim() == 2, "expect 2D");
  int64_t R = x.size(0), n64 = x.size(1);
  TORCH_CHECK(n64 < INT32_MAX && k_ >= 1, "bad sizes");
  at::Tensor xc = x.is_contiguous() ? x : x.contiguous();
  const int n = (int)n64;
  const int kfull = (int)k_;
  TORCH_CHECK(kfull <= n, "k>n");

  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  at::Tensor out_v(at::detail::empty_cuda({R, k_}, at::kFloat, x.device(), std::nullopt));
  at::Tensor out_i(at::detail::empty_cuda({R, k_}, at::kLong, x.device(), std::nullopt));

  float* xv = xc.data_ptr<float>();
  float* vp = out_v.data_ptr<float>();
  long long* ip = (long long*)out_i.data_ptr<int64_t>();

  const bool vec = (n % 4 == 0) &&
      ((reinterpret_cast<uintptr_t>(xv) & 0xF) == 0);

  // ---------------- policy ----------------
  int W, C;
  static const char* envc = getenv("TOPK_FORCE_CFG");
  static int envW = 0, envC = 0;
  static const bool env_ok = (envc != nullptr && std::strlen(envc) > 0 && sscanf(envc, "%d,%d", &envW, &envC) == 2);
  if (env_ok) {
    W = envW;
    C = envC;
  } else if (R < 8 && n > 4096) {
    // tiny batch, big rows: split for DRAM parallelism (pow2-ish, <=32 hits
    // the sorted-tournament final merge)
    C = std::min(std::min(cdiv(148, (int)R), std::max(1, n / 8192)), 32);
    W = 16;
  } else {
    C = 1;
    W = 16;
  }
  const int nv = (n + 3) >> 2;
  const int chunkv = cdiv(nv, C);
  const int chunk_len = chunkv << 2;

  const bool use_thresh = (kfull <= 128);
  const int kw0 = std::min(kfull, 128);

  unsigned long long* scratch = nullptr;
  unsigned long long* counters = nullptr;
  if (C > 1) {
    int64_t need = (int64_t)R * C * std::max(kw0, 1);
    if (!g_scratch.defined() || g_scratch.numel() < need || g_scratch.device() != x.device()) {
      g_scratch = at::empty({need}, at::TensorOptions().dtype(at::kLong).device(x.device()));
    }
    if (!g_counters.defined() || g_counters.numel() < R || g_counters.device() != x.device()) {
      g_counters = at::zeros({R}, at::TensorOptions().dtype(at::kLong).device(x.device()));
    }
    scratch = (unsigned long long*)g_scratch.data_ptr<int64_t>();
    counters = (unsigned long long*)g_counters.data_ptr<int64_t>();
  }

  dim3 grid(C, (unsigned)R);

  if (kfull == 1) {
#define GO(WV, VV) argmax_kernel<WV, VV><<<grid, (WV)*32, 0, stream>>>(xv, n, C, chunk_len, vp, ip, scratch, counters)
    if (vec) { if (W == 4) GO(4, true); else if (W == 8) GO(8, true); else if (W == 16) GO(16, true); else GO(32, true); }
    else     { if (W == 4) GO(4, false); else if (W == 8) GO(8, false); else if (W == 16) GO(16, false); else GO(32, false); }
#undef GO
    return {out_v, out_i};
  }

  if (use_thresh) {
#define GO(WV, VV) topk_thresh_kernel<WV, VV><<<grid, (WV)*32, 0, stream>>>(xv, n, kfull, C, chunk_len, vp, ip, scratch, counters)
    if (vec) { if (W == 4) GO(4, true); else if (W == 8) GO(8, true); else GO(16, true); }
    else     { if (W == 4) GO(4, false); else if (W == 8) GO(8, false); else GO(16, false); }
#undef GO
    return {out_v, out_i};
  }

  // ---------------- k > 128: multi-round exact bitonic fallback -----------
  unsigned long long* roundbuf = nullptr;
  if (!g_roundbuf.defined() || g_roundbuf.numel() < 2 * R || g_roundbuf.device() != x.device()) {
    g_roundbuf = at::zeros({2 * R}, at::TensorOptions().dtype(at::kLong).device(x.device()));
  }
  roundbuf = (unsigned long long*)g_roundbuf.data_ptr<int64_t>();

  for (int base_k = 0; base_k < kfull; base_k += 128) {
    const int kw = std::min(128, kfull - base_k);
    float* vo = vp + (size_t)base_k;
    long long* io = ip + (size_t)base_k;
    const unsigned long long* cur_excl = (base_k == 0) ? nullptr : roundbuf + ((base_k / 128 + 1) % 2) * R;
    unsigned long long* next_excl = (base_k + 128 < kfull) ? roundbuf + ((base_k / 128) % 2) * R : nullptr;
    if (vec) {
      if (W == 8) topk_kernel<8, true><<<grid, W * 32, 0, stream>>>(xv, n, kw, C, kfull, vo, io, scratch, counters, cur_excl, next_excl);
      else topk_kernel<16, true><<<grid, W * 32, 0, stream>>>(xv, n, kw, C, kfull, vo, io, scratch, counters, cur_excl, next_excl);
    } else {
      if (W == 8) topk_kernel<8, false><<<grid, W * 32, 0, stream>>>(xv, n, kw, C, kfull, vo, io, scratch, counters, cur_excl, next_excl);
      else topk_kernel<16, false><<<grid, W * 32, 0, stream>>>(xv, n, kw, C, kfull, vo, io, scratch, counters, cur_excl, next_excl);
    }
  }
  return {out_v, out_i};
}
"""

CPP_SRC = r"""
#include <torch/extension.h>
#include <torch/csrc/autograd/python_variable.h>
#include <torch/csrc/Exceptions.h>
#include <vector>
std::vector<at::Tensor> topk_forward(const at::Tensor& x, int64_t k);

static PyObject* py_topk_raw(PyObject*, PyObject* const* args, Py_ssize_t nargs) {
  HANDLE_TH_ERRORS
  if (nargs != 2 || !THPVariable_Check(args[0])) {
    PyErr_SetString(PyExc_TypeError, "topk_raw(x: Tensor, k: int)");
    return nullptr;
  }
  const at::Tensor& x = THPVariable_Unpack(args[0]);
  int64_t k = (int64_t)PyLong_AsLongLong(args[1]);
  auto out = topk_forward(x, k);
  PyObject* tup = PyTuple_New(2);
  PyTuple_SET_ITEM(tup, 0, THPVariable_Wrap(std::move(out[0])));
  PyTuple_SET_ITEM(tup, 1, THPVariable_Wrap(std::move(out[1])));
  return tup;
  END_HANDLE_TH_ERRORS
}

static PyMethodDef _topk_methods[] = {
    {"topk_raw", (PyCFunction)(void*)py_topk_raw, METH_FASTCALL, nullptr},
    {nullptr, nullptr, 0, nullptr}};

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  PyModule_AddFunctions(m.ptr(), _topk_methods);
}
"""

_ext = load_inline(
    name="topk_bitonic_v5",
    cpp_sources=CPP_SRC,
    cuda_sources=CUDA_SRC,
    extra_cuda_cflags=["-O3", "--use_fast_math", "-gencode=arch=compute_100,code=sm_100"],
    verbose=False,
)


class Model(nn.Module):
    def __init__(self, batch: int, n: int, k: int):
        super().__init__()
        self.batch, self.n, self.k = batch, n, k
        self._k = int(k)
        self.register_buffer("_dummy", torch.zeros(1))

    def forward(self, x: torch.Tensor):
        return _ext.topk_raw(x, self._k)

    __call__ = forward


# Module-level shims rebuilt by check.py / benchmark.py per shape.
batch = 64
n = 8192
k = 8


def get_inputs():
    x = torch.randn(batch, n, dtype=torch.float32)
    return [x]


def get_init_inputs():
    return [batch, n, k]