TopK Bitonic Kimi K3 (256k)

"""Fused top-k kernel: register run-merge selection, single launch.

Design notes (RTX PRO 6000 / sm_120):
  - u64 keys pack (order-preserving float bits map << 32) | element index, so
    descending u64 order == descending value order with unique index
    tie-break; selection is exact.
  - Selection uses TRUNCATING BITONIC RUN-MERGES in registers: a sorted run
    is spread 8-registers-per-lane across a lane group; merging two sibling
    runs is rev-cross shuffles + a bitonic clean, dropping the lower half
    once past k. A per-warp merge tree (no smem, no block barriers) reduces
    32 lane runs to one warp top-k run.
  - Phase 1: per block of 256 threads, scan a 2048-element chunk with
    vectorized loads, per-thread sort-8, warp merge tree, then a 3-level
    smem block merge -> chunk top-k candidates.
  - Phase 2: the last block per row (atomic ticket) merges the per-chunk
    candidate runs with the same primitive in macro-rounds + accumulator.
  - k == 1 has a dedicated argmax reduction path.
  - Counters self-clean; workspaces persist and are never returned.
"""
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline


def _emit_merge_fns():
    L = []
    L.append("template <int G> DEVINL void warp_grow(ull* r, int lane, uint mask);\n"
             "template <int G> DEVINL void warp_trunc(ull* r, int lane, uint mask, int sp, int km);\n")
    for G in (1, 2, 4, 8):
        L.append(f"template <> DEVINL void warp_grow<{G}>(ull* r, int lane, uint mask) {{")
        L.append(f"    ull o_[8];")
        L.append(f"    #pragma unroll")
        L.append(f"    for (int s = 0; s < 8; ++s) o_[s] = __shfl_xor_sync(mask, r[7 - s], {2*G-1});")
        L.append(f"    #pragma unroll")
        L.append(f"    for (int s = 0; s < 8; ++s) {{")
        L.append(f"        r[s] = ((lane & {G}) == 0) ? (r[s] > o_[s] ? r[s] : o_[s]) : (r[s] < o_[s] ? r[s] : o_[s]);")
        L.append(f"    }}")
        LL = 8 * G
        j = 16 * G // 2
        while j >= 1:
            if j >= 8:
                p = j // 8
                L.append(f"    {{ bool hi_ = ((lane & {p}) == 0);")
                L.append(f"    #pragma unroll")
                L.append(f"    for (int s = 0; s < 8; ++s) {{")
                L.append(f"        ull o_ = __shfl_xor_sync(mask, r[s], {p});")
                L.append(f"        r[s] = hi_ ? (r[s] > o_ ? r[s] : o_) : (r[s] < o_ ? r[s] : o_);")
                L.append(f"    }} }}")
            else:
                L.append(f"    #pragma unroll")
                L.append(f"    for (int s = 0; s < 8; ++s) if (!((s & {j}) != 0)) {{")
                L.append(f"        ull a_ = r[s], b_ = r[s | {j}];")
                L.append(f"        r[s] = a_ > b_ ? a_ : b_; r[s | {j}] = a_ > b_ ? b_ : a_;")
                L.append(f"    }}")
            j >>= 1
        L.append("}")
    for G in (1, 2, 4, 8):
        L.append(f"template <> DEVINL void warp_trunc<{G}>(ull* r, int lane, uint mask, int sp, int km) {{")
        L.append(f"    ull o_[8];")
        L.append(f"    #pragma unroll")
        L.append(f"    for (int s = 0; s < 8; ++s) o_[s] = __shfl_xor_sync(mask, r[7 - s], sp);")
        L.append(f"    #pragma unroll")
        L.append(f"    for (int s = 0; s < 8; ++s) {{")
        L.append(f"        if ((lane & km) == 0) r[s] = r[s] > o_[s] ? r[s] : o_[s];")
        L.append(f"    }}")
        LL = 8 * G
        j = LL // 2
        while j >= 8:
            p = j // 8
            L.append(f"    {{ bool hi_ = ((lane & {p}) == 0);")
            L.append(f"    #pragma unroll")
            L.append(f"    for (int s = 0; s < 8; ++s) {{")
            L.append(f"        ull o_ = __shfl_xor_sync(mask, r[s], {p});")
            L.append(f"        r[s] = hi_ ? (r[s] > o_ ? r[s] : o_) : (r[s] < o_ ? r[s] : o_);")
            L.append(f"    }} }}")
            j >>= 1
        while j >= 1:
            L.append(f"    #pragma unroll")
            L.append(f"    for (int s = 0; s < 8; ++s) if (!((s & {j}) != 0)) {{")
            L.append(f"        ull a_ = r[s], b_ = r[s | {j}];")
            L.append(f"        r[s] = a_ > b_ ? a_ : b_; r[s | {j}] = a_ > b_ ? b_ : a_;")
            L.append(f"    }}")
            j >>= 1
        L.append("}")
    return "\n".join(L)


_MERGE_FNS = _emit_merge_fns()

_CUDA_TEMPLATE = r"""
#include <torch/extension.h>
#include <ATen/cuda/CUDAContext.h>
#include <cuda_runtime.h>

using ull = unsigned long long;
using uint = unsigned int;

#define DEVINL __device__ __forceinline__

// Order-preserving map float bits -> uint (total order; NaNs above +inf,
// -0.0 below +0.0). keyf inverts it.
DEVINL uint fkey(float f) {
    uint b = __float_as_uint(f);
    return (b & 0x80000000u) ? ~b : (b | 0x80000000u);
}
DEVINL float keyf(uint u) {
    uint b = (u & 0x80000000u) ? (u & 0x7fffffffu) : ~u;
    return __uint_as_float(b);
}
DEVINL ull mkkey(float f, uint idx) { return ((ull)fkey(f) << 32) | (ull)idx; }
DEVINL ull umax64(ull a, ull b) { return a > b ? a : b; }

#define CSD(i, j) do { ull a_ = r[i], b_ = r[j]; \
    r[i] = a_ > b_ ? a_ : b_; r[j] = a_ > b_ ? b_ : a_; } while (0)

__MERGE_FNS__

// per-thread 8-element sorting network (Batcher odd-even, descending)
DEVINL void sort8(ull* r) {
    CSD(0, 1);
    CSD(2, 3);
    CSD(0, 2);
    CSD(1, 3);
    CSD(1, 2);
    CSD(4, 5);
    CSD(6, 7);
    CSD(4, 6);
    CSD(5, 7);
    CSD(5, 6);
    CSD(0, 4);
    CSD(2, 6);
    CSD(2, 4);
    CSD(1, 5);
    CSD(3, 7);
    CSD(3, 5);
    CSD(1, 2);
    CSD(3, 4);
    CSD(5, 6);
}

// Merge two sorted-desc runs of 8G u64 each at A, B; keep desc top 8G in out.
// Executed by lanes 0..2G-1 of one warp (mask covers exactly those lanes).
__device__ void smem_merge_pair(const ull* __restrict__ A, const ull* __restrict__ B,
                                ull* __restrict__ out, int G, int lane, uint mask) {
    ull r[8];
    int gl = lane % (2 * G);
    bool upper = gl >= G;
    int p = upper ? gl - G : gl;
    const ull* base = upper ? B : A;
#pragma unroll
    for (int i = 0; i < 8; ++i) r[i] = base[p * 8 + i];
    if (G == 1)      warp_trunc<1>(r, lane, mask, 1, 1);
    else if (G == 2) warp_trunc<2>(r, lane, mask, 3, 2);
    else if (G == 4) warp_trunc<4>(r, lane, mask, 7, 4);
    else             warp_trunc<8>(r, lane, mask, 15, 8);
    if (!upper) {
#pragma unroll
        for (int i = 0; i < 8; ++i) out[p * 8 + i] = r[i];
    }
}

DEVINL int np2d(int x) {
    x--;
    x |= x >> 1; x |= x >> 2; x |= x >> 4; x |= x >> 8; x |= x >> 16;
    return x + 1;
}

// M1: per-thread element count (pow2). E: elems scanned per thread.
template <int M1, int E>
__global__ void topk_kernel(const float* __restrict__ x,
                            ull* __restrict__ cand, int* __restrict__ cnt, int* __restrict__ cnt2,
                            float* __restrict__ outv, long long* __restrict__ outi,
                            int n, int k, int K2, int C, int CH) {
    const int T = 256;
    const int tid = threadIdx.x;
    const int bid = blockIdx.x;
    const int row = bid / C;
    const int cid = bid - row * C;
    const float* xrow = x + (long long)row * n;
    const int cstart = cid * CH;
    const int clen = min(CH, n - cstart);
    const int warp = tid >> 5;
    const int lane = tid & 31;
    const int RS = K2 < 8 ? 8 : K2;      // run length in merge machinery
    const int G = RS >> 3;               // lane group size in merge prims

    __shared__ ull wsm[640];
    __shared__ int s_flag;

    ull r[8];
    bool vec = (clen == CH) && ((n & 3) == 0);

    // ---------------- phase 1 -------------------------------------------
    if (M1 == 1) {
        // argmax path (k == 1)
        ull best = 0ull;
        if (vec) {
            const float4* xc = reinterpret_cast<const float4*>(xrow + cstart);
#pragma unroll
            for (int j = 0; j < E / 4; ++j) {
                float4 v = xc[j * T + tid];
                int base = cstart + j * 4 * T + tid * 4;
                best = umax64(best, mkkey(v.x, base + 0));
                best = umax64(best, mkkey(v.y, base + 1));
                best = umax64(best, mkkey(v.z, base + 2));
                best = umax64(best, mkkey(v.w, base + 3));
            }
        } else {
#pragma unroll
            for (int e = 0; e < E; ++e) {
                int local = e * T + tid;
                if (local < clen) best = umax64(best, mkkey(xrow[cstart + local], cstart + local));
            }
        }
#pragma unroll
        for (int s = 16; s >= 1; s >>= 1) {
            ull o = __shfl_xor_sync(0xffffffffu, best, s);
            best = umax64(best, o);
        }
        if (lane == 0) wsm[warp] = best;
        __syncthreads();
        if (tid < 32) {
            ull b = tid < 8 ? wsm[tid] : 0ull;
#pragma unroll
            for (int s = 4; s >= 1; s >>= 1) {
                ull o = __shfl_xor_sync(0xffffffffu, b, s);
                b = umax64(b, o);
            }
            if (tid == 0) wsm[512] = b;
        }
        __syncthreads();
        if (C == 1) {
            if (tid == 0) {
                ull u = wsm[512];
                outv[(long long)row * k] = keyf((uint)(u >> 32));
                outi[(long long)row * k] = (long long)(uint)(u & 0xffffffffu);
            }
            return;
        }
        if (tid == 0) cand[(long long)bid * k] = wsm[512];
        __threadfence();
        __syncthreads();
        if (tid == 0) { int old = atomicAdd(&cnt[row], 1); s_flag = (old == C - 1) ? 1 : 0; }
        __syncthreads();
        if (s_flag == 0) return;
        if (tid == 0) cnt[row] = 0;
        __threadfence();
        const ull* candrow = cand + (long long)row * C;
        if (tid < 32) {
            ull b = 0ull;
            for (int q = tid; q < C; q += 32) b = umax64(b, candrow[q]);
#pragma unroll
            for (int s = 16; s >= 1; s >>= 1) {
                ull o = __shfl_xor_sync(0xffffffffu, b, s);
                b = umax64(b, o);
            }
            if (tid == 0) {
                outv[(long long)row * k] = keyf((uint)(b >> 32));
                outi[(long long)row * k] = (long long)(uint)(b & 0xffffffffu);
            }
        }
        return;
    }

#pragma unroll
    for (int i = 0; i < 8; ++i) r[i] = 0ull;
    if (M1 == E) {
        if (vec) {
            const float4* xc = reinterpret_cast<const float4*>(xrow + cstart);
#pragma unroll
            for (int j = 0; j < E / 4; ++j) {
                float4 v = xc[j * T + tid];
                int base = cstart + j * 4 * T + tid * 4;
                r[4 * j + 0] = mkkey(v.x, base + 0);
                r[4 * j + 1] = mkkey(v.y, base + 1);
                r[4 * j + 2] = mkkey(v.z, base + 2);
                r[4 * j + 3] = mkkey(v.w, base + 3);
            }
        } else {
#pragma unroll
            for (int e = 0; e < E; ++e) {
                int local = e * T + tid;
                r[e] = (local < clen) ? mkkey(xrow[cstart + local], cstart + local) : 0ull;
            }
        }
    } else {
        // threshold filter keeps per-thread top-M1 in r[0..M1), rest zeros
        ull rmin = 0ull;
#pragma unroll
        for (int e = 0; e < E; ++e) {
            int local = e * T + tid;
            if (local < clen) {
                ull key = mkkey(xrow[cstart + local], cstart + local);
                if (key > rmin) {
                    int mp = 0;
#pragma unroll
                    for (int i = 1; i < M1; ++i) if (r[i] < r[mp]) mp = i;
                    r[mp] = key;
                    ull m = r[0];
#pragma unroll
                    for (int i = 1; i < M1; ++i) m = m < r[i] ? m : r[i];
                    rmin = m;
                }
            }
        }
    }
    sort8(r);

    // register merge tree: 32 lane-runs -> 1 warp run of RS (zeros padded)
    if (K2 <= 8) {
        warp_trunc<1>(r, lane, 0xffffffffu, 1, 1);
        warp_trunc<1>(r, lane, 0xffffffffu, 2, 2);
        warp_trunc<1>(r, lane, 0xffffffffu, 4, 4);
        warp_trunc<1>(r, lane, 0xffffffffu, 8, 8);
        warp_trunc<1>(r, lane, 0xffffffffu, 16, 16);
    } else if (K2 == 16) {
        warp_grow<1>(r, lane, 0xffffffffu);
        warp_trunc<2>(r, lane, 0xffffffffu, 3, 2);
        warp_trunc<2>(r, lane, 0xffffffffu, 5, 4);
        warp_trunc<2>(r, lane, 0xffffffffu, 9, 8);
        warp_trunc<2>(r, lane, 0xffffffffu, 17, 16);
    } else if (K2 == 32) {
        warp_grow<1>(r, lane, 0xffffffffu);
        warp_grow<2>(r, lane, 0xffffffffu);
        warp_trunc<4>(r, lane, 0xffffffffu, 7, 4);
        warp_trunc<4>(r, lane, 0xffffffffu, 11, 8);
        warp_trunc<4>(r, lane, 0xffffffffu, 19, 16);
    } else {  // K2 == 64
        warp_grow<1>(r, lane, 0xffffffffu);
        warp_grow<2>(r, lane, 0xffffffffu);
        warp_grow<4>(r, lane, 0xffffffffu);
        warp_trunc<8>(r, lane, 0xffffffffu, 15, 8);
        warp_trunc<8>(r, lane, 0xffffffffu, 23, 16);
    }

    // warp winners -> wsm (runs at stride RS)
#pragma unroll
    for (int i = 0; i < 8; ++i) {
        int v = lane * 8 + i;
        if (v < RS) wsm[warp * RS + v] = r[i];
    }
    __syncthreads();
    // block tree: 8 runs -> 1
    if (warp < 4) smem_merge_pair(wsm + (2 * warp) * RS, wsm + (2 * warp + 1) * RS,
                                  wsm + 256 + warp * RS, G, lane, (1u << (2 * G)) - 1);
    __syncthreads();
    if (warp < 2) smem_merge_pair(wsm + 256 + (2 * warp) * RS, wsm + 256 + (2 * warp + 1) * RS,
                                  wsm + 384 + warp * RS, G, lane, (1u << (2 * G)) - 1);
    __syncthreads();
    if (warp == 0) smem_merge_pair(wsm + 384, wsm + 384 + RS,
                                   wsm + 512, G, lane, (1u << (2 * G)) - 1);
    __syncthreads();

    if (C == 1) {
        for (int w = tid; w < k; w += T) {
            ull u = wsm[512 + w];
            outv[(long long)row * k + w] = keyf((uint)(u >> 32));
            outi[(long long)row * k + w] = (long long)(uint)(u & 0xffffffffu);
        }
        return;
    }
    {
        int SEC = C < 64 ? C : 64;
        long long coff = (long long)row * (C + SEC) * RS + (long long)cid * RS;
        for (int w = tid; w < RS; w += T) cand[coff + w] = wsm[512 + w];
    }

    __threadfence();

    // ---------------- phase 2 -------------------------------------------
    // Hierarchical count-tree cascade for ALL C > 1: phase-1 blocks climb
    // pairwise as their ticket rights allow (block-level pipelining, no
    // big single-block serial tree). Node counters self-clean.
    const int SEC = C < 64 ? C : 64;          // scratch runs per row
    const ull* candrow = cand + (long long)row * (C + SEC) * RS;
    const ull* scrrow = candrow + (long long)C * RS;
    ull* ACC = wsm + 576;
    {
        // Cascade runs on warp 0 only of every block (merge work is
        // warp-internal); other warps exit immediately after publication.
        if (warp != 0) return;
        __syncwarp();
        // level run offsets inside scrrow: base(0)=0; base(l) = sum of
        // ceil(C/2^i) for i<l. Node id for (lvl, pair) = lvl*64 + pair.
        int C_lvl = C;
        int wid = cid;
        int lvl = 0;
        int lbase = -C;                // run-offset of current level (level 0 = candrow)
        int lnext = 0;                 // run-offset of next level (in scrrow)
        bool i_am_root = false;
        while (true) {
            int pair = wid >> 1;
            int target = (2 * pair + 1 < C_lvl) ? 2 : 1;
            if (tid == 0) {
                s_flag = atomicAdd(&cnt2[row * (8 * 64) + lvl * 64 + pair], 1) + 1;
            }
            __syncwarp();
            if (s_flag < target) break;
            if (tid == 0) cnt2[row * (8 * 64) + lvl * 64 + pair] = 0;  // self-clean
            __threadfence();
            const ull* A = scrrow + ((long long)lbase + 2 * pair) * RS;
            ull* outp = const_cast<ull*>(scrrow) + ((long long)lnext + pair) * RS;
            if (target == 2) {
                smem_merge_pair(A, A + RS, outp, G, lane, (1u << (2 * G)) - 1);
            } else {
                for (int w = lane; w < RS; w += 32) outp[w] = A[w];
            }
            __syncwarp();
            wid = pair;
            int newC = (C_lvl + 1) >> 1;
            lvl++;
            lbase = lnext;
            lnext += newC;
            C_lvl = newC;
            if (C_lvl == 1) { i_am_root = true; break; }
        }
        if (i_am_root) {
            const ull* res = scrrow + (long long)lbase * RS;
            for (int w = lane; w < k; w += 32) {
                ull u = res[w];
                outv[(long long)row * k + w] = keyf((uint)(u >> 32));
                outi[(long long)row * k + w] = (long long)(uint)(u & 0xffffffffu);
            }
        }
    }
}

// ------------------------------ host side ---------------------------------
static torch::Tensor g_cand, g_cnt, g_cnt2;

template <int M1, int E>
static void launch(const float* xp, ull* candp, int* cntp, int* cnt2p, float* vp,
                   long long* ip, int B, int n, int k, int K2, int C, int CH,
                   cudaStream_t stream) {
    int RS = K2 < 8 ? 8 : K2;
    (void)RS;
    dim3 grid((unsigned)(B * (int64_t)C));
    topk_kernel<M1, E><<<grid, 256, 0, stream>>>(xp, candp, cntp, cnt2p, vp, ip,
                                                    n, k, K2, C, CH);
}

std::tuple<torch::Tensor, torch::Tensor> topk_run(const torch::Tensor& x,
                                                  int64_t k_) {
    TORCH_CHECK(x.is_cuda() && x.dim() == 2 && x.scalar_type() == at::kFloat, "bad x");
    auto xc = x.is_contiguous() ? x : x.contiguous();
    int64_t B = xc.size(0), n = xc.size(1);
    int k = (int)k_;
    TORCH_CHECK(k >= 1 && (int64_t)k <= n, "bad k");
    int K2 = 1; while (K2 < k) K2 <<= 1;
    TORCH_CHECK(K2 <= 64, "k too large for fast path");

    auto opts = xc.options();
    auto vals = at::empty({B, k}, opts);
    auto idxs = at::empty({B, k}, opts.dtype(at::kLong));

    static thread_local int64_t cb = -1, cn = -1, ck = -1;
    static thread_local int cE, cC, cM1;
    int64_t sig = (B << 32) ^ (n << 16) ^ k;
    bool hit = (sig == ((cb << 32) ^ (cn << 16) ^ ck));
    int E, CH, C, M1;
    if (hit) {
        E = cE; C = cC; M1 = cM1; CH = E * 256;
    } else {
        E = 8; CH = 2048;
        if (n <= 4096 && K2 == 1) { E = 16; CH = 4096; }
        C = (int)((n + CH - 1) / CH);
        M1 = CH / 256 < K2 ? CH / 256 : K2;  // min(E, K2)
        M1 = M1 < 2 ? M1 : (M1 > 8 ? 8 : M1);    // cap run at 8
        cb = B; cn = n; ck = k; cE = E; cC = C; cM1 = M1;
    }

    int RS_ = K2 < 8 ? 8 : K2;
    int SEC_ = C < 64 ? C : 64;
    int64_t need = (C > 1) ? (int64_t)B * (C + SEC_) * RS_ : 64;
    if (!g_cand.defined() || g_cand.numel() < need) {
        g_cand = at::empty({std::max<int64_t>(need, 1 << 20)}, opts.dtype(at::kLong));
    }
    if (!g_cnt.defined() || g_cnt.numel() < B) {
        g_cnt = at::zeros({std::max<int64_t>(B, 4096)}, opts.dtype(at::kInt));
    }
    if (!g_cnt2.defined()) {
        g_cnt2 = at::zeros({512 * 128}, opts.dtype(at::kInt));
    }

    const float* xp = xc.data_ptr<float>();
    ull* candp = reinterpret_cast<ull*>(g_cand.data_ptr<int64_t>());
    int* cntp = g_cnt.data_ptr<int>();
    int* cnt2p = g_cnt2.data_ptr<int>();
    float* vp = vals.data_ptr<float>();
    long long* ip = reinterpret_cast<long long*>(idxs.data_ptr<int64_t>());
    auto stream = at::cuda::getCurrentCUDAStream();

#define LAUNCH(M1_, E_) launch<M1_, E_>(xp, candp, cntp, cnt2p, vp, ip, (int)B, \
                                        (int)n, k, K2, C, CH, stream)
    switch (E * 100 + M1) {
        case 801:  LAUNCH(1, 8);  break;
        case 802:  LAUNCH(2, 8);  break;
        case 804:  LAUNCH(4, 8);  break;
        case 808:  LAUNCH(8, 8);  break;
        case 1601: LAUNCH(1, 16); break;
        case 1602: LAUNCH(2, 16); break;
        case 1604: LAUNCH(4, 16); break;
        default: TORCH_CHECK(false, "no kernel for config");
    }
#undef LAUNCH
    return {vals, idxs};
}
"""

CUDA_SRC = _CUDA_TEMPLATE.replace("__MERGE_FNS__", _MERGE_FNS)

CPP_SRC = "#include <torch/extension.h>\nstd::tuple<torch::Tensor, torch::Tensor> topk_run(const torch::Tensor& x, int64_t k_);"

_mod = load_inline(
    name="topk_bitonic_v5",
    cpp_sources=CPP_SRC,
    cuda_sources=CUDA_SRC,
    functions=["topk_run"],
    extra_cuda_cflags=["-O3"],
    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.register_buffer("_dummy", torch.zeros(1))
        self._fn = _mod.topk_run

    def forward(self, x: torch.Tensor):
        return self._fn(x, self.k)


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]