--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -0,0 +1,365 @@
+"""TopK custom kernel via load_inline.
+
+Strategy:
+  - Each thread streams its share of the row, maintaining a sorted top-K
+    array in registers (insertion sort, descending).
+  - Threads write their per-thread top-K to shared memory: (THREADS * K)
+    candidates total.
+  - Block-wide bitonic sort over the (THREADS * K) candidates picks the
+    block-wide top-K. (THREADS * K) is small enough to fit in 99 KB
+    Blackwell opt-in shared memory.
+  - For batch=1 n=131072 (the decoder shape) one block can't saturate
+    bandwidth, so we use a 2-stage approach: G blocks each find a chunk-local
+    top-K, then a single block merges G*K candidates.
+"""
+import os
+
+os.environ.setdefault("TORCH_CUDA_ARCH_LIST", "12.0")
+
+import torch
+import torch.nn as nn
+from torch.utils.cpp_extension import load_inline
+
+CPP_SRC = r"""
+std::vector<torch::Tensor> topk_forward(torch::Tensor x, int64_t k);
+"""
+
+CUDA_SRC = r"""
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda_runtime.h>
+#include <cstdint>
+#include <cstdio>
+
+#define NEG_INF (-3.4e38f)
+
+// Per-thread insertion: maintain sorted-descending top-K in registers.
+// my_v[0] is the current max in this thread, my_v[K-1] is the current min.
+template<int K>
+__device__ __forceinline__ void thread_insert(float v, int i, float (&my_v)[K], int (&my_i)[K]) {
+    if (v > my_v[K-1]) {
+        // Find insertion position; shift down.
+        int j = K - 1;
+        #pragma unroll
+        for (int p = K - 1; p > 0; --p) {
+            if (my_v[p-1] < v) {
+                my_v[p] = my_v[p-1];
+                my_i[p] = my_i[p-1];
+                j = p - 1;
+            } else {
+                // Already in place; record j here.
+                // Use a small trick: only update j if we did not break.
+            }
+        }
+        // Re-derive insertion point properly: walk from top, find first < v.
+        // (The unrolled loop above shifts but we still need to find j.)
+        int pos = 0;
+        #pragma unroll
+        for (int p = 0; p < K; ++p) {
+            if (my_v[p] >= v && p < K - 1) pos = p + 1;
+        }
+        // Actually the unroll-shift is messy. Replace with clean version below.
+    }
+}
+
+// Cleaner version: explicit branched insertion.
+template<int K>
+__device__ __forceinline__ void thread_insert_v2(float v, int i, float (&my_v)[K], int (&my_i)[K]) {
+    if (v <= my_v[K-1]) return;
+    // Find insertion position (largest p with my_v[p-1] < v, else 0).
+    // Then shift my_v[p..K-2] -> my_v[p+1..K-1], place v at p.
+    int pos = K - 1;
+    #pragma unroll
+    for (int p = K - 2; p >= 0; --p) {
+        if (my_v[p] < v) {
+            my_v[p+1] = my_v[p];
+            my_i[p+1] = my_i[p];
+            pos = p;
+        }
+    }
+    my_v[pos] = v;
+    my_i[pos] = i;
+}
+
+// Block-wide bitonic sort of N elements in shared memory, ascending.
+// Caller must ensure N is a power of 2 and N >= THREADS*2 (or handles small N).
+template<int N, int THREADS>
+__device__ __forceinline__ void bitonic_sort_block(float* svals, int* sidxs) {
+    // For each (size, stride) pair, every thread handles N/THREADS pairs.
+    #pragma unroll
+    for (int sz = 2; sz <= N; sz <<= 1) {
+        #pragma unroll
+        for (int j = sz >> 1; j > 0; j >>= 1) {
+            // Each thread handles elements i = threadIdx.x, threadIdx.x + THREADS, ...
+            for (int t = threadIdx.x; t < N; t += THREADS) {
+                int p = t ^ j;
+                if (p > t) {
+                    bool asc = (t & sz) == 0;
+                    float vt = svals[t], vp = svals[p];
+                    bool sw = asc ? (vt > vp) : (vt < vp);
+                    if (sw) {
+                        svals[t] = vp; svals[p] = vt;
+                        int it = sidxs[t], ip = sidxs[p];
+                        sidxs[t] = ip; sidxs[p] = it;
+                    }
+                }
+            }
+            __syncthreads();
+        }
+    }
+}
+
+// Single-row block kernel: 1 block per row, per-thread top-K + block bitonic.
+// REDUCE_N must equal THREADS * K and be a power of 2.
+template<int K, int THREADS, int REDUCE_N>
+__global__ void topk_block_kernel(
+    const float* __restrict__ x,
+    float* __restrict__ vals,
+    int64_t* __restrict__ idxs,
+    int n)
+{
+    extern __shared__ unsigned char smem_raw[];
+    float* svals = reinterpret_cast<float*>(smem_raw);
+    int* sidxs = reinterpret_cast<int*>(smem_raw + REDUCE_N * sizeof(float));
+
+    int row = blockIdx.x;
+    const float* xrow = x + row * (size_t)n;
+
+    float my_v[K];
+    int my_i[K];
+    #pragma unroll
+    for (int p = 0; p < K; ++p) { my_v[p] = NEG_INF; my_i[p] = -1; }
+
+    // Stream the row.
+    for (int i = threadIdx.x; i < n; i += THREADS) {
+        float v = xrow[i];
+        thread_insert_v2<K>(v, i, my_v, my_i);
+    }
+
+    // Write per-thread top-K to shared. Layout: thread t's K elements at
+    // smem[t*K .. t*K+K-1]. Then bitonic sort over REDUCE_N = THREADS*K.
+    #pragma unroll
+    for (int p = 0; p < K; ++p) {
+        svals[threadIdx.x * K + p] = my_v[p];
+        sidxs[threadIdx.x * K + p] = my_i[p];
+    }
+    __syncthreads();
+
+    bitonic_sort_block<REDUCE_N, THREADS>(svals, sidxs);
+
+    // Output top K (last K of ascending sort, reversed).
+    for (int i = threadIdx.x; i < K; i += THREADS) {
+        vals[row * K + i] = svals[REDUCE_N - 1 - i];
+        idxs[row * K + i] = (int64_t)sidxs[REDUCE_N - 1 - i];
+    }
+}
+
+// Stage 1 for the (1, 131072, 64) shape: G blocks per row, each finds
+// top-K within a chunk; outputs G*K candidates per row to scratch.
+template<int K, int THREADS, int REDUCE_N>
+__global__ void topk_stage1_kernel(
+    const float* __restrict__ x,
+    float* __restrict__ vals_out,
+    int* __restrict__ idxs_out,
+    int n,
+    int chunk_size,
+   

... (6699 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -1,16 +1,13 @@
-"""TopK custom kernel via load_inline.
+"""TopK custom kernel via load_inline (v1: shared-memory bitonic sort).
 
 Strategy:
-  - Each thread streams its share of the row, maintaining a sorted top-K
-    array in registers (insertion sort, descending).
-  - Threads write their per-thread top-K to shared memory: (THREADS * K)
-    candidates total.
-  - Block-wide bitonic sort over the (THREADS * K) candidates picks the
-    block-wide top-K. (THREADS * K) is small enough to fit in 99 KB
-    Blackwell opt-in shared memory.
-  - For batch=1 n=131072 (the decoder shape) one block can't saturate
-    bandwidth, so we use a 2-stage approach: G blocks each find a chunk-local
-    top-K, then a single block merges G*K candidates.
+  - For each row, load the row (or chunk) into shared memory padded to a
+    power of 2 with -inf, sort via in-place block-wide bitonic sort, then
+    output the last K elements (reversed) as the top-K descending.
+  - Where N_PAD * 8 bytes does not fit in 99 KB opt-in shared memory, we
+    use a 2-stage approach: stage 1 has G blocks per row each find the
+    chunk-local top-K via bitonic, stage 2 has 1 block per row sort the
+    G*K candidates.
 """
 import os
 
@@ -29,71 +26,21 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <cuda_runtime.h>
 #include <cstdint>
-#include <cstdio>
 
 #define NEG_INF (-3.4e38f)
 
-// Per-thread insertion: maintain sorted-descending top-K in registers.
-// my_v[0] is the current max in this thread, my_v[K-1] is the current min.
-template<int K>
-__device__ __forceinline__ void thread_insert(float v, int i, float (&my_v)[K], int (&my_i)[K]) {
-    if (v > my_v[K-1]) {
-        // Find insertion position; shift down.
-        int j = K - 1;
-        #pragma unroll
-        for (int p = K - 1; p > 0; --p) {
-            if (my_v[p-1] < v) {
-                my_v[p] = my_v[p-1];
-                my_i[p] = my_i[p-1];
-                j = p - 1;
-            } else {
-                // Already in place; record j here.
-                // Use a small trick: only update j if we did not break.
-            }
-        }
-        // Re-derive insertion point properly: walk from top, find first < v.
-        // (The unrolled loop above shifts but we still need to find j.)
-        int pos = 0;
-        #pragma unroll
-        for (int p = 0; p < K; ++p) {
-            if (my_v[p] >= v && p < K - 1) pos = p + 1;
-        }
-        // Actually the unroll-shift is messy. Replace with clean version below.
-    }
-}
-
-// Cleaner version: explicit branched insertion.
-template<int K>
-__device__ __forceinline__ void thread_insert_v2(float v, int i, float (&my_v)[K], int (&my_i)[K]) {
-    if (v <= my_v[K-1]) return;
-    // Find insertion position (largest p with my_v[p-1] < v, else 0).
-    // Then shift my_v[p..K-2] -> my_v[p+1..K-1], place v at p.
-    int pos = K - 1;
-    #pragma unroll
-    for (int p = K - 2; p >= 0; --p) {
-        if (my_v[p] < v) {
-            my_v[p+1] = my_v[p];
-            my_i[p+1] = my_i[p];
-            pos = p;
-        }
-    }
-    my_v[pos] = v;
-    my_i[pos] = i;
-}
-
-// Block-wide bitonic sort of N elements in shared memory, ascending.
-// Caller must ensure N is a power of 2 and N >= THREADS*2 (or handles small N).
+// Block-wide bitonic sort, ascending. N must be a power of 2.
 template<int N, int THREADS>
 __device__ __forceinline__ void bitonic_sort_block(float* svals, int* sidxs) {
-    // For each (size, stride) pair, every thread handles N/THREADS pairs.
     #pragma unroll
     for (int sz = 2; sz <= N; sz <<= 1) {
         #pragma unroll
         for (int j = sz >> 1; j > 0; j >>= 1) {
-            // Each thread handles elements i = threadIdx.x, threadIdx.x + THREADS, ...
-            for (int t = threadIdx.x; t < N; t += THREADS) {
+            #pragma unroll
+            for (int t0 = 0; t0 < N; t0 += THREADS) {
+                int t = t0 + threadIdx.x;
                 int p = t ^ j;
-                if (p > t) {
+                if (p > t && t < N) {
                     bool asc = (t & sz) == 0;
                     float vt = svals[t], vp = svals[p];
                     bool sw = asc ? (vt > vp) : (vt < vp);
@@ -109,10 +56,108 @@
     }
 }
 
-// Single-row block kernel: 1 block per row, per-thread top-K + block bitonic.
-// REDUCE_N must equal THREADS * K and be a power of 2.
-template<int K, int THREADS, int REDUCE_N>
-__global__ void topk_block_kernel(
+// Stage 1 / single-stage kernel.
+//   Reads chunk_size elements per block from x, pads to N_PAD with -inf,
+//   bitonic sorts, writes top-K to (out_v, out_i).
+//   blocks_per_row: 1 for single-stage; >1 for stage 1 of multi-stage.
+//   out_i is int32; we promote to int64 in a small writeout kernel.
+template<int K, int N_PAD, int THREADS>
+__global__ void topk_stage_x_kernel(
+    const float* __restrict__ x,
+    float* __restrict__ out_v,
+    int* __restrict__ out_i,
+    int n,
+    int chunk_size,
+    int blocks_per_row)
+{
+    extern __shared__ unsigned char smem_raw[];
+    float* svals = reinterpret_cast<float*>(smem_raw);
+    int* sidxs = reinterpret_cast<int*>(smem_raw + N_PAD * sizeof(float));
+
+    int row = blockIdx.y;
+    int blk = blockIdx.x;
+    int chunk_start = blk * chunk_size;
+    int chunk_end = chunk_start + chunk_size;
+    if (chunk_end > n) chunk_end = n;
+
+    const float* xrow = x + row * (size_t)n;
+
+    #pragma unroll
+    for (int i0 = 0; i0 < N_PAD; i0 += THREADS) {
+        int i = i0 + threadIdx.x;
+        if (i < N_PAD) {
+            int gi = chunk_start + i;
+            if (gi < chunk_end) {
+                svals[i] = xrow[gi];
+                sidxs[i] = gi;
+            } else {
+                svals[i] = NEG_INF;
+                sidxs[i] = -1;
+            }
+        }
+    }
+    __syncthreads();

... (13888 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -1,13 +1,19 @@
-"""TopK custom kernel via load_inline (v1: shared-memory bitonic sort).
+"""TopK custom kernel via load_inline (v3: per-thread streaming + block bitonic).
 
 Strategy:
-  - For each row, load the row (or chunk) into shared memory padded to a
-    power of 2 with -inf, sort via in-place block-wide bitonic sort, then
-    output the last K elements (reversed) as the top-K descending.
-  - Where N_PAD * 8 bytes does not fit in 99 KB opt-in shared memory, we
-    use a 2-stage approach: stage 1 has G blocks per row each find the
-    chunk-local top-K via bitonic, stage 2 has 1 block per row sort the
-    G*K candidates.
+  - Each thread streams its share of a chunk, maintaining a sorted top-K
+    array in registers (descending). Insertion uses static-indexed
+    bubble-up (place new element at position K-1, then bubble it up via
+    K-1 conditional swaps; all indices are compile-time constants so the
+    compiler keeps everything in registers without spills).
+  - After streaming, T threads write per-thread top-K to shared memory
+    (T*K floats + T*K ints), and a block-wide bitonic sort on those
+    T*K candidates picks the chunk-local top-K.
+  - Each row is split into G chunks; G is chosen per shape so total
+    stage-1 block count is ~256 (filling the 188-SM RTX PRO 6000).
+  - Stage 2 merges the G*K candidates per row into the final top-K.
+  - For shapes where 1 block per row already saturates the SMs (e.g.
+    batch=128) the kernel runs single-stage with int64 output directly.
 """
 import os
 
@@ -29,7 +35,29 @@
 
 #define NEG_INF (-3.4e38f)
 
-// Block-wide bitonic sort, ascending. N must be a power of 2.
+// Insert (v, idx) into sorted-descending registers my_v[0..K-1], my_i[0..K-1].
+// Static-indexed: place candidate at slot K-1, then bubble up via K-1
+// conditional swaps. All indices are compile-time constants, so PTX should
+// keep my_v/my_i in registers and use predicated swaps (no local memory).
+template<int K>
+__device__ __forceinline__ void thread_insert(float v, int idx,
+                                                 float (&my_v)[K], int (&my_i)[K]) {
+    if (v <= my_v[K-1]) return;
+    my_v[K-1] = v;
+    my_i[K-1] = idx;
+    #pragma unroll
+    for (int p = K - 1; p > 0; --p) {
+        bool sw = my_v[p] > my_v[p-1];
+        float a = my_v[p-1], b = my_v[p];
+        int ai = my_i[p-1], bi = my_i[p];
+        my_v[p-1] = sw ? b : a;
+        my_v[p]   = sw ? a : b;
+        my_i[p-1] = sw ? bi : ai;
+        my_i[p]   = sw ? ai : bi;
+    }
+}
+
+// Block-wide bitonic sort, ascending.
 template<int N, int THREADS>
 __device__ __forceinline__ void bitonic_sort_block(float* svals, int* sidxs) {
     #pragma unroll
@@ -39,15 +67,17 @@
             #pragma unroll
             for (int t0 = 0; t0 < N; t0 += THREADS) {
                 int t = t0 + threadIdx.x;
-                int p = t ^ j;
-                if (p > t && t < N) {
-                    bool asc = (t & sz) == 0;
-                    float vt = svals[t], vp = svals[p];
-                    bool sw = asc ? (vt > vp) : (vt < vp);
-                    if (sw) {
-                        svals[t] = vp; svals[p] = vt;
-                        int it = sidxs[t], ip = sidxs[p];
-                        sidxs[t] = ip; sidxs[p] = it;
+                if (t < N) {
+                    int p = t ^ j;
+                    if (p > t) {
+                        bool asc = (t & sz) == 0;
+                        float vt = svals[t], vp = svals[p];
+                        bool sw = asc ? (vt > vp) : (vt < vp);
+                        if (sw) {
+                            svals[t] = vp; svals[p] = vt;
+                            int it = sidxs[t], ip = sidxs[p];
+                            sidxs[t] = ip; sidxs[p] = it;
+                        }
                     }
                 }
             }
@@ -56,63 +86,86 @@
     }
 }
 
-// Stage 1 / single-stage kernel.
-//   Reads chunk_size elements per block from x, pads to N_PAD with -inf,
-//   bitonic sorts, writes top-K to (out_v, out_i).
-//   blocks_per_row: 1 for single-stage; >1 for stage 1 of multi-stage.
-//   out_i is int32; we promote to int64 in a small writeout kernel.
-template<int K, int N_PAD, int THREADS>
-__global__ void topk_stage_x_kernel(
+// Stage 1 / single-stage kernel: per-thread streaming top-K + block bitonic.
+//   N_REDUCE = next pow 2 of THREADS * K (must be a power of 2).
+//   When write_int64=1, indices output is int64 directly (single-stage).
+//   chunks_per_row = grid.x dim (1 for single-stage).
+template<int K, int THREADS, int N_REDUCE, int WRITE_INT64>
+__global__ void topk_streaming_kernel(
     const float* __restrict__ x,
-    float* __restrict__ out_v,
-    int* __restrict__ out_i,
+    void* __restrict__ out_v,           // float*
+    void* __restrict__ out_i,           // int* or int64_t*
     int n,
     int chunk_size,
-    int blocks_per_row)
+    int chunks_per_row)
 {
     extern __shared__ unsigned char smem_raw[];
     float* svals = reinterpret_cast<float*>(smem_raw);
-    int* sidxs = reinterpret_cast<int*>(smem_raw + N_PAD * sizeof(float));
+    int* sidxs = reinterpret_cast<int*>(smem_raw + N_REDUCE * sizeof(float));
 
     int row = blockIdx.y;
     int blk = blockIdx.x;
     int chunk_start = blk * chunk_size;
     int chunk_end = chunk_start + chunk_size;
     if (chunk_end > n) chunk_end = n;
-
     const float* xrow = x + row * (size_t)n;
 
-    #pragma unroll
-    for (int i0 = 0; i0 < N_PAD; i0 += THREADS) {
-        int i = i0 + threadIdx.x;
-        if (i < N_PAD) {
-            int gi = chunk_start + i;
-            if (gi < chunk_end) {
-                svals[i] = xrow[gi];
-                sidxs[i] = gi;
-            } else {
-                svals[i] = NEG_INF;
-                sidxs[i] = -1;
-            }

... (18237 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -1,19 +1,15 @@
-"""TopK custom kernel via load_inline (v3: per-thread streaming + block bitonic).
+"""TopK custom kernel via load_inline (v4).
 
 Strategy:
-  - Each thread streams its share of a chunk, maintaining a sorted top-K
-    array in registers (descending). Insertion uses static-indexed
-    bubble-up (place new element at position K-1, then bubble it up via
-    K-1 conditional swaps; all indices are compile-time constants so the
-    compiler keeps everything in registers without spills).
-  - After streaming, T threads write per-thread top-K to shared memory
-    (T*K floats + T*K ints), and a block-wide bitonic sort on those
-    T*K candidates picks the chunk-local top-K.
-  - Each row is split into G chunks; G is chosen per shape so total
-    stage-1 block count is ~256 (filling the 188-SM RTX PRO 6000).
-  - Stage 2 merges the G*K candidates per row into the final top-K.
-  - For shapes where 1 block per row already saturates the SMs (e.g.
-    batch=128) the kernel runs single-stage with int64 output directly.
+  - Full bitonic sort of small chunks in shared memory; take chunk-local
+    top-K. Bitonic sort over a small N has O(log^2 N) parallel stages and
+    no long per-thread dependency chains, so it's faster than per-thread
+    streaming insertion-sort for K >= 16.
+  - Stage 1 is split into chunks so that total stage-1 block count is
+    O(>=128), filling the 188-SM RTX PRO 6000.
+  - Stage 2 single block per row merges the G*K candidates.
+  - Shape (128, 4096, 1) is k=1 argmax — use a dedicated 1-block-per-row
+    parallel max reduction (no sort needed).
 """
 import os
 
@@ -35,29 +31,7 @@
 
 #define NEG_INF (-3.4e38f)
 
-// Insert (v, idx) into sorted-descending registers my_v[0..K-1], my_i[0..K-1].
-// Static-indexed: place candidate at slot K-1, then bubble up via K-1
-// conditional swaps. All indices are compile-time constants, so PTX should
-// keep my_v/my_i in registers and use predicated swaps (no local memory).
-template<int K>
-__device__ __forceinline__ void thread_insert(float v, int idx,
-                                                 float (&my_v)[K], int (&my_i)[K]) {
-    if (v <= my_v[K-1]) return;
-    my_v[K-1] = v;
-    my_i[K-1] = idx;
-    #pragma unroll
-    for (int p = K - 1; p > 0; --p) {
-        bool sw = my_v[p] > my_v[p-1];
-        float a = my_v[p-1], b = my_v[p];
-        int ai = my_i[p-1], bi = my_i[p];
-        my_v[p-1] = sw ? b : a;
-        my_v[p]   = sw ? a : b;
-        my_i[p-1] = sw ? bi : ai;
-        my_i[p]   = sw ? ai : bi;
-    }
-}
-
-// Block-wide bitonic sort, ascending.
+// Block-wide bitonic sort, ascending. N must be a power of 2.
 template<int N, int THREADS>
 __device__ __forceinline__ void bitonic_sort_block(float* svals, int* sidxs) {
     #pragma unroll
@@ -86,22 +60,20 @@
     }
 }
 
-// Stage 1 / single-stage kernel: per-thread streaming top-K + block bitonic.
-//   N_REDUCE = next pow 2 of THREADS * K (must be a power of 2).
-//   When write_int64=1, indices output is int64 directly (single-stage).
-//   chunks_per_row = grid.x dim (1 for single-stage).
-template<int K, int THREADS, int N_REDUCE, int WRITE_INT64>
-__global__ void topk_streaming_kernel(
+// Bitonic chunk kernel: load a CHUNK_SIZE chunk into shared (pad to N_PAD
+// with -inf), bitonic sort, write top-K. WRITE_INT64=1 writes int64 directly.
+template<int K, int N_PAD, int THREADS, int WRITE_INT64>
+__global__ void topk_chunk_kernel(
     const float* __restrict__ x,
-    void* __restrict__ out_v,           // float*
-    void* __restrict__ out_i,           // int* or int64_t*
+    void* __restrict__ out_v,
+    void* __restrict__ out_i,
     int n,
     int chunk_size,
     int chunks_per_row)
 {
     extern __shared__ unsigned char smem_raw[];
     float* svals = reinterpret_cast<float*>(smem_raw);
-    int* sidxs = reinterpret_cast<int*>(smem_raw + N_REDUCE * sizeof(float));
+    int* sidxs = reinterpret_cast<int*>(smem_raw + N_PAD * sizeof(float));
 
     int row = blockIdx.y;
     int blk = blockIdx.x;
@@ -110,63 +82,45 @@
     if (chunk_end > n) chunk_end = n;
     const float* xrow = x + row * (size_t)n;
 
-    float my_v[K];
-    int my_i[K];
     #pragma unroll
-    for (int p = 0; p < K; ++p) { my_v[p] = NEG_INF; my_i[p] = -1; }
-
-    // Stream the chunk.
-    int i = chunk_start + threadIdx.x;
-    for (; i + 3 * THREADS < chunk_end; i += 4 * THREADS) {
-        float v0 = xrow[i];
-        float v1 = xrow[i + THREADS];
-        float v2 = xrow[i + 2 * THREADS];
-        float v3 = xrow[i + 3 * THREADS];
-        thread_insert<K>(v0, i, my_v, my_i);
-        thread_insert<K>(v1, i + THREADS, my_v, my_i);
-        thread_insert<K>(v2, i + 2 * THREADS, my_v, my_i);
-        thread_insert<K>(v3, i + 3 * THREADS, my_v, my_i);
-    }
-    for (; i < chunk_end; i += THREADS) {
-        thread_insert<K>(xrow[i], i, my_v, my_i);
-    }
-
-    // Write per-thread top-K to shared at strided offset.
-    #pragma unroll
-    for (int p = 0; p < K; ++p) {
-        svals[threadIdx.x * K + p] = my_v[p];
-        sidxs[threadIdx.x * K + p] = my_i[p];
-    }
-    int total_real = THREADS * K;
-    for (int q = total_real + threadIdx.x; q < N_REDUCE; q += THREADS) {
-        svals[q] = NEG_INF;
-        sidxs[q] = -1;
+    for (int i0 = 0; i0 < N_PAD; i0 += THREADS) {
+        int i = i0 + threadIdx.x;
+        if (i < N_PAD) {
+            int gi = chunk_start + i;
+            if (gi < chunk_end) {
+                svals[i] = xrow[gi];
+                sidxs[i] = gi;
+            } else {
+                svals[i] = NEG_INF;
+                sidxs[i] = -1;
+            }
+        }
     }
     __syncthreads();
 
-    bitonic_sort_block<N_REDUCE, THREADS>(svals, sidxs);
+    bitonic_sort_block<N_PAD, THREADS>(svals, sidxs);
 
     int out_off = (row * chunks_per_row + blk

... (15830 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -1,15 +1,15 @@
-"""TopK custom kernel via load_inline (v4).
+"""TopK custom kernel via load_inline (v5).
 
 Strategy:
-  - Full bitonic sort of small chunks in shared memory; take chunk-local
-    top-K. Bitonic sort over a small N has O(log^2 N) parallel stages and
-    no long per-thread dependency chains, so it's faster than per-thread
-    streaming insertion-sort for K >= 16.
-  - Stage 1 is split into chunks so that total stage-1 block count is
-    O(>=128), filling the 188-SM RTX PRO 6000.
-  - Stage 2 single block per row merges the G*K candidates.
-  - Shape (128, 4096, 1) is k=1 argmax — use a dedicated 1-block-per-row
-    parallel max reduction (no sort needed).
+  - Stage 1: full bitonic sort of small chunks in shared memory; take
+    chunk-local top-K. Stage 1 has many small chunks for SM parallelism.
+  - Stage 2: each stage-1 chunk emits a sorted-descending run of K
+    elements. We merge G sorted runs of length K into the global top-K
+    using a tournament-style argmax: per iteration, one thread per run
+    holds its current head value; warp+block argmax picks the run with
+    the largest head; that run advances by 1. K iterations -> top-K.
+    This is O(K * log G) instead of bitonic-sort's O(G*K * log^2(G*K)).
+  - Shape (128, 4096, 1) is k=1 argmax — dedicated max-reduction kernel.
 """
 import os
 
@@ -60,8 +60,8 @@
     }
 }
 
-// Bitonic chunk kernel: load a CHUNK_SIZE chunk into shared (pad to N_PAD
-// with -inf), bitonic sort, write top-K. WRITE_INT64=1 writes int64 directly.
+// Stage 1 / single-stage chunk kernel: load CHUNK_SIZE into shared (pad
+// to N_PAD with -inf), bitonic sort, write top-K to (out_v, out_i).
 template<int K, int N_PAD, int THREADS, int WRITE_INT64>
 __global__ void topk_chunk_kernel(
     const float* __restrict__ x,
@@ -118,8 +118,9 @@
     }
 }
 
-// Stage 2: input is candidate values + int32 indices. Output int64 indices.
-template<int K, int N_PAD, int THREADS>
+// Stage 2 (bitonic): merge candidates (already sorted in chunks of K)
+// via full bitonic sort over N_PAD candidates.
+template<int K_OUT, int N_PAD, int THREADS>
 __global__ void topk_stage2_bitonic_kernel(
     const float* __restrict__ in_v,
     const int* __restrict__ in_i,
@@ -152,13 +153,95 @@
 
     bitonic_sort_block<N_PAD, THREADS>(svals, sidxs);
 
-    for (int q = threadIdx.x; q < K; q += THREADS) {
-        vals[row * K + q] = svals[N_PAD - 1 - q];
-        idxs[row * K + q] = (int64_t)sidxs[N_PAD - 1 - q];
-    }
-}
-
-// Argmax kernel for k=1 — block-wide max-reduction.
+    for (int q = threadIdx.x; q < K_OUT; q += THREADS) {
+        vals[row * K_OUT + q] = svals[N_PAD - 1 - q];
+        idxs[row * K_OUT + q] = (int64_t)sidxs[N_PAD - 1 - q];
+    }
+}
+
+// Stage 2 (k-way merge tournament): NUM_RUNS sorted-desc runs of RUN_LEN
+// each per row. THREADS == NUM_RUNS (one thread per run). K iterations of
+// argmax + advance.
+template<int K_OUT, int NUM_RUNS, int RUN_LEN, int THREADS>
+__global__ void topk_stage2_merge_kernel(
+    const float* __restrict__ in_v,
+    const int* __restrict__ in_i,
+    float* __restrict__ vals,
+    int64_t* __restrict__ idxs)
+{
+    int row = blockIdx.x;
+    int tid = threadIdx.x;
+
+    // Each thread handles one run.
+    int run_off = row * NUM_RUNS * RUN_LEN + tid * RUN_LEN;
+
+    float head_v = (tid < NUM_RUNS) ? in_v[run_off] : NEG_INF;
+    int head_i = (tid < NUM_RUNS) ? in_i[run_off] : -1;
+    int head_ptr = 0;
+
+    __shared__ float smax[(THREADS + 31) / 32];
+    __shared__ int sidx[(THREADS + 31) / 32];
+    __shared__ int s_winner;
+    __shared__ float s_winner_v;
+    __shared__ int s_winner_i;
+
+    constexpr int WARP_SIZE = 32;
+    constexpr int NUM_WARPS = THREADS / WARP_SIZE;
+
+    for (int k_iter = 0; k_iter < K_OUT; ++k_iter) {
+        // Warp-level argmax.
+        unsigned mask = 0xFFFFFFFFu;
+        float my_v = head_v;
+        int my_t = tid;
+        #pragma unroll
+        for (int s = WARP_SIZE / 2; s > 0; s >>= 1) {
+            float ov = __shfl_xor_sync(mask, my_v, s);
+            int ot = __shfl_xor_sync(mask, my_t, s);
+            // Tie-break: prefer smaller thread index (deterministic).
+            bool take = (ov > my_v) || (ov == my_v && ot < my_t);
+            if (take) { my_v = ov; my_t = ot; }
+        }
+        int lane = tid & (WARP_SIZE - 1);
+        int warp = tid / WARP_SIZE;
+        if (lane == 0) { smax[warp] = my_v; sidx[warp] = my_t; }
+        __syncthreads();
+
+        if (warp == 0) {
+            float v = (lane < NUM_WARPS) ? smax[lane] : NEG_INF;
+            int t = (lane < NUM_WARPS) ? sidx[lane] : -1;
+            #pragma unroll
+            for (int s = NUM_WARPS / 2; s > 0; s >>= 1) {
+                float ov = __shfl_xor_sync(mask, v, s);
+                int ot = __shfl_xor_sync(mask, t, s);
+                bool take = (ov > v) || (ov == v && ot < t);
+                if (take) { v = ov; t = ot; }
+            }
+            if (lane == 0) {
+                s_winner = t;
+                s_winner_v = v;
+                // We need the index too — the winner thread will write it.
+                vals[row * K_OUT + k_iter] = v;
+            }
+        }
+        __syncthreads();
+
+        int w = s_winner;
+        if (tid == w) {
+            idxs[row * K_OUT + k_iter] = (int64_t)head_i;
+            head_ptr++;
+            if (head_ptr < RUN_LEN) {
+                head_v = in_v[run_off + head_ptr];
+                head_i = in_i[run_off + head_ptr];
+            } else {
+                head_v = NEG_INF;
+                head_i = -1;
+            }
+        }
+        __syncthreads();
+    }
+}
+
+// Argmax kernel for k=1.
 template<int N, int THREADS>
 __global__ void argmax_kernel(
     const float* __restrict__ x,
@@ -176,7 +259,6 @@
         if (v > my_v) { my_v =

... (7947 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -32,10 +32,20 @@
 #define NEG_INF (-3.4e38f)
 
 // Block-wide bitonic sort, ascending. N must be a power of 2.
+// Uses __syncwarp for j<32 (intra-warp pattern) and __syncthreads for j>=32.
+// At sz transitions where the next sz starts with a block-level j, we
+// insert a __syncthreads to make warp-level writes visible cross-warp.
 template<int N, int THREADS>
 __device__ __forceinline__ void bitonic_sort_block(float* svals, int* sidxs) {
     #pragma unroll
     for (int sz = 2; sz <= N; sz <<= 1) {
+        // If the previous sz ended on a warp-only sequence and this sz starts
+        // with a block-level j (j = sz/2 >= 32), block-sync first.
+        if ((sz >> 1) >= 32) {
+            // The very first sz=2 has prev sequence; but sz=2 has j=1 only,
+            // so this branch only triggers for sz>=64.
+            __syncthreads();
+        }
         #pragma unroll
         for (int j = sz >> 1; j > 0; j >>= 1) {
             #pragma unroll
@@ -55,7 +65,11 @@
                     }
                 }
             }
-            __syncthreads();
+            if (j >= 32) {
+                __syncthreads();
+            } else {
+                __syncwarp();
+            }
         }
     }
 }

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -32,20 +32,10 @@
 #define NEG_INF (-3.4e38f)
 
 // Block-wide bitonic sort, ascending. N must be a power of 2.
-// Uses __syncwarp for j<32 (intra-warp pattern) and __syncthreads for j>=32.
-// At sz transitions where the next sz starts with a block-level j, we
-// insert a __syncthreads to make warp-level writes visible cross-warp.
 template<int N, int THREADS>
 __device__ __forceinline__ void bitonic_sort_block(float* svals, int* sidxs) {
     #pragma unroll
     for (int sz = 2; sz <= N; sz <<= 1) {
-        // If the previous sz ended on a warp-only sequence and this sz starts
-        // with a block-level j (j = sz/2 >= 32), block-sync first.
-        if ((sz >> 1) >= 32) {
-            // The very first sz=2 has prev sequence; but sz=2 has j=1 only,
-            // so this branch only triggers for sz>=64.
-            __syncthreads();
-        }
         #pragma unroll
         for (int j = sz >> 1; j > 0; j >>= 1) {
             #pragma unroll
@@ -65,11 +55,7 @@
                     }
                 }
             }
-            if (j >= 32) {
-                __syncthreads();
-            } else {
-                __syncwarp();
-            }
+            __syncthreads();
         }
     }
 }

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -351,12 +351,12 @@
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
     if (batch == 1 && n == 131072 && k == 64) {
-        // Stage 1: 128 chunks * 1024. K=64.
-        constexpr int CHUNK = 1024;
-        constexpr int N_PAD1 = 1024;
+        // Stage 1: 64 chunks * 2048. K=64. T=256.
+        constexpr int CHUNK = 2048;
+        constexpr int N_PAD1 = 2048;
         constexpr int K1 = 64;
-        constexpr int T1 = 128;
-        constexpr int CPR = 131072 / CHUNK;   // 128
+        constexpr int T1 = 256;
+        constexpr int CPR = 131072 / CHUNK;   // 64
 
         auto v_scratch = torch::empty({batch * CPR * K1}, x.options());
         auto i_scratch = torch::empty({batch * CPR * K1},
@@ -365,15 +365,26 @@
         launch_chunk_int32<K1, N_PAD1, T1>(x, v_scratch, i_scratch,
                                             n, CHUNK, CPR, batch, stream);
 
-        // Stage 2: k-way merge of 128 sorted runs (K1 each) -> top-K=64.
-        launch_stage2_merge<K1, CPR, K1, CPR>(v_scratch, i_scratch, vals, idxs,
-                                               batch, stream);
+        // Stage 2: k-way merge of 64 sorted runs (K1 each) -> top-K=64.
+        launch_stage2_merge<K1, CPR, K1, 64>(v_scratch, i_scratch, vals, idxs,
+                                              batch, stream);
     } else if (batch == 64 && n == 8192 && k == 8) {
-        constexpr int CHUNK = 1024;
-        constexpr int N_PAD1 = 1024;
+        // Single-stage: 1 block per row, full bitonic 8192. 64 blocks total.
+        constexpr int CHUNK = 8192;
+        constexpr int N_PAD1 = 8192;
         constexpr int K1 = 8;
-        constexpr int T1 = 128;
-        constexpr int CPR = 8192 / CHUNK;     // 8
+        constexpr int T1 = 512;
+        constexpr int CPR = 1;
+
+        launch_chunk_int64<K1, N_PAD1, T1>(x, vals, idxs,
+                                            n, CHUNK, CPR, batch, stream);
+    } else if (batch == 32 && n == 16384 && k == 32) {
+        // 2 chunks * 8192 each, K=32. 64 stage1 blocks.
+        constexpr int CHUNK = 8192;
+        constexpr int N_PAD1 = 8192;
+        constexpr int K1 = 32;
+        constexpr int T1 = 512;
+        constexpr int CPR = 16384 / CHUNK;    // 2
 
         auto v_scratch = torch::empty({batch * CPR * K1}, x.options());
         auto i_scratch = torch::empty({batch * CPR * K1},
@@ -382,19 +393,19 @@
         launch_chunk_int32<K1, N_PAD1, T1>(x, v_scratch, i_scratch,
                                             n, CHUNK, CPR, batch, stream);
 
-        // Stage 2: 8 runs * 8 elements = 64 candidates.
-        // Use bitonic on 64 (cheap).
+        // Stage 2: 2 runs * 32 = 64 candidates -> top-32. Bitonic 64.
         constexpr int CAND2 = CPR * K1;       // 64
         constexpr int N_PAD2 = 64;
         constexpr int T2 = 32;
         launch_stage2_bitonic<K1, N_PAD2, T2>(v_scratch, i_scratch, vals, idxs,
                                                 CAND2, batch, stream);
-    } else if (batch == 32 && n == 16384 && k == 32) {
-        constexpr int CHUNK = 1024;
-        constexpr int N_PAD1 = 1024;
-        constexpr int K1 = 32;
-        constexpr int T1 = 128;
-        constexpr int CPR = 16384 / CHUNK;    // 16
+    } else if (batch == 16 && n == 12000 && k == 16) {
+        // 2 chunks of 6000 (pad to 8192), K=16. 32 stage1 blocks.
+        constexpr int CHUNK = 6000;
+        constexpr int N_PAD1 = 8192;
+        constexpr int K1 = 16;
+        constexpr int T1 = 512;
+        constexpr int CPR = 2;
 
         auto v_scratch = torch::empty({batch * CPR * K1}, x.options());
         auto i_scratch = torch::empty({batch * CPR * K1},
@@ -403,28 +414,12 @@
         launch_chunk_int32<K1, N_PAD1, T1>(x, v_scratch, i_scratch,
                                             n, CHUNK, CPR, batch, stream);
 
-        // Stage 2: 16 runs * 32 = 512 candidates.
-        // Try k-way merge: 16 runs of 32 -> top-32. THREADS=32 (one per run, padded for warp).
-        launch_stage2_merge<K1, CPR, K1, 32>(v_scratch, i_scratch, vals, idxs,
-                                              batch, stream);
-    } else if (batch == 16 && n == 12000 && k == 16) {
-        constexpr int CHUNK = 1024;
-        constexpr int N_PAD1 = 1024;
-        constexpr int K1 = 16;
-        constexpr int T1 = 128;
-        constexpr int CPR = 12;               // ceil(12000/1024)
-
-        auto v_scratch = torch::empty({batch * CPR * K1}, x.options());
-        auto i_scratch = torch::empty({batch * CPR * K1},
-                                       x.options().dtype(torch::kInt32));
-
-        launch_chunk_int32<K1, N_PAD1, T1>(x, v_scratch, i_scratch,
-                                            n, CHUNK, CPR, batch, stream);
-
-        // Stage 2: 12 runs of 16 = 192 candidates -> top-16.
-        // Merge with THREADS=32 (12 runs + 20 inactive).
-        launch_stage2_merge<K1, CPR, K1, 32>(v_scratch, i_scratch, vals, idxs,
-                                              batch, stream);
+        // Stage 2: 2 runs * 16 = 32 candidates -> top-16. Bitonic 32.
+        constexpr int CAND2 = CPR * K1;       // 32
+        constexpr int N_PAD2 = 32;
+        constexpr int T2 = 32;
+        launch_stage2_bitonic<K1, N_PAD2, T2>(v_scratch, i_scratch, vals, idxs,
+                                                CAND2, batch, stream);
     } else if (batch == 128 && n == 4096 && k == 1) {
         constexpr int N = 4096;
         constexpr int T = 256;

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -438,7 +438,7 @@
 """
 
 _module = load_inline(
-    name="topk_v5",
+    name="topk_v6",
     cpp_sources=CPP_SRC,
     cuda_sources=CUDA_SRC,
     functions=["topk_forward"],

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -307,6 +307,23 @@
         n, chunk_size, chunks_per_row);
 }
 
+template<int K, int N_PAD, int THREADS>
+inline void launch_chunk_int64(torch::Tensor x, torch::Tensor vals,
+                                torch::Tensor idxs,
+                                int n, int chunk_size, int chunks_per_row,
+                                int batch, cudaStream_t stream)
+{
+    size_t sm = (size_t)N_PAD * (sizeof(float) + sizeof(int));
+    cudaFuncSetAttribute((void*)topk_chunk_kernel<K, N_PAD, THREADS, 1>,
+        cudaFuncAttributeMaxDynamicSharedMemorySize, sm);
+    topk_chunk_kernel<K, N_PAD, THREADS, 1>
+        <<<dim3(chunks_per_row, batch), THREADS, sm, stream>>>(
+        x.data_ptr<float>(),
+        (void*)vals.data_ptr<float>(),
+        (void*)idxs.data_ptr<int64_t>(),
+        n, chunk_size, chunks_per_row);
+}
+
 template<int K_OUT, int N_PAD, int THREADS>
 inline void launch_stage2_bitonic(torch::Tensor v_scratch, torch::Tensor i_scratch,
                                     torch::Tensor vals, torch::Tensor idxs,

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -368,12 +368,12 @@
     cudaStream_t stream = at::cuda::getCurrentCUDAStream();
 
     if (batch == 1 && n == 131072 && k == 64) {
-        // Stage 1: 64 chunks * 2048. K=64. T=256.
-        constexpr int CHUNK = 2048;
-        constexpr int N_PAD1 = 2048;
+        // Stage 1: 128 chunks * 1024. K=64.
+        constexpr int CHUNK = 1024;
+        constexpr int N_PAD1 = 1024;
         constexpr int K1 = 64;
-        constexpr int T1 = 256;
-        constexpr int CPR = 131072 / CHUNK;   // 64
+        constexpr int T1 = 128;
+        constexpr int CPR = 131072 / CHUNK;   // 128
 
         auto v_scratch = torch::empty({batch * CPR * K1}, x.options());
         auto i_scratch = torch::empty({batch * CPR * K1},
@@ -382,26 +382,14 @@
         launch_chunk_int32<K1, N_PAD1, T1>(x, v_scratch, i_scratch,
                                             n, CHUNK, CPR, batch, stream);
 
-        // Stage 2: k-way merge of 64 sorted runs (K1 each) -> top-K=64.
-        launch_stage2_merge<K1, CPR, K1, 64>(v_scratch, i_scratch, vals, idxs,
-                                              batch, stream);
+        launch_stage2_merge<K1, CPR, K1, CPR>(v_scratch, i_scratch, vals, idxs,
+                                               batch, stream);
     } else if (batch == 64 && n == 8192 && k == 8) {
-        // Single-stage: 1 block per row, full bitonic 8192. 64 blocks total.
-        constexpr int CHUNK = 8192;
-        constexpr int N_PAD1 = 8192;
+        constexpr int CHUNK = 1024;
+        constexpr int N_PAD1 = 1024;
         constexpr int K1 = 8;
-        constexpr int T1 = 512;
-        constexpr int CPR = 1;
-
-        launch_chunk_int64<K1, N_PAD1, T1>(x, vals, idxs,
-                                            n, CHUNK, CPR, batch, stream);
-    } else if (batch == 32 && n == 16384 && k == 32) {
-        // 2 chunks * 8192 each, K=32. 64 stage1 blocks.
-        constexpr int CHUNK = 8192;
-        constexpr int N_PAD1 = 8192;
-        constexpr int K1 = 32;
-        constexpr int T1 = 512;
-        constexpr int CPR = 16384 / CHUNK;    // 2
+        constexpr int T1 = 128;
+        constexpr int CPR = 8192 / CHUNK;     // 8
 
         auto v_scratch = torch::empty({batch * CPR * K1}, x.options());
         auto i_scratch = torch::empty({batch * CPR * K1},
@@ -410,19 +398,17 @@
         launch_chunk_int32<K1, N_PAD1, T1>(x, v_scratch, i_scratch,
                                             n, CHUNK, CPR, batch, stream);
 
-        // Stage 2: 2 runs * 32 = 64 candidates -> top-32. Bitonic 64.
         constexpr int CAND2 = CPR * K1;       // 64
         constexpr int N_PAD2 = 64;
         constexpr int T2 = 32;
         launch_stage2_bitonic<K1, N_PAD2, T2>(v_scratch, i_scratch, vals, idxs,
                                                 CAND2, batch, stream);
-    } else if (batch == 16 && n == 12000 && k == 16) {
-        // 2 chunks of 6000 (pad to 8192), K=16. 32 stage1 blocks.
-        constexpr int CHUNK = 6000;
-        constexpr int N_PAD1 = 8192;
-        constexpr int K1 = 16;
-        constexpr int T1 = 512;
-        constexpr int CPR = 2;
+    } else if (batch == 32 && n == 16384 && k == 32) {
+        constexpr int CHUNK = 1024;
+        constexpr int N_PAD1 = 1024;
+        constexpr int K1 = 32;
+        constexpr int T1 = 128;
+        constexpr int CPR = 16384 / CHUNK;    // 16
 
         auto v_scratch = torch::empty({batch * CPR * K1}, x.options());
         auto i_scratch = torch::empty({batch * CPR * K1},
@@ -431,12 +417,24 @@
         launch_chunk_int32<K1, N_PAD1, T1>(x, v_scratch, i_scratch,
                                             n, CHUNK, CPR, batch, stream);
 
-        // Stage 2: 2 runs * 16 = 32 candidates -> top-16. Bitonic 32.
-        constexpr int CAND2 = CPR * K1;       // 32
-        constexpr int N_PAD2 = 32;
-        constexpr int T2 = 32;
-        launch_stage2_bitonic<K1, N_PAD2, T2>(v_scratch, i_scratch, vals, idxs,
-                                                CAND2, batch, stream);
+        launch_stage2_merge<K1, CPR, K1, 32>(v_scratch, i_scratch, vals, idxs,
+                                              batch, stream);
+    } else if (batch == 16 && n == 12000 && k == 16) {
+        constexpr int CHUNK = 1024;
+        constexpr int N_PAD1 = 1024;
+        constexpr int K1 = 16;
+        constexpr int T1 = 128;
+        constexpr int CPR = 12;
+
+        auto v_scratch = torch::empty({batch * CPR * K1}, x.options());
+        auto i_scratch = torch::empty({batch * CPR * K1},
+                                       x.options().dtype(torch::kInt32));
+
+        launch_chunk_int32<K1, N_PAD1, T1>(x, v_scratch, i_scratch,
+                                            n, CHUNK, CPR, batch, stream);
+
+        launch_stage2_merge<K1, CPR, K1, 32>(v_scratch, i_scratch, vals, idxs,
+                                              batch, stream);
     } else if (batch == 128 && n == 4096 && k == 1) {
         constexpr int N = 4096;
         constexpr int T = 256;