--- /dev/null
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -0,0 +1,193 @@
+"""Custom top-k kernel for RTX PRO 6000 (SM120 Blackwell).
+
+Memory-bound top-k. Strategy: single-pass streaming bitonic top-k. Each block
+keeps a shared-memory buffer of size BUF (= 2*nextpow2(k), >=64); it streams the
+input in tiles of BUF/2, sorting BUF elements descending each tile and keeping
+the best BUF/2. Rows are split across multiple blocks (phase 1); a small second
+kernel merges the per-block partial top-k (phase 2).
+"""
+import torch
+import torch.nn as nn
+from torch.utils.cpp_extension import load_inline
+
+_CUDA = r'''
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+#include <cuda_fp32.h>
+
+// Bitonic sort of BUF elements DESCENDING, blockDim.x == BUF threads.
+__device__ __forceinline__ void bitonic_desc(float* sv, int* si, int BUF) {
+    const int tid = threadIdx.x;
+    for (int k = 2; k <= BUF; k <<= 1) {
+        for (int j = k >> 1; j > 0; j >>= 1) {
+            int ixj = tid ^ j;
+            if (ixj > tid) {
+                bool up = ((tid & k) == 0);
+                float a = sv[tid], b = sv[ixj];
+                // descending: smallest pushed toward high index
+                bool swap = up ? (a < b) : (a > b);
+                if (swap) {
+                    sv[tid] = b; sv[ixj] = a;
+                    int ta = si[tid]; si[tid] = si[ixj]; si[ixj] = ta;
+                }
+            }
+            __syncthreads();
+        }
+    }
+}
+
+// Phase 1: each block processes a contiguous slice of its row.
+__global__ void topk_phase1(const float* __restrict__ x, float* __restrict__ pv,
+                            int* __restrict__ pi, int n, int k, int BUF, int blocksPerRow) {
+    extern __shared__ char smem[];
+    float* sv = (float*)smem;
+    int* si = (int*)(sv + BUF);
+    const int tid = threadIdx.x;
+    const int row = blockIdx.y;
+    const int b = blockIdx.x;
+    const int TILE = BUF >> 1;
+
+    long epb = (long)(n + blocksPerRow - 1) / blocksPerRow;
+    long start = (long)b * epb;
+    long end = start + epb; if (end > n) end = n;
+
+    sv[tid] = -CUDART_INF_F; si[tid] = 0;
+    __syncthreads();
+
+    const float* row_x = x + (long)row * n;
+    for (long base = start; base < end; base += TILE) {
+        if (tid < TILE) {
+            long gi = base + tid;
+            if (gi < end) { sv[TILE + tid] = row_x[gi]; si[TILE + tid] = (int)gi; }
+            else { sv[TILE + tid] = -CUDART_INF_F; si[TILE + tid] = 0; }
+        }
+        __syncthreads();
+        bitonic_desc(sv, si, BUF);
+    }
+    if (tid < k) {
+        long o = ((long)row * blocksPerRow + b) * k + tid;
+        pv[o] = sv[tid]; pi[o] = si[tid];
+    }
+}
+
+// Phase 2: merge blocksPerRow*k partials per row -> final sorted top-k.
+__global__ void topk_phase2(const float* __restrict__ pv, const int* __restrict__ pi,
+                            float* __restrict__ outv, long* __restrict__ outi,
+                            int k, int BUF, int blocksPerRow) {
+    extern __shared__ char smem[];
+    float* sv = (float*)smem;
+    int* si = (int*)(sv + BUF);
+    const int tid = threadIdx.x;
+    const int row = blockIdx.x;
+    const int TILE = BUF >> 1;
+    const int m = blocksPerRow * k;
+
+    sv[tid] = -CUDART_INF_F; si[tid] = 0;
+    __syncthreads();
+
+    const float* base_v = pv + (long)row * m;
+    const int* base_i = pi + (long)row * m;
+    for (int base = 0; base < m; base += TILE) {
+        if (tid < TILE) {
+            int gi = base + tid;
+            if (gi < m) { sv[TILE + tid] = base_v[gi]; si[TILE + tid] = base_i[gi]; }
+            else { sv[TILE + tid] = -CUDART_INF_F; si[TILE + tid] = 0; }
+        }
+        __syncthreads();
+        bitonic_desc(sv, si, BUF);
+    }
+    if (tid < k) {
+        outv[(long)row * k + tid] = sv[tid];
+        outi[(long)row * k + tid] = (long)si[tid];
+    }
+}
+
+std::vector<torch::Tensor> topk_cuda(torch::Tensor x, int k, int BUF, int blocksPerRow) {
+    int batch = x.size(0);
+    int n = x.size(1);
+    auto opt_v = torch::TensorOptions().dtype(torch::kFloat32).device(x.device());
+    auto opt_i = torch::TensorOptions().dtype(torch::kInt64).device(x.device());
+    auto outv = torch::empty({batch, k}, opt_v);
+    auto outi = torch::empty({batch, k}, opt_i);
+
+    size_t shmem = (size_t)BUF * (sizeof(float) + sizeof(int));
+    auto stream = at::cuda::getCurrentCUDAStream();
+
+    if (blocksPerRow == 1) {
+        // phase1 writes directly to output (but output idx must be int64)
+        auto pv = torch::empty({batch, k}, opt_v);
+        auto pi = torch::empty({batch, k}, torch::TensorOptions().dtype(torch::kInt32).device(x.device()));
+        dim3 grid(1, batch);
+        topk_phase1<<<grid, BUF, shmem, stream>>>(x.data_ptr<float>(), pv.data_ptr<float>(),
+            pi.data_ptr<int>(), n, k, BUF, 1);
+        outv.copy_(pv);
+        outi.copy_(pi.to(torch::kInt64));
+        return {outv, outi};
+    }
+
+    auto opt_i32 = torch::TensorOptions().dtype(torch::kInt32).device(x.device());
+    auto pv = torch::empty({(long)batch * blocksPerRow * k}, opt_v);
+    auto pi = torch::empty({(long)batch * blocksPerRow * k}, opt_i32);
+    dim3 grid1(blocksPerRow, batch);
+    topk_phase1<<<grid1, BUF, shmem, stream>>>(x.data_ptr<float>(), pv.data_ptr<float>(),
+        pi.data_ptr<int>(), n, k, BUF, blocksPerRow);
+    topk_phase2<<<batch, BUF, shmem, stream>>>(pv.data_ptr<float>(), pi.data_ptr<int>(),
+        outv.data_ptr<float>(), outi.data_ptr<long>(), k, BUF, blocksPerRow);
+    return {outv, outi};
+}
+'''
+
+_CPP = "std::vector<torch::Tensor> topk_cuda(torch::Tensor x, int k, int BUF, int blocksPerRow);"
+
+_mod = load_inline(
+    name="topk_bitonic",
+    cpp_sources=_CPP,
+    cuda_sources=_CUDA,
+    functions=["topk_cuda"],
+    extra_cuda_cflags=["-O3", "--use_fast_math"],
+    verbose=False,
+)
+
+
+def _nextpow2(v):
+    p = 1
+    while p < v:
+        p <<= 1
+    return p
+
+
+class Model

... (949 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -13,7 +13,10 @@
 _CUDA = r'''
 #include <torch/extension.h>
 #include <cuda_runtime.h>
-#include <cuda_fp32.h>
+#include <math_constants.h>
+#include <c10/cuda/CUDAStream.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <vector>
 
 // Bitonic sort of BUF elements DESCENDING, blockDim.x == BUF threads.
 __device__ __forceinline__ void bitonic_desc(float* sv, int* si, int BUF) {

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,10 +1,9 @@
 """Custom top-k kernel for RTX PRO 6000 (SM120 Blackwell).
 
-Memory-bound top-k. Strategy: single-pass streaming bitonic top-k. Each block
-keeps a shared-memory buffer of size BUF (= 2*nextpow2(k), >=64); it streams the
-input in tiles of BUF/2, sorting BUF elements descending each tile and keeping
-the best BUF/2. Rows are split across multiple blocks (phase 1); a small second
-kernel merges the per-block partial top-k (phase 2).
+Memory-bound top-k on tiny inputs (0.5-2MB) => latency bound, ~8us floor.
+Strategy: ~188 total blocks (one wave). Each block loads a contiguous chunk of
+its row into shared memory, bitonic-sorts it descending, keeps top-k. A small
+phase-2 kernel merges the per-block partials. k=1 uses a dedicated argmax.
 """
 import torch
 import torch.nn as nn
@@ -18,20 +17,20 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <vector>
 
-// Bitonic sort of BUF elements DESCENDING, blockDim.x == BUF threads.
-__device__ __forceinline__ void bitonic_desc(float* sv, int* si, int BUF) {
-    const int tid = threadIdx.x;
-    for (int k = 2; k <= BUF; k <<= 1) {
+// Bitonic sort of N (pow2) elements DESCENDING. blockDim.x threads, N>=blockDim.
+__device__ __forceinline__ void bitonic_desc(float* sv, int* si, int N) {
+    for (int k = 2; k <= N; k <<= 1) {
         for (int j = k >> 1; j > 0; j >>= 1) {
-            int ixj = tid ^ j;
-            if (ixj > tid) {
-                bool up = ((tid & k) == 0);
-                float a = sv[tid], b = sv[ixj];
-                // descending: smallest pushed toward high index
-                bool swap = up ? (a < b) : (a > b);
-                if (swap) {
-                    sv[tid] = b; sv[ixj] = a;
-                    int ta = si[tid]; si[tid] = si[ixj]; si[ixj] = ta;
+            for (int i = threadIdx.x; i < N; i += blockDim.x) {
+                int ixj = i ^ j;
+                if (ixj > i) {
+                    bool up = ((i & k) == 0);
+                    float a = sv[i], b = sv[ixj];
+                    bool sw = up ? (a < b) : (a > b);
+                    if (sw) {
+                        sv[i] = b; sv[ixj] = a;
+                        int t = si[i]; si[i] = si[ixj]; si[ixj] = t;
+                    }
                 }
             }
             __syncthreads();
@@ -39,111 +38,168 @@
     }
 }
 
-// Phase 1: each block processes a contiguous slice of its row.
-__global__ void topk_phase1(const float* __restrict__ x, float* __restrict__ pv,
-                            int* __restrict__ pi, int n, int k, int BUF, int blocksPerRow) {
+// Phase 1: each block sorts a contiguous chunk of its row, emits top-k.
+__global__ void phase1(const float* __restrict__ x,
+                       float* __restrict__ pv, int* __restrict__ pi,
+                       float* __restrict__ outv, long* __restrict__ outi,
+                       int n, int k, int CHUNK, int bpr, int final_direct) {
     extern __shared__ char smem[];
     float* sv = (float*)smem;
-    int* si = (int*)(sv + BUF);
-    const int tid = threadIdx.x;
+    int* si = (int*)(sv + CHUNK);
     const int row = blockIdx.y;
     const int b = blockIdx.x;
-    const int TILE = BUF >> 1;
-
-    long epb = (long)(n + blocksPerRow - 1) / blocksPerRow;
-    long start = (long)b * epb;
-    long end = start + epb; if (end > n) end = n;
-
-    sv[tid] = -CUDART_INF_F; si[tid] = 0;
+    long start = (long)b * CHUNK;          // chunks are power-of-2 sized & contiguous
+    const float* rx = x + (long)row * n;
+
+    for (int i = threadIdx.x; i < CHUNK; i += blockDim.x) {
+        long gi = start + i;
+        if (gi < n) { sv[i] = rx[gi]; si[i] = (int)gi; }
+        else { sv[i] = -CUDART_INF_F; si[i] = 0; }
+    }
     __syncthreads();
-
-    const float* row_x = x + (long)row * n;
-    for (long base = start; base < end; base += TILE) {
-        if (tid < TILE) {
-            long gi = base + tid;
-            if (gi < end) { sv[TILE + tid] = row_x[gi]; si[TILE + tid] = (int)gi; }
-            else { sv[TILE + tid] = -CUDART_INF_F; si[TILE + tid] = 0; }
-        }
-        __syncthreads();
-        bitonic_desc(sv, si, BUF);
-    }
-    if (tid < k) {
-        long o = ((long)row * blocksPerRow + b) * k + tid;
-        pv[o] = sv[tid]; pi[o] = si[tid];
-    }
-}
-
-// Phase 2: merge blocksPerRow*k partials per row -> final sorted top-k.
-__global__ void topk_phase2(const float* __restrict__ pv, const int* __restrict__ pi,
-                            float* __restrict__ outv, long* __restrict__ outi,
-                            int k, int BUF, int blocksPerRow) {
+    bitonic_desc(sv, si, CHUNK);
+
+    if (final_direct) {
+        for (int i = threadIdx.x; i < k; i += blockDim.x) {
+            outv[(long)row * k + i] = sv[i];
+            outi[(long)row * k + i] = (long)si[i];
+        }
+    } else {
+        for (int i = threadIdx.x; i < k; i += blockDim.x) {
+            long o = ((long)row * bpr + b) * k + i;
+            pv[o] = sv[i]; pi[o] = si[i];
+        }
+    }
+}
+
+// Phase 2: merge bpr*k partials per row -> sorted top-k.
+__global__ void phase2(const float* __restrict__ pv, const int* __restrict__ pi,
+                       float* __restrict__ outv, long* __restrict__ outi,
+                       int k, int CHUNK, int m) {
     extern __shared__ char smem[];
     float* sv = (float*)smem;
-    int* si = (int*)(sv + BUF);
-    const int tid = threadIdx.x;
+    int* si = (int*)(sv + CHUNK);
     const int row = blockIdx.x;
-    const int TILE = BUF >> 1;
-    const int m = blocksPerRow * k;
-
-    sv[tid] = -CUDART_INF_F; si[tid] = 0;
+    const float* bv = pv + (long)row * m;
+    const int* bi = pi + (long)row * m;
+    for (int i = threadIdx.x; i < CHUNK; i += blockDim.x) {
+        if (i < m) { sv[i] = bv[i]; si[i] = bi[i]; }
+        else { sv[i] = -CUDART_INF_F; si[i] = 0; }
+    }
     __syncthreads();
-
-    const fl

... (9028 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -149,16 +149,6 @@
 
     if (k == 1) {
         int chunk = (n + bpr - 1) / bpr;
-        if (bpr == 1) {
-            dim3 grid(1, batch);
-            auto pv = torch::empty({batch, 1}, fopt);
-            auto pi = torch::empty({batch, 1}, i32opt);
-            // direct: write to scratch then argmax2 trivial; instead reduce directly
-            argmax1<<<grid, 1024, 0, stream>>>(x.data_ptr<float>(), outv.data_ptr<float>(),
-                (int*)nullptr, n, chunk, 1);
-            // argmax1 writes float to pv path; but we passed outv as pv and need idx too.
-            return {outv, outi};
-        }
         auto pv = torch::empty({(long)batch * bpr}, fopt);
         auto pi = torch::empty({(long)batch * bpr}, i32opt);
         dim3 grid(bpr, batch);

--- /dev/null
+++ b//home/agent/.claude/projects/-workspace-problems-05-topk-bitonic/memory/MEMORY.md
@@ -0,0 +1,3 @@
+# Memory Index
+
+- [topk-bitonic-problem](topk-bitonic-problem.md) — key facts about the 05_topk_bitonic kernel optimization problem

--- /dev/null
+++ b//home/agent/.claude/projects/-workspace-problems-05-topk-bitonic/memory/topk-bitonic-problem.md
@@ -0,0 +1,18 @@
+---
+name: topk-bitonic-problem
+description: Key facts about the 05_topk_bitonic kernel optimization (RTX PRO 6000)
+metadata:
+  type: project
+---
+
+Problem: custom top-k kernel, scored on geomean of achieved DRAM bandwidth fraction (vs 1800 GB/s) over 5 shapes. solution.py, interface matches reference.py. Forbidden: torch.topk/sort/kthvalue/argsort.
+
+Shapes: (1,131072,64) decoder; (64,8192,8); (32,16384,32); (16,12000,16) non-pow2; (128,4096,1) argmax.
+
+KEY INSIGHT: inputs are tiny (0.5-2MB) so this is LATENCY/launch bound, NOT bandwidth bound. Pure read+reduce floor measured ~8-10us per shape at gridDim≈188 (one wave). More blocks = slower launch. Achievable geomean ceiling ≈ 0.08.
+
+torch.topk (sota) ms: s0=0.055, s1=0.062, s2=0.054, s3=0.046, s4=0.022. Beat these.
+
+Run: `uv run python check.py` (must print PASS), `uv run python benchmark.py`. With baselines: prefix KBH_BENCHMARK_BASELINES=1.
+
+Progress: v2 (load-chunk-to-shared + bitonic sort, dedicated argmax for k=1) → geomean 0.0123. shape4 argmax near floor (13us). Full bitonic sort wastes work for small k; single-block phase2 sort hurts shape0. Moving to per-thread register top-k + tree-merge.

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,9 +1,10 @@
 """Custom top-k kernel for RTX PRO 6000 (SM120 Blackwell).
 
-Memory-bound top-k on tiny inputs (0.5-2MB) => latency bound, ~8us floor.
-Strategy: ~188 total blocks (one wave). Each block loads a contiguous chunk of
-its row into shared memory, bitonic-sorts it descending, keeps top-k. A small
-phase-2 kernel merges the per-block partials. k=1 uses a dedicated argmax.
+Tiny inputs (0.5-2MB) => latency bound, ~8us read floor. Approach: per-thread
+register top-k (threshold-gated insertion) over a coalesced single pass, then a
+pairwise tree-merge of the per-thread sorted lists in shared memory. Rows are
+split across blocks (phase 1) and merged by a second pass (phase 2, same kernel
+with bpr=1 over the partials). k=1 uses a dedicated argmax reduction.
 """
 import torch
 import torch.nn as nn
@@ -17,127 +18,108 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <vector>
 
-// Bitonic sort of N (pow2) elements DESCENDING. blockDim.x threads, N>=blockDim.
-__device__ __forceinline__ void bitonic_desc(float* sv, int* si, int N) {
-    for (int k = 2; k <= N; k <<= 1) {
-        for (int j = k >> 1; j > 0; j >>= 1) {
-            for (int i = threadIdx.x; i < N; i += blockDim.x) {
-                int ixj = i ^ j;
-                if (ixj > i) {
-                    bool up = ((i & k) == 0);
-                    float a = sv[i], b = sv[ixj];
-                    bool sw = up ? (a < b) : (a > b);
-                    if (sw) {
-                        sv[i] = b; sv[ixj] = a;
-                        int t = si[i]; si[i] = si[ixj]; si[ixj] = t;
-                    }
-                }
-            }
-            __syncthreads();
-        }
-    }
-}
-
-// Phase 1: each block sorts a contiguous chunk of its row, emits top-k.
-__global__ void phase1(const float* __restrict__ x,
-                       float* __restrict__ pv, int* __restrict__ pi,
-                       float* __restrict__ outv, long* __restrict__ outi,
-                       int n, int k, int CHUNK, int bpr, int final_direct) {
+template<int K>
+__device__ __forceinline__ void insert(float* rv, int* ri, float v, int idx) {
+    if (v <= rv[K-1]) return;
+    #pragma unroll
+    for (int i = K-1; i > 0; --i) {
+        if (rv[i-1] < v) { rv[i] = rv[i-1]; ri[i] = ri[i-1]; }
+        else { rv[i] = v; ri[i] = idx; return; }
+    }
+    rv[0] = v; ri[0] = idx;
+}
+
+// Generic top-k kernel. Reads a contiguous chunk of `src` row, builds per-thread
+// register top-K, tree-merges in shared, writes top-K (to out* if final else p*).
+template<int K>
+__global__ void topk_kernel(const float* __restrict__ src, const int* __restrict__ srcidx,
+                            float* __restrict__ pv, int* __restrict__ pi,
+                            float* __restrict__ outv, long* __restrict__ outi,
+                            int nrow, int bpr, int final_direct) {
     extern __shared__ char smem[];
     float* sv = (float*)smem;
-    int* si = (int*)(sv + CHUNK);
+    int* si = (int*)(sv + blockDim.x * K);
+    const int tid = threadIdx.x;
     const int row = blockIdx.y;
     const int b = blockIdx.x;
-    long start = (long)b * CHUNK;          // chunks are power-of-2 sized & contiguous
-    const float* rx = x + (long)row * n;
-
-    for (int i = threadIdx.x; i < CHUNK; i += blockDim.x) {
-        long gi = start + i;
-        if (gi < n) { sv[i] = rx[gi]; si[i] = (int)gi; }
-        else { sv[i] = -CUDART_INF_F; si[i] = 0; }
-    }
+    const int BD = blockDim.x;
+
+    int chunk = (nrow + bpr - 1) / bpr;
+    long start = (long)b * chunk;
+    long end = start + chunk; if (end > nrow) end = nrow;
+
+    float rv[K]; int ri[K];
+    #pragma unroll
+    for (int i = 0; i < K; ++i) { rv[i] = -CUDART_INF_F; ri[i] = 0; }
+
+    const float* rx = src + (long)row * nrow;
+    if (srcidx == nullptr) {
+        for (long i = start + tid; i < end; i += BD) {
+            insert<K>(rv, ri, rx[i], (int)i);
+        }
+    } else {
+        const int* rxi = srcidx + (long)row * nrow;
+        for (long i = start + tid; i < end; i += BD) {
+            insert<K>(rv, ri, rx[i], rxi[i]);
+        }
+    }
+
+    #pragma unroll
+    for (int i = 0; i < K; ++i) { sv[tid*K + i] = rv[i]; si[tid*K + i] = ri[i]; }
     __syncthreads();
-    bitonic_desc(sv, si, CHUNK);
+
+    for (int stride = BD >> 1; stride > 0; stride >>= 1) {
+        if (tid < stride) {
+            int a = tid * K, bb = (tid + stride) * K;
+            float tv[K]; int ti[K];
+            int ia = 0, ib = 0;
+            #pragma unroll
+            for (int o = 0; o < K; ++o) {
+                float av = sv[a+ia], bv = sv[bb+ib];
+                if (av >= bv) { tv[o] = av; ti[o] = si[a+ia]; ia++; }
+                else { tv[o] = bv; ti[o] = si[bb+ib]; ib++; }
+            }
+            #pragma unroll
+            for (int o = 0; o < K; ++o) { sv[a+o] = tv[o]; si[a+o] = ti[o]; }
+        }
+        __syncthreads();
+    }
 
     if (final_direct) {
-        for (int i = threadIdx.x; i < k; i += blockDim.x) {
-            outv[(long)row * k + i] = sv[i];
-            outi[(long)row * k + i] = (long)si[i];
-        }
+        if (tid < K) { outv[(long)row*K + tid] = sv[tid]; outi[(long)row*K + tid] = (long)si[tid]; }
     } else {
-        for (int i = threadIdx.x; i < k; i += blockDim.x) {
-            long o = ((long)row * bpr + b) * k + i;
-            pv[o] = sv[i]; pi[o] = si[i];
-        }
-    }
-}
-
-// Phase 2: merge bpr*k partials per row -> sorted top-k.
-__global__ void phase2(const float* __restrict__ pv, const int* __restrict__ pi,
-                       float* __restrict__ outv, long* __restrict__ outi,
-                       int k, int CHUNK, int m) {
-    extern __shared__ char smem[];
-    float* sv = (float*)smem;
-    int* si = (int*)(sv + CHUNK);
-    const int row = blockIdx.x;
-    const float* bv = pv + (long)row * m;
-    const int* 

... (9608 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -86,9 +86,9 @@
     }
 
     if (final_direct) {
-        if (tid < K) { outv[(long)row*K + tid] = sv[tid]; outi[(long)row*K + tid] = (long)si[tid]; }
+        for (int i = tid; i < K; i += BD) { outv[(long)row*K + i] = sv[i]; outi[(long)row*K + i] = (long)si[i]; }
     } else {
-        if (tid < K) { long o = ((long)row*bpr + b)*K + tid; pv[o] = sv[tid]; pi[o] = si[tid]; }
+        for (int i = tid; i < K; i += BD) { long o = ((long)row*bpr + b)*K + i; pv[o] = sv[i]; pi[o] = si[i]; }
     }
 }
 

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -155,7 +155,8 @@
             nullptr, nullptr, n, bpr, 0);
         int m = bpr * K;
         size_t sh2 = (size_t)BD2 * K * (sizeof(float) + sizeof(int));
-        topk_kernel<K><<<batch, BD2, sh2, stream>>>(pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr, nullptr,
+        dim3 grid2(1, batch);
+        topk_kernel<K><<<grid2, BD2, sh2, stream>>>(pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr, nullptr,
             outv.data_ptr<float>(), outi.data_ptr<long>(), m, 1, 1);
     };
 

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,10 +1,11 @@
 """Custom top-k kernel for RTX PRO 6000 (SM120 Blackwell).
 
-Tiny inputs (0.5-2MB) => latency bound, ~8us read floor. Approach: per-thread
-register top-k (threshold-gated insertion) over a coalesced single pass, then a
-pairwise tree-merge of the per-thread sorted lists in shared memory. Rows are
-split across blocks (phase 1) and merged by a second pass (phase 2, same kernel
-with bpr=1 over the partials). k=1 uses a dedicated argmax reduction.
+Tiny inputs (0.5-2MB) => latency bound, ~8us read floor. Hybrid by k:
+  k==1 : block argmax reduction.
+  k<=16: per-thread register top-k (threshold-gated) + pairwise tree-merge.
+  k>=32: load chunk to shared, bitonic sort descending, keep top-k
+         (register arrays of size k spill to local mem for large k).
+Rows split across blocks (phase1); phase2 merges per-block partials.
 """
 import torch
 import torch.nn as nn
@@ -18,6 +19,7 @@
 #include <ATen/cuda/CUDAContext.h>
 #include <vector>
 
+// ===================== register top-k (small k) =====================
 template<int K>
 __device__ __forceinline__ void insert(float* rv, int* ri, float v, int idx) {
     if (v <= rv[K-1]) return;
@@ -29,13 +31,11 @@
     rv[0] = v; ri[0] = idx;
 }
 
-// Generic top-k kernel. Reads a contiguous chunk of `src` row, builds per-thread
-// register top-K, tree-merges in shared, writes top-K (to out* if final else p*).
 template<int K>
-__global__ void topk_kernel(const float* __restrict__ src, const int* __restrict__ srcidx,
-                            float* __restrict__ pv, int* __restrict__ pi,
-                            float* __restrict__ outv, long* __restrict__ outi,
-                            int nrow, int bpr, int final_direct) {
+__global__ void regmerge(const float* __restrict__ src, const int* __restrict__ srcidx,
+                         float* __restrict__ pv, int* __restrict__ pi,
+                         float* __restrict__ outv, long* __restrict__ outi,
+                         int nrow, int bpr, int final_direct) {
     extern __shared__ char smem[];
     float* sv = (float*)smem;
     int* si = (int*)(sv + blockDim.x * K);
@@ -54,14 +54,10 @@
 
     const float* rx = src + (long)row * nrow;
     if (srcidx == nullptr) {
-        for (long i = start + tid; i < end; i += BD) {
-            insert<K>(rv, ri, rx[i], (int)i);
-        }
+        for (long i = start + tid; i < end; i += BD) insert<K>(rv, ri, rx[i], (int)i);
     } else {
         const int* rxi = srcidx + (long)row * nrow;
-        for (long i = start + tid; i < end; i += BD) {
-            insert<K>(rv, ri, rx[i], rxi[i]);
-        }
+        for (long i = start + tid; i < end; i += BD) insert<K>(rv, ri, rx[i], rxi[i]);
     }
 
     #pragma unroll
@@ -92,34 +88,84 @@
     }
 }
 
+// ===================== bitonic (large k) =====================
+__device__ __forceinline__ void bitonic_desc(float* sv, int* si, int N) {
+    for (int k = 2; k <= N; k <<= 1) {
+        for (int j = k >> 1; j > 0; j >>= 1) {
+            for (int i = threadIdx.x; i < N; i += blockDim.x) {
+                int ixj = i ^ j;
+                if (ixj > i) {
+                    bool up = ((i & k) == 0);
+                    float a = sv[i], b = sv[ixj];
+                    bool sw = up ? (a < b) : (a > b);
+                    if (sw) { sv[i]=b; sv[ixj]=a; int t=si[i]; si[i]=si[ixj]; si[ixj]=t; }
+                }
+            }
+            __syncthreads();
+        }
+    }
+}
+
+__global__ void bitonic_kernel(const float* __restrict__ src, const int* __restrict__ srcidx,
+                               float* __restrict__ pv, int* __restrict__ pi,
+                               float* __restrict__ outv, long* __restrict__ outi,
+                               int nrow, int k, int CHUNK, int bpr, int final_direct) {
+    extern __shared__ char smem[];
+    float* sv = (float*)smem;
+    int* si = (int*)(sv + CHUNK);
+    const int row = blockIdx.y, b = blockIdx.x;
+    long start = (long)b * CHUNK;
+    const float* rx = src + (long)row * nrow;
+    if (srcidx == nullptr) {
+        for (int i = threadIdx.x; i < CHUNK; i += blockDim.x) {
+            long gi = start + i;
+            if (gi < nrow) { sv[i]=rx[gi]; si[i]=(int)gi; } else { sv[i]=-CUDART_INF_F; si[i]=0; }
+        }
+    } else {
+        const int* rxi = srcidx + (long)row * nrow;
+        for (int i = threadIdx.x; i < CHUNK; i += blockDim.x) {
+            long gi = start + i;
+            if (gi < nrow) { sv[i]=rx[gi]; si[i]=rxi[gi]; } else { sv[i]=-CUDART_INF_F; si[i]=0; }
+        }
+    }
+    __syncthreads();
+    bitonic_desc(sv, si, CHUNK);
+    if (final_direct) {
+        for (int i = threadIdx.x; i < k; i += blockDim.x) { outv[(long)row*k+i]=sv[i]; outi[(long)row*k+i]=(long)si[i]; }
+    } else {
+        for (int i = threadIdx.x; i < k; i += blockDim.x) { long o=((long)row*bpr+b)*k+i; pv[o]=sv[i]; pi[o]=si[i]; }
+    }
+}
+
+// ===================== argmax (k=1) =====================
 __global__ void argmax1(const float* __restrict__ x, float* __restrict__ pv,
                         int* __restrict__ pi, int n, int chunk, int bpr) {
     __shared__ float sv[1024]; __shared__ int si[1024];
     const int row = blockIdx.y, b = blockIdx.x, t = threadIdx.x;
-    long start = (long)b * chunk, end = start + chunk; if (end > n) end = n;
-    const float* rx = x + (long)row * n;
-    float best = -CUDART_INF_F; int bidx = 0;
-    for (long i = start + t; i < end; i += blockDim.x) { float v = rx[i]; if (v > best) { best = v; bidx = (int)i; } }
-    sv[t] = best; si[t] = bidx; __syncthreads();
-    for (int s = blockDim.x >> 1; s > 0; s >>= 1) { if (t < s && sv[t+s] > sv[t]) { sv[t]=sv[t+s]; si[t]=si[t+s]; } __syncthreads(); }
-    if (t == 0) { long o = (long)row*bpr + b; pv[o] = sv[0]; pi[o] = si[0]; }
+    long start = (long)b*chunk, end = start+chunk; if (end>n) end=n;
+    const f

... (9413 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -86,6 +86,54 @@
     } else {
         for (int i = tid; i < K; i += BD) { long o = ((long)row*bpr + b)*K + i; pv[o] = sv[i]; pi[o] = si[i]; }
     }
+}
+
+// ===== register top-k + scalar double-buffer tree-merge (large k, no spill) =====
+template<int K>
+__global__ void regmerge_db(const float* __restrict__ src, const int* __restrict__ srcidx,
+                            float* __restrict__ pv, int* __restrict__ pi,
+                            float* __restrict__ outv, long* __restrict__ outi,
+                            int nrow, int bpr, int final_direct) {
+    extern __shared__ char smem[];
+    const int BD = blockDim.x;
+    float* av = (float*)smem;
+    int* ai = (int*)(av + BD*K);
+    float* bv = (float*)(ai + BD*K);
+    int* bi = (int*)(bv + BD*K);
+    const int tid = threadIdx.x;
+    const int row = blockIdx.y, b = blockIdx.x;
+
+    int chunk = (nrow + bpr - 1) / bpr;
+    long start = (long)b*chunk, end = start+chunk; if (end>nrow) end=nrow;
+
+    float rv[K]; int ri[K];
+    #pragma unroll
+    for (int i=0;i<K;++i){ rv[i]=-CUDART_INF_F; ri[i]=0; }
+    const float* rx = src + (long)row*nrow;
+    if (srcidx==nullptr) { for(long i=start+tid;i<end;i+=BD) insert<K>(rv,ri,rx[i],(int)i); }
+    else { const int* rxi=srcidx+(long)row*nrow; for(long i=start+tid;i<end;i+=BD) insert<K>(rv,ri,rx[i],rxi[i]); }
+
+    #pragma unroll
+    for (int i=0;i<K;++i){ av[tid*K+i]=rv[i]; ai[tid*K+i]=ri[i]; }
+    __syncthreads();
+
+    float* sv=av; int* si=ai; float* dv=bv; int* di=bi;
+    for (int nL=BD; nL>1; nL>>=1) {
+        int half=nL>>1;
+        if (tid<half) {
+            int a=2*tid*K, bb=(2*tid+1)*K, d=tid*K, ia=0, ib=0;
+            #pragma unroll
+            for (int o=0;o<K;++o){
+                float x=sv[a+ia], y=sv[bb+ib];
+                if (x>=y){ dv[d+o]=x; di[d+o]=si[a+ia]; ia++; }
+                else { dv[d+o]=y; di[d+o]=si[bb+ib]; ib++; }
+            }
+        }
+        __syncthreads();
+        float* t1=sv; sv=dv; dv=t1; int* t2=si; si=di; di=t2;
+    }
+    if (final_direct) { for(int i=tid;i<K;i+=BD){ outv[(long)row*K+i]=sv[i]; outi[(long)row*K+i]=(long)si[i]; } }
+    else { for(int i=tid;i<K;i+=BD){ long o=((long)row*bpr+b)*K+i; pv[o]=sv[i]; pi[o]=si[i]; } }
 }
 
 // ===================== bitonic (large k) =====================

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -233,50 +233,37 @@
         return {outv, outi};
     }
 
-    if (k <= 16) {
-        auto run = [&](auto kc) {
-            constexpr int K = decltype(kc)::value;
-            size_t sh1 = (size_t)BD * K * (sizeof(float)+sizeof(int));
-            if (bpr == 1) {
-                dim3 grid(1, batch);
-                regmerge<K><<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), n, 1, 1);
-                return;
-            }
-            auto pv = torch::empty({(long)batch*bpr*K}, fopt);
-            auto pi = torch::empty({(long)batch*bpr*K}, i32opt);
-            dim3 grid(bpr, batch);
-            regmerge<K><<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, n, bpr, 0);
-            int m = bpr*K;
-            size_t sh2 = (size_t)BD2 * K * (sizeof(float)+sizeof(int));
-            dim3 grid2(1, batch);
-            regmerge<K><<<grid2,BD2,sh2,stream>>>(pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), m, 1, 1);
+    auto run = [&](auto kc) {
+        constexpr int K = decltype(kc)::value;
+        const bool db = (K >= 32);
+        size_t sh1 = (size_t)BD * K * (sizeof(float)+sizeof(int)) * (db?2:1);
+        size_t sh2 = (size_t)BD2 * K * (sizeof(float)+sizeof(int)) * (db?2:1);
+        if (db && !attr_set) {
+            cudaFuncSetAttribute(regmerge_db<K>, cudaFuncAttributeMaxDynamicSharedMemorySize, 101376);
+        }
+        auto launch1 = [&](dim3 g, int bd, size_t sh, const float* s, const int* sidx,
+                           float* p_v, int* p_i, float* o_v, long* o_i, int nr, int br, int fd) {
+            if (db) regmerge_db<K><<<g,bd,sh,stream>>>(s,sidx,p_v,p_i,o_v,o_i,nr,br,fd);
+            else regmerge<K><<<g,bd,sh,stream>>>(s,sidx,p_v,p_i,o_v,o_i,nr,br,fd);
         };
-        if (k==8) run(std::integral_constant<int,8>{});
-        else if (k==16) run(std::integral_constant<int,16>{});
-        else TORCH_CHECK(false,"bad k");
-        return {outv, outi};
-    }
-
-    // large k: bitonic
-    if (!attr_set) {
-        cudaFuncSetAttribute(bitonic_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 101376);
-        attr_set = true;
-    }
-    size_t sh1 = (size_t)CHUNK * (sizeof(float)+sizeof(int));
-    if (bpr == 1) {
-        dim3 grid(1, batch);
-        bitonic_kernel<<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), n, k, CHUNK, 1, 1);
-        return {outv, outi};
-    }
-    auto pv = torch::empty({(long)batch*bpr*k}, fopt);
-    auto pi = torch::empty({(long)batch*bpr*k}, i32opt);
-    dim3 grid(bpr, batch);
-    bitonic_kernel<<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, n, k, CHUNK, bpr, 0);
-    int m = bpr*k;
-    size_t sh2 = (size_t)CHUNK2 * (sizeof(float)+sizeof(int));
-    int bd2 = BD2 < CHUNK2 ? BD2 : CHUNK2;
-    dim3 grid2(1, batch);
-    bitonic_kernel<<<grid2,bd2,sh2,stream>>>(pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), m, k, CHUNK2, 1, 1);
+        if (bpr == 1) {
+            dim3 grid(1, batch);
+            launch1(grid,BD,sh1, x.data_ptr<float>(), nullptr, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), n, 1, 1);
+            return;
+        }
+        auto pv = torch::empty({(long)batch*bpr*K}, fopt);
+        auto pi = torch::empty({(long)batch*bpr*K}, i32opt);
+        dim3 grid(bpr, batch);
+        launch1(grid,BD,sh1, x.data_ptr<float>(), nullptr, pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, n, bpr, 0);
+        int m = bpr*K;
+        dim3 grid2(1, batch);
+        launch1(grid2,BD2,sh2, pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), m, 1, 1);
+    };
+    if (k==8) run(std::integral_constant<int,8>{});
+    else if (k==16) run(std::integral_constant<int,16>{});
+    else if (k==32) { run(std::integral_constant<int,32>{}); attr_set=true; }
+    else if (k==64) { run(std::integral_constant<int,64>{}); attr_set=true; }
+    else TORCH_CHECK(false,"bad k");
     return {outv, outi};
 }
 '''

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -305,33 +305,17 @@
             self.bpr = max(1, min(round(188 / batch), (n + 1023) // 1024))
             return
 
-        if k <= 16:
-            BD = _floorpow2(max(32, min(1024, 4096 // k)))
-            bpr = max(1, round(188 / batch))
-            while bpr * k > 2048 and bpr > 1:
-                bpr -= 1
-            while bpr > 1 and (n + bpr - 1) // bpr < BD:
-                bpr -= 1
-            m = bpr * k
-            self.BD = BD
-            self.bpr = bpr
-            self.BD2 = _floorpow2(max(32, min(1024, (m + k - 1) // k)))
-            return
-
-        # large k: bitonic
+        # k in {8,16,32,64}: register top-k + tree merge
+        BD = _floorpow2(max(32, min(1024, 4096 // k)))
         bpr = max(1, round(188 / batch))
-        chunk = _nextpow2((n + bpr - 1) // bpr)
-        while chunk > 8192:
-            bpr += 1
-            chunk = _nextpow2((n + bpr - 1) // bpr)
-        while bpr * k > 4096 and bpr > 1:
+        while bpr * k > 2048 and bpr > 1:
             bpr -= 1
-            chunk = _nextpow2((n + bpr - 1) // bpr)
+        while bpr > 1 and (n + bpr - 1) // bpr < BD:
+            bpr -= 1
+        m = bpr * k
+        self.BD = BD
         self.bpr = bpr
-        self.CHUNK = chunk
-        self.BD = min(1024, chunk)
-        self.CHUNK2 = _nextpow2(bpr * k) if bpr > 1 else 0
-        self.BD2 = min(1024, self.CHUNK2) if self.CHUNK2 else 0
+        self.BD2 = _floorpow2(max(32, min(1024, (m + k - 1) // k)))
 
     def forward(self, x: torch.Tensor):
         return tuple(_mod.topk_cuda(x, self.k, self.bpr, self.BD, self.BD2,

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -238,7 +238,7 @@
         const bool db = (K >= 32);
         size_t sh1 = (size_t)BD * K * (sizeof(float)+sizeof(int)) * (db?2:1);
         size_t sh2 = (size_t)BD2 * K * (sizeof(float)+sizeof(int)) * (db?2:1);
-        if (db && !attr_set) {
+        if (db) {
             cudaFuncSetAttribute(regmerge_db<K>, cudaFuncAttributeMaxDynamicSharedMemorySize, 101376);
         }
         auto launch1 = [&](dim3 g, int bd, size_t sh, const float* s, const int* sidx,

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -233,37 +233,46 @@
         return {outv, outi};
     }
 
-    auto run = [&](auto kc) {
-        constexpr int K = decltype(kc)::value;
-        const bool db = (K >= 32);
-        size_t sh1 = (size_t)BD * K * (sizeof(float)+sizeof(int)) * (db?2:1);
-        size_t sh2 = (size_t)BD2 * K * (sizeof(float)+sizeof(int)) * (db?2:1);
-        if (db) {
-            cudaFuncSetAttribute(regmerge_db<K>, cudaFuncAttributeMaxDynamicSharedMemorySize, 101376);
-        }
-        auto launch1 = [&](dim3 g, int bd, size_t sh, const float* s, const int* sidx,
-                           float* p_v, int* p_i, float* o_v, long* o_i, int nr, int br, int fd) {
-            if (db) regmerge_db<K><<<g,bd,sh,stream>>>(s,sidx,p_v,p_i,o_v,o_i,nr,br,fd);
-            else regmerge<K><<<g,bd,sh,stream>>>(s,sidx,p_v,p_i,o_v,o_i,nr,br,fd);
+    if (k <= 16) {
+        auto run = [&](auto kc) {
+            constexpr int K = decltype(kc)::value;
+            size_t sh1 = (size_t)BD * K * (sizeof(float)+sizeof(int));
+            size_t sh2 = (size_t)BD2 * K * (sizeof(float)+sizeof(int));
+            if (bpr == 1) {
+                dim3 grid(1, batch);
+                regmerge<K><<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), n, 1, 1);
+                return;
+            }
+            auto pv = torch::empty({(long)batch*bpr*K}, fopt);
+            auto pi = torch::empty({(long)batch*bpr*K}, i32opt);
+            dim3 grid(bpr, batch);
+            regmerge<K><<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, n, bpr, 0);
+            int m = bpr*K;
+            dim3 grid2(1, batch);
+            regmerge<K><<<grid2,BD2,sh2,stream>>>(pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), m, 1, 1);
         };
-        if (bpr == 1) {
-            dim3 grid(1, batch);
-            launch1(grid,BD,sh1, x.data_ptr<float>(), nullptr, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), n, 1, 1);
-            return;
-        }
-        auto pv = torch::empty({(long)batch*bpr*K}, fopt);
-        auto pi = torch::empty({(long)batch*bpr*K}, i32opt);
-        dim3 grid(bpr, batch);
-        launch1(grid,BD,sh1, x.data_ptr<float>(), nullptr, pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, n, bpr, 0);
-        int m = bpr*K;
-        dim3 grid2(1, batch);
-        launch1(grid2,BD2,sh2, pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), m, 1, 1);
-    };
-    if (k==8) run(std::integral_constant<int,8>{});
-    else if (k==16) run(std::integral_constant<int,16>{});
-    else if (k==32) { run(std::integral_constant<int,32>{}); attr_set=true; }
-    else if (k==64) { run(std::integral_constant<int,64>{}); attr_set=true; }
-    else TORCH_CHECK(false,"bad k");
+        if (k==8) run(std::integral_constant<int,8>{});
+        else run(std::integral_constant<int,16>{});
+        return {outv, outi};
+    }
+
+    // large k: bitonic load-to-shared
+    cudaFuncSetAttribute(bitonic_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 101376);
+    size_t sh1 = (size_t)CHUNK * (sizeof(float)+sizeof(int));
+    if (bpr == 1) {
+        dim3 grid(1, batch);
+        bitonic_kernel<<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), n, k, CHUNK, 1, 1);
+        return {outv, outi};
+    }
+    auto pv = torch::empty({(long)batch*bpr*k}, fopt);
+    auto pi = torch::empty({(long)batch*bpr*k}, i32opt);
+    dim3 grid(bpr, batch);
+    bitonic_kernel<<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, n, k, CHUNK, bpr, 0);
+    int m = bpr*k;
+    size_t sh2 = (size_t)CHUNK2 * (sizeof(float)+sizeof(int));
+    int bd2 = BD2 < CHUNK2 ? BD2 : CHUNK2;
+    dim3 grid2(1, batch);
+    bitonic_kernel<<<grid2,bd2,sh2,stream>>>(pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), m, k, CHUNK2, 1, 1);
     return {outv, outi};
 }
 '''

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -314,17 +314,34 @@
             self.bpr = max(1, min(round(188 / batch), (n + 1023) // 1024))
             return
 
-        # k in {8,16,32,64}: register top-k + tree merge
-        BD = _floorpow2(max(32, min(1024, 4096 // k)))
+        if k <= 16:
+            # register top-k + tree merge
+            BD = _floorpow2(max(32, min(1024, 4096 // k)))
+            bpr = max(1, round(188 / batch))
+            while bpr * k > 2048 and bpr > 1:
+                bpr -= 1
+            while bpr > 1 and (n + bpr - 1) // bpr < BD:
+                bpr -= 1
+            m = bpr * k
+            self.BD = BD
+            self.bpr = bpr
+            self.BD2 = _floorpow2(max(32, min(1024, (m + k - 1) // k)))
+            return
+
+        # large k: bitonic load-to-shared
         bpr = max(1, round(188 / batch))
-        while bpr * k > 2048 and bpr > 1:
+        chunk = _nextpow2((n + bpr - 1) // bpr)
+        while chunk > 8192:
+            bpr += 1
+            chunk = _nextpow2((n + bpr - 1) // bpr)
+        while bpr * k > 4096 and bpr > 1:
             bpr -= 1
-        while bpr > 1 and (n + bpr - 1) // bpr < BD:
-            bpr -= 1
-        m = bpr * k
-        self.BD = BD
+            chunk = _nextpow2((n + bpr - 1) // bpr)
         self.bpr = bpr
-        self.BD2 = _floorpow2(max(32, min(1024, (m + k - 1) // k)))
+        self.CHUNK = chunk
+        self.BD = min(1024, chunk)
+        self.CHUNK2 = _nextpow2(bpr * k) if bpr > 1 else 0
+        self.BD2 = min(1024, self.CHUNK2) if self.CHUNK2 else 0
 
     def forward(self, x: torch.Tensor):
         return tuple(_mod.topk_cuda(x, self.k, self.bpr, self.BD, self.BD2,

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -136,17 +136,32 @@
     else { for(int i=tid;i<K;i+=BD){ long o=((long)row*bpr+b)*K+i; pv[o]=sv[i]; pi[o]=si[i]; } }
 }
 
-// ===================== bitonic (large k) =====================
-__device__ __forceinline__ void bitonic_desc(float* sv, int* si, int N) {
+// ===================== bitonic (large k), packed uint64 keys =====================
+// Map float to order-preserving uint32, then pack key = (ord<<32)|idx so a single
+// uint64 descending sort orders by value (ties by idx, arbitrary -> fine).
+__device__ __forceinline__ unsigned f2ord(float f) {
+    unsigned b = __float_as_uint(f);
+    return (b & 0x80000000u) ? ~b : (b | 0x80000000u);
+}
+__device__ __forceinline__ float ord2f(unsigned o) {
+    unsigned b = (o & 0x80000000u) ? (o & 0x7fffffffu) : ~o;
+    return __uint_as_float(b);
+}
+__device__ __forceinline__ unsigned long long pack(float v, int idx) {
+    return ((unsigned long long)f2ord(v) << 32) | (unsigned)idx;
+}
+#define NEG_KEY ((unsigned long long)0)  // f2ord(-inf)=0 packs to smallest
+
+__device__ __forceinline__ void bitonic_desc(unsigned long long* s, int N) {
     for (int k = 2; k <= N; k <<= 1) {
         for (int j = k >> 1; j > 0; j >>= 1) {
             for (int i = threadIdx.x; i < N; i += blockDim.x) {
                 int ixj = i ^ j;
                 if (ixj > i) {
                     bool up = ((i & k) == 0);
-                    float a = sv[i], b = sv[ixj];
+                    unsigned long long a = s[i], b = s[ixj];
                     bool sw = up ? (a < b) : (a > b);
-                    if (sw) { sv[i]=b; sv[ixj]=a; int t=si[i]; si[i]=si[ixj]; si[ixj]=t; }
+                    if (sw) { s[i]=b; s[ixj]=a; }
                 }
             }
             __syncthreads();
@@ -159,29 +174,28 @@
                                float* __restrict__ outv, long* __restrict__ outi,
                                int nrow, int k, int CHUNK, int bpr, int final_direct) {
     extern __shared__ char smem[];
-    float* sv = (float*)smem;
-    int* si = (int*)(sv + CHUNK);
+    unsigned long long* s = (unsigned long long*)smem;
     const int row = blockIdx.y, b = blockIdx.x;
     long start = (long)b * CHUNK;
     const float* rx = src + (long)row * nrow;
     if (srcidx == nullptr) {
         for (int i = threadIdx.x; i < CHUNK; i += blockDim.x) {
             long gi = start + i;
-            if (gi < nrow) { sv[i]=rx[gi]; si[i]=(int)gi; } else { sv[i]=-CUDART_INF_F; si[i]=0; }
+            s[i] = (gi < nrow) ? pack(rx[gi], (int)gi) : NEG_KEY;
         }
     } else {
         const int* rxi = srcidx + (long)row * nrow;
         for (int i = threadIdx.x; i < CHUNK; i += blockDim.x) {
             long gi = start + i;
-            if (gi < nrow) { sv[i]=rx[gi]; si[i]=rxi[gi]; } else { sv[i]=-CUDART_INF_F; si[i]=0; }
+            s[i] = (gi < nrow) ? pack(rx[gi], rxi[gi]) : NEG_KEY;
         }
     }
     __syncthreads();
-    bitonic_desc(sv, si, CHUNK);
+    bitonic_desc(s, CHUNK);
     if (final_direct) {
-        for (int i = threadIdx.x; i < k; i += blockDim.x) { outv[(long)row*k+i]=sv[i]; outi[(long)row*k+i]=(long)si[i]; }
+        for (int i = threadIdx.x; i < k; i += blockDim.x) { outv[(long)row*k+i]=ord2f((unsigned)(s[i]>>32)); outi[(long)row*k+i]=(long)(unsigned)s[i]; }
     } else {
-        for (int i = threadIdx.x; i < k; i += blockDim.x) { long o=((long)row*bpr+b)*k+i; pv[o]=sv[i]; pi[o]=si[i]; }
+        for (int i = threadIdx.x; i < k; i += blockDim.x) { long o=((long)row*bpr+b)*k+i; pv[o]=ord2f((unsigned)(s[i]>>32)); pi[o]=(int)(unsigned)s[i]; }
     }
 }
 

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -342,15 +342,13 @@
             self.BD2 = _floorpow2(max(32, min(1024, (m + k - 1) // k)))
             return
 
-        # large k: bitonic load-to-shared
-        bpr = max(1, round(188 / batch))
+        # large k: bitonic load-to-shared, target CHUNK~4096, C2<=2048
+        bpr = max(1, (n + 4095) // 4096)
         chunk = _nextpow2((n + bpr - 1) // bpr)
-        while chunk > 8192:
-            bpr += 1
-            chunk = _nextpow2((n + bpr - 1) // bpr)
-        while bpr * k > 4096 and bpr > 1:
-            bpr -= 1
-            chunk = _nextpow2((n + bpr - 1) // bpr)
+        # keep phase2 candidate sort bounded
+        while bpr * k > 2048 and (n + (bpr - 1)) // bpr < 8192:
+            # cannot reduce bpr without growing chunk past 8192; accept
+            break
         self.bpr = bpr
         self.CHUNK = chunk
         self.BD = min(1024, chunk)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -329,9 +329,9 @@
             return
 
         if k <= 16:
-            # register top-k + tree merge
-            BD = _floorpow2(max(32, min(1024, 4096 // k)))
-            bpr = max(1, round(188 / batch))
+            # register top-k + tree merge; BD=128 sweet spot, target chunk~1536
+            BD = 128
+            bpr = max(1, round(n / 1536))
             while bpr * k > 2048 and bpr > 1:
                 bpr -= 1
             while bpr > 1 and (n + bpr - 1) // bpr < BD:

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -270,23 +270,52 @@
         return {outv, outi};
     }
 
-    // large k: bitonic load-to-shared
+    // large k: bitonic load-to-shared, multi-pass reduction.
+    // params: bpr=phase1 blocks/row, CHUNK=phase1 chunk, CHUNK2=reduce-chunk target.
     cudaFuncSetAttribute(bitonic_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 101376);
-    size_t sh1 = (size_t)CHUNK * (sizeof(float)+sizeof(int));
+    auto np2 = [](int v){ int p=1; while(p<v) p<<=1; return p; };
+
     if (bpr == 1) {
+        size_t sh = (size_t)CHUNK * 8;
         dim3 grid(1, batch);
-        bitonic_kernel<<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), n, k, CHUNK, 1, 1);
+        bitonic_kernel<<<grid,BD,sh,stream>>>(x.data_ptr<float>(), nullptr, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), n, k, CHUNK, 1, 1);
         return {outv, outi};
     }
-    auto pv = torch::empty({(long)batch*bpr*k}, fopt);
-    auto pi = torch::empty({(long)batch*bpr*k}, i32opt);
-    dim3 grid(bpr, batch);
-    bitonic_kernel<<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, n, k, CHUNK, bpr, 0);
-    int m = bpr*k;
-    size_t sh2 = (size_t)CHUNK2 * (sizeof(float)+sizeof(int));
-    int bd2 = BD2 < CHUNK2 ? BD2 : CHUNK2;
-    dim3 grid2(1, batch);
-    bitonic_kernel<<<grid2,bd2,sh2,stream>>>(pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), m, k, CHUNK2, 1, 1);
+
+    // phase1: x -> scratch A
+    int cap = bpr * k;
+    auto Av = torch::empty({(long)batch*cap}, fopt);
+    auto Ai = torch::empty({(long)batch*cap}, i32opt);
+    {
+        size_t sh = (size_t)CHUNK * 8;
+        dim3 grid(bpr, batch);
+        bitonic_kernel<<<grid,BD,sh,stream>>>(x.data_ptr<float>(), nullptr, Av.data_ptr<float>(), Ai.data_ptr<int>(), nullptr,nullptr, n, k, CHUNK, bpr, 0);
+    }
+    int m = bpr * k;
+    const int RC = CHUNK2;        // reduce chunk target
+    auto Bv = torch::empty({(long)batch*cap}, fopt);
+    auto Bi = torch::empty({(long)batch*cap}, i32opt);
+    float *curV = Av.data_ptr<float>(); int *curI = Ai.data_ptr<int>();
+    float *altV = Bv.data_ptr<float>(); int *altI = Bi.data_ptr<int>();
+    // reduce while too many candidates for a single-block final sort
+    while (m > RC) {
+        int g = (m + RC - 1) / RC;
+        int chunk = np2((m + g - 1) / g);
+        int bd = chunk < 1024 ? chunk : 1024;
+        size_t sh = (size_t)chunk * 8;
+        dim3 grid(g, batch);
+        bitonic_kernel<<<grid,bd,sh,stream>>>(curV, curI, altV, altI, nullptr,nullptr, m, k, chunk, g, 0);
+        m = g * k;
+        std::swap(curV, altV); std::swap(curI, altI);
+    }
+    // final: 1 block per row -> output
+    {
+        int chunk = np2(m);
+        int bd = chunk < 1024 ? chunk : 1024;
+        size_t sh = (size_t)chunk * 8;
+        dim3 grid(1, batch);
+        bitonic_kernel<<<grid,bd,sh,stream>>>(curV, curI, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), m, k, chunk, 1, 1);
+    }
     return {outv, outi};
 }
 '''

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -371,18 +371,15 @@
             self.BD2 = _floorpow2(max(32, min(1024, (m + k - 1) // k)))
             return
 
-        # large k: bitonic load-to-shared, target CHUNK~4096, C2<=2048
-        bpr = max(1, (n + 4095) // 4096)
-        chunk = _nextpow2((n + bpr - 1) // bpr)
-        # keep phase2 candidate sort bounded
-        while bpr * k > 2048 and (n + (bpr - 1)) // bpr < 8192:
-            # cannot reduce bpr without growing chunk past 8192; accept
-            break
+        # large k: bitonic, multi-pass reduction. Small phase1 chunk => many
+        # blocks (fill SMs) + less total sort work; reduce passes shrink cands.
+        CHUNK1 = 1024
+        bpr = max(1, (n + CHUNK1 - 1) // CHUNK1)
         self.bpr = bpr
-        self.CHUNK = chunk
-        self.BD = min(1024, chunk)
-        self.CHUNK2 = _nextpow2(bpr * k) if bpr > 1 else 0
-        self.BD2 = min(1024, self.CHUNK2) if self.CHUNK2 else 0
+        self.CHUNK = CHUNK1
+        self.BD = min(1024, CHUNK1)
+        self.CHUNK2 = 2048   # reduce-chunk target (RC)
+        self.BD2 = 0
 
     def forward(self, x: torch.Tensor):
         return tuple(_mod.topk_cuda(x, self.k, self.bpr, self.BD, self.BD2,

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -373,12 +373,12 @@
 
         # large k: bitonic, multi-pass reduction. Small phase1 chunk => many
         # blocks (fill SMs) + less total sort work; reduce passes shrink cands.
-        CHUNK1 = 1024
+        CHUNK1 = 512
         bpr = max(1, (n + CHUNK1 - 1) // CHUNK1)
         self.bpr = bpr
         self.CHUNK = CHUNK1
         self.BD = min(1024, CHUNK1)
-        self.CHUNK2 = 2048   # reduce-chunk target (RC)
+        self.CHUNK2 = 1024   # reduce-chunk target (RC)
         self.BD2 = 0
 
     def forward(self, x: torch.Tensor):

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -247,7 +247,7 @@
         return {outv, outi};
     }
 
-    if (k <= 16) {
+    if (false) {
         auto run = [&](auto kc) {
             constexpr int K = decltype(kc)::value;
             size_t sh1 = (size_t)BD * K * (sizeof(float)+sizeof(int));

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -357,21 +357,7 @@
             self.bpr = max(1, min(round(188 / batch), (n + 1023) // 1024))
             return
 
-        if k <= 16:
-            # register top-k + tree merge; BD=128 sweet spot, target chunk~1536
-            BD = 128
-            bpr = max(1, round(n / 1536))
-            while bpr * k > 2048 and bpr > 1:
-                bpr -= 1
-            while bpr > 1 and (n + bpr - 1) // bpr < BD:
-                bpr -= 1
-            m = bpr * k
-            self.BD = BD
-            self.bpr = bpr
-            self.BD2 = _floorpow2(max(32, min(1024, (m + k - 1) // k)))
-            return
-
-        # large k: bitonic, multi-pass reduction. Small phase1 chunk => many
+        # bitonic, multi-pass reduction. Small phase1 chunk => many
         # blocks (fill SMs) + less total sort work; reduce passes shrink cands.
         CHUNK1 = 512
         bpr = max(1, (n + CHUNK1 - 1) // CHUNK1)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -247,7 +247,7 @@
         return {outv, outi};
     }
 
-    if (false) {
+    if (k <= 8) {
         auto run = [&](auto kc) {
             constexpr int K = decltype(kc)::value;
             size_t sh1 = (size_t)BD * K * (sizeof(float)+sizeof(int));

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -357,6 +357,20 @@
             self.bpr = max(1, min(round(188 / batch), (n + 1023) // 1024))
             return
 
+        if k <= 8:
+            # register top-k + tree merge; BD=128 sweet spot
+            BD = 128
+            bpr = max(1, round(n / 1536))
+            while bpr * k > 2048 and bpr > 1:
+                bpr -= 1
+            while bpr > 1 and (n + bpr - 1) // bpr < BD:
+                bpr -= 1
+            m = bpr * k
+            self.BD = BD
+            self.bpr = bpr
+            self.BD2 = _floorpow2(max(32, min(1024, (m + k - 1) // k)))
+            return
+
         # bitonic, multi-pass reduction. Small phase1 chunk => many
         # blocks (fill SMs) + less total sort work; reduce passes shrink cands.
         CHUNK1 = 512

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -226,25 +226,23 @@
 
 static bool attr_set = false;
 
-std::vector<torch::Tensor> topk_cuda(torch::Tensor x, int k, int bpr, int BD,
-                                     int BD2, int CHUNK, int CHUNK2) {
+void topk_cuda(torch::Tensor x, int k, int bpr, int BD, int BD2, int CHUNK, int CHUNK2,
+               torch::Tensor outv, torch::Tensor outi,
+               torch::Tensor sAv, torch::Tensor sAi, torch::Tensor sBv, torch::Tensor sBi) {
     int batch = x.size(0);
     int n = x.size(1);
-    auto fopt = torch::TensorOptions().dtype(torch::kFloat32).device(x.device());
-    auto lopt = torch::TensorOptions().dtype(torch::kInt64).device(x.device());
-    auto i32opt = torch::TensorOptions().dtype(torch::kInt32).device(x.device());
-    auto outv = torch::empty({batch, k}, fopt);
-    auto outi = torch::empty({batch, k}, lopt);
     auto stream = at::cuda::getCurrentCUDAStream();
+    float* OV = outv.data_ptr<float>(); long* OI = outi.data_ptr<long>();
+    float* Av = sAv.data_ptr<float>(); int* Ai = sAi.data_ptr<int>();
+    float* Bv = sBv.data_ptr<float>(); int* Bi = sBi.data_ptr<int>();
+    const float* X = x.data_ptr<float>();
 
     if (k == 1) {
         int chunk = (n + bpr - 1) / bpr;
-        auto pv = torch::empty({(long)batch*bpr}, fopt);
-        auto pi = torch::empty({(long)batch*bpr}, i32opt);
         dim3 grid(bpr, batch);
-        argmax1<<<grid,1024,0,stream>>>(x.data_ptr<float>(), pv.data_ptr<float>(), pi.data_ptr<int>(), n, chunk, bpr);
-        argmax2<<<batch,1024,0,stream>>>(pv.data_ptr<float>(), pi.data_ptr<int>(), outv.data_ptr<float>(), outi.data_ptr<long>(), bpr);
-        return {outv, outi};
+        argmax1<<<grid,1024,0,stream>>>(X, Av, Ai, n, chunk, bpr);
+        argmax2<<<batch,1024,0,stream>>>(Av, Ai, OV, OI, bpr);
+        return;
     }
 
     if (k <= 8) {
@@ -254,50 +252,40 @@
             size_t sh2 = (size_t)BD2 * K * (sizeof(float)+sizeof(int));
             if (bpr == 1) {
                 dim3 grid(1, batch);
-                regmerge<K><<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), n, 1, 1);
+                regmerge<K><<<grid,BD,sh1,stream>>>(X, nullptr, nullptr,nullptr, OV, OI, n, 1, 1);
                 return;
             }
-            auto pv = torch::empty({(long)batch*bpr*K}, fopt);
-            auto pi = torch::empty({(long)batch*bpr*K}, i32opt);
             dim3 grid(bpr, batch);
-            regmerge<K><<<grid,BD,sh1,stream>>>(x.data_ptr<float>(), nullptr, pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, n, bpr, 0);
+            regmerge<K><<<grid,BD,sh1,stream>>>(X, nullptr, Av, Ai, nullptr,nullptr, n, bpr, 0);
             int m = bpr*K;
             dim3 grid2(1, batch);
-            regmerge<K><<<grid2,BD2,sh2,stream>>>(pv.data_ptr<float>(), pi.data_ptr<int>(), nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), m, 1, 1);
+            regmerge<K><<<grid2,BD2,sh2,stream>>>(Av, Ai, nullptr,nullptr, OV, OI, m, 1, 1);
         };
         if (k==8) run(std::integral_constant<int,8>{});
         else run(std::integral_constant<int,16>{});
-        return {outv, outi};
+        return;
     }
 
     // large k: bitonic load-to-shared, multi-pass reduction.
-    // params: bpr=phase1 blocks/row, CHUNK=phase1 chunk, CHUNK2=reduce-chunk target.
     cudaFuncSetAttribute(bitonic_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, 101376);
     auto np2 = [](int v){ int p=1; while(p<v) p<<=1; return p; };
 
     if (bpr == 1) {
         size_t sh = (size_t)CHUNK * 8;
         dim3 grid(1, batch);
-        bitonic_kernel<<<grid,BD,sh,stream>>>(x.data_ptr<float>(), nullptr, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), n, k, CHUNK, 1, 1);
-        return {outv, outi};
-    }
-
-    // phase1: x -> scratch A
-    int cap = bpr * k;
-    auto Av = torch::empty({(long)batch*cap}, fopt);
-    auto Ai = torch::empty({(long)batch*cap}, i32opt);
+        bitonic_kernel<<<grid,BD,sh,stream>>>(X, nullptr, nullptr,nullptr, OV, OI, n, k, CHUNK, 1, 1);
+        return;
+    }
+
     {
         size_t sh = (size_t)CHUNK * 8;
         dim3 grid(bpr, batch);
-        bitonic_kernel<<<grid,BD,sh,stream>>>(x.data_ptr<float>(), nullptr, Av.data_ptr<float>(), Ai.data_ptr<int>(), nullptr,nullptr, n, k, CHUNK, bpr, 0);
+        bitonic_kernel<<<grid,BD,sh,stream>>>(X, nullptr, Av, Ai, nullptr,nullptr, n, k, CHUNK, bpr, 0);
     }
     int m = bpr * k;
-    const int RC = CHUNK2;        // reduce chunk target
-    auto Bv = torch::empty({(long)batch*cap}, fopt);
-    auto Bi = torch::empty({(long)batch*cap}, i32opt);
-    float *curV = Av.data_ptr<float>(); int *curI = Ai.data_ptr<int>();
-    float *altV = Bv.data_ptr<float>(); int *altI = Bi.data_ptr<int>();
-    // reduce while too many candidates for a single-block final sort
+    const int RC = CHUNK2;
+    float *curV = Av; int *curI = Ai;
+    float *altV = Bv; int *altI = Bi;
     while (m > RC) {
         int g = (m + RC - 1) / RC;
         int chunk = np2((m + g - 1) / g);
@@ -308,15 +296,13 @@
         m = g * k;
         std::swap(curV, altV); std::swap(curI, altI);
     }
-    // final: 1 block per row -> output
     {
         int chunk = np2(m);
         int bd = chunk < 1024 ? chunk : 1024;
         size_t sh = (size_t)chunk * 8;
         dim3 grid(1, batch);
-        bitonic_kernel<<<grid,bd,sh,stream>>>(curV, curI, nullptr,nullptr, outv.data_ptr<float>(), outi.data_ptr<long>(), m, k, chunk, 1, 1);
-    }
-    return {outv, outi};
+        bitonic_kernel<<<grid,bd,sh,stream>>>(curV, curI, nullptr,nullptr, OV, OI, m, k, chunk, 1, 1);
+    }
 }
 '''
 

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -306,7 +306,9 @@
 }
 '''
 
-_CPP = "std::vector<torch::Tensor> topk_cuda(torch::Tensor x, int k, int bpr, int BD, int BD2, int CHUNK, int CHUNK2);"
+_CPP = ("void topk_cuda(torch::Tensor x, int k, int bpr, int BD, int BD2, int CHUNK, int CHUNK2, "
+        "torch::Tensor outv, torch::Tensor outi, torch::Tensor sAv, torch::Tensor sAi, "
+        "torch::Tensor sBv, torch::Tensor sBi);")
 
 _mod = load_inline(
     name="topk_hybrid",

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -359,6 +359,9 @@
             self.BD2 = _floorpow2(max(32, min(1024, (m + k - 1) // k)))
             return
 
+            self._alloc()
+            return
+
         # bitonic, multi-pass reduction. Small phase1 chunk => many
         # blocks (fill SMs) + less total sort work; reduce passes shrink cands.
         CHUNK1 = 512
@@ -368,10 +371,23 @@
         self.BD = min(1024, CHUNK1)
         self.CHUNK2 = 1024   # reduce-chunk target (RC)
         self.BD2 = 0
+        self._alloc()
+
+    def _alloc(self):
+        # preallocate persistent scratch + output buffers (avoid per-call alloc)
+        dev = torch.device("cuda")
+        cap = max(self.batch * self.bpr * self.k, self.batch * self.k)
+        self.sAv = torch.empty(cap, dtype=torch.float32, device=dev)
+        self.sAi = torch.empty(cap, dtype=torch.int32, device=dev)
+        self.sBv = torch.empty(cap, dtype=torch.float32, device=dev)
+        self.sBi = torch.empty(cap, dtype=torch.int32, device=dev)
+        self.outv = torch.empty(self.batch, self.k, dtype=torch.float32, device=dev)
+        self.outi = torch.empty(self.batch, self.k, dtype=torch.int64, device=dev)
 
     def forward(self, x: torch.Tensor):
-        return tuple(_mod.topk_cuda(x, self.k, self.bpr, self.BD, self.BD2,
-                                    self.CHUNK, self.CHUNK2))
+        _mod.topk_cuda(x, self.k, self.bpr, self.BD, self.BD2, self.CHUNK, self.CHUNK2,
+                       self.outv, self.outi, self.sAv, self.sAi, self.sBv, self.sBi)
+        return self.outv, self.outi
 
 
 def get_inputs():

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -343,6 +343,7 @@
 
         if k == 1:
             self.bpr = max(1, min(round(188 / batch), (n + 1023) // 1024))
+            self._alloc()
             return
 
         if k <= 8:

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -358,8 +358,6 @@
             self.BD = BD
             self.bpr = bpr
             self.BD2 = _floorpow2(max(32, min(1024, (m + k - 1) // k)))
-            return
-
             self._alloc()
             return
 

--- a//home/agent/.claude/projects/-workspace-problems-05-topk-bitonic/memory/topk-bitonic-problem.md
+++ b//home/agent/.claude/projects/-workspace-problems-05-topk-bitonic/memory/topk-bitonic-problem.md
@@ -15,4 +15,9 @@
 
 Run: `uv run python check.py` (must print PASS), `uv run python benchmark.py`. With baselines: prefix KBH_BENCHMARK_BASELINES=1.
 
-Progress: v2 (load-chunk-to-shared + bitonic sort, dedicated argmax for k=1) → geomean 0.0123. shape4 argmax near floor (13us). Full bitonic sort wastes work for small k; single-block phase2 sort hurts shape0. Moving to per-thread register top-k + tree-merge.
+WINNING DESIGN (geomean ~0.031, beats sota ~0.016 on every shape):
+- k==1: dedicated argmax (2-kernel reduction). near floor ~12us.
+- k<=8: per-thread register top-k (threshold-gated insert, templated K, unrolled) + pairwise tree-merge in shared. BD=128 sweet spot. 2 kernels.
+- k>=16: bitonic sort in shared with PACKED uint64 keys (f2ord(val)<<32|idx) — single 64-bit compare-swap halves shared traffic, ~1.5x speedup. MULTI-PASS reduction: small phase1 chunk (CHUNK1=512, many blocks fill SMs, less total sort work) -> reduce passes (RC=1024) -> final 1-block sort.
+KEY LESSONS: (1) ~188 blocks (one wave) best, more=slower launch. (2) phase1 bitonic was occupancy-bound (waves 0.17) not compute/mem — smaller chunks = more blocks AND less log^2 work. (3) per-thread rv[K] register array SPILLS for K>=32 -> bitonic better for large k. (4) PREALLOCATE scratch+output buffers in __init__ (not per forward) - torch.empty per call costs ~6us CPU/GPU-idle, huge for these tiny kernels. (5) single-row shape0 tail (reduce+final) is serial, inherently limited ~34us.
+Tuning: shape0 prefers CHUNK1=1024 (smaller tail), shape2 prefers 512 (more blocks).

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -363,7 +363,7 @@
 
         # bitonic, multi-pass reduction. Small phase1 chunk => many
         # blocks (fill SMs) + less total sort work; reduce passes shrink cands.
-        CHUNK1 = 512
+        CHUNK1 = 1024 if batch <= 2 else 512
         bpr = max(1, (n + CHUNK1 - 1) // CHUNK1)
         self.bpr = bpr
         self.CHUNK = CHUNK1

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -212,6 +212,18 @@
     for(int s=blockDim.x>>1;s>0;s>>=1){ if(t<s && sv[t+s]>sv[t]){sv[t]=sv[t+s];si[t]=si[t+s];} __syncthreads(); }
     if(t==0){ long o=(long)row*bpr+b; pv[o]=sv[0]; pi[o]=si[0]; }
 }
+// one block per row, full reduction, write int64 directly
+__global__ void argmax_single(const float* __restrict__ x, float* __restrict__ outv,
+                              long* __restrict__ outi, int n) {
+    __shared__ float sv[1024]; __shared__ int si[1024];
+    const int row = blockIdx.x, t = threadIdx.x;
+    const float* rx = x + (long)row*n;
+    float best=-CUDART_INF_F; int bidx=0;
+    for (long i=t; i<n; i+=blockDim.x){ float v=rx[i]; if(v>best){best=v;bidx=(int)i;} }
+    sv[t]=best; si[t]=bidx; __syncthreads();
+    for(int s=blockDim.x>>1;s>0;s>>=1){ if(t<s && sv[t+s]>sv[t]){sv[t]=sv[t+s];si[t]=si[t+s];} __syncthreads(); }
+    if(t==0){ outv[row]=sv[0]; outi[row]=(long)si[0]; }
+}
 __global__ void argmax2(const float* __restrict__ pv, const int* __restrict__ pi,
                         float* __restrict__ outv, long* __restrict__ outi, int bpr) {
     __shared__ float sv[1024]; __shared__ int si[1024];

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -250,6 +250,10 @@
     const float* X = x.data_ptr<float>();
 
     if (k == 1) {
+        if (bpr == 1) {
+            argmax_single<<<batch,1024,0,stream>>>(X, OV, OI, n);
+            return;
+        }
         int chunk = (n + bpr - 1) / bpr;
         dim3 grid(bpr, batch);
         argmax1<<<grid,1024,0,stream>>>(X, Av, Ai, n, chunk, bpr);

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -379,12 +379,12 @@
 
         # bitonic, multi-pass reduction. Small phase1 chunk => many
         # blocks (fill SMs) + less total sort work; reduce passes shrink cands.
-        CHUNK1 = 1024 if batch <= 2 else 512
+        CHUNK1 = 256
         bpr = max(1, (n + CHUNK1 - 1) // CHUNK1)
         self.bpr = bpr
         self.CHUNK = CHUNK1
         self.BD = min(1024, CHUNK1)
-        self.CHUNK2 = 1024   # reduce-chunk target (RC)
+        self.CHUNK2 = 512   # reduce-chunk target (RC)
         self.BD2 = 0
         self._alloc()
 

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -14,6 +14,7 @@
 _CUDA = r'''
 #include <torch/extension.h>
 #include <cuda_runtime.h>
+#include <cooperative_groups.h>
 #include <math_constants.h>
 #include <c10/cuda/CUDAStream.h>
 #include <ATen/cuda/CUDAContext.h>

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -197,6 +197,67 @@
         for (int i = threadIdx.x; i < k; i += blockDim.x) { outv[(long)row*k+i]=ord2f((unsigned)(s[i]>>32)); outi[(long)row*k+i]=(long)(unsigned)s[i]; }
     } else {
         for (int i = threadIdx.x; i < k; i += blockDim.x) { long o=((long)row*bpr+b)*k+i; pv[o]=ord2f((unsigned)(s[i]>>32)); pi[o]=(int)(unsigned)s[i]; }
+    }
+}
+
+// ===================== cooperative single-launch (single row, large k) =====
+namespace cg = cooperative_groups;
+__device__ __forceinline__ int dpow2(int v){ int p=1; while(p<v) p<<=1; return p; }
+
+// One launch: phase A (each block sorts its chunk -> top k) then grid-sync'd
+// reduce passes, all in-kernel. SHP = shared capacity (ull). Row = blockIdx.y.
+__global__ void coop_topk(const float* __restrict__ x, int n, int k, int RC,
+                          unsigned long long* __restrict__ bufA,
+                          unsigned long long* __restrict__ bufB,
+                          float* __restrict__ outv, long* __restrict__ outi) {
+    cg::grid_group grid = cg::this_grid();
+    extern __shared__ unsigned long long sm[];
+    const int G = gridDim.x;
+    const int row = blockIdx.y;
+    const int t = threadIdx.x, BD = blockDim.x;
+    unsigned long long* A = bufA + (long)row * G * k;
+    unsigned long long* B = bufB + (long)row * G * k;
+    const float* rx = x + (long)row * n;
+
+    // phase A
+    int chunk = (n + G - 1) / G;
+    int CP = dpow2(chunk);
+    long start = (long)blockIdx.x * chunk;
+    for (int i = t; i < CP; i += BD) {
+        long gi = start + i;
+        sm[i] = (i < chunk && gi < n) ? pack(rx[gi], (int)gi) : 0ULL;
+    }
+    __syncthreads();
+    bitonic_desc(sm, CP);
+    for (int i = t; i < k; i += BD) A[(long)blockIdx.x * k + i] = sm[i];
+    grid.sync();
+
+    // reduce passes
+    int m = G * k;
+    unsigned long long* cur = A; unsigned long long* alt = B;
+    while (m > k) {
+        int groups = (m + RC - 1) / RC;
+        int cchunk = (m + groups - 1) / groups;
+        int CP2 = dpow2(cchunk);
+        if (blockIdx.x < groups) {
+            long st = (long)blockIdx.x * cchunk;
+            for (int i = t; i < CP2; i += BD) {
+                long gi = st + i;
+                sm[i] = (i < cchunk && gi < m) ? cur[gi] : 0ULL;
+            }
+            __syncthreads();
+            bitonic_desc(sm, CP2);
+            for (int i = t; i < k; i += BD) alt[(long)blockIdx.x * k + i] = sm[i];
+        }
+        grid.sync();
+        m = groups * k;
+        unsigned long long* tmp = cur; cur = alt; alt = tmp;
+    }
+    if (blockIdx.x == 0) {
+        for (int i = t; i < k; i += BD) {
+            outv[(long)row*k + i] = ord2f((unsigned)(cur[i] >> 32));
+            outi[(long)row*k + i] = (long)(unsigned)cur[i];
+        }
     }
 }
 

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -299,6 +299,50 @@
 }
 
 static bool attr_set = false;
+
+static int g_numSM = -1;
+
+// Cooperative single-launch top-k for a single row (batch==1), large k.
+// Returns 0 on success, nonzero (cuda error) so caller can fall back.
+int topk_coop(torch::Tensor x, int k, int RC, torch::Tensor outv, torch::Tensor outi,
+              torch::Tensor bufA, torch::Tensor bufB) {
+    int batch = x.size(0);
+    int n = x.size(1);
+    auto stream = at::cuda::getCurrentCUDAStream();
+    if (g_numSM < 0) {
+        cudaDeviceGetAttribute(&g_numSM, cudaDevAttrMultiProcessorCount, 0);
+    }
+    const int BD = 256;
+    auto np2 = [](int v){ int p=1; while(p<v) p<<=1; return p; };
+    int target = (n + 255) / 256;          // aim chunkA ~256
+
+    // occupancy with a provisional shared size, then finalize
+    int SHP = RC > 512 ? RC : 512;
+    int maxbpsm = 1;
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&maxbpsm, coop_topk, BD, (size_t)SHP*8);
+    int maxG = maxbpsm * g_numSM / batch;
+    int G = target < maxG ? target : maxG;
+    if (G < 1) G = 1;
+    int chunkA = (n + G - 1) / G;
+    int CP = np2(chunkA);
+    SHP = CP > RC ? CP : RC;
+    // re-check occupancy with final shared, clamp G if needed
+    cudaFuncSetAttribute(coop_topk, cudaFuncAttributeMaxDynamicSharedMemorySize, (size_t)SHP*8);
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&maxbpsm, coop_topk, BD, (size_t)SHP*8);
+    maxG = maxbpsm * g_numSM / batch;
+    if (G > maxG) { G = maxG; if (G < 1) G = 1; chunkA = (n + G - 1) / G; CP = np2(chunkA); SHP = CP > RC ? CP : RC; cudaFuncSetAttribute(coop_topk, cudaFuncAttributeMaxDynamicSharedMemorySize, (size_t)SHP*8); }
+
+    void* args[] = { (void*)&x, nullptr };  // placeholder, filled below
+    const float* X = x.data_ptr<float>();
+    float* OV = outv.data_ptr<float>(); long* OI = outi.data_ptr<long>();
+    unsigned long long* A = (unsigned long long*)bufA.data_ptr<long>();
+    unsigned long long* B = (unsigned long long*)bufB.data_ptr<long>();
+    void* kargs[] = { (void*)&X, (void*)&n, (void*)&k, (void*)&RC, (void*)&A, (void*)&B, (void*)&OV, (void*)&OI };
+    dim3 grid(G, batch);
+    dim3 block(BD);
+    cudaError_t err = cudaLaunchCooperativeKernel((void*)coop_topk, grid, block, kargs, (size_t)SHP*8, stream);
+    return (int)err;
+}
 
 void topk_cuda(torch::Tensor x, int k, int bpr, int BD, int BD2, int CHUNK, int CHUNK2,
                torch::Tensor outv, torch::Tensor outi,

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -333,6 +333,47 @@
     if (G > maxG) { G = maxG; if (G < 1) G = 1; chunkA = (n + G - 1) / G; CP = np2(chunkA); SHP = CP > RC ? CP : RC; cudaFuncSetAttribute(coop_topk, cudaFuncAttributeMaxDynamicSharedMemorySize, (size_t)SHP*8); }
 
     void* args[] = { (void*)&x, nullptr };  // placeholder, filled below
+    const float* X = x.data_ptr<float>();
+    float* OV = outv.data_ptr<float>(); long* OI = outi.data_ptr<long>();
+    unsigned long long* A = (unsigned long long*)bufA.data_ptr<long>();
+    unsigned long long* B = (unsigned long long*)bufB.data_ptr<long>();
+    void* kargs[] = { (void*)&X, (void*)&n, (void*)&k, (void*)&RC, (void*)&A, (void*)&B, (void*)&OV, (void*)&OI };
+    dim3 grid(G, batch);
+    dim3 block(BD);
+    cudaError_t err = cudaLaunchCooperativeKernel((void*)coop_topk, grid, block, kargs, (size_t)SHP*8, stream);
+    return (int)err;
+}
+
+static int g_numSM = -1;
+
+// Cooperative single-launch top-k for a single row (batch==1), large k.
+// Returns 0 on success, nonzero (cuda error) so caller can fall back.
+int topk_coop(torch::Tensor x, int k, int RC, torch::Tensor outv, torch::Tensor outi,
+              torch::Tensor bufA, torch::Tensor bufB) {
+    int batch = x.size(0);
+    int n = x.size(1);
+    auto stream = at::cuda::getCurrentCUDAStream();
+    if (g_numSM < 0) {
+        cudaDeviceGetAttribute(&g_numSM, cudaDevAttrMultiProcessorCount, 0);
+    }
+    const int BD = 256;
+    auto np2 = [](int v){ int p=1; while(p<v) p<<=1; return p; };
+    int target = (n + 255) / 256;
+
+    int SHP = RC > 512 ? RC : 512;
+    int maxbpsm = 1;
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&maxbpsm, coop_topk, BD, (size_t)SHP*8);
+    int maxG = maxbpsm * g_numSM / batch;
+    int G = target < maxG ? target : maxG;
+    if (G < 1) G = 1;
+    int chunkA = (n + G - 1) / G;
+    int CP = np2(chunkA);
+    SHP = CP > RC ? CP : RC;
+    cudaFuncSetAttribute(coop_topk, cudaFuncAttributeMaxDynamicSharedMemorySize, (size_t)SHP*8);
+    cudaOccupancyMaxActiveBlocksPerMultiprocessor(&maxbpsm, coop_topk, BD, (size_t)SHP*8);
+    maxG = maxbpsm * g_numSM / batch;
+    if (G > maxG) { G = maxG; if (G < 1) G = 1; chunkA = (n + G - 1) / G; CP = np2(chunkA); SHP = CP > RC ? CP : RC; cudaFuncSetAttribute(coop_topk, cudaFuncAttributeMaxDynamicSharedMemorySize, (size_t)SHP*8); }
+
     const float* X = x.data_ptr<float>();
     float* OV = outv.data_ptr<float>(); long* OI = outi.data_ptr<long>();
     unsigned long long* A = (unsigned long long*)bufA.data_ptr<long>();

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -471,13 +471,15 @@
 
 _CPP = ("void topk_cuda(torch::Tensor x, int k, int bpr, int BD, int BD2, int CHUNK, int CHUNK2, "
         "torch::Tensor outv, torch::Tensor outi, torch::Tensor sAv, torch::Tensor sAi, "
-        "torch::Tensor sBv, torch::Tensor sBi);")
+        "torch::Tensor sBv, torch::Tensor sBi);\n"
+        "int topk_coop(torch::Tensor x, int k, int RC, torch::Tensor outv, torch::Tensor outi, "
+        "torch::Tensor bufA, torch::Tensor bufB);")
 
 _mod = load_inline(
     name="topk_hybrid",
     cpp_sources=_CPP,
     cuda_sources=_CUDA,
-    functions=["topk_cuda"],
+    functions=["topk_cuda", "topk_coop"],
     extra_cuda_cflags=["-O3", "--use_fast_math"],
     verbose=False,
 )

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -536,6 +536,15 @@
         self.CHUNK2 = 512   # reduce-chunk target (RC)
         self.BD2 = 0
         self._alloc()
+        # single-row large-k: try cooperative single-launch kernel
+        self.use_coop = (batch == 1)
+        if self.use_coop:
+            self.RC = 512
+            Gmax = (n + 255) // 256
+            cap = Gmax * k  # ull buffer per row (batch==1)
+            self.coopA = torch.empty(cap, dtype=torch.int64, device="cuda")
+            self.coopB = torch.empty(cap, dtype=torch.int64, device="cuda")
+            self._coop_ok = True
 
     def _alloc(self):
         # preallocate persistent scratch + output buffers (avoid per-call alloc)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -505,6 +505,7 @@
         self.batch, self.n, self.k = batch, n, k
         self.register_buffer("_dummy", torch.zeros(1))
         self.BD = self.BD2 = self.CHUNK = self.CHUNK2 = 0
+        self.use_coop = False
 
         if k == 1:
             self.bpr = max(1, min(round(188 / batch), (n + 1023) // 1024))

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -559,6 +559,12 @@
         self.outi = torch.empty(self.batch, self.k, dtype=torch.int64, device=dev)
 
     def forward(self, x: torch.Tensor):
+        if self.use_coop and self._coop_ok:
+            err = _mod.topk_coop(x, self.k, self.RC, self.outv, self.outi,
+                                 self.coopA, self.coopB)
+            if err == 0:
+                return self.outv, self.outi
+            self._coop_ok = False  # fall back permanently
         _mod.topk_cuda(x, self.k, self.bpr, self.BD, self.BD2, self.CHUNK, self.CHUNK2,
                        self.outv, self.outi, self.sAv, self.sAi, self.sBv, self.sBi)
         return self.outv, self.outi

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -537,8 +537,9 @@
         self.CHUNK2 = 512   # reduce-chunk target (RC)
         self.BD2 = 0
         self._alloc()
-        # single-row large-k: try cooperative single-launch kernel
-        self.use_coop = (batch == 1)
+        # single-row large-k: cooperative single-launch kernel was slower than
+        # multi-pass relaunch (grid.sync barrier across many idle blocks), disabled.
+        self.use_coop = False
         if self.use_coop:
             self.RC = 512
             Gmax = (n + 255) // 256

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -515,7 +515,7 @@
         if k <= 8:
             # register top-k + tree merge; BD=128 sweet spot
             BD = 128
-            bpr = max(1, round(n / 1536))
+            bpr = max(1, round(n / 1365))
             while bpr * k > 2048 and bpr > 1:
                 bpr -= 1
             while bpr > 1 and (n + bpr - 1) // bpr < BD:

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -515,7 +515,7 @@
         if k <= 8:
             # register top-k + tree merge; BD=128 sweet spot
             BD = 128
-            bpr = max(1, round(n / 1365))
+            bpr = max(1, round(n / 1536))
             while bpr * k > 2048 and bpr > 1:
                 bpr -= 1
             while bpr > 1 and (n + bpr - 1) // bpr < BD:

--- a//home/agent/.claude/projects/-workspace-problems-05-topk-bitonic/memory/topk-bitonic-problem.md
+++ b//home/agent/.claude/projects/-workspace-problems-05-topk-bitonic/memory/topk-bitonic-problem.md
@@ -20,4 +20,6 @@
 - k<=8: per-thread register top-k (threshold-gated insert, templated K, unrolled) + pairwise tree-merge in shared. BD=128 sweet spot. 2 kernels.
 - k>=16: bitonic sort in shared with PACKED uint64 keys (f2ord(val)<<32|idx) — single 64-bit compare-swap halves shared traffic, ~1.5x speedup. MULTI-PASS reduction: small phase1 chunk (CHUNK1=512, many blocks fill SMs, less total sort work) -> reduce passes (RC=1024) -> final 1-block sort.
 KEY LESSONS: (1) ~188 blocks (one wave) best, more=slower launch. (2) phase1 bitonic was occupancy-bound (waves 0.17) not compute/mem — smaller chunks = more blocks AND less log^2 work. (3) per-thread rv[K] register array SPILLS for K>=32 -> bitonic better for large k. (4) PREALLOCATE scratch+output buffers in __init__ (not per forward) - torch.empty per call costs ~6us CPU/GPU-idle, huge for these tiny kernels. (5) single-row shape0 tail (reduce+final) is serial, inherently limited ~34us.
-Tuning: shape0 prefers CHUNK1=1024 (smaller tail), shape2 prefers 512 (more blocks).
+Tuning (FINAL, geomean ~0.034): large-k bitonic CHUNK1=256, RC(CHUNK2)=512 uniformly best. k<=8 regmerge BD=128 bpr=round(n/1536). k=1 bpr=1 -> single argmax kernel (no phase2). Per-shape fractions: s0~0.009(30us), s1~0.049, s2~0.037, s3~0.019, s4~0.14.
+FAILED experiments (don't retry): (a) regmerge_db scalar double-buffer merge for large k — read-pass rv[K] still spills, 722us. (b) cooperative single-launch kernel (grid.sync) for shape0 — SLOWER than multi-pass relaunch (0.0071 vs 0.009) because grid.sync barrier across many idle blocks is expensive. (c) bpr=6 for s1 worse than bpr=5 post-prealloc.
+shape0 (single row k=64) is the floor-limiting min ~30us via 4-pass bitonic; tail is launch-overhead bound but multi-pass still beats alternatives.