--- /dev/null
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -0,0 +1,655 @@
+"""Top-k via streaming CUDA kernel with warp-shuffle merge.
+
+Strategy:
+  - One thread block per batch element (or multiple chunks per row for large n).
+  - Each thread streams through its slice, maintaining a sorted top-k in registers.
+  - Within-warp reduction via __shfl_down_sync.
+  - Cross-warp merge via shared memory.
+  - For large n with batch=1, use multiple chunks + a merge pass.
+"""
+
+import math
+import os
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.utils.cpp_extension
+
+OP_TYPE = "topk"
+SUPPORTED_PRECISIONS = ["fp32"]
+HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
+
+# ---------------------------------------------------------------------------
+# CUDA kernel source
+# ---------------------------------------------------------------------------
+
+_kernel_src = """
+#include <cuda_runtime.h>
+#include <cstdint>
+#include <cfloat>
+
+// ---------------------------------------------------------------------------
+// Device helpers: sorted (descending) top-k insertion and merge
+// ---------------------------------------------------------------------------
+
+// Insert (val, idx) into a sorted-descending array of size k.
+// Returns true if the value was inserted.
+__device__ __forceinline__ bool insert_topk(
+    float val, int idx,
+    float* vals, int* indices, int k
+) {
+    if (val <= vals[k - 1]) return false;
+
+    // Binary search for insertion position in descending array
+    int lo = 0, hi = k - 1;
+    #pragma unroll
+    while (lo < hi) {
+        int mid = (lo + hi) >> 1;
+        if (val > vals[mid])
+            hi = mid;
+        else
+            lo = mid + 1;
+    }
+
+    // Shift tail down and insert
+    #pragma unroll
+    for (int i = k - 1; i > lo; --i) {
+        vals[i] = vals[i - 1];
+        indices[i] = indices[i - 1];
+    }
+    vals[lo] = val;
+    indices[lo] = idx;
+    return true;
+}
+
+// Merge two sorted-descending lists a[0..k-1] and b[0..k-1] into dst[0..k-1],
+// keeping only the largest k elements.
+__device__ __forceinline__ void merge_two_k(
+    const float* __restrict__ a_vals, const int* __restrict__ a_idx,
+    const float* __restrict__ b_vals, const int* __restrict__ b_idx,
+    float* __restrict__ dst_vals, int* __restrict__ dst_idx,
+    int k
+) {
+    int ia = 0, ib = 0, io = 0;
+    while (io < k) {
+        bool take_a;
+        if (ia >= k)      take_a = false;
+        else if (ib >= k) take_a = true;
+        else              take_a = (a_vals[ia] >= b_vals[ib]);
+        dst_vals[io] = take_a ? a_vals[ia] : b_vals[ib];
+        dst_idx[io] = take_a ? a_idx[ia] : b_idx[ib];
+        ia += take_a ? 1 : 0;
+        ib += take_a ? 0 : 1;
+        io++;
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Kernel: streaming top-k with warp-shuffle + shared-memory reduction
+//
+// Grid:   (batch * num_chunks,)
+// Block:  N_THREADS (must be multiple of 32)
+//
+// Each block processes 'chunk_size' elements of one row, producing a sorted
+// top-k result in global memory.
+// ---------------------------------------------------------------------------
+
+template <int K, int N_THREADS, int N_WARPS>
+__global__ void topk_stream_kernel(
+    const float* __restrict__ x,      // (batch, n)
+    float* __restrict__ out_vals,     // (batch * num_chunks, K)
+    int64_t* __restrict__ out_idx,    // (batch * num_chunks, K)
+    int n,
+    int chunk_size,
+    int n_chunks_per_row
+) {
+    int global_bid = blockIdx.x;
+    int batch_idx = global_bid / n_chunks_per_row;
+    int chunk_idx = global_bid % n_chunks_per_row;
+    int tid = threadIdx.x;
+
+    const float* row_ptr = x + (int64_t)batch_idx * n;
+    int start = chunk_idx * chunk_size;
+    int end = min(start + chunk_size, n);
+    if (start >= end) return; // empty chunk (shouldn't happen)
+
+    // ---- Phase 1: streaming with per-thread local top-k ----
+    float local_vals[K];
+    int   local_idx[K];
+
+    // Init to -inf
+    #pragma unroll
+    for (int i = 0; i < K; ++i) {
+        local_vals[i] = -CUDART_INF_F;
+        local_idx[i] = -1;
+    }
+
+    // Each thread processes its strided slice
+    for (int pos = start + tid; pos < end; pos += N_THREADS) {
+        float v = row_ptr[pos];
+        // Only try to insert — the function checks the threshold
+        insert_topk(v, pos, local_vals, local_idx, K);
+    }
+
+    // ---- Phase 2: within-warp reduction via __shfl_down_sync ----
+    // Tree reduction: merge pairs of threads within each warp
+    const int warp_id = tid / 32;
+    const int lane_id = tid & 31;
+    unsigned mask = __activemask();
+
+    #pragma unroll
+    for (int offset = 16; offset > 0; offset >>= 1) {
+        // Shuffle partner's top-k values and indices
+        float partner_vals[K];
+        int   partner_idx[K];
+        #pragma unroll
+        for (int i = 0; i < K; ++i) {
+            partner_vals[i] = __shfl_down_sync(mask, local_vals[i], offset);
+            partner_idx[i] = __shfl_down_sync(mask, local_idx[i], offset);
+        }
+        // Merge if we're the receiving lane
+        if (lane_id < offset) {
+            float merged_vals[K];
+            int   merged_idx[K];
+            merge_two_k(local_vals, local_idx, partner_vals, partner_idx,
+                        merged_vals, merged_idx, K);
+            #pragma unroll
+            for (int i = 0; i < K; ++i) {
+                local_vals[i] = merged_vals[i];
+                local_idx[i] = merged_idx[i];
+            }
+        }
+    }
+
+    // After warp reduction, lane 0 of each warp has the warp's top-k
+    // ---- Phase 3: cross-warp reduction via shared memory ----
+    __shared__ float smem_vals[N_WARPS * K];
+    __shared__ int   smem_idx[N_WARPS * K];
+
+    if (lane_id == 0) {
+        #pragma unroll
+        for (int i = 0; i < K; ++i) {
+            smem_vals[warp_id * K

... (17194 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -402,7 +402,9 @@
 _extra_cflags = [
     "-O3",
     "--use_fast_math",
-    f"-arch={_arch}",
+    f"-arch=compute_{_sm_major}{_sm_minor}a",
+    "-code",
+    f"sm_{_sm_major}{_sm_minor}a",
     "--expt-relaxed-constexpr",
     "-std=c++17",
     "-DNDEBUG",

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -402,9 +402,7 @@
 _extra_cflags = [
     "-O3",
     "--use_fast_math",
-    f"-arch=compute_{_sm_major}{_sm_minor}a",
-    "-code",
-    f"sm_{_sm_major}{_sm_minor}a",
+    f"-arch=sm_{_sm_major}{_sm_minor}",
     "--expt-relaxed-constexpr",
     "-std=c++17",
     "-DNDEBUG",

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -9,7 +9,6 @@
 """
 
 import math
-import os
 from typing import List
 
 import torch
@@ -21,30 +20,32 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 # ---------------------------------------------------------------------------
-# CUDA kernel source
-# ---------------------------------------------------------------------------
-
-_kernel_src = """
+# CUDA kernel source — main top-k streaming + merge
+# ---------------------------------------------------------------------------
+
+_kernel_src = r"""
 #include <cuda_runtime.h>
 #include <cstdint>
 #include <cfloat>
+#include <math.h>
 
 // ---------------------------------------------------------------------------
 // Device helpers: sorted (descending) top-k insertion and merge
 // ---------------------------------------------------------------------------
 
-// Insert (val, idx) into a sorted-descending array of size k.
-// Returns true if the value was inserted.
+// Insert (val, idx) into a sorted-descending array of size k (k is a template
+// parameter, known at compile time).
 __device__ __forceinline__ bool insert_topk(
     float val, int idx,
     float* vals, int* indices, int k
 ) {
     if (val <= vals[k - 1]) return false;
 
-    // Binary search for insertion position in descending array
+    // Binary search for insertion position in descending array.
+    // k is small (<= 64) so this is fast.
     int lo = 0, hi = k - 1;
-    #pragma unroll
-    while (lo < hi) {
+    for (int step = 0; step < 10; ++step) {  // enough for k <= 64 (log2 64 = 6)
+        if (lo >= hi) break;
         int mid = (lo + hi) >> 1;
         if (val > vals[mid])
             hi = mid;
@@ -53,7 +54,6 @@
     }
 
     // Shift tail down and insert
-    #pragma unroll
     for (int i = k - 1; i > lo; --i) {
         vals[i] = vals[i - 1];
         indices[i] = indices[i - 1];
@@ -91,7 +91,7 @@
 // Grid:   (batch * num_chunks,)
 // Block:  N_THREADS (must be multiple of 32)
 //
-// Each block processes 'chunk_size' elements of one row, producing a sorted
+// Each block processes 'chunk_size' elements of one row, producing its sorted
 // top-k result in global memory.
 // ---------------------------------------------------------------------------
 
@@ -111,8 +111,8 @@
 
     const float* row_ptr = x + (int64_t)batch_idx * n;
     int start = chunk_idx * chunk_size;
-    int end = min(start + chunk_size, n);
-    if (start >= end) return; // empty chunk (shouldn't happen)
+    int end = (start + chunk_size < n) ? (start + chunk_size) : n;
+    if (start >= end) return;
 
     // ---- Phase 1: streaming with per-thread local top-k ----
     float local_vals[K];
@@ -121,26 +121,22 @@
     // Init to -inf
     #pragma unroll
     for (int i = 0; i < K; ++i) {
-        local_vals[i] = -CUDART_INF_F;
+        local_vals[i] = -INFINITY;
         local_idx[i] = -1;
     }
 
     // Each thread processes its strided slice
     for (int pos = start + tid; pos < end; pos += N_THREADS) {
         float v = row_ptr[pos];
-        // Only try to insert — the function checks the threshold
         insert_topk(v, pos, local_vals, local_idx, K);
     }
 
     // ---- Phase 2: within-warp reduction via __shfl_down_sync ----
-    // Tree reduction: merge pairs of threads within each warp
-    const int warp_id = tid / 32;
     const int lane_id = tid & 31;
     unsigned mask = __activemask();
 
     #pragma unroll
     for (int offset = 16; offset > 0; offset >>= 1) {
-        // Shuffle partner's top-k values and indices
         float partner_vals[K];
         int   partner_idx[K];
         #pragma unroll
@@ -148,7 +144,6 @@
             partner_vals[i] = __shfl_down_sync(mask, local_vals[i], offset);
             partner_idx[i] = __shfl_down_sync(mask, local_idx[i], offset);
         }
-        // Merge if we're the receiving lane
         if (lane_id < offset) {
             float merged_vals[K];
             int   merged_idx[K];
@@ -162,11 +157,11 @@
         }
     }
 
-    // After warp reduction, lane 0 of each warp has the warp's top-k
     // ---- Phase 3: cross-warp reduction via shared memory ----
     __shared__ float smem_vals[N_WARPS * K];
     __shared__ int   smem_idx[N_WARPS * K];
 
+    const int warp_id = tid / 32;
     if (lane_id == 0) {
         #pragma unroll
         for (int i = 0; i < K; ++i) {
@@ -176,24 +171,31 @@
     }
     __syncthreads();
 
-    // Tree merge across warps
-    #pragma unroll
-    for (int stride = N_WARPS / 2; stride > 0; stride >>= 1) {
-        if (warp_id < stride && lane_id == 0) {
-            float merged_vals[K];
-            int   merged_idx[K];
-            merge_two_k(
-                &smem_vals[warp_id * K], &smem_idx[warp_id * K],
-                &smem_vals[(warp_id + stride) * K], &smem_idx[(warp_id + stride) * K],
-                merged_vals, merged_idx, K
-            );
-            #pragma unroll
-            for (int i = 0; i < K; ++i) {
-                smem_vals[warp_id * K + i] = merged_vals[i];
-                smem_idx[warp_id * K + i] = merged_idx[i];
+    // Tree merge across warps (lane 0 of each warp participates)
+    if (lane_id == 0) {
+        #pragma unroll
+        for (int stride = N_WARPS / 2; stride > 0; stride >>= 1) {
+            if (warp_id < stride) {
+                float merged_vals[K];
+                int   merged_idx[K];
+                merge_two_k(
+                    &smem_vals[warp_id * K], &smem_idx[warp_id * K],
+                    &smem_vals[(warp_id + stride) * K], &smem_idx[(warp_id + stride) * K],
+                    merged_vals, merged_idx, K
+                );
+                #pragma unroll
+                for (int i = 0; i < K; ++i) {
+                    smem_vals[warp_id * K + i] = merged_vals[i];
+                    smem_idx[warp_id * K + i] = merged_idx[i];
+                }
             }
-        }
-        __syncthreads();
+ 

... (17152 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -370,6 +370,21 @@
 } // extern "C"
 """
 
+# C++ header with forward declarations for pybind11
+_cpp_src = """
+#include <cuda_runtime.h>
+#include <cstdint>
+
+extern "C" {
+void launch_topk_stream_64(const float*, float*, int64_t*, int, int, int, int, cudaStream_t);
+void launch_topk_stream_32(const float*, float*, int64_t*, int, int, int, int, cudaStream_t);
+void launch_topk_stream_16(const float*, float*, int64_t*, int, int, int, int, cudaStream_t);
+void launch_topk_stream_8(const float*, float*, int64_t*, int, int, int, int, cudaStream_t);
+void launch_topk_stream_1(const float*, float*, int64_t*, int, int, int, int, cudaStream_t);
+void launch_topk_merge_64_64(const float*, const int64_t*, float*, int64_t*, int, cudaStream_t);
+}
+"""
+
 # ---------------------------------------------------------------------------
 # Generic merge kernel (for arbitrary chunk counts)
 # ---------------------------------------------------------------------------

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -491,7 +491,7 @@
     if _topk_module is None:
         _topk_module = torch.utils.cpp_extension.load_inline(
             name="topk_bitonic_kernel",
-            cpp_sources="",
+            cpp_sources=_cpp_src,
             cuda_sources=_kernel_src,
             functions=[
                 "launch_topk_stream_64",

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -465,6 +465,15 @@
 }
 """
 
+_merge_cpp_src = """
+#include <cuda_runtime.h>
+#include <cstdint>
+
+extern "C" {
+void launch_topk_merge_generic(const float*, const int64_t*, float*, int64_t*, int, int, int, cudaStream_t);
+}
+"""
+
 # ---------------------------------------------------------------------------
 # Compile the CUDA extensions
 # ---------------------------------------------------------------------------

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -521,7 +521,7 @@
     if _merge_module is None:
         _merge_module = torch.utils.cpp_extension.load_inline(
             name="topk_merge_generic",
-            cpp_sources="",
+            cpp_sources=_merge_cpp_src,
             cuda_sources=_merge_generic_src,
             functions=["launch_topk_merge_generic"],
             extra_cuda_cflags=_extra_cflags,

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -376,12 +376,12 @@
 #include <cstdint>
 
 extern "C" {
-void launch_topk_stream_64(const float*, float*, int64_t*, int, int, int, int, cudaStream_t);
-void launch_topk_stream_32(const float*, float*, int64_t*, int, int, int, int, cudaStream_t);
-void launch_topk_stream_16(const float*, float*, int64_t*, int, int, int, int, cudaStream_t);
-void launch_topk_stream_8(const float*, float*, int64_t*, int, int, int, int, cudaStream_t);
-void launch_topk_stream_1(const float*, float*, int64_t*, int, int, int, int, cudaStream_t);
-void launch_topk_merge_64_64(const float*, const int64_t*, float*, int64_t*, int, cudaStream_t);
+void launch_topk_stream_64(const float*, float*, int64_t*, int, int, int, int, void*);
+void launch_topk_stream_32(const float*, float*, int64_t*, int, int, int, int, void*);
+void launch_topk_stream_16(const float*, float*, int64_t*, int, int, int, int, void*);
+void launch_topk_stream_8(const float*, float*, int64_t*, int, int, int, int, void*);
+void launch_topk_stream_1(const float*, float*, int64_t*, int, int, int, int, void*);
+void launch_topk_merge_64_64(const float*, const int64_t*, float*, int64_t*, int, void*);
 }
 """
 

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -376,12 +376,12 @@
 #include <cstdint>
 
 extern "C" {
-void launch_topk_stream_64(const float*, float*, int64_t*, int, int, int, int, void*);
-void launch_topk_stream_32(const float*, float*, int64_t*, int, int, int, int, void*);
-void launch_topk_stream_16(const float*, float*, int64_t*, int, int, int, int, void*);
-void launch_topk_stream_8(const float*, float*, int64_t*, int, int, int, int, void*);
-void launch_topk_stream_1(const float*, float*, int64_t*, int, int, int, int, void*);
-void launch_topk_merge_64_64(const float*, const int64_t*, float*, int64_t*, int, void*);
+void launch_topk_stream_64(const float*, float*, int64_t*, int, int, int, int, int64_t);
+void launch_topk_stream_32(const float*, float*, int64_t*, int, int, int, int, int64_t);
+void launch_topk_stream_16(const float*, float*, int64_t*, int, int, int, int, int64_t);
+void launch_topk_stream_8(const float*, float*, int64_t*, int, int, int, int, int64_t);
+void launch_topk_stream_1(const float*, float*, int64_t*, int, int, int, int, int64_t);
+void launch_topk_merge_64_64(const float*, const int64_t*, float*, int64_t*, int, int64_t);
 }
 """
 

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -19,9 +19,9 @@
 SUPPORTED_PRECISIONS = ["fp32"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
-# ---------------------------------------------------------------------------
+# ===========================================================================
 # CUDA kernel source — main top-k streaming + merge
-# ---------------------------------------------------------------------------
+# ===========================================================================
 
 _kernel_src = r"""
 #include <cuda_runtime.h>
@@ -30,22 +30,19 @@
 #include <math.h>
 
 // ---------------------------------------------------------------------------
-// Device helpers: sorted (descending) top-k insertion and merge
-// ---------------------------------------------------------------------------
-
-// Insert (val, idx) into a sorted-descending array of size k (k is a template
-// parameter, known at compile time).
+// Device helpers
+// ---------------------------------------------------------------------------
+
+// Insert (val, idx) into a sorted-descending array of size k.
 __device__ __forceinline__ bool insert_topk(
     float val, int idx,
     float* vals, int* indices, int k
 ) {
     if (val <= vals[k - 1]) return false;
 
-    // Binary search for insertion position in descending array.
-    // k is small (<= 64) so this is fast.
+    // Binary search for insertion position in descending array
     int lo = 0, hi = k - 1;
-    for (int step = 0; step < 10; ++step) {  // enough for k <= 64 (log2 64 = 6)
-        if (lo >= hi) break;
+    while (lo < hi) {
         int mid = (lo + hi) >> 1;
         if (val > vals[mid])
             hi = mid;
@@ -63,8 +60,7 @@
     return true;
 }
 
-// Merge two sorted-descending lists a[0..k-1] and b[0..k-1] into dst[0..k-1],
-// keeping only the largest k elements.
+// Merge two sorted-descending lists a[0..k-1] and b[0..k-1] into dst[0..k-1].
 __device__ __forceinline__ void merge_two_k(
     const float* __restrict__ a_vals, const int* __restrict__ a_idx,
     const float* __restrict__ b_vals, const int* __restrict__ b_idx,
@@ -87,19 +83,13 @@
 
 // ---------------------------------------------------------------------------
 // Kernel: streaming top-k with warp-shuffle + shared-memory reduction
-//
-// Grid:   (batch * num_chunks,)
-// Block:  N_THREADS (must be multiple of 32)
-//
-// Each block processes 'chunk_size' elements of one row, producing its sorted
-// top-k result in global memory.
 // ---------------------------------------------------------------------------
 
 template <int K, int N_THREADS, int N_WARPS>
 __global__ void topk_stream_kernel(
-    const float* __restrict__ x,      // (batch, n)
-    float* __restrict__ out_vals,     // (batch * num_chunks, K)
-    int64_t* __restrict__ out_idx,    // (batch * num_chunks, K)
+    const float* __restrict__ x,
+    float* __restrict__ out_vals,
+    int64_t* __restrict__ out_idx,
     int n,
     int chunk_size,
     int n_chunks_per_row
@@ -114,18 +104,16 @@
     int end = (start + chunk_size < n) ? (start + chunk_size) : n;
     if (start >= end) return;
 
-    // ---- Phase 1: streaming with per-thread local top-k ----
+    // ---- Phase 1: per-thread local top-k ----
     float local_vals[K];
     int   local_idx[K];
 
-    // Init to -inf
     #pragma unroll
     for (int i = 0; i < K; ++i) {
         local_vals[i] = -INFINITY;
         local_idx[i] = -1;
     }
 
-    // Each thread processes its strided slice
     for (int pos = start + tid; pos < end; pos += N_THREADS) {
         float v = row_ptr[pos];
         insert_topk(v, pos, local_vals, local_idx, K);
@@ -171,7 +159,7 @@
     }
     __syncthreads();
 
-    // Tree merge across warps (lane 0 of each warp participates)
+    // Lane 0 of each warp does the tree merge across warps
     if (lane_id == 0) {
         #pragma unroll
         for (int stride = N_WARPS / 2; stride > 0; stride >>= 1) {
@@ -192,7 +180,6 @@
             __syncthreads();
         }
     } else {
-        // Other lanes just sync
         for (int stride = N_WARPS / 2; stride > 0; stride >>= 1) {
             __syncthreads();
         }
@@ -211,15 +198,15 @@
 
 
 // ---------------------------------------------------------------------------
-// Kernel: merge chunk results into final top-k via bitonic sort in shared mem
+// Kernel: merge chunk results via bitonic sort in shared memory
 // ---------------------------------------------------------------------------
 
 template <int K, int N_CHUNKS>
 __global__ void topk_merge_kernel(
-    const float* __restrict__ chunk_vals,   // (batch, N_CHUNKS, K)
-    const int64_t* __restrict__ chunk_idx,  // (batch, N_CHUNKS, K)
-    float* __restrict__ out_vals,           // (batch, K)
-    int64_t* __restrict__ out_idx           // (batch, K)
+    const float* __restrict__ chunk_vals,
+    const int64_t* __restrict__ chunk_idx,
+    float* __restrict__ out_vals,
+    int64_t* __restrict__ out_idx
 ) {
     int batch_idx = blockIdx.x;
     int tid = threadIdx.x;
@@ -235,7 +222,7 @@
     }
     __syncthreads();
 
-    // Bitonic sort the shared memory buffer (total elements), descending.
+    // Bitonic sort (descending)
     int n_pow2 = 1;
     while (n_pow2 < total) n_pow2 <<= 1;
 
@@ -253,8 +240,6 @@
                         swap_needed = smem_vals[idx_i] > smem_vals[idx_j];
                     else
                         swap_needed = smem_vals[idx_i] < smem_vals[idx_j];
-                } else if (idx_i < total && idx_j >= total) {
-                    swap_needed = false;  // padded element is -inf for descending
                 }
                 if (swap_needed) {
                     float tv = smem_vals[idx_i];
@@ -278,7 +263,7 @@
 
 
 // ---------------------------------------------------------------------------
-// Host launchers — extern "C" so pybind11 finds them
+// Host launchers — extern "C" with 

... (9208 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,14 +1,14 @@
-"""Top-k via streaming CUDA kernel with warp-shuffle merge.
-
-Strategy:
-  - One thread block per batch element (or multiple chunks per row for large n).
-  - Each thread streams through its slice, maintaining a sorted top-k in registers.
+"""Top-k via CUDA kernel with warp-shuffle merge.
+
+Approach:
+  - For batch >= 1: one thread block per batch element.
+  - Each thread streams through its slice, maintaining sorted top-k in registers.
   - Within-warp reduction via __shfl_down_sync.
   - Cross-warp merge via shared memory.
-  - For large n with batch=1, use multiple chunks + a merge pass.
+  - For the large (1, 131072, 64) case: use multiple chunks + merge pass.
 """
 
-import math
+import os
 from typing import List
 
 import torch
@@ -20,37 +20,43 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 # ===========================================================================
-# CUDA kernel source — main top-k streaming + merge
-# ===========================================================================
-
-_kernel_src = r"""
+# Combined C++/CUDA source (cpp_sources provides declarations for pybind11)
+# ===========================================================================
+
+_cpp_src = r"""
+#include <cstdint>
+
+// Forward declarations for pybind11 — no stream needed (use default stream)
+void launch_topk_stream_64(const float*, float*, int64_t*, int, int, int, int);
+void launch_topk_stream_32(const float*, float*, int64_t*, int, int, int, int);
+void launch_topk_stream_16(const float*, float*, int64_t*, int, int, int, int);
+void launch_topk_stream_8(const float*, float*, int64_t*, int, int, int, int);
+void launch_topk_stream_1(const float*, float*, int64_t*, int, int, int, int);
+void launch_topk_merge_64_64(const float*, const int64_t*, float*, int64_t*, int);
+void launch_topk_merge_generic(const float*, const int64_t*, float*, int64_t*, int, int, int);
+"""
+
+_cuda_src = r"""
 #include <cuda_runtime.h>
 #include <cstdint>
 #include <cfloat>
 #include <math.h>
 
-// ---------------------------------------------------------------------------
+// ===========================================================================
 // Device helpers
-// ---------------------------------------------------------------------------
+// ===========================================================================
 
 // Insert (val, idx) into a sorted-descending array of size k.
 __device__ __forceinline__ bool insert_topk(
-    float val, int idx,
-    float* vals, int* indices, int k
+    float val, int idx, float* vals, int* indices, int k
 ) {
     if (val <= vals[k - 1]) return false;
-
-    // Binary search for insertion position in descending array
     int lo = 0, hi = k - 1;
     while (lo < hi) {
         int mid = (lo + hi) >> 1;
-        if (val > vals[mid])
-            hi = mid;
-        else
-            lo = mid + 1;
-    }
-
-    // Shift tail down and insert
+        if (val > vals[mid]) hi = mid;
+        else                 lo = mid + 1;
+    }
     for (int i = k - 1; i > lo; --i) {
         vals[i] = vals[i - 1];
         indices[i] = indices[i - 1];
@@ -60,469 +66,289 @@
     return true;
 }
 
-// Merge two sorted-descending lists a[0..k-1] and b[0..k-1] into dst[0..k-1].
+// Merge two sorted-descending lists of size k, keeping the largest k.
 __device__ __forceinline__ void merge_two_k(
-    const float* __restrict__ a_vals, const int* __restrict__ a_idx,
-    const float* __restrict__ b_vals, const int* __restrict__ b_idx,
-    float* __restrict__ dst_vals, int* __restrict__ dst_idx,
-    int k
+    const float* a_vals, const int* a_idx,
+    const float* b_vals, const int* b_idx,
+    float* dst_vals, int* dst_idx, int k
 ) {
     int ia = 0, ib = 0, io = 0;
     while (io < k) {
-        bool take_a;
-        if (ia >= k)      take_a = false;
-        else if (ib >= k) take_a = true;
-        else              take_a = (a_vals[ia] >= b_vals[ib]);
-        dst_vals[io] = take_a ? a_vals[ia] : b_vals[ib];
-        dst_idx[io] = take_a ? a_idx[ia] : b_idx[ib];
-        ia += take_a ? 1 : 0;
-        ib += take_a ? 0 : 1;
+        bool ta;
+        if (ia >= k)      ta = false;
+        else if (ib >= k) ta = true;
+        else              ta = (a_vals[ia] >= b_vals[ib]);
+        dst_vals[io] = ta ? a_vals[ia] : b_vals[ib];
+        dst_idx[io]  = ta ? a_idx[ia]  : b_idx[ib];
+        ia += ta ? 1 : 0;
+        ib += ta ? 0 : 1;
         io++;
     }
 }
 
-// ---------------------------------------------------------------------------
-// Kernel: streaming top-k with warp-shuffle + shared-memory reduction
-// ---------------------------------------------------------------------------
+// ===========================================================================
+// Streaming top-k kernel
+// ===========================================================================
 
 template <int K, int N_THREADS, int N_WARPS>
 __global__ void topk_stream_kernel(
     const float* __restrict__ x,
     float* __restrict__ out_vals,
     int64_t* __restrict__ out_idx,
-    int n,
-    int chunk_size,
-    int n_chunks_per_row
+    int n, int chunk_size, int n_chunks_per_row
 ) {
     int global_bid = blockIdx.x;
-    int batch_idx = global_bid / n_chunks_per_row;
-    int chunk_idx = global_bid % n_chunks_per_row;
-    int tid = threadIdx.x;
-
-    const float* row_ptr = x + (int64_t)batch_idx * n;
+    int batch_idx   = global_bid / n_chunks_per_row;
+    int chunk_idx   = global_bid % n_chunks_per_row;
+    int tid         = threadIdx.x;
+
+    int64_t row_off = (int64_t)batch_idx * n;
     int start = chunk_idx * chunk_size;
-    int end = (start + chunk_size < n) ? (start + chunk_size) : n;
+    int end   = min(start + chunk_size, n);
     if (start >= end) return;
 
     // ---- Phase 1: per-thread local top-k ----
     float local_vals[

... (26345 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -26,7 +26,7 @@
 _cpp_src = r"""
 #include <cstdint>
 
-// Forward declarations for pybind11 — no stream needed (use default stream)
+extern "C" {
 void launch_topk_stream_64(const float*, float*, int64_t*, int, int, int, int);
 void launch_topk_stream_32(const float*, float*, int64_t*, int, int, int, int);
 void launch_topk_stream_16(const float*, float*, int64_t*, int, int, int, int);
@@ -34,6 +34,7 @@
 void launch_topk_stream_1(const float*, float*, int64_t*, int, int, int, int);
 void launch_topk_merge_64_64(const float*, const int64_t*, float*, int64_t*, int);
 void launch_topk_merge_generic(const float*, const int64_t*, float*, int64_t*, int, int, int);
+}
 """
 
 _cuda_src = r"""

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -251,27 +251,27 @@
 
 void launch_topk_64(const float* x, float* ov, int64_t* oi,
                     int B, int n, int cs, int nc) {
-    constexpr int K=64, NT=128, NE=NT*K; // 8192
+    constexpr int K=64, NT=64, NE=NT*K; // 4096
     topk_kernel<K,NT,NE><<<B*nc, NT, 0, 0>>>(x, ov, oi, n, cs, nc);
 }
 void launch_topk_32(const float* x, float* ov, int64_t* oi,
                     int B, int n, int cs, int nc) {
-    constexpr int K=32, NT=128, NE=NT*K; // 4096
+    constexpr int K=32, NT=64, NE=NT*K; // 2048
     topk_kernel<K,NT,NE><<<B*nc, NT, 0, 0>>>(x, ov, oi, n, cs, nc);
 }
 void launch_topk_16(const float* x, float* ov, int64_t* oi,
                     int B, int n, int cs, int nc) {
-    constexpr int K=16, NT=128, NE=NT*K; // 2048
+    constexpr int K=16, NT=64, NE=NT*K; // 1024
     topk_kernel<K,NT,NE><<<B*nc, NT, 0, 0>>>(x, ov, oi, n, cs, nc);
 }
 void launch_topk_8(const float* x, float* ov, int64_t* oi,
                    int B, int n, int cs, int nc) {
-    constexpr int K=8, NT=128, NE=NT*K; // 1024
+    constexpr int K=8, NT=64, NE=NT*K; // 512
     topk_kernel<K,NT,NE><<<B*nc, NT, 0, 0>>>(x, ov, oi, n, cs, nc);
 }
 void launch_topk_1(const float* x, float* ov, int64_t* oi,
                    int B, int n, int cs, int nc) {
-    constexpr int K=1, NT=128, NE=NT*K; // 128
+    constexpr int K=1, NT=64, NE=NT*K; // 64
     topk_kernel<K,NT,NE><<<B*nc, NT, 0, 0>>>(x, ov, oi, n, cs, nc);
 }
 void launch_merge_64_64(const float* cv, const int64_t* ci,

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -72,30 +72,29 @@
 
 // ---------------------------------------------------------------------------
 // Bitonic sort in shared memory (descending)
-// Helper macro-style function: all threads must participate, with
-// __syncthreads() between stages.
-// ---------------------------------------------------------------------------
-template <int N>
-__device__ void bitonic_sort_desc(__shared__ float* sv, __shared__ int* si,
-                                   int n_actual, int tid, int n_threads) {
-    int n = 1; while (n < N) n <<= 1;  // padded size
+// All threads participate; __syncthreads() between stages.
+// ---------------------------------------------------------------------------
+template <int N_ELEM>
+__device__ void bitonic_sort_desc(int* sv_as_int, int* si, int tid, int n_threads) {
+    // sv_as_int is float values bitcast to int, si are indices
+    int n = N_ELEM;
     for (int stage = 2; stage <= n; stage <<= 1) {
         for (int step = stage >> 1; step > 0; step >>= 1) {
-            int np = n >> 1;
-            for (int p = tid; p < np; p += n_threads) {
-                int i = p, j = p + step;
-                bool asc = ((p / stage) & 1) == 0;
-                bool sw = false;
-                if (i < n_actual && j < n_actual)
-                    sw = asc ? (sv[i] > sv[j]) : (sv[i] < sv[j]);
-                // Padded elements: treat as -inf for descending
-                else if (i < n_actual && j >= n_actual)
-                    sw = false;
-                else if (i >= n_actual && j < n_actual)
-                    sw = true;
-                if (sw) {
-                    float tv = sv[i]; sv[i] = sv[j]; sv[j] = tv;
-                    int   ti = si[i]; si[i] = si[j]; si[j] = ti;
+            for (int i = tid; i < n; i += n_threads) {
+                int ixj = i ^ step;
+                if (ixj > i) {
+                    bool asc = ((i & stage) == 0);
+                    float vi = __int_as_float(sv_as_int[i]);
+                    float vj = __int_as_float(sv_as_int[ixj]);
+                    bool swap = asc ? (vi > vj) : (vi < vj);
+                    if (swap) {
+                        int ti = sv_as_int[i];
+                        sv_as_int[i] = sv_as_int[ixj];
+                        sv_as_int[ixj] = ti;
+                        int tii = si[i];
+                        si[i] = si[ixj];
+                        si[ixj] = tii;
+                    }
                 }
             }
             __syncthreads();

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,14 +1,13 @@
-"""Top-k via CUDA kernel with shared-memory merge.
+"""Top-k via CUDA kernel with shared-memory bitonic merge.
 
 Approach:
-  - One thread block per batch element (or chunk).
+  - Each thread block processes one chunk of a batch element.
   - Each thread streams through its slice, maintaining sorted top-k in registers.
-  - All threads write local top-k to shared memory.
+  - All threads write local top-k to shared memory (N_THREADS * K elements).
   - Parallel bitonic sort in shared memory to find global top-k.
   - For large n with batch=1: multi-chunk + merge pass.
 """
 
-import os
 from typing import List
 
 import torch
@@ -48,9 +47,8 @@
 #include <math.h>
 
 // ---------------------------------------------------------------------------
-// Device helpers
-// ---------------------------------------------------------------------------
-
+// Insert into sorted-descending top-k buffer
+// ---------------------------------------------------------------------------
 __device__ __forceinline__ bool insert_topk(
     float val, int idx, float* vals, int* indices, int k
 ) {
@@ -71,29 +69,25 @@
 }
 
 // ---------------------------------------------------------------------------
-// Bitonic sort in shared memory (descending)
-// All threads participate; __syncthreads() between stages.
-// ---------------------------------------------------------------------------
-template <int N_ELEM>
-__device__ void bitonic_sort_desc(int* sv_as_int, int* si, int tid, int n_threads) {
-    // sv_as_int is float values bitcast to int, si are indices
-    int n = N_ELEM;
-    for (int stage = 2; stage <= n; stage <<= 1) {
+// Bitonic sort (descending) on shared memory arrays.
+// N = total element count (must be a power of 2).
+// All n_threads participate; __syncthreads() called between stages.
+// ---------------------------------------------------------------------------
+template <int N>
+__device__ void bitonic_sort_shared(volatile float* sv, volatile int* si,
+                                     int tid, int n_threads) {
+    for (int stage = 2; stage <= N; stage <<= 1) {
         for (int step = stage >> 1; step > 0; step >>= 1) {
-            for (int i = tid; i < n; i += n_threads) {
+            for (int i = tid; i < N; i += n_threads) {
                 int ixj = i ^ step;
                 if (ixj > i) {
+                    float vi = sv[i];
+                    float vj = sv[ixj];
                     bool asc = ((i & stage) == 0);
-                    float vi = __int_as_float(sv_as_int[i]);
-                    float vj = __int_as_float(sv_as_int[ixj]);
-                    bool swap = asc ? (vi > vj) : (vi < vj);
-                    if (swap) {
-                        int ti = sv_as_int[i];
-                        sv_as_int[i] = sv_as_int[ixj];
-                        sv_as_int[ixj] = ti;
-                        int tii = si[i];
-                        si[i] = si[ixj];
-                        si[ixj] = tii;
+                    bool sw = asc ? (vi > vj) : (vi < vj);
+                    if (sw) {
+                        float tv = sv[i]; sv[i] = sv[ixj]; sv[ixj] = tv;
+                        int   ti = si[i]; si[i] = si[ixj]; si[ixj] = ti;
                     }
                 }
             }
@@ -104,12 +98,10 @@
 
 // ---------------------------------------------------------------------------
 // Main streaming top-k kernel
-//
-// K        — top-k value (compile-time)
-// N_THREADS — threads per block (compile-time)
-// N_ELEM   — N_THREADS * K (shared memory elements, compile-time)
-// ---------------------------------------------------------------------------
-
+// K         — top-k value
+// N_THREADS — threads per block (must be power of 2)
+// N_ELEM    — N_THREADS * K (shared memory size, must be power of 2)
+// ---------------------------------------------------------------------------
 template <int K, int N_THREADS, int N_ELEM>
 __global__ void topk_kernel(
     const float* __restrict__ x,
@@ -127,23 +119,18 @@
     int end = (start + chunk_size < n) ? (start + chunk_size) : n;
     if (start >= end) return;
 
-    // ---- Phase 1: streaming per-thread local top-k ----
+    // ---- Phase 1: per-thread local top-k ----
     float lv[K];
     int   li[K];
     #pragma unroll
-    for (int i = 0; i < K; ++i) {
-        lv[i] = -INFINITY;
-        li[i] = -1;
-    }
+    for (int i = 0; i < K; ++i) { lv[i] = -INFINITY; li[i] = -1; }
     for (int pos = start + tid; pos < end; pos += N_THREADS) {
-        float v = x[row_off + pos];
-        insert_topk(v, pos, lv, li, K);
-    }
-
-    // ---- Phase 2: write local top-k to shared memory ----
+        insert_topk(x[row_off + pos], pos, lv, li, K);
+    }
+
+    // ---- Phase 2: write to shared memory ----
     __shared__ float sv[N_ELEM];
     __shared__ int   si[N_ELEM];
-
     #pragma unroll
     for (int i = 0; i < K; ++i) {
         sv[tid * K + i] = lv[i];
@@ -151,10 +138,10 @@
     }
     __syncthreads();
 
-    // ---- Phase 3: bitonic sort in shared memory ----
-    bitonic_sort_desc<N_ELEM>(sv, si, N_ELEM, tid, N_THREADS);
-
-    // ---- Phase 4: write top K (first K elements are the largest, descending) ----
+    // ---- Phase 3: bitonic sort ----
+    bitonic_sort_shared<N_ELEM>(sv, si, tid, N_THREADS);
+
+    // ---- Phase 4: write top K (now sorted descending) ----
     if (tid < K) {
         int64_t off = (int64_t)gid * K;
         out_vals[off + tid] = sv[tid];
@@ -163,9 +150,8 @@
 }
 
 // ---------------------------------------------------------------------------
-// Merge kernel: bitonic sort N_CHUNKS * K elements in shared memory
-// ---------------------------------------------------------------------------
-
+// Merge kernel (for 64-chunk case)
+// ---------------------------------------------------------------------------
 template <int K, int N_CHUNKS>
 __global__ void merg

... (4968 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -74,8 +74,7 @@
 // All n_threads participate; __syncthreads() called between stages.
 // ---------------------------------------------------------------------------
 template <int N>
-__device__ void bitonic_sort_shared(volatile float* sv, volatile int* si,
-                                     int tid, int n_threads) {
+__device__ void bitonic_sort_shared(float* sv, int* si, int tid, int n_threads) {
     for (int stage = 2; stage <= N; stage <<= 1) {
         for (int step = stage >> 1; step > 0; step >>= 1) {
             for (int i = tid; i < N; i += n_threads) {

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,11 +1,6 @@
 """Top-k via CUDA kernel with shared-memory bitonic merge.
 
-Approach:
-  - Each thread block processes one chunk of a batch element.
-  - Each thread streams through its slice, maintaining sorted top-k in registers.
-  - All threads write local top-k to shared memory (N_THREADS * K elements).
-  - Parallel bitonic sort in shared memory to find global top-k.
-  - For large n with batch=1: multi-chunk + merge pass.
+pybind11 workaround: all pointers passed as int64_t then cast internally.
 """
 
 from typing import List
@@ -19,20 +14,20 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 # ===========================================================================
-# C++ forward declarations for pybind11
+# C++ forward declarations for pybind11 — all pointers as int64_t
 # ===========================================================================
 
 _cpp_src = r"""
 #include <cstdint>
 
 extern "C" {
-void launch_topk_64(const float*, float*, int64_t*, int, int, int, int);
-void launch_topk_32(const float*, float*, int64_t*, int, int, int, int);
-void launch_topk_16(const float*, float*, int64_t*, int, int, int, int);
-void launch_topk_8(const float*, float*, int64_t*, int, int, int, int);
-void launch_topk_1(const float*, float*, int64_t*, int, int, int, int);
-void launch_merge_64_64(const float*, const int64_t*, float*, int64_t*, int);
-void launch_merge_generic(const float*, const int64_t*, float*, int64_t*, int, int, int);
+void launch_topk_64(int64_t, int64_t, int64_t, int, int, int, int);
+void launch_topk_32(int64_t, int64_t, int64_t, int, int, int, int);
+void launch_topk_16(int64_t, int64_t, int64_t, int, int, int, int);
+void launch_topk_8(int64_t, int64_t, int64_t, int, int, int, int);
+void launch_topk_1(int64_t, int64_t, int64_t, int, int, int, int);
+void launch_merge_64_64(int64_t, int64_t, int64_t, int64_t, int);
+void launch_merge_generic(int64_t, int64_t, int64_t, int64_t, int, int, int);
 }
 """
 
@@ -49,9 +44,7 @@
 // ---------------------------------------------------------------------------
 // Insert into sorted-descending top-k buffer
 // ---------------------------------------------------------------------------
-__device__ __forceinline__ bool insert_topk(
-    float val, int idx, float* vals, int* indices, int k
-) {
+__device__ bool insert_topk(float val, int idx, float* vals, int* indices, int k) {
     if (val <= vals[k - 1]) return false;
     int lo = 0, hi = k - 1;
     while (lo < hi) {
@@ -69,9 +62,8 @@
 }
 
 // ---------------------------------------------------------------------------
-// Bitonic sort (descending) on shared memory arrays.
-// N = total element count (must be a power of 2).
-// All n_threads participate; __syncthreads() called between stages.
+// Bitonic sort (descending) on shared memory arrays
+// N = total element count (must be power of 2)
 // ---------------------------------------------------------------------------
 template <int N>
 __device__ void bitonic_sort_shared(float* sv, int* si, int tid, int n_threads) {
@@ -80,8 +72,7 @@
             for (int i = tid; i < N; i += n_threads) {
                 int ixj = i ^ step;
                 if (ixj > i) {
-                    float vi = sv[i];
-                    float vj = sv[ixj];
+                    float vi = sv[i], vj = sv[ixj];
                     bool asc = ((i & stage) == 0);
                     bool sw = asc ? (vi > vj) : (vi < vj);
                     if (sw) {
@@ -97,9 +88,6 @@
 
 // ---------------------------------------------------------------------------
 // Main streaming top-k kernel
-// K         — top-k value
-// N_THREADS — threads per block (must be power of 2)
-// N_ELEM    — N_THREADS * K (shared memory size, must be power of 2)
 // ---------------------------------------------------------------------------
 template <int K, int N_THREADS, int N_ELEM>
 __global__ void topk_kernel(
@@ -118,29 +106,26 @@
     int end = (start + chunk_size < n) ? (start + chunk_size) : n;
     if (start >= end) return;
 
-    // ---- Phase 1: per-thread local top-k ----
-    float lv[K];
-    int   li[K];
-    #pragma unroll
+    // Phase 1: per-thread local top-k
+    float lv[K]; int li[K];
     for (int i = 0; i < K; ++i) { lv[i] = -INFINITY; li[i] = -1; }
     for (int pos = start + tid; pos < end; pos += N_THREADS) {
         insert_topk(x[row_off + pos], pos, lv, li, K);
     }
 
-    // ---- Phase 2: write to shared memory ----
+    // Phase 2: write to shared memory
     __shared__ float sv[N_ELEM];
     __shared__ int   si[N_ELEM];
-    #pragma unroll
     for (int i = 0; i < K; ++i) {
         sv[tid * K + i] = lv[i];
         si[tid * K + i] = li[i];
     }
     __syncthreads();
 
-    // ---- Phase 3: bitonic sort ----
+    // Phase 3: bitonic sort
     bitonic_sort_shared<N_ELEM>(sv, si, tid, N_THREADS);
 
-    // ---- Phase 4: write top K (now sorted descending) ----
+    // Phase 4: write top K
     if (tid < K) {
         int64_t off = (int64_t)gid * K;
         out_vals[off + tid] = sv[tid];
@@ -149,7 +134,7 @@
 }
 
 // ---------------------------------------------------------------------------
-// Merge kernel (for 64-chunk case)
+// Merge kernel (64 chunks of 64)
 // ---------------------------------------------------------------------------
 template <int K, int N_CHUNKS>
 __global__ void merge_kernel(
@@ -178,7 +163,7 @@
 }
 
 // ---------------------------------------------------------------------------
-// Generic merge kernel with dynamic shared memory + runtime-size bitonic sort
+// Generic merge kernel
 // ---------------------------------------------------------------------------
 __global__ void merge_generic_kernel(
     const float* __restrict__ cv, const int64_t* __restrict__ ci,
@@ -199,7 +184,7 @@
     }
     __syncthreads();
 
-    // Bitonic sort: pad to next power of 2
+    // Bitonic sort with padding


... (6049 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -73,8 +73,10 @@
                 int ixj = i ^ step;
                 if (ixj > i) {
                     float vi = sv[i], vj = sv[ixj];
+                    // Standard bitonic ascending: asc ? (vi>vj) : (vi<vj)
+                    // For DESCENDING: flip both conditions
                     bool asc = ((i & stage) == 0);
-                    bool sw = asc ? (vi > vj) : (vi < vj);
+                    bool sw = asc ? (vi < vj) : (vi > vj);
                     if (sw) {
                         float tv = sv[i]; sv[i] = sv[ixj]; sv[ixj] = tv;
                         int   ti = si[i]; si[i] = si[ixj]; si[ixj] = ti;

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -196,17 +196,21 @@
                     float vi = (i < total) ? sv[i] : -INFINITY;
                     float vj = (ixj < total) ? sv[ixj] : -INFINITY;
                     bool asc = ((i & stage) == 0);
-                    bool sw = asc ? (vi > vj) : (vi < vj);
+                    bool sw = asc ? (vi < vj) : (vi > vj);
                     if (sw) {
-                        float tv = vi; int tii = (i < total) ? si[i] : -1;
                         if (i < total && ixj < total) {
+                            float tv = sv[i]; sv[i] = sv[ixj]; sv[ixj] = tv;
+                            int   ti = si[i]; si[i] = si[ixj]; si[ixj] = ti;
+                        } else if (i < total) {
+                            // ixj is padding (-inf); move padding to ixj if needed
+                            // Since vj is -inf and vi is real, sw = asc ? (vi < -inf) : (vi > -inf)
+                            // vi < -inf is never true, vi > -inf is always true
+                            // For asc block: no swap. For desc block: swap.
+                            // Just leave sv[i] as is — it's a real value
+                        } else if (ixj < total) {
+                            // i is padding, ixj is real: move real value to i
                             sv[i] = sv[ixj]; si[i] = si[ixj];
-                            sv[ixj] = tv; si[ixj] = tii;
-                        } else if (i < total) {
-                            sv[i] = vj; si[i] = (ixj < total) ? si[ixj] : -1;
-                            // ixj is padding, no need to store back
-                        } else if (ixj < total) {
-                            sv[ixj] = tv; si[ixj] = tii;
+                            sv[ixj] = -INFINITY; si[ixj] = -1;
                         }
                     }
                 }

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,6 +1,9 @@
-"""Top-k via CUDA kernel with shared-memory bitonic merge.
-
-pybind11 workaround: all pointers passed as int64_t then cast internally.
+"""Top-k via CUDA kernel with warp-shuffle merge.
+
+Fast path:
+  - 32 threads (1 warp) per block: shuffle-only merge, zero shared memory.
+  - 64 threads (2 warps): shuffle within warp + tiny shared-memory merge.
+  - For (1, 131072, 64): multi-chunk with shuffle + shared + merge pass.
 """
 
 from typing import List
@@ -14,25 +17,23 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 # ===========================================================================
-# C++ forward declarations for pybind11 — all pointers as int64_t
+# C++ forward declarations
 # ===========================================================================
 
 _cpp_src = r"""
 #include <cstdint>
-
 extern "C" {
-void launch_topk_64(int64_t, int64_t, int64_t, int, int, int, int);
-void launch_topk_32(int64_t, int64_t, int64_t, int, int, int, int);
-void launch_topk_16(int64_t, int64_t, int64_t, int, int, int, int);
-void launch_topk_8(int64_t, int64_t, int64_t, int, int, int, int);
-void launch_topk_1(int64_t, int64_t, int64_t, int, int, int, int);
-void launch_merge_64_64(int64_t, int64_t, int64_t, int64_t, int);
-void launch_merge_generic(int64_t, int64_t, int64_t, int64_t, int, int, int);
+void launch_topk_64(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_topk_32(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_topk_16(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_topk_8(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_topk_1(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_merge(int64_t,int64_t,int64_t,int64_t,int,int,int);
 }
 """
 
 # ===========================================================================
-# CUDA kernel source
+# CUDA kernels
 # ===========================================================================
 
 _cuda_src = r"""
@@ -42,175 +43,147 @@
 #include <math.h>
 
 // ---------------------------------------------------------------------------
-// Insert into sorted-descending top-k buffer
+// Helpers
 // ---------------------------------------------------------------------------
 __device__ bool insert_topk(float val, int idx, float* vals, int* indices, int k) {
     if (val <= vals[k - 1]) return false;
     int lo = 0, hi = k - 1;
     while (lo < hi) {
         int mid = (lo + hi) >> 1;
-        if (val > vals[mid]) hi = mid;
-        else                 lo = mid + 1;
+        if (val > vals[mid]) hi = mid; else lo = mid + 1;
     }
     for (int i = k - 1; i > lo; --i) {
-        vals[i] = vals[i - 1];
-        indices[i] = indices[i - 1];
-    }
-    vals[lo] = val;
-    indices[lo] = idx;
+        vals[i] = vals[i - 1]; indices[i] = indices[i - 1];
+    }
+    vals[lo] = val; indices[lo] = idx;
     return true;
 }
 
-// ---------------------------------------------------------------------------
-// Bitonic sort (descending) on shared memory arrays
-// N = total element count (must be power of 2)
-// ---------------------------------------------------------------------------
-template <int N>
-__device__ void bitonic_sort_shared(float* sv, int* si, int tid, int n_threads) {
-    for (int stage = 2; stage <= N; stage <<= 1) {
-        for (int step = stage >> 1; step > 0; step >>= 1) {
-            for (int i = tid; i < N; i += n_threads) {
-                int ixj = i ^ step;
-                if (ixj > i) {
-                    float vi = sv[i], vj = sv[ixj];
-                    // Standard bitonic ascending: asc ? (vi>vj) : (vi<vj)
-                    // For DESCENDING: flip both conditions
-                    bool asc = ((i & stage) == 0);
-                    bool sw = asc ? (vi < vj) : (vi > vj);
-                    if (sw) {
-                        float tv = sv[i]; sv[i] = sv[ixj]; sv[ixj] = tv;
-                        int   ti = si[i]; si[i] = si[ixj]; si[ixj] = ti;
+__device__ void merge_two(const float* av, const int* ai,
+                          const float* bv, const int* bi,
+                          float* dv, int* di, int k) {
+    int ia=0,ib=0,io=0;
+    while (io < k) {
+        bool ta = (ia<k) && (ib>=k || av[ia]>=bv[ib]);
+        dv[io] = ta ? av[ia] : bv[ib];
+        di[io] = ta ? ai[ia] : bi[ib];
+        ia += ta; ib += !ta; io++;
+    }
+}
+
+// ---------------------------------------------------------------------------
+// Kernel: streaming + within-warp shuffle + optional cross-warp shared merge
+// K = top-k, NT = threads (32 or 64), NW = NT/32 warps, NS = NW*K
+// ---------------------------------------------------------------------------
+template <int K, int NT, int NW, int NS>
+__global__ void topk_kernel(
+    const float* __restrict__ x,
+    float* __restrict__ ov, int64_t* __restrict__ oi,
+    int n, int cs, int nc
+) {
+    int gid = blockIdx.x;
+    int batch_idx = gid / nc;
+    int chunk_idx = gid % nc;
+    int tid = threadIdx.x;
+    int64_t row_off = (int64_t)batch_idx * n;
+    int start = chunk_idx * cs;
+    int end = min(start + cs, n);
+    if (start >= end) return;
+
+    // ---- Phase 1: streaming per-thread local top-k ----
+    float lv[K]; int li[K];
+    for (int i=0;i<K;++i) { lv[i]=-INFINITY; li[i]=-1; }
+    for (int pos=start+tid; pos<end; pos+=NT)
+        insert_topk(x[row_off+pos], pos, lv, li, K);
+
+    // ---- Phase 2: within-warp shuffle merge ----
+    const int lane = tid & 31;
+    for (int off=16; off>0; off>>=1) {
+        float pv[K]; int pi[K];
+        for (int i=0;i<K;++i) {
+            pv[i]=__shfl_down_sync(0xffffffff, lv[i], off);
+            pi[i]=__shfl_down_sync(0xffffffff, li[i], off);
+        }
+        if (lane < off) {
+            float mv[K]; int mi[K];
+            merge_two(lv,li, pv,pi, mv,mi, K);
+            for (int i=0;i<K;++i) { lv[i]=mv[i]; li[i]=mi[i]; }
+        }
+    

... (17389 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,9 +1,12 @@
-"""Top-k via CUDA kernel with warp-shuffle merge.
+"""Top-k via CUDA kernel with warp-shuffle + shared-memory reduction.
 
-Fast path:
-  - 32 threads (1 warp) per block: shuffle-only merge, zero shared memory.
-  - 64 threads (2 warps): shuffle within warp + tiny shared-memory merge.
-  - For (1, 131072, 64): multi-chunk with shuffle + shared + merge pass.
+Strategy:
+  - Multiple blocks per row for parallelism.
+  - Within each block: 256 threads (8 warps).
+  - Per-thread streaming top-k in registers.
+  - Within-warp shuffle merge (5 rounds).
+  - Cross-warp tree merge in shared memory (3 levels).
+  - Merge pass to combine chunk results.
 """
 
 from typing import List
@@ -17,9 +20,6 @@
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
 # ===========================================================================
-# C++ forward declarations
-# ===========================================================================
-
 _cpp_src = r"""
 #include <cstdint>
 extern "C" {
@@ -32,206 +32,114 @@
 }
 """
 
-# ===========================================================================
-# CUDA kernels
-# ===========================================================================
-
 _cuda_src = r"""
 #include <cuda_runtime.h>
 #include <cstdint>
 #include <cfloat>
 #include <math.h>
 
-// ---------------------------------------------------------------------------
-// Helpers
-// ---------------------------------------------------------------------------
+// ---- helpers ----
 __device__ bool insert_topk(float val, int idx, float* vals, int* indices, int k) {
-    if (val <= vals[k - 1]) return false;
-    int lo = 0, hi = k - 1;
-    while (lo < hi) {
-        int mid = (lo + hi) >> 1;
-        if (val > vals[mid]) hi = mid; else lo = mid + 1;
-    }
-    for (int i = k - 1; i > lo; --i) {
-        vals[i] = vals[i - 1]; indices[i] = indices[i - 1];
-    }
-    vals[lo] = val; indices[lo] = idx;
-    return true;
+    if (val <= vals[k-1]) return false;
+    int lo=0, hi=k-1;
+    while(lo<hi){ int mid=(lo+hi)>>1; if(val>vals[mid])hi=mid; else lo=mid+1; }
+    for(int i=k-1;i>lo;--i){ vals[i]=vals[i-1]; indices[i]=indices[i-1]; }
+    vals[lo]=val; indices[lo]=idx; return true;
+}
+__device__ void merge_two(const float* av,const int* ai,const float* bv,const int* bi,
+                          float* dv,int* di,int k){
+    int ia=0,ib=0,io=0;
+    while(io<k){ bool ta=(ia<k)&&(ib>=k||av[ia]>=bv[ib]);
+        dv[io]=ta?av[ia]:bv[ib]; di[io]=ta?ai[ia]:bi[ib]; ia+=ta; ib+=!ta; io++; }
 }
 
-__device__ void merge_two(const float* av, const int* ai,
-                          const float* bv, const int* bi,
-                          float* dv, int* di, int k) {
-    int ia=0,ib=0,io=0;
-    while (io < k) {
-        bool ta = (ia<k) && (ib>=k || av[ia]>=bv[ib]);
-        dv[io] = ta ? av[ia] : bv[ib];
-        di[io] = ta ? ai[ia] : bi[ib];
-        ia += ta; ib += !ta; io++;
+// ---- streaming kernel: 256 threads (8 warps) ----
+template <int K, int NT, int NW, int NS>
+__global__ void topk_kernel(
+    const float* __restrict__ x, float* __restrict__ ov, int64_t* __restrict__ oi,
+    int n, int cs, int nc
+) {
+    int gid=blockIdx.x, bid=gid/nc, cid=gid%nc, tid=threadIdx.x;
+    int64_t ro = (int64_t)bid*n;
+    int start=cid*cs, end=min(start+cs,n);
+    if(start>=end) return;
+
+    // Phase 1: streaming
+    float lv[K]; int li[K];
+    for(int i=0;i<K;++i){ lv[i]=-INFINITY; li[i]=-1; }
+    for(int pos=start+tid; pos<end; pos+=NT) insert_topk(x[ro+pos],pos,lv,li,K);
+
+    // Phase 2: within-warp shuffle merge
+    const int lane=tid&31;
+    for(int off=16;off>0;off>>=1){
+        float pv[K]; int pi[K];
+        for(int i=0;i<K;++i){ pv[i]=__shfl_down_sync(0xffffffff,lv[i],off); pi[i]=__shfl_down_sync(0xffffffff,li[i],off); }
+        if(lane<off){ float mv[K];int mi[K]; merge_two(lv,li,pv,pi,mv,mi,K); for(int i=0;i<K;++i){lv[i]=mv[i];li[i]=mi[i];} }
     }
+
+    // Phase 3: cross-warp tree merge in shared memory
+    __shared__ float sv[NS];
+    __shared__ int   si[NS];
+    int wid=tid/32;
+    if(lane==0){ for(int i=0;i<K;++i){ sv[wid*K+i]=lv[i]; si[wid*K+i]=li[i]; } }
+    __syncthreads();
+    // Tree reduction: log2(NW) levels
+    for(int stride=NW/2;stride>0;stride>>=1){
+        if(lane==0 && wid<stride){
+            float mv[K];int mi[K];
+            merge_two(&sv[wid*K],&si[wid*K],&sv[(wid+stride)*K],&si[(wid+stride)*K],mv,mi,K);
+            for(int i=0;i<K;++i){ sv[wid*K+i]=mv[i]; si[wid*K+i]=mi[i]; }
+        }
+        __syncthreads();
+    }
+    if(tid==0){ int64_t off=(int64_t)gid*K; for(int i=0;i<K;++i){ ov[off+i]=sv[i]; oi[off+i]=(int64_t)si[i]; } }
 }
 
-// ---------------------------------------------------------------------------
-// Kernel: streaming + within-warp shuffle + optional cross-warp shared merge
-// K = top-k, NT = threads (32 or 64), NW = NT/32 warps, NS = NW*K
-// ---------------------------------------------------------------------------
-template <int K, int NT, int NW, int NS>
-__global__ void topk_kernel(
-    const float* __restrict__ x,
-    float* __restrict__ ov, int64_t* __restrict__ oi,
-    int n, int cs, int nc
-) {
-    int gid = blockIdx.x;
-    int batch_idx = gid / nc;
-    int chunk_idx = gid % nc;
-    int tid = threadIdx.x;
-    int64_t row_off = (int64_t)batch_idx * n;
-    int start = chunk_idx * cs;
-    int end = min(start + cs, n);
-    if (start >= end) return;
-
-    // ---- Phase 1: streaming per-thread local top-k ----
-    float lv[K]; int li[K];
-    for (int i=0;i<K;++i) { lv[i]=-INFINITY; li[i]=-1; }
-    for (int pos=start+tid; pos<end; pos+=NT)
-        insert_topk(x[row_off+pos], pos, lv, li, K);
-
-    // ---- Phase 2: within-warp shuffle merge ----
-    const int lane = tid & 31;
-    for (int off=16; off>0; off>>=1) {
-        float pv[K]; int pi[K];
-        for (int i=0;i<K

... (9722 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -140,6 +140,10 @@
 void launch_merge(int64_t cvp,int64_t cip,int64_t ovp,int64_t oip,int B,int nc,int k){
     auto*cv=(const float*)cvp;auto*ci=(const int64_t*)cip;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;
     int total=nc*k,block=total<1024?total:1024; size_t sm=total*(sizeof(float)+sizeof(int));
+    // Opt-in to larger shared memory if needed (max 99 KB on this GPU)
+    if(sm > 49152) {
+        cudaFuncSetAttribute(merge_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, sm);
+    }
     merge_kernel<<<B,block,sm>>>(cv,ci,ov,oi,nc,k);
 }
 

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -177,23 +177,29 @@
 
     def forward(self,x):
         B,N,K=self.batch,self.n,self.k
-        # Chunk for parallelism: target at least 128 total blocks
-        target_blocks = max(B, 128)
-        cs = max(256, (N + (target_blocks//B) - 1) // (target_blocks//B)) if B>0 else N
+        # Chunking: optimize for latency.
+        # Single row → single block (no merge overhead).
+        # Batched → multiple blocks per row for SM utilization, then merge.
+        if B == 1:
+            cs = N  # one block processes the entire row
+        else:
+            # Each row processed by ~4 chunks = good balance of parallelism vs merge
+            target_chunks_per_row = 4
+            cs = max(256, (N + target_chunks_per_row - 1) // target_chunks_per_row)
         cs = min(cs, N)
-        nc = (N+cs-1)//cs
+        nc = (N + cs - 1) // cs
         m = _get_mod()
         launch = {64:m.launch_topk_64,32:m.launch_topk_32,16:m.launch_topk_16,8:m.launch_topk_8,1:m.launch_topk_1}[K]
-        if nc==1:
-            ov=torch.empty(B,K,dtype=torch.float32,device=x.device)
-            oi=torch.empty(B,K,dtype=torch.int64,device=x.device)
-            launch(x.data_ptr(),ov.data_ptr(),oi.data_ptr(),B,N,cs,nc)
-            return ov,oi
-        total=B*nc
-        cv=torch.empty(total,K,dtype=torch.float32,device=x.device)
-        ci=torch.empty(total,K,dtype=torch.int64,device=x.device)
-        launch(x.data_ptr(),cv.data_ptr(),ci.data_ptr(),B,N,cs,nc)
-        ov=torch.empty(B,K,dtype=torch.float32,device=x.device)
-        oi=torch.empty(B,K,dtype=torch.int64,device=x.device)
-        m.launch_merge(cv.data_ptr(),ci.data_ptr(),ov.data_ptr(),oi.data_ptr(),B,nc,K)
-        return ov,oi
+        if nc == 1:
+            ov = torch.empty(B,K,dtype=torch.float32,device=x.device)
+            oi = torch.empty(B,K,dtype=torch.int64,device=x.device)
+            launch(x.data_ptr(), ov.data_ptr(), oi.data_ptr(), B, N, cs, nc)
+            return ov, oi
+        total = B * nc
+        cv = torch.empty(total,K,dtype=torch.float32,device=x.device)
+        ci = torch.empty(total,K,dtype=torch.int64,device=x.device)
+        launch(x.data_ptr(), cv.data_ptr(), ci.data_ptr(), B, N, cs, nc)
+        ov = torch.empty(B,K,dtype=torch.float32,device=x.device)
+        oi = torch.empty(B,K,dtype=torch.int64,device=x.device)
+        m.launch_merge(cv.data_ptr(), ci.data_ptr(), ov.data_ptr(), oi.data_ptr(), B, nc, K)
+        return ov, oi

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,12 +1,8 @@
-"""Top-k via CUDA kernel with warp-shuffle + shared-memory reduction.
-
-Strategy:
-  - Multiple blocks per row for parallelism.
-  - Within each block: 256 threads (8 warps).
-  - Per-thread streaming top-k in registers.
-  - Within-warp shuffle merge (5 rounds).
-  - Cross-warp tree merge in shared memory (3 levels).
-  - Merge pass to combine chunk results.
+"""Top-k via multi-pass radix-select CUDA kernel.
+
+For each row, scan bits from MSB to LSB, narrowing down candidates
+until exactly k remain. One kernel launch per batch element with
+multiple blocks cooperating via global atomics.
 """
 
 from typing import List
@@ -23,12 +19,7 @@
 _cpp_src = r"""
 #include <cstdint>
 extern "C" {
-void launch_topk_64(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_topk_32(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_topk_16(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_topk_8(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_topk_1(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_merge(int64_t,int64_t,int64_t,int64_t,int,int,int);
+void launch_radix_topk(int64_t, int64_t, int64_t, int, int);
 }
 """
 
@@ -36,115 +27,152 @@
 #include <cuda_runtime.h>
 #include <cstdint>
 #include <cfloat>
-#include <math.h>
-
-// ---- helpers ----
-__device__ bool insert_topk(float val, int idx, float* vals, int* indices, int k) {
-    if (val <= vals[k-1]) return false;
-    int lo=0, hi=k-1;
-    while(lo<hi){ int mid=(lo+hi)>>1; if(val>vals[mid])hi=mid; else lo=mid+1; }
-    for(int i=k-1;i>lo;--i){ vals[i]=vals[i-1]; indices[i]=indices[i-1]; }
-    vals[lo]=val; indices[lo]=idx; return true;
-}
-__device__ void merge_two(const float* av,const int* ai,const float* bv,const int* bi,
-                          float* dv,int* di,int k){
-    int ia=0,ib=0,io=0;
-    while(io<k){ bool ta=(ia<k)&&(ib>=k||av[ia]>=bv[ib]);
-        dv[io]=ta?av[ia]:bv[ib]; di[io]=ta?ai[ia]:bi[ib]; ia+=ta; ib+=!ta; io++; }
-}
-
-// ---- streaming kernel: 256 threads (8 warps) ----
-template <int K, int NT, int NW, int NS>
-__global__ void topk_kernel(
-    const float* __restrict__ x, float* __restrict__ ov, int64_t* __restrict__ oi,
-    int n, int cs, int nc
+
+// Convert fp32 to sortable uint32: flip sign bit, then if sign was 1, flip all bits
+__device__ __forceinline__ unsigned int float_to_sortable(float f) {
+    unsigned int u = __float_as_uint(f);
+    // If sign bit is 1 (negative), flip all bits
+    // If sign bit is 0 (positive), flip sign bit only
+    unsigned int mask = (int)u >> 31;  // 0xFFFFFFFF if negative, 0 if positive
+    return u ^ (mask | 0x80000000u);
+}
+
+// Convert back
+__device__ __forceinline__ float sortable_to_float(unsigned int u) {
+    unsigned int mask = (int)u >> 31;
+    return __uint_as_float(u ^ ((~mask) & 0x80000000u));
+}
+
+// ---------------------------------------------------------------------------
+// Multi-pass radix select: each pass counts+partitions by a range of bits.
+// Uses global atomics for cross-block counting.
+// ---------------------------------------------------------------------------
+
+__global__ void radix_select_kernel(
+    const float* __restrict__ x,      // (batch, n)
+    float* __restrict__ out_vals,     // (batch, k)
+    int64_t* __restrict__ out_idx,    // (batch, k)
+    int n, int k,
+    int* __restrict__ global_counters // 2 ints: [count_0, count_1]
 ) {
-    int gid=blockIdx.x, bid=gid/nc, cid=gid%nc, tid=threadIdx.x;
-    int64_t ro = (int64_t)bid*n;
-    int start=cid*cs, end=min(start+cs,n);
-    if(start>=end) return;
-
-    // Phase 1: streaming
+    __shared__ int s_count[2];
+    int tid = threadIdx.x;
+    int bid = blockIdx.x;
+
+    // Phase 1: count how many elements in this block have bit=0 vs bit=1
+    // at the current bit position (passed via global_counter[2] implicitly
+    // through the iteration protocol — we'll iterate bits in the launcher).
+
+    // This kernel is called multiple times from Python, each time with
+    // a specific bit position and a mask of "active" elements.
+
+    // ... This is getting complex for a single kernel.
+    // Let me use a different approach: streaming top-k per row, but optimized.
+}
+
+// ---------------------------------------------------------------------------
+// Simplified: single-pass streaming top-k per block, then merge via shared mem.
+// Optimized for speed: each block handles one row, no inter-block communication.
+// For batch=1 with large n, use multiple warps within one block for parallel scan.
+// ---------------------------------------------------------------------------
+
+template <int K, int NT>
+__global__ void streaming_topk_kernel(
+    const float* __restrict__ x,
+    float* __restrict__ ov, int64_t* __restrict__ oi,
+    int n
+) {
+    // NT threads per block, one block launched. All threads scan the row.
+    int tid = threadIdx.x;
+    int total_warps = NT / 32;
+
+    // Each thread maintains local top-K
     float lv[K]; int li[K];
-    for(int i=0;i<K;++i){ lv[i]=-INFINITY; li[i]=-1; }
-    for(int pos=start+tid; pos<end; pos+=NT) insert_topk(x[ro+pos],pos,lv,li,K);
-
-    // Phase 2: within-warp shuffle merge
-    const int lane=tid&31;
-    for(int off=16;off>0;off>>=1){
+    for (int i=0;i<K;++i){ lv[i]=-INFINITY; li[i]=-1; }
+
+    // Strided scan over the entire row
+    for (int pos=tid; pos<n; pos+=NT) {
+        float v = x[pos];
+        // inline insertion (avoid function call overhead)
+        if (v > lv[K-1]) {
+            int lo=0, hi=K-1;
+            while(lo<hi){ int mid=(lo+hi)>>1; if(v>lv[mid])hi=mid; else lo=mid+1; }
+            for(int i=K-1;i>lo;--i){ lv[i]=lv[i-1]; li[i]=li[i-1]; }
+            lv[lo]=v; li[lo]=pos;
+        }
+    }
+
+    // Within-warp shuffle reduction
+    const int lane = tid & 31;
+    for (int off=16; off>0; off>>=1) {
         float pv[K]; int pi[K];


... (8387 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,8 +1,7 @@
-"""Top-k via multi-pass radix-select CUDA kernel.
+"""Top-k via streaming CUDA kernel with warp-shuffle + shared-memory merge.
 
-For each row, scan bits from MSB to LSB, narrowing down candidates
-until exactly k remain. One kernel launch per batch element with
-multiple blocks cooperating via global atomics.
+One kernel launch per batch element, or one launch for the whole batch.
+Uses 256-1024 threads per block for bandwidth saturation.
 """
 
 from typing import List
@@ -15,11 +14,14 @@
 SUPPORTED_PRECISIONS = ["fp32"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
-# ===========================================================================
 _cpp_src = r"""
 #include <cstdint>
 extern "C" {
-void launch_radix_topk(int64_t, int64_t, int64_t, int, int);
+void launch_topk_64_1024(int64_t,int64_t,int64_t,int,int);
+void launch_topk_32_256(int64_t,int64_t,int64_t,int,int);
+void launch_topk_16_128(int64_t,int64_t,int64_t,int,int);
+void launch_topk_8_128(int64_t,int64_t,int64_t,int,int);
+void launch_topk_1_128(int64_t,int64_t,int64_t,int,int);
 }
 """
 
@@ -27,158 +29,88 @@
 #include <cuda_runtime.h>
 #include <cstdint>
 #include <cfloat>
+#include <math.h>
 
-// Convert fp32 to sortable uint32: flip sign bit, then if sign was 1, flip all bits
-__device__ __forceinline__ unsigned int float_to_sortable(float f) {
-    unsigned int u = __float_as_uint(f);
-    // If sign bit is 1 (negative), flip all bits
-    // If sign bit is 0 (positive), flip sign bit only
-    unsigned int mask = (int)u >> 31;  // 0xFFFFFFFF if negative, 0 if positive
-    return u ^ (mask | 0x80000000u);
-}
+template <int K, int NT, int NW>
+__global__ void topk_kernel(
+    const float* __restrict__ x, float* __restrict__ ov, int64_t* __restrict__ oi,
+    int n, int stride  // stride = batch stride (n)
+) {
+    int bid = blockIdx.x;      // which batch element
+    int tid = threadIdx.x;
+    const float* row = x + (int64_t)bid * stride;
 
-// Convert back
-__device__ __forceinline__ float sortable_to_float(unsigned int u) {
-    unsigned int mask = (int)u >> 31;
-    return __uint_as_float(u ^ ((~mask) & 0x80000000u));
-}
-
-// ---------------------------------------------------------------------------
-// Multi-pass radix select: each pass counts+partitions by a range of bits.
-// Uses global atomics for cross-block counting.
-// ---------------------------------------------------------------------------
-
-__global__ void radix_select_kernel(
-    const float* __restrict__ x,      // (batch, n)
-    float* __restrict__ out_vals,     // (batch, k)
-    int64_t* __restrict__ out_idx,    // (batch, k)
-    int n, int k,
-    int* __restrict__ global_counters // 2 ints: [count_0, count_1]
-) {
-    __shared__ int s_count[2];
-    int tid = threadIdx.x;
-    int bid = blockIdx.x;
-
-    // Phase 1: count how many elements in this block have bit=0 vs bit=1
-    // at the current bit position (passed via global_counter[2] implicitly
-    // through the iteration protocol — we'll iterate bits in the launcher).
-
-    // This kernel is called multiple times from Python, each time with
-    // a specific bit position and a mask of "active" elements.
-
-    // ... This is getting complex for a single kernel.
-    // Let me use a different approach: streaming top-k per row, but optimized.
-}
-
-// ---------------------------------------------------------------------------
-// Simplified: single-pass streaming top-k per block, then merge via shared mem.
-// Optimized for speed: each block handles one row, no inter-block communication.
-// For batch=1 with large n, use multiple warps within one block for parallel scan.
-// ---------------------------------------------------------------------------
-
-template <int K, int NT>
-__global__ void streaming_topk_kernel(
-    const float* __restrict__ x,
-    float* __restrict__ ov, int64_t* __restrict__ oi,
-    int n
-) {
-    // NT threads per block, one block launched. All threads scan the row.
-    int tid = threadIdx.x;
-    int total_warps = NT / 32;
-
-    // Each thread maintains local top-K
+    // Phase 1: streaming per-thread top-K
     float lv[K]; int li[K];
-    for (int i=0;i<K;++i){ lv[i]=-INFINITY; li[i]=-1; }
-
-    // Strided scan over the entire row
-    for (int pos=tid; pos<n; pos+=NT) {
-        float v = x[pos];
-        // inline insertion (avoid function call overhead)
-        if (v > lv[K-1]) {
-            int lo=0, hi=K-1;
+    for(int i=0;i<K;++i){ lv[i]=-INFINITY; li[i]=-1; }
+    for(int pos=tid; pos<n; pos+=NT) {
+        float v = row[pos];
+        if(v > lv[K-1]) {
+            int lo=0,hi=K-1;
             while(lo<hi){ int mid=(lo+hi)>>1; if(v>lv[mid])hi=mid; else lo=mid+1; }
             for(int i=K-1;i>lo;--i){ lv[i]=lv[i-1]; li[i]=li[i-1]; }
             lv[lo]=v; li[lo]=pos;
         }
     }
 
-    // Within-warp shuffle reduction
-    const int lane = tid & 31;
-    for (int off=16; off>0; off>>=1) {
+    // Phase 2: within-warp shuffle merge
+    const int lane=tid&31;
+    for(int off=16;off>0;off>>=1){
         float pv[K]; int pi[K];
         for(int i=0;i<K;++i){ pv[i]=__shfl_down_sync(0xffffffff,lv[i],off); pi[i]=__shfl_down_sync(0xffffffff,li[i],off); }
         if(lane<off){
-            // merge two sorted lists of size K
-            float mv[K]; int mi[K];
-            int ia=0,ib=0,io=0;
-            while(io<K){
-                bool ta=(ia<K)&&(ib>=K||lv[ia]>=pv[ib]);
-                mv[io]=ta?lv[ia]:pv[ib]; mi[io]=ta?li[ia]:pi[ib]; ia+=ta; ib+=!ta; io++;
-            }
+            float mv[K];int mi[K]; int ia=0,ib=0,io=0;
+            while(io<K){ bool ta=(ia<K)&&(ib>=K||lv[ia]>=pv[ib]); mv[io]=ta?lv[ia]:pv[ib]; mi[io]=ta?li[ia]:pi[ib]; ia+=ta; ib+=!ta; io++; }
             for(int i=0;i<K;++i){ lv[i]=mv[i]; li[i]=mi[i]; }
         }
     }
 
-    // Cross-warp tree merge in shared mem

... (6008 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -90,7 +90,7 @@
 extern "C" {
 void launch_topk_64_1024(int64_t xp,int64_t ovp,int64_t oip,int B,int n){
     auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;
-    topk_kernel<64,1024,32><<<B,1024>>>(x,ov,oi,n,n);
+    topk_kernel<64,256,8><<<B,256>>>(x,ov,oi,n,n);
 }
 void launch_topk_32_256(int64_t xp,int64_t ovp,int64_t oip,int B,int n){
     auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,152 +1,58 @@
-"""Top-k via streaming CUDA kernel with warp-shuffle + shared-memory merge.
+"""Top-k via iterative batched argmax + masking.
 
-One kernel launch per batch element, or one launch for the whole batch.
-Uses 256-1024 threads per block for bandwidth saturation.
+Uses torch.argmax (not banned) in a loop over k. For small k this
+approaches memory-bandwidth-limited performance since each argmax
+reads the entire input.
 """
 
 from typing import List
 
 import torch
 import torch.nn as nn
-import torch.utils.cpp_extension
 
 OP_TYPE = "topk"
 SUPPORTED_PRECISIONS = ["fp32"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
-_cpp_src = r"""
-#include <cstdint>
-extern "C" {
-void launch_topk_64_1024(int64_t,int64_t,int64_t,int,int);
-void launch_topk_32_256(int64_t,int64_t,int64_t,int,int);
-void launch_topk_16_128(int64_t,int64_t,int64_t,int,int);
-void launch_topk_8_128(int64_t,int64_t,int64_t,int,int);
-void launch_topk_1_128(int64_t,int64_t,int64_t,int,int);
-}
-"""
+batch = 64
+n = 8192
+k = 8
 
-_cuda_src = r"""
-#include <cuda_runtime.h>
-#include <cstdint>
-#include <cfloat>
-#include <math.h>
+def get_inputs():
+    return [torch.randn(batch, n, dtype=torch.float32)]
 
-template <int K, int NT, int NW>
-__global__ void topk_kernel(
-    const float* __restrict__ x, float* __restrict__ ov, int64_t* __restrict__ oi,
-    int n, int stride  // stride = batch stride (n)
-) {
-    int bid = blockIdx.x;      // which batch element
-    int tid = threadIdx.x;
-    const float* row = x + (int64_t)bid * stride;
-
-    // Phase 1: streaming per-thread top-K
-    float lv[K]; int li[K];
-    for(int i=0;i<K;++i){ lv[i]=-INFINITY; li[i]=-1; }
-    for(int pos=tid; pos<n; pos+=NT) {
-        float v = row[pos];
-        if(v > lv[K-1]) {
-            int lo=0,hi=K-1;
-            while(lo<hi){ int mid=(lo+hi)>>1; if(v>lv[mid])hi=mid; else lo=mid+1; }
-            for(int i=K-1;i>lo;--i){ lv[i]=lv[i-1]; li[i]=li[i-1]; }
-            lv[lo]=v; li[lo]=pos;
-        }
-    }
-
-    // Phase 2: within-warp shuffle merge
-    const int lane=tid&31;
-    for(int off=16;off>0;off>>=1){
-        float pv[K]; int pi[K];
-        for(int i=0;i<K;++i){ pv[i]=__shfl_down_sync(0xffffffff,lv[i],off); pi[i]=__shfl_down_sync(0xffffffff,li[i],off); }
-        if(lane<off){
-            float mv[K];int mi[K]; int ia=0,ib=0,io=0;
-            while(io<K){ bool ta=(ia<K)&&(ib>=K||lv[ia]>=pv[ib]); mv[io]=ta?lv[ia]:pv[ib]; mi[io]=ta?li[ia]:pi[ib]; ia+=ta; ib+=!ta; io++; }
-            for(int i=0;i<K;++i){ lv[i]=mv[i]; li[i]=mi[i]; }
-        }
-    }
-
-    // Phase 3: cross-warp tree merge in shared memory
-    constexpr int NS = NW * K;
-    __shared__ float sv[NS];
-    __shared__ int   si[NS];
-    int wid = tid/32;
-    if(lane==0){ for(int i=0;i<K;++i){ sv[wid*K+i]=lv[i]; si[wid*K+i]=li[i]; } }
-    __syncthreads();
-
-    for(int stride=NW/2; stride>0; stride>>=1){
-        if(lane==0 && wid<stride){
-            float mv[K];int mi[K]; int ia=0,ib=0,io=0;
-            float* av=&sv[wid*K]; int* ai=&si[wid*K];
-            float* bv=&sv[(wid+stride)*K]; int* bi=&si[(wid+stride)*K];
-            while(io<K){ bool ta=(ia<K)&&(ib>=K||av[ia]>=bv[ib]); mv[io]=ta?av[ia]:bv[ib]; mi[io]=ta?ai[ia]:bi[ib]; ia+=ta; ib+=!ta; io++; }
-            for(int i=0;i<K;++i){ sv[wid*K+i]=mv[i]; si[wid*K+i]=mi[i]; }
-        }
-        __syncthreads();
-    }
-
-    if(tid==0){ for(int i=0;i<K;++i){ ov[(int64_t)bid*K+i]=sv[i]; oi[(int64_t)bid*K+i]=(int64_t)si[i]; } }
-}
-
-extern "C" {
-void launch_topk_64_1024(int64_t xp,int64_t ovp,int64_t oip,int B,int n){
-    auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;
-    topk_kernel<64,256,8><<<B,256>>>(x,ov,oi,n,n);
-}
-void launch_topk_32_256(int64_t xp,int64_t ovp,int64_t oip,int B,int n){
-    auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;
-    topk_kernel<32,256,8><<<B,256>>>(x,ov,oi,n,n);
-}
-void launch_topk_16_128(int64_t xp,int64_t ovp,int64_t oip,int B,int n){
-    auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;
-    topk_kernel<16,128,4><<<B,128>>>(x,ov,oi,n,n);
-}
-void launch_topk_8_128(int64_t xp,int64_t ovp,int64_t oip,int B,int n){
-    auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;
-    topk_kernel<8,128,4><<<B,128>>>(x,ov,oi,n,n);
-}
-void launch_topk_1_128(int64_t xp,int64_t ovp,int64_t oip,int B,int n){
-    auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;
-    topk_kernel<1,128,4><<<B,128>>>(x,ov,oi,n,n);
-}
-} // extern "C"
-"""
-
-_props = torch.cuda.get_device_properties(0)
-_arch = f"sm_{_props.major}{_props.minor}"
-_cflags = ["-O3","--use_fast_math",f"-arch={_arch}","--expt-relaxed-constexpr","-std=c++17","-DNDEBUG"]
-_mod = None
-
-def _get_mod():
-    global _mod
-    if _mod is None:
-        _mod = torch.utils.cpp_extension.load_inline(
-            name="topk_kern",cpp_sources=_cpp_src,cuda_sources=_cuda_src,
-            functions=["launch_topk_64_1024","launch_topk_32_256","launch_topk_16_128","launch_topk_8_128","launch_topk_1_128"],
-            extra_cuda_cflags=_cflags,verbose=False)
-    return _mod
-
-batch=64; n=8192; k=8
-def get_inputs(): return [torch.randn(batch,n,dtype=torch.float32)]
-def get_init_inputs(): return [batch,n,k]
+def get_init_inputs():
+    return [batch, n, k]
 
 class Model(nn.Module):
-    def __init__(self,batch:int,n:int,k:int):
-        super().__init__(); self.batch,self.n,self.k=batch,n,k
-        self.register_buffer("_dummy",torch.zeros(1))
+    def __init__(self, batch: int, n: int, k: int):
+        super().__init__()
+        self.batch, self.n, self.k = batch, n, k
+        self.register_buffer("_dummy", torch.zeros(1))
 
-    def forward(self,x):
-        B,N,K = self.batch,self.n,self.k
-        m = _get_mod()
-        ov = torch.empty(B,K,dtype=torch.float32,device=x.device)
-        oi

... (1414 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -52,7 +52,7 @@
             values[:, i] = val
             indices[:, i] = idx
 
-            # Mask out the found max so next argmax finds the next largest
-            work.scatter_(1, idx.unsqueeze(-1), neg_inf)
+            # Mask out the found max
+            work = work.scatter(1, idx.unsqueeze(-1), neg_inf.expand_as(idx.unsqueeze(-1)))
 
         return values, indices

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -47,12 +47,12 @@
         for i in range(K):
             # Find the max element in each row
             idx = torch.argmax(work, dim=-1)  # (B,)
-            val = work.gather(1, idx.unsqueeze(-1)).squeeze(-1)  # (B,)
+            val = work[torch.arange(B, device=x.device), idx]  # (B,)
 
             values[:, i] = val
             indices[:, i] = idx
 
             # Mask out the found max
-            work = work.scatter(1, idx.unsqueeze(-1), neg_inf.expand_as(idx.unsqueeze(-1)))
+            work[torch.arange(B, device=x.device), idx] = float('-inf')
 
         return values, indices

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -42,7 +42,6 @@
 
         # Work on a copy since we need to mask found elements
         work = x.clone()
-        neg_inf = torch.tensor(float('-inf'), dtype=torch.float32, device=x.device)
 
         for i in range(K):
             # Find the max element in each row

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,57 +1,203 @@
-"""Top-k via iterative batched argmax + masking.
-
-Uses torch.argmax (not banned) in a loop over k. For small k this
-approaches memory-bandwidth-limited performance since each argmax
-reads the entire input.
+"""Top-k via CUDA streaming + chunked merge.
+
+Each block processes one chunk with warp-shuffle + shared-memory reduction.
+A merge pass combines chunk results via bitonic sort.
+Chunk sizes chosen so merge fits in 48 KB default shared memory.
 """
 
 from typing import List
 
 import torch
 import torch.nn as nn
+import torch.utils.cpp_extension
 
 OP_TYPE = "topk"
 SUPPORTED_PRECISIONS = ["fp32"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
-batch = 64
-n = 8192
-k = 8
-
-def get_inputs():
-    return [torch.randn(batch, n, dtype=torch.float32)]
-
-def get_init_inputs():
-    return [batch, n, k]
+_cpp_src = r"""
+#include <cstdint>
+extern "C" {
+void launch_topk_64(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_topk_32(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_topk_16(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_topk_8(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_topk_1(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_merge(int64_t,int64_t,int64_t,int64_t,int,int,int);
+}
+"""
+
+_cuda_src = r"""
+#include <cuda_runtime.h>
+#include <cstdint>
+#include <cfloat>
+#include <math.h>
+
+// ---- helpers ----
+__device__ bool insert_topk(float val, int idx, float* vals, int* indices, int k) {
+    if (val <= vals[k-1]) return false;
+    int lo=0, hi=k-1;
+    while(lo<hi){ int mid=(lo+hi)>>1; if(val>vals[mid])hi=mid; else lo=mid+1; }
+    for(int i=k-1;i>lo;--i){ vals[i]=vals[i-1]; indices[i]=indices[i-1]; }
+    vals[lo]=val; indices[lo]=idx; return true;
+}
+
+// ---- streaming kernel ----
+template <int K, int NT>
+__global__ void topk_kernel(
+    const float* __restrict__ x, float* __restrict__ ov, int64_t* __restrict__ oi,
+    int n, int cs, int nc
+) {
+    int gid=blockIdx.x, bid=gid/nc, cid=gid%nc, tid=threadIdx.x;
+    int64_t ro = (int64_t)bid*n;
+    int start=cid*cs, end=min(start+cs,n);
+    if(start>=end) return;
+
+    float lv[K]; int li[K];
+    for(int i=0;i<K;++i){ lv[i]=-INFINITY; li[i]=-1; }
+    for(int pos=start+tid; pos<end; pos+=NT) insert_topk(x[ro+pos],pos,lv,li,K);
+
+    // within-warp shuffle merge
+    const int lane=tid&31;
+    for(int off=16;off>0;off>>=1){
+        float pv[K]; int pi[K];
+        for(int i=0;i<K;++i){ pv[i]=__shfl_down_sync(0xffffffff,lv[i],off); pi[i]=__shfl_down_sync(0xffffffff,li[i],off); }
+        if(lane<off){
+            float mv[K];int mi[K]; int ia=0,ib=0,io=0;
+            while(io<K){ bool ta=(ia<K)&&(ib>=K||lv[ia]>=pv[ib]); mv[io]=ta?lv[ia]:pv[ib]; mi[io]=ta?li[ia]:pi[ib]; ia+=ta; ib+=!ta; io++; }
+            for(int i=0;i<K;++i){ lv[i]=mv[i]; li[i]=mi[i]; }
+        }
+    }
+
+    // cross-warp tree merge in shared memory
+    constexpr int NW=NT/32, NS=NW*K;
+    __shared__ float sv[NS];
+    __shared__ int   si[NS];
+    int wid=tid/32;
+    if(lane==0){ for(int i=0;i<K;++i){ sv[wid*K+i]=lv[i]; si[wid*K+i]=li[i]; } }
+    __syncthreads();
+    for(int stride=NW/2;stride>0;stride>>=1){
+        if(lane==0 && wid<stride){
+            float mv[K];int mi[K]; int ia=0,ib=0,io=0;
+            float* av=&sv[wid*K];int* ai=&si[wid*K];
+            float* bv=&sv[(wid+stride)*K];int* bi=&si[(wid+stride)*K];
+            while(io<K){ bool ta=(ia<K)&&(ib>=K||av[ia]>=bv[ib]); mv[io]=ta?av[ia]:bv[ib]; mi[io]=ta?ai[ia]:bi[ib]; ia+=ta; ib+=!ta; io++; }
+            for(int i=0;i<K;++i){ sv[wid*K+i]=mv[i]; si[wid*K+i]=mi[i]; }
+        }
+        __syncthreads();
+    }
+    if(tid==0){ int64_t off=(int64_t)gid*K; for(int i=0;i<K;++i){ ov[off+i]=sv[i]; oi[off+i]=(int64_t)si[i]; } }
+}
+
+// ---- merge kernel: bitonic sort in shared memory ----
+__global__ void merge_kernel(
+    const float* __restrict__ cv, const int64_t* __restrict__ ci,
+    float* __restrict__ ov, int64_t* __restrict__ oi,
+    int nc, int k
+) {
+    int bid=blockIdx.x, tid=threadIdx.x, total=nc*k;
+    extern __shared__ char smem[];
+    float* sv=(float*)smem;
+    int*   si=(int*)(smem+total*sizeof(float));
+    for(int i=tid;i<total;i+=blockDim.x){ int64_t src=(int64_t)bid*total+i; sv[i]=cv[src]; si[i]=(int)ci[src]; }
+    __syncthreads();
+    int n2=1; while(n2<total) n2<<=1;
+    for(int stage=2;stage<=n2;stage<<=1){
+        for(int step=stage>>1;step>0;step>>=1){
+            for(int i=tid;i<n2;i+=blockDim.x){
+                int ixj=i^step; if(ixj<=i) continue;
+                float vi=(i<total)?sv[i]:-INFINITY, vj=(ixj<total)?sv[ixj]:-INFINITY;
+                bool asc=((i&stage)==0), sw=asc?(vi<vj):(vi>vj);
+                if(sw){
+                    if(i<total&&ixj<total){ float tv=sv[i];sv[i]=sv[ixj];sv[ixj]=tv; int ti=si[i];si[i]=si[ixj];si[ixj]=ti; }
+                    else if(i<total){ sv[i]=vj; si[i]=(ixj<total)?si[ixj]:-1; }
+                    else if(ixj<total){ sv[i]=vi; si[i]=(i<total)?si[i]:-1; }
+                }
+            }
+            __syncthreads();
+        }
+    }
+    for(int i=tid;i<k;i+=blockDim.x){ ov[(int64_t)bid*k+i]=sv[i]; oi[(int64_t)bid*k+i]=(int64_t)si[i]; }
+}
+
+// ---- launchers ----
+extern "C" {
+void launch_topk_64(int64_t xp,int64_t ovp,int64_t oip,int B,int n,int cs,int nc){
+    auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;
+    topk_kernel<64,256><<<B*nc,256>>>(x,ov,oi,n,cs,nc);
+}
+void launch_topk_32(int64_t xp,int64_t ovp,int64_t oip,int B,int n,int cs,int nc){
+    auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;
+    topk_kernel<32,128><<<B*nc,128>>>(x,ov,oi,n,cs,nc);
+}
+void launch_topk_16(int64_t xp,int64_t ovp,int64_t oip,int B,int n,int cs,int nc){
+    auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;
+    topk_kernel<16,128><<<B*nc,128>>>(x

... (4268 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -40,6 +40,16 @@
     while(lo<hi){ int mid=(lo+hi)>>1; if(val>vals[mid])hi=mid; else lo=mid+1; }
     for(int i=k-1;i>lo;--i){ vals[i]=vals[i-1]; indices[i]=indices[i-1]; }
     vals[lo]=val; indices[lo]=idx; return true;
+}
+
+// Non-inlined merge to reduce register pressure in caller
+template<int K>
+__device__ __noinline__ void merge_lists(
+    const float* av, const int* ai, const float* bv, const int* bi,
+    float* dv, int* di
+) {
+    int ia=0,ib=0,io=0;
+    while(io<K){ bool ta=(ia<K)&&(ib>=K||av[ia]>=bv[ib]); dv[io]=ta?av[ia]:bv[ib]; di[io]=ta?ai[ia]:bi[ib]; ia+=ta; ib+=!ta; io++; }
 }
 
 // ---- streaming kernel ----

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -67,14 +67,14 @@
     for(int i=0;i<K;++i){ lv[i]=-INFINITY; li[i]=-1; }
     for(int pos=start+tid; pos<end; pos+=NT) insert_topk(x[ro+pos],pos,lv,li,K);
 
-    // within-warp shuffle merge
+    // within-warp shuffle merge using non-inlined merge_lists
     const int lane=tid&31;
     for(int off=16;off>0;off>>=1){
         float pv[K]; int pi[K];
         for(int i=0;i<K;++i){ pv[i]=__shfl_down_sync(0xffffffff,lv[i],off); pi[i]=__shfl_down_sync(0xffffffff,li[i],off); }
         if(lane<off){
-            float mv[K];int mi[K]; int ia=0,ib=0,io=0;
-            while(io<K){ bool ta=(ia<K)&&(ib>=K||lv[ia]>=pv[ib]); mv[io]=ta?lv[ia]:pv[ib]; mi[io]=ta?li[ia]:pi[ib]; ia+=ta; ib+=!ta; io++; }
+            float mv[K]; int mi[K];
+            merge_lists<K>(lv,li, pv,pi, mv,mi);
             for(int i=0;i<K;++i){ lv[i]=mv[i]; li[i]=mi[i]; }
         }
     }
@@ -88,10 +88,8 @@
     __syncthreads();
     for(int stride=NW/2;stride>0;stride>>=1){
         if(lane==0 && wid<stride){
-            float mv[K];int mi[K]; int ia=0,ib=0,io=0;
-            float* av=&sv[wid*K];int* ai=&si[wid*K];
-            float* bv=&sv[(wid+stride)*K];int* bi=&si[(wid+stride)*K];
-            while(io<K){ bool ta=(ia<K)&&(ib>=K||av[ia]>=bv[ib]); mv[io]=ta?av[ia]:bv[ib]; mi[io]=ta?ai[ia]:bi[ib]; ia+=ta; ib+=!ta; io++; }
+            float mv[K]; int mi[K];
+            merge_lists<K>(&sv[wid*K],&si[wid*K], &sv[(wid+stride)*K],&si[(wid+stride)*K], mv,mi);
             for(int i=0;i<K;++i){ sv[wid*K+i]=mv[i]; si[wid*K+i]=mi[i]; }
         }
         __syncthreads();

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1 +1,2 @@
-            max_nc = min(max_nc, 64)+            max_nc = min(max_nc, 64)
+        if K==64: max_nc = min(max_nc, 32)  // fewer chunks for K=64

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -132,7 +132,7 @@
 extern "C" {
 void launch_topk_64(int64_t xp,int64_t ovp,int64_t oip,int B,int n,int cs,int nc){
     auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;
-    topk_kernel<64,256><<<B*nc,256>>>(x,ov,oi,n,cs,nc);
+    topk_kernel<64,128><<<B*nc,128>>>(x,ov,oi,n,cs,nc);
 }
 void launch_topk_32(int64_t xp,int64_t ovp,int64_t oip,int B,int n,int cs,int nc){
     auto*x=(const float*)xp;auto*ov=(float*)ovp;auto*oi=(int64_t*)oip;

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -188,9 +188,9 @@
         # max_chunks = 4096 // K  (but keep merge efficient)
         max_merge = 4096  # max elements in merge (fits in 32 KB shared)
         max_nc = max_merge // K
-        if K==64: max_nc = min(max_nc, 64)
+        if K==64: max_nc = min(max_nc, 32)  # fewer chunks = faster merge for large K
         # Choose chunk_size to get reasonable number of chunks
-        target_nc = min(max_nc, max(1, 128 // max(B, 1)))
+        target_nc = min(max_nc, max(1, 64 // max(B, 1) if K==64 else 128 // max(B, 1)))
         cs = max(256, (N + target_nc - 1) // target_nc)
         cs = min(cs, N)
         nc = (N + cs - 1) // cs

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,8 +1,9 @@
-"""Top-k via CUDA streaming + chunked merge.
-
-Each block processes one chunk with warp-shuffle + shared-memory reduction.
-A merge pass combines chunk results via bitonic sort.
-Chunk sizes chosen so merge fits in 48 KB default shared memory.
+"""Top-k via radix-select CUDA kernel.
+
+Two-pass approach:
+  Pass 1: histogram top 11 bits of fp32 values across all blocks.
+          Determine which bucket(s) contain the top-k threshold.
+  Pass 2: within the target bucket(s), select top-k via register streaming.
 """
 
 from typing import List
@@ -18,12 +19,11 @@
 _cpp_src = r"""
 #include <cstdint>
 extern "C" {
-void launch_topk_64(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_topk_32(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_topk_16(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_topk_8(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_topk_1(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_merge(int64_t,int64_t,int64_t,int64_t,int,int,int);
+void launch_topk_1(int64_t,int64_t,int64_t,int,int);
+void launch_topk_8(int64_t,int64_t,int64_t,int,int);
+void launch_topk_16(int64_t,int64_t,int64_t,int,int);
+void launch_topk_32(int64_t,int64_t,int64_t,int,int);
+void launch_topk_64(int64_t,int64_t,int64_t,int,int);
 }
 """
 
@@ -32,129 +32,186 @@
 #include <cstdint>
 #include <cfloat>
 #include <math.h>
-
-// ---- helpers ----
-__device__ bool insert_topk(float val, int idx, float* vals, int* indices, int k) {
-    if (val <= vals[k-1]) return false;
-    int lo=0, hi=k-1;
-    while(lo<hi){ int mid=(lo+hi)>>1; if(val>vals[mid])hi=mid; else lo=mid+1; }
-    for(int i=k-1;i>lo;--i){ vals[i]=vals[i-1]; indices[i]=indices[i-1]; }
-    vals[lo]=val; indices[lo]=idx; return true;
-}
-
-// Non-inlined merge to reduce register pressure in caller
-template<int K>
-__device__ __noinline__ void merge_lists(
-    const float* av, const int* ai, const float* bv, const int* bi,
-    float* dv, int* di
+#include <stdio.h>
+
+// ---- radix select for single-row top-k ----
+// Uses two passes:
+// Pass 1: histogram bits [31:21] (top 11 bits) → 2048 buckets
+// Pass 2: within selected buckets, histogram bits [20:10] → refine
+// Final: collect top-k
+
+#define HIST_BINS 2048
+#define HIST_BITS 11
+
+__global__ void radix_histogram_kernel(
+    const float* __restrict__ x,
+    int* __restrict__ histogram,  // [HIST_BINS] per row
+    int n
 ) {
-    int ia=0,ib=0,io=0;
-    while(io<K){ bool ta=(ia<K)&&(ib>=K||av[ia]>=bv[ib]); dv[io]=ta?av[ia]:bv[ib]; di[io]=ta?ai[ia]:bi[ib]; ia+=ta; ib+=!ta; io++; }
-}
-
-// ---- streaming kernel ----
-template <int K, int NT>
-__global__ void topk_kernel(
-    const float* __restrict__ x, float* __restrict__ ov, int64_t* __restrict__ oi,
-    int n, int cs, int nc
+    __shared__ int shist[HIST_BINS];
+    int tid = threadIdx.x;
+    // Init shared histogram
+    for(int i=tid; i<HIST_BINS; i+=blockDim.x) shist[i] = 0;
+    __syncthreads();
+
+    // Count elements by top HIST_BITS bits
+    for(int pos = blockIdx.x * blockDim.x + tid; pos < n; pos += gridDim.x * blockDim.x) {
+        unsigned int u = __float_as_uint(x[pos]);
+        // Flip sign bit for correct ordering: pos → 0x80000000, neg → flip all
+        unsigned int mask = (int)u >> 31;
+        u ^= (mask | 0x80000000u);
+        int bin = (int)(u >> (32 - HIST_BITS));
+        atomicAdd(&shist[bin], 1);
+    }
+    __syncthreads();
+
+    // Write to global histogram
+    for(int i=tid; i<HIST_BINS; i+=blockDim.x) {
+        if(shist[i] > 0) atomicAdd(&histogram[i], shist[i]);
+    }
+}
+
+// Global: after histogram, CPU determines which bins contain top-k.
+// Then launch a collection kernel for those bins.
+
+__global__ void radix_collect_kernel(
+    const float* __restrict__ x,
+    float* __restrict__ out_vals,
+    int64_t* __restrict__ out_idx,
+    int n, int k,
+    int bin_start, int bin_end  // inclusive range of histogram bins to scan
 ) {
-    int gid=blockIdx.x, bid=gid/nc, cid=gid%nc, tid=threadIdx.x;
-    int64_t ro = (int64_t)bid*n;
-    int start=cid*cs, end=min(start+cs,n);
-    if(start>=end) return;
-
+    int tid = threadIdx.x;
+    // Each thread maintains local top-K (same as streaming approach)
+    float lv[64]; int li[64]; // max K=64
+    if(k > 64) return;
+    for(int i=0;i<k;++i){ lv[i]=-INFINITY; li[i]=-1; }
+
+    // Scan elements that fall into the target bin range
+    for(int pos = blockIdx.x * blockDim.x + tid; pos < n; pos += gridDim.x * blockDim.x) {
+        unsigned int u = __float_as_uint(x[pos]);
+        unsigned int mask = (int)u >> 31;
+        u ^= (mask | 0x80000000u);
+        int bin = (int)(u >> (32 - HIST_BITS));
+        if(bin >= bin_start && bin <= bin_end) {
+            float v = x[pos];
+            // Insert into top-k
+            if(v > lv[k-1]) {
+                int lo=0, hi=k-1;
+                while(lo<hi){ int mid=(lo+hi)>>1; if(v>lv[mid])hi=mid; else lo=mid+1; }
+                for(int i=k-1;i>lo;--i){ lv[i]=lv[i-1]; li[i]=li[i-1]; }
+                lv[lo]=v; li[lo]=pos;
+            }
+        }
+    }
+
+    // Write to global, use atomics to merge across blocks
+    // Simple approach: write all local top-k to a scratch buffer, merge in a second pass
+    // For now, write to shared memory and merge within block
+    __shared__ float sv[4096];  // max 1024 threads * 64 = 65536... too big.
+    // Use a different approach: write to global and merge in a separate kernel
+    int64_t gid = (int64_t)blockIdx.x * k;
+    for(int i=0;i<k;++i){ out_vals[gid+i] = lv[i]; out_idx[gid+i] = (int64_t)li[i]; }
+}
+"""
+
+# --- Actually, the radix approach is getting too complex with multiple passes.
+# Let me go back to the simple streaming approach but optimize it differently.
+# The key insight: we just need to be competitive, not beat torch by 10x.
+
+# For the required shapes: use a si

... (9732 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,13 +1,10 @@
-"""Top-k via radix-select CUDA kernel.
-
-Two-pass approach:
-  Pass 1: histogram top 11 bits of fp32 values across all blocks.
-          Determine which bucket(s) contain the top-k threshold.
-  Pass 2: within the target bucket(s), select top-k via register streaming.
+"""Top-k via CUDA streaming + chunked merge with noinline helpers.
+
+Each block processes one chunk with warp-shuffle + shared-memory reduction.
+A merge pass combines chunk results via bitonic sort in shared memory.
 """
 
 from typing import List
-
 import torch
 import torch.nn as nn
 import torch.utils.cpp_extension
@@ -19,11 +16,12 @@
 _cpp_src = r"""
 #include <cstdint>
 extern "C" {
-void launch_topk_1(int64_t,int64_t,int64_t,int,int);
-void launch_topk_8(int64_t,int64_t,int64_t,int,int);
-void launch_topk_16(int64_t,int64_t,int64_t,int,int);
-void launch_topk_32(int64_t,int64_t,int64_t,int,int);
-void launch_topk_64(int64_t,int64_t,int64_t,int,int);
+void launch_stream_64(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_stream_32(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_stream_16(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_stream_8(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_stream_1(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_merge(int64_t,int64_t,int64_t,int64_t,int,int,int);
 }
 """
 
@@ -32,186 +30,121 @@
 #include <cstdint>
 #include <cfloat>
 #include <math.h>
-#include <stdio.h>
-
-// ---- radix select for single-row top-k ----
-// Uses two passes:
-// Pass 1: histogram bits [31:21] (top 11 bits) → 2048 buckets
-// Pass 2: within selected buckets, histogram bits [20:10] → refine
-// Final: collect top-k
-
-#define HIST_BINS 2048
-#define HIST_BITS 11
-
-__global__ void radix_histogram_kernel(
-    const float* __restrict__ x,
-    int* __restrict__ histogram,  // [HIST_BINS] per row
-    int n
+
+// ------- Helpers (noinline to reduce register pressure) -------
+
+__device__ bool insert_topk(float val, int idx, float* vals, int* indices, int k) {
+    if (val <= vals[k-1]) return false;
+    int lo=0, hi=k-1;
+    while(lo<hi){ int mid=(lo+hi)>>1; if(val>vals[mid])hi=mid; else lo=mid+1; }
+    for(int i=k-1;i>lo;--i){ vals[i]=vals[i-1]; indices[i]=indices[i-1]; }
+    vals[lo]=val; indices[lo]=idx; return true;
+}
+
+template<int K>
+__device__ __noinline__ void merge_lists(
+    const float* av, const int* ai, const float* bv, const int* bi,
+    float* dv, int* di
 ) {
-    __shared__ int shist[HIST_BINS];
-    int tid = threadIdx.x;
-    // Init shared histogram
-    for(int i=tid; i<HIST_BINS; i+=blockDim.x) shist[i] = 0;
-    __syncthreads();
-
-    // Count elements by top HIST_BITS bits
-    for(int pos = blockIdx.x * blockDim.x + tid; pos < n; pos += gridDim.x * blockDim.x) {
-        unsigned int u = __float_as_uint(x[pos]);
-        // Flip sign bit for correct ordering: pos → 0x80000000, neg → flip all
-        unsigned int mask = (int)u >> 31;
-        u ^= (mask | 0x80000000u);
-        int bin = (int)(u >> (32 - HIST_BITS));
-        atomicAdd(&shist[bin], 1);
-    }
-    __syncthreads();
-
-    // Write to global histogram
-    for(int i=tid; i<HIST_BINS; i+=blockDim.x) {
-        if(shist[i] > 0) atomicAdd(&histogram[i], shist[i]);
-    }
-}
-
-// Global: after histogram, CPU determines which bins contain top-k.
-// Then launch a collection kernel for those bins.
-
-__global__ void radix_collect_kernel(
-    const float* __restrict__ x,
-    float* __restrict__ out_vals,
-    int64_t* __restrict__ out_idx,
-    int n, int k,
-    int bin_start, int bin_end  // inclusive range of histogram bins to scan
+    int ia=0,ib=0,io=0;
+    while(io<K){ bool ta=(ia<K)&&(ib>=K||av[ia]>=bv[ib]); dv[io]=ta?av[ia]:bv[ib]; di[io]=ta?ai[ia]:bi[ib]; ia+=ta; ib+=!ta; io++; }
+}
+
+// ------- Streaming kernel -------
+
+template <int K, int NT>
+__global__ void stream_kernel(
+    const float* __restrict__ x, float* __restrict__ ov, int64_t* __restrict__ oi,
+    int n, int cs, int nc
 ) {
-    int tid = threadIdx.x;
-    // Each thread maintains local top-K (same as streaming approach)
-    float lv[64]; int li[64]; // max K=64
-    if(k > 64) return;
-    for(int i=0;i<k;++i){ lv[i]=-INFINITY; li[i]=-1; }
-
-    // Scan elements that fall into the target bin range
-    for(int pos = blockIdx.x * blockDim.x + tid; pos < n; pos += gridDim.x * blockDim.x) {
-        unsigned int u = __float_as_uint(x[pos]);
-        unsigned int mask = (int)u >> 31;
-        u ^= (mask | 0x80000000u);
-        int bin = (int)(u >> (32 - HIST_BITS));
-        if(bin >= bin_start && bin <= bin_end) {
-            float v = x[pos];
-            // Insert into top-k
-            if(v > lv[k-1]) {
-                int lo=0, hi=k-1;
-                while(lo<hi){ int mid=(lo+hi)>>1; if(v>lv[mid])hi=mid; else lo=mid+1; }
-                for(int i=k-1;i>lo;--i){ lv[i]=lv[i-1]; li[i]=li[i-1]; }
-                lv[lo]=v; li[lo]=pos;
-            }
-        }
-    }
-
-    // Write to global, use atomics to merge across blocks
-    // Simple approach: write all local top-k to a scratch buffer, merge in a second pass
-    // For now, write to shared memory and merge within block
-    __shared__ float sv[4096];  // max 1024 threads * 64 = 65536... too big.
-    // Use a different approach: write to global and merge in a separate kernel
-    int64_t gid = (int64_t)blockIdx.x * k;
-    for(int i=0;i<k;++i){ out_vals[gid+i] = lv[i]; out_idx[gid+i] = (int64_t)li[i]; }
-}
-"""
-
-# --- Actually, the radix approach is getting too complex with multiple passes.
-# Let me go back to the simple streaming approach but optimize it differently.
-# The key insight: we just need to be competitive, not beat torch by 10x.
-
-# For the required shapes: use a simple streaming kernel that's memory-bandwidth optimal
-# and minimize register pressure by keeping k small in registers

... (9894 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -174,15 +174,19 @@
 
     def forward(self,x):
         B,N,K = self.batch,self.n,self.k
-        # Chunk: keep merge buffer <= 32KB (4096 elements shared, 0 elements = 0 bytes)
-        max_merge = 4096 // K  # max chunks for shared memory budget
-        # Target ~2-8 chunks per row for parallelism vs merge tradeoff
-        tgt = 4 if B < 32 else 2
-        tgt = min(max_merge, tgt, (N+255)//256)  # at least 256 elem per chunk
-        tgt = max(1, tgt)
-        cs = (N + tgt - 1) // tgt
-        cs = max(256, cs)
-        nc = (N + cs - 1) // cs
+        # Optimized chunk sizes per shape
+        if B==1 and N==131072 and K==64:
+            cs, nc = 2048, 64
+        elif B==64 and N==8192 and K==8:
+            cs, nc = N, 1  # no merge needed
+        elif B==32 and N==16384 and K==32:
+            cs, nc = 4096, 4
+        elif B==16 and N==12000 and K==16:
+            cs, nc = 1500, 8
+        elif B==128 and N==4096 and K==1:
+            cs, nc = N, 1
+        else:
+            cs, nc = N, 1
 
         m = _get_mod()
         launch = {64:m.launch_stream_64,32:m.launch_stream_32,16:m.launch_stream_16,8:m.launch_stream_8,1:m.launch_stream_1}[K]
@@ -190,7 +194,7 @@
         if nc == 1:
             ov = torch.empty(B,K,dtype=torch.float32,device=x.device)
             oi = torch.empty(B,K,dtype=torch.int64,device=x.device)
-            launch(x.data_ptr(),ov.data_ptr(),oi.data_ptr(),B,N,N,nc)
+            launch(x.data_ptr(),ov.data_ptr(),oi.data_ptr(),B,N,cs,nc)
             return ov,oi
 
         total = B * nc

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -86,8 +86,10 @@
     if(tid==0){ int64_t off=(int64_t)gid*K; for(int i=0;i<K;++i){ ov[off+i]=sv[i]; oi[off+i]=(int64_t)si[i]; } }
 }
 
-// ------- Merge kernel -------
-
+// ------- Merge kernel: tournament selection -------
+// Load all nc*k elements into shared memory.
+// Then: repeatedly find global max (via reduction), extract it, write to output.
+// Uses warp-shuffle for fast max reduction.
 __global__ void merge_kernel(
     const float* __restrict__ cv, const int64_t* __restrict__ ci,
     float* __restrict__ ov, int64_t* __restrict__ oi,
@@ -99,23 +101,58 @@
     int*   si=(int*)(smem+total*sizeof(float));
     for(int i=tid;i<total;i+=blockDim.x){ int64_t src=(int64_t)bid*total+i; sv[i]=cv[src]; si[i]=(int)ci[src]; }
     __syncthreads();
-    int n2=1; while(n2<total) n2<<=1;
-    for(int stage=2;stage<=n2;stage<<=1){
-        for(int step=stage>>1;step>0;step>>=1){
-            for(int i=tid;i<n2;i+=blockDim.x){
-                int ixj=i^step; if(ixj<=i) continue;
-                float vi=(i<total)?sv[i]:-INFINITY, vj=(ixj<total)?sv[ixj]:-INFINITY;
-                bool asc=((i&stage)==0), sw=asc?(vi<vj):(vi>vj);
-                if(sw){
-                    if(i<total&&ixj<total){ float tv=sv[i];sv[i]=sv[ixj];sv[ixj]=tv; int ti=si[i];si[i]=si[ixj];si[ixj]=ti; }
-                    else if(i<total){ sv[i]=vj; si[i]=(ixj<total)?si[ixj]:-1; }
-                    else if(ixj<total){ sv[i]=vi; si[i]=(i<total)?si[i]:-1; }
-                }
+
+    for(int out_idx=0; out_idx<k; out_idx++) {
+        // Each thread finds max among its assigned elements
+        float best_val = -INFINITY;
+        int   best_pos = -1;
+        int elems_per_thread = (total + blockDim.x - 1) / blockDim.x;
+        int start = tid * elems_per_thread;
+        int end = min(start + elems_per_thread, total);
+        for(int i=start; i<end; i++) {
+            if(sv[i] > best_val) { best_val = sv[i]; best_pos = i; }
+        }
+
+        // Warp-level reduction to find global max
+        const int lane = tid & 31;
+        for(int off=16; off>0; off>>=1) {
+            float other_val = __shfl_down_sync(0xffffffff, best_val, off);
+            int   other_pos = __shfl_down_sync(0xffffffff, best_pos, off);
+            if(other_val > best_val || (other_val == best_val && other_pos < best_pos)) {
+                best_val = other_val;
+                best_pos = other_pos;
             }
-            __syncthreads();
-        }
+        }
+
+        // Cross-warp reduction using shared memory (small: 32 warps max)
+        __shared__ float wv[32];
+        __shared__ int   wp[32];
+        int wid = tid / 32;
+        if(lane == 0) { wv[wid] = best_val; wp[wid] = best_pos; }
+        __syncthreads();
+
+        // First warp does final reduction
+        if(wid == 0) {
+            if(lane == 0) { best_val = wv[0]; best_pos = wp[0]; }
+            for(int i=1; i<blockDim.x/32; i++) {
+                if(lane == 0 && wv[i] > best_val) { best_val = wv[i]; best_pos = wp[i]; }
+            }
+            // Shared the result with all threads via shared mem
+            if(lane == 0) { wv[0] = best_val; wp[0] = best_pos; }
+        }
+        __syncthreads();
+        best_val = wv[0]; best_pos = wp[0];
+        __syncthreads();
+
+        // Write the found max to output
+        if(tid == 0) {
+            ov[(int64_t)bid*k + out_idx] = best_val;
+            oi[(int64_t)bid*k + out_idx] = (int64_t)si[best_pos];
+            // Remove this element from consideration
+            sv[best_pos] = -INFINITY;
+        }
+        __syncthreads();
     }
-    for(int i=tid;i<k;i+=blockDim.x){ ov[(int64_t)bid*k+i]=sv[i]; oi[(int64_t)bid*k+i]=(int64_t)si[i]; }
 }
 
 // ------- Launchers -------

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -213,7 +213,7 @@
         B,N,K = self.batch,self.n,self.k
         # Optimized chunk sizes per shape
         if B==1 and N==131072 and K==64:
-            cs, nc = 2048, 64
+            cs, nc = N, 1  # single block, no merge needed
         elif B==64 and N==8192 and K==8:
             cs, nc = N, 1  # no merge needed
         elif B==32 and N==16384 and K==32:

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -213,7 +213,7 @@
         B,N,K = self.batch,self.n,self.k
         # Optimized chunk sizes per shape
         if B==1 and N==131072 and K==64:
-            cs, nc = N, 1  # single block, no merge needed
+            cs, nc = 2048, 64  # 64 chunks, merge 4096 elems
         elif B==64 and N==8192 and K==8:
             cs, nc = N, 1  # no merge needed
         elif B==32 and N==16384 and K==32:

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -64,7 +64,15 @@
 
     float lv[K]; int li[K];
     for(int i=0;i<K;++i){ lv[i]=-INFINITY; li[i]=-1; }
-    for(int pos=start+tid; pos<end; pos+=NT) insert_topk(x[ro+pos],pos,lv,li,K);
+    for(int pos=start+tid; pos<end; pos+=NT) {
+        float v = __ldg(x + ro + pos);
+        if(v > lv[K-1]) {
+            int lo=0,hi=K-1;
+            while(lo<hi){ int mid=(lo+hi)>>1; if(v>lv[mid])hi=mid; else lo=mid+1; }
+            for(int i=K-1;i>lo;--i){ lv[i]=lv[i-1]; li[i]=li[i-1]; }
+            lv[lo]=v; li[lo]=pos;
+        }
+    }
 
     const int lane=tid&31;
     for(int off=16;off>0;off>>=1){

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -94,10 +94,7 @@
     if(tid==0){ int64_t off=(int64_t)gid*K; for(int i=0;i<K;++i){ ov[off+i]=sv[i]; oi[off+i]=(int64_t)si[i]; } }
 }
 
-// ------- Merge kernel: tournament selection -------
-// Load all nc*k elements into shared memory.
-// Then: repeatedly find global max (via reduction), extract it, write to output.
-// Uses warp-shuffle for fast max reduction.
+// ------- Merge kernel: bitonic sort -------
 __global__ void merge_kernel(
     const float* __restrict__ cv, const int64_t* __restrict__ ci,
     float* __restrict__ ov, int64_t* __restrict__ oi,
@@ -109,58 +106,23 @@
     int*   si=(int*)(smem+total*sizeof(float));
     for(int i=tid;i<total;i+=blockDim.x){ int64_t src=(int64_t)bid*total+i; sv[i]=cv[src]; si[i]=(int)ci[src]; }
     __syncthreads();
-
-    for(int out_idx=0; out_idx<k; out_idx++) {
-        // Each thread finds max among its assigned elements
-        float best_val = -INFINITY;
-        int   best_pos = -1;
-        int elems_per_thread = (total + blockDim.x - 1) / blockDim.x;
-        int start = tid * elems_per_thread;
-        int end = min(start + elems_per_thread, total);
-        for(int i=start; i<end; i++) {
-            if(sv[i] > best_val) { best_val = sv[i]; best_pos = i; }
+    int n2=1; while(n2<total) n2<<=1;
+    for(int stage=2;stage<=n2;stage<<=1){
+        for(int step=stage>>1;step>0;step>>=1){
+            for(int i=tid;i<n2;i+=blockDim.x){
+                int ixj=i^step; if(ixj<=i) continue;
+                float vi=(i<total)?sv[i]:-INFINITY, vj=(ixj<total)?sv[ixj]:-INFINITY;
+                bool asc=((i&stage)==0), sw=asc?(vi<vj):(vi>vj);
+                if(sw){
+                    if(i<total&&ixj<total){ float tv=sv[i];sv[i]=sv[ixj];sv[ixj]=tv; int ti=si[i];si[i]=si[ixj];si[ixj]=ti; }
+                    else if(i<total){ sv[i]=vj; si[i]=(ixj<total)?si[ixj]:-1; }
+                    else if(ixj<total){ sv[i]=vi; si[i]=(i<total)?si[i]:-1; }
+                }
+            }
+            __syncthreads();
         }
-
-        // Warp-level reduction to find global max
-        const int lane = tid & 31;
-        for(int off=16; off>0; off>>=1) {
-            float other_val = __shfl_down_sync(0xffffffff, best_val, off);
-            int   other_pos = __shfl_down_sync(0xffffffff, best_pos, off);
-            if(other_val > best_val || (other_val == best_val && other_pos < best_pos)) {
-                best_val = other_val;
-                best_pos = other_pos;
-            }
-        }
-
-        // Cross-warp reduction using shared memory (small: 32 warps max)
-        __shared__ float wv[32];
-        __shared__ int   wp[32];
-        int wid = tid / 32;
-        if(lane == 0) { wv[wid] = best_val; wp[wid] = best_pos; }
-        __syncthreads();
-
-        // First warp does final reduction
-        if(wid == 0) {
-            if(lane == 0) { best_val = wv[0]; best_pos = wp[0]; }
-            for(int i=1; i<blockDim.x/32; i++) {
-                if(lane == 0 && wv[i] > best_val) { best_val = wv[i]; best_pos = wp[i]; }
-            }
-            // Shared the result with all threads via shared mem
-            if(lane == 0) { wv[0] = best_val; wp[0] = best_pos; }
-        }
-        __syncthreads();
-        best_val = wv[0]; best_pos = wp[0];
-        __syncthreads();
-
-        // Write the found max to output
-        if(tid == 0) {
-            ov[(int64_t)bid*k + out_idx] = best_val;
-            oi[(int64_t)bid*k + out_idx] = (int64_t)si[best_pos];
-            // Remove this element from consideration
-            sv[best_pos] = -INFINITY;
-        }
-        __syncthreads();
-    }
+    }
+    for(int i=tid;i<k;i+=blockDim.x){ ov[(int64_t)bid*k+i]=sv[i]; oi[(int64_t)bid*k+i]=(int64_t)si[i]; }
 }
 
 // ------- Launchers -------

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,7 +1,7 @@
-"""Top-k via CUDA streaming + chunked merge with noinline helpers.
-
-Each block processes one chunk with warp-shuffle + shared-memory reduction.
-A merge pass combines chunk results via bitonic sort in shared memory.
+"""Top-k via unified CUDA kernel with last-block merge.
+
+Single kernel launch: all blocks process chunks in parallel.
+The last finishing block merges all chunk results via bitonic sort.
 """
 
 from typing import List
@@ -16,12 +16,11 @@
 _cpp_src = r"""
 #include <cstdint>
 extern "C" {
-void launch_stream_64(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_stream_32(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_stream_16(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_stream_8(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_stream_1(int64_t,int64_t,int64_t,int,int,int,int);
-void launch_merge(int64_t,int64_t,int64_t,int64_t,int,int,int);
+void launch_topk_64(int64_t,int64_t,int64_t,int64_t,int64_t,int,int,int,int,int);
+void launch_topk_32(int64_t,int64_t,int64_t,int64_t,int64_t,int,int,int,int,int);
+void launch_topk_16(int64_t,int64_t,int64_t,int64_t,int64_t,int,int,int,int,int);
+void launch_topk_8(int64_t,int64_t,int64_t,int64_t,int64_t,int,int,int,int,int);
+void launch_topk_1(int64_t,int64_t,int64_t,int64_t,int64_t,int,int,int,int,int);
 }
 """
 
@@ -31,49 +30,45 @@
 #include <cfloat>
 #include <math.h>
 
-// ------- Helpers (noinline to reduce register pressure) -------
-
-__device__ bool insert_topk(float val, int idx, float* vals, int* indices, int k) {
-    if (val <= vals[k-1]) return false;
-    int lo=0, hi=k-1;
-    while(lo<hi){ int mid=(lo+hi)>>1; if(val>vals[mid])hi=mid; else lo=mid+1; }
-    for(int i=k-1;i>lo;--i){ vals[i]=vals[i-1]; indices[i]=indices[i-1]; }
-    vals[lo]=val; indices[lo]=idx; return true;
-}
-
+// ---- noinline merge helper ----
 template<int K>
 __device__ __noinline__ void merge_lists(
-    const float* av, const int* ai, const float* bv, const int* bi,
-    float* dv, int* di
+    const float* av, const int* ai, const float* bv, const int* bi, float* dv, int* di
 ) {
     int ia=0,ib=0,io=0;
     while(io<K){ bool ta=(ia<K)&&(ib>=K||av[ia]>=bv[ib]); dv[io]=ta?av[ia]:bv[ib]; di[io]=ta?ai[ia]:bi[ib]; ia+=ta; ib+=!ta; io++; }
 }
 
-// ------- Streaming kernel -------
-
+// ---- Unified kernel: streaming + last-block merge ----
 template <int K, int NT>
-__global__ void stream_kernel(
-    const float* __restrict__ x, float* __restrict__ ov, int64_t* __restrict__ oi,
+__global__ void unified_kernel(
+    const float* __restrict__ x,
+    float* __restrict__ out_vals, int64_t* __restrict__ out_idx,
+    float* __restrict__ chunk_vals, int64_t* __restrict__ chunk_idx,  // scratch buffers
+    int* __restrict__ counter,  // atomic counter for last-block detection
     int n, int cs, int nc
 ) {
     int gid=blockIdx.x, bid=gid/nc, cid=gid%nc, tid=threadIdx.x;
     int64_t ro=(int64_t)bid*n;
     int start=cid*cs, end=min(start+cs,n);
-    if(start>=end) return;
-
+    constexpr int NW=NT/32, NS=NW*K;
+
+    // ---- Phase 1: streaming top-K for this chunk ----
     float lv[K]; int li[K];
     for(int i=0;i<K;++i){ lv[i]=-INFINITY; li[i]=-1; }
-    for(int pos=start+tid; pos<end; pos+=NT) {
-        float v = __ldg(x + ro + pos);
-        if(v > lv[K-1]) {
-            int lo=0,hi=K-1;
-            while(lo<hi){ int mid=(lo+hi)>>1; if(v>lv[mid])hi=mid; else lo=mid+1; }
-            for(int i=K-1;i>lo;--i){ lv[i]=lv[i-1]; li[i]=li[i-1]; }
-            lv[lo]=v; li[lo]=pos;
+    if(start<end) {
+        for(int pos=start+tid; pos<end; pos+=NT) {
+            float v = x[ro+pos];
+            if(v > lv[K-1]) {
+                int lo=0,hi=K-1;
+                while(lo<hi){ int mid=(lo+hi)>>1; if(v>lv[mid])hi=mid; else lo=mid+1; }
+                for(int i=K-1;i>lo;--i){ lv[i]=lv[i-1]; li[i]=li[i-1]; }
+                lv[lo]=v; li[lo]=pos;
+            }
         }
     }
 
+    // Within-warp shuffle merge
     const int lane=tid&31;
     for(int off=16;off>0;off>>=1){
         float pv[K]; int pi[K];
@@ -81,7 +76,7 @@
         if(lane<off){ float mv[K];int mi[K]; merge_lists<K>(lv,li,pv,pi,mv,mi); for(int i=0;i<K;++i){lv[i]=mv[i];li[i]=mi[i];} }
     }
 
-    constexpr int NW=NT/32, NS=NW*K;
+    // Cross-warp tree merge
     __shared__ float sv[NS];
     __shared__ int   si[NS];
     int wid=tid/32;
@@ -91,69 +86,82 @@
         if(lane==0&&wid<stride){ float mv[K];int mi[K]; merge_lists<K>(&sv[wid*K],&si[wid*K],&sv[(wid+stride)*K],&si[(wid+stride)*K],mv,mi); for(int i=0;i<K;++i){sv[wid*K+i]=mv[i];si[wid*K+i]=mi[i];} }
         __syncthreads();
     }
-    if(tid==0){ int64_t off=(int64_t)gid*K; for(int i=0;i<K;++i){ ov[off+i]=sv[i]; oi[off+i]=(int64_t)si[i]; } }
-}
-
-// ------- Merge kernel: bitonic sort -------
-__global__ void merge_kernel(
-    const float* __restrict__ cv, const int64_t* __restrict__ ci,
-    float* __restrict__ ov, int64_t* __restrict__ oi,
-    int nc, int k
-) {
-    int bid=blockIdx.x, tid=threadIdx.x, total=nc*k;
-    extern __shared__ char smem[];
-    float* sv=(float*)smem;
-    int*   si=(int*)(smem+total*sizeof(float));
-    for(int i=tid;i<total;i+=blockDim.x){ int64_t src=(int64_t)bid*total+i; sv[i]=cv[src]; si[i]=(int)ci[src]; }
+
+    // Write chunk result to global memory
+    if(tid==0){ int64_t off=(int64_t)gid*K; for(int i=0;i<K;++i){ chunk_vals[off+i]=sv[i]; chunk_idx[off+i]=(int64_t)si[i]; } }
+
+    // ---- Phase 2: last-block merge ----
+    __threadfence();
+    int prev = atomicAdd(counter, 1);
+
+    if(prev == gid && gid == gridDim.x - 1) {  // Last block: do merge  (use gid check as fallback)
+        // Wait for all other blocks: spin until counter reaches gridDim.x
+        // But counter was just incremented by this block. Need ordering.
+        // Simple approach: use a second counter
+    }
+    // Spin until we know all blocks

... (7448 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,7 +1,10 @@
-"""Top-k via unified CUDA kernel with last-block merge.
+"""Top-k via CUDA streaming + bitonic merge.
 
-Single kernel launch: all blocks process chunks in parallel.
-The last finishing block merges all chunk results via bitonic sort.
+Two-kernel approach:
+  - Stream kernel: each block processes one chunk with warp-shuffle + shared-memory reduction.
+  - Merge kernel: bitonic sort combines chunk results.
+
+Optimized chunk sizes per shape for best parallelism/merge tradeoff.
 """
 
 from typing import List
@@ -16,11 +19,12 @@
 _cpp_src = r"""
 #include <cstdint>
 extern "C" {
-void launch_topk_64(int64_t,int64_t,int64_t,int64_t,int64_t,int,int,int,int,int);
-void launch_topk_32(int64_t,int64_t,int64_t,int64_t,int64_t,int,int,int,int,int);
-void launch_topk_16(int64_t,int64_t,int64_t,int64_t,int64_t,int,int,int,int,int);
-void launch_topk_8(int64_t,int64_t,int64_t,int64_t,int64_t,int,int,int,int,int);
-void launch_topk_1(int64_t,int64_t,int64_t,int64_t,int64_t,int,int,int,int,int);
+void launch_stream_64(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_stream_32(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_stream_16(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_stream_8(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_stream_1(int64_t,int64_t,int64_t,int,int,int,int);
+void launch_merge(int64_t,int64_t,int64_t,int64_t,int,int,int);
 }
 """
 
@@ -30,137 +34,73 @@
 #include <cfloat>
 #include <math.h>
 
-// ---- noinline merge helper ----
 template<int K>
 __device__ __noinline__ void merge_lists(
     const float* av, const int* ai, const float* bv, const int* bi, float* dv, int* di
-) {
-    int ia=0,ib=0,io=0;
-    while(io<K){ bool ta=(ia<K)&&(ib>=K||av[ia]>=bv[ib]); dv[io]=ta?av[ia]:bv[ib]; di[io]=ta?ai[ia]:bi[ib]; ia+=ta; ib+=!ta; io++; }
-}
+) { int ia=0,ib=0,io=0; while(io<K){ bool ta=(ia<K)&&(ib>=K||av[ia]>=bv[ib]); dv[io]=ta?av[ia]:bv[ib]; di[io]=ta?ai[ia]:bi[ib]; ia+=ta; ib+=!ta; io++; } }
 
-// ---- Unified kernel: streaming + last-block merge ----
 template <int K, int NT>
-__global__ void unified_kernel(
-    const float* __restrict__ x,
-    float* __restrict__ out_vals, int64_t* __restrict__ out_idx,
-    float* __restrict__ chunk_vals, int64_t* __restrict__ chunk_idx,  // scratch buffers
-    int* __restrict__ counter,  // atomic counter for last-block detection
+__global__ void stream_kernel(
+    const float* __restrict__ x, float* __restrict__ ov, int64_t* __restrict__ oi,
     int n, int cs, int nc
 ) {
     int gid=blockIdx.x, bid=gid/nc, cid=gid%nc, tid=threadIdx.x;
     int64_t ro=(int64_t)bid*n;
     int start=cid*cs, end=min(start+cs,n);
-    constexpr int NW=NT/32, NS=NW*K;
-
-    // ---- Phase 1: streaming top-K for this chunk ----
+    if(start>=end) return;
     float lv[K]; int li[K];
     for(int i=0;i<K;++i){ lv[i]=-INFINITY; li[i]=-1; }
-    if(start<end) {
-        for(int pos=start+tid; pos<end; pos+=NT) {
-            float v = x[ro+pos];
-            if(v > lv[K-1]) {
-                int lo=0,hi=K-1;
-                while(lo<hi){ int mid=(lo+hi)>>1; if(v>lv[mid])hi=mid; else lo=mid+1; }
-                for(int i=K-1;i>lo;--i){ lv[i]=lv[i-1]; li[i]=li[i-1]; }
-                lv[lo]=v; li[lo]=pos;
-            }
+    for(int pos=start+tid; pos<end; pos+=NT) {
+        float v = x[ro+pos];
+        if(v > lv[K-1]) {
+            int lo=0,hi=K-1; while(lo<hi){ int mid=(lo+hi)>>1; if(v>lv[mid])hi=mid; else lo=mid+1; }
+            for(int i=K-1;i>lo;--i){ lv[i]=lv[i-1]; li[i]=li[i-1]; } lv[lo]=v; li[lo]=pos;
         }
     }
-
-    // Within-warp shuffle merge
     const int lane=tid&31;
-    for(int off=16;off>0;off>>=1){
-        float pv[K]; int pi[K];
+    for(int off=16;off>0;off>>=1){ float pv[K];int pi[K];
         for(int i=0;i<K;++i){ pv[i]=__shfl_down_sync(0xffffffff,lv[i],off); pi[i]=__shfl_down_sync(0xffffffff,li[i],off); }
         if(lane<off){ float mv[K];int mi[K]; merge_lists<K>(lv,li,pv,pi,mv,mi); for(int i=0;i<K;++i){lv[i]=mv[i];li[i]=mi[i];} }
     }
-
-    // Cross-warp tree merge
-    __shared__ float sv[NS];
-    __shared__ int   si[NS];
-    int wid=tid/32;
+    constexpr int NW=NT/32, NS=NW*K;
+    __shared__ float sv[NS]; __shared__ int si[NS]; int wid=tid/32;
     if(lane==0){ for(int i=0;i<K;++i){ sv[wid*K+i]=lv[i]; si[wid*K+i]=li[i]; } }
     __syncthreads();
-    for(int stride=NW/2;stride>0;stride>>=1){
-        if(lane==0&&wid<stride){ float mv[K];int mi[K]; merge_lists<K>(&sv[wid*K],&si[wid*K],&sv[(wid+stride)*K],&si[(wid+stride)*K],mv,mi); for(int i=0;i<K;++i){sv[wid*K+i]=mv[i];si[wid*K+i]=mi[i];} }
-        __syncthreads();
-    }
-
-    // Write chunk result to global memory
-    if(tid==0){ int64_t off=(int64_t)gid*K; for(int i=0;i<K;++i){ chunk_vals[off+i]=sv[i]; chunk_idx[off+i]=(int64_t)si[i]; } }
-
-    // ---- Phase 2: last-block merge ----
-    __threadfence();
-    int prev = atomicAdd(counter, 1);
-
-    if(prev == gid && gid == gridDim.x - 1) {  // Last block: do merge  (use gid check as fallback)
-        // Wait for all other blocks: spin until counter reaches gridDim.x
-        // But counter was just incremented by this block. Need ordering.
-        // Simple approach: use a second counter
-    }
-    // Spin until we know all blocks are done
-    // atomicAdd returns the OLD value, so the last block increments counter from (gridDim.x-1) to gridDim.x
-    __threadfence();
-
-    if(prev < gridDim.x - 1) return;  // Not last, exit
-
-    // ---- Only last block continues: merge all chunk results ----
-    int total = nc * K;  // total elements from all chunks
-    // Use shared memory for merge (up to 4096 elements)
-    if(total > 4096) { if(tid==0) *counter = -2; return; } // too many, shouldn't happen
-
-    // Load all chunk results into shared memory
-    __shared__ float msv[4096];
-    __shared__ int   msi[4096];
-    for(int i=tid; i<total; 

... (7194 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -130,8 +130,8 @@
     def forward(self,x):
         B,N,K=self.batch,self.n,self.k
         # Best chunk sizes per shape (tuned)
-        cfg = {(1,131072,64):(2048,64), (64,8192,8):(N,1), (32,16384,32):(4096,4),
-               (16,12000,16):(1500,8), (128,4096,1):(N,1)}
+        cfg = {(1,131072,64):(2048,64), (64,8192,8):(8192,1), (32,16384,32):(4096,4),
+               (16,12000,16):(1500,8), (128,4096,1):(4096,1)}
         cs,nc = cfg.get((B,N,K), (N,1))
         m=_get_mod()
         launch={64:m.launch_stream_64,32:m.launch_stream_32,16:m.launch_stream_16,8:m.launch_stream_8,1:m.launch_stream_1}[K]