--- /dev/null
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -0,0 +1,45 @@
+"""Correct-but-slow baseline using only allowed ops (torch.max).
+
+This is a placeholder to verify the check.py / benchmark.py plumbing and to
+measure the torch.topk SOTA bar. It will be replaced by a custom CUDA kernel.
+"""
+import torch
+import torch.nn as nn
+
+OP_TYPE = "topk"
+SUPPORTED_PRECISIONS = ["fp32"]
+HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
+
+
+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 forward(self, x: torch.Tensor):
+        b, n, k = x.shape[0], x.shape[1], self.k
+        vals = torch.empty(b, k, dtype=torch.float32, device=x.device)
+        idxs = torch.empty(b, k, dtype=torch.int64, device=x.device)
+        work = x.clone()
+        for i in range(k):
+            v, idx = work.max(dim=-1)
+            vals[:, i] = v
+            idxs[:, i] = idx
+            work[torch.arange(b, device=x.device), idx] = float("-inf")
+        return vals, idxs
+
+
+# Module-level shims rebuilt by check.py / benchmark.py per shape.
+batch = 64
+n = 8192
+k = 8
+
+
+def get_inputs():
+    x = torch.randn(batch, n, dtype=torch.float32)
+    return [x]
+
+
+def get_init_inputs():
+    return [batch, n, k]

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,7 +1,7 @@
 """Correct-but-slow baseline using only allowed ops (torch.max).
 
-This is a placeholder to verify the check.py / benchmark.py plumbing and to
-measure the torch.topk SOTA bar. It will be replaced by a custom CUDA kernel.
+Placeholder to verify the check.py / benchmark.py plumbing and to measure the
+SOTA bar. Replaced later by a custom CUDA kernel.
 """
 import torch
 import torch.nn as nn

--- /dev/null
+++ b//workspace/problems/05_topk_bitonic/scratch_overhead.py
@@ -0,0 +1,118 @@
+"""Microbench: measure the overhead floor on this box.
+- raw kernel launch (load_inline no-op kernel)
+- CUDA graph replay of a no-op kernel
+- nn.Module call overhead (empty forward)
+- a pure copy kernel reading N bytes (bandwidth achievable from a single launch)
+"""
+import time
+import torch
+import torch.nn as nn
+from torch.utils.cpp_extension import load_inline
+
+SRC = r"""
+#include <torch/extension.h>
+#include <cuda_runtime.h>
+
+__global__ void noop_kernel() {}
+__global__ void copy_kernel(const float* __restrict__ in, float* __restrict__ out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = in[i] + 1.0f;
+}
+
+void launch_noop() {
+    noop_kernel<<<1, 32>>>();
+}
+
+void launch_copy(torch::Tensor in, torch::Tensor out) {
+    int n = in.numel();
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    copy_kernel<<<grid, block, 0, c10::cuda::getCurrentCUDAStream()>>>(
+        in.data_ptr<float>(), out.data_ptr<float>(), n);
+}
+"""
+
+mod = load_inline(name="microbench_mod", cpp_sources=[SRC], functions=["launch_noop", "launch_copy"], verbose=False)
+
+
+def time_loop(fn, iters=200, warmup=20, flush=True):
+    # mirror harness: L2 flush before each timed call
+    flush_buf = torch.empty(128 * 1024 * 1024 // 4, dtype=torch.float32, device="cuda") if flush else None
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    times = []
+    for _ in range(iters):
+        if flush:
+            flush_buf.zero_()
+            torch.cuda.synchronize()
+        s = torch.cuda.Event(enable_timing=True)
+        e = torch.cuda.Event(enable_timing=True)
+        s.record()
+        fn()
+        e.record()
+        torch.cuda.synchronize()
+        times.append(s.elapsed_time(e))
+    times.sort()
+    return times[len(times)//2]  # median ms
+
+
+class EmptyModule(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("_dummy", torch.zeros(1))
+    def forward(self, x):
+        return x
+
+
+def main():
+    dev = "cuda"
+    x = torch.randn(1, 131072, device=dev)
+    xs = [x]
+
+    # 1. raw no-op kernel launch (python -> C++ -> 1 kernel)
+    print(f"raw noop kernel launch:        {time_loop(lambda: mod.launch_noop()):.4f} ms")
+
+    # 2. nn.Module empty forward (python dispatch overhead only)
+    em = EmptyModule()
+    print(f"nn.Module empty forward:       {time_loop(lambda: em(x)):.4f} ms")
+
+    # 3. CUDA graph replay (capture the module forward)
+    em_warm = EmptyModule().to(dev)
+    # static input/output
+    g = torch.cuda.CUDAGraph()
+    # warmup capture stream
+    s = torch.cuda.Stream()
+    s.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s):
+        for _ in range(3):
+            y = em_warm(x)
+    torch.cuda.current_stream().wait_stream(s)
+    with torch.cuda.graph(g):
+        y = em_warm(x)
+    print(f"CUDAGraph replay (nn.Module):  {time_loop(lambda: g.replay()):.4f} ms")
+
+    # 4. graph replay wrapping just the noop C++ launch
+    g2 = torch.cuda.CUDAGraph()
+    s2 = torch.cuda.Stream()
+    s2.wait_stream(torch.cuda.current_stream())
+    with torch.cuda.stream(s2):
+        for _ in range(3):
+            mod.launch_noop()
+    torch.cuda.current_stream().wait_stream(s2)
+    with torch.cuda.graph(g2):
+        mod.launch_noop()
+    print(f"CUDAGraph replay (noop C++):   {time_loop(lambda: g2.replay()):.4f} ms")
+
+    # 5. single copy kernel reading 0.5MB and 2MB (bandwidth from one launch)
+    for nbytes_mb in [0.5, 2.0, 0.75]:
+        n = int(nbytes_mb * 1024 * 1024 / 4)
+        xi = torch.randn(n, device=dev)
+        xo = torch.empty(n, device=dev)
+        ms = time_loop(lambda: mod.launch_copy(xi, xo))
+        gbps = nbytes_mb*1e6 / (ms*1e-3) / 1e9
+        print(f"single copy {nbytes_mb:.2f}MB:       {ms:.4f} ms  -> {gbps:.1f} GB/s  (peak_frac {gbps/1800:.3f})")
+
+
+if __name__ == "__main__":
+    main()

--- a//workspace/problems/05_topk_bitonic/scratch_overhead.py
+++ b//workspace/problems/05_topk_bitonic/scratch_overhead.py
@@ -32,7 +32,8 @@
 }
 """
 
-mod = load_inline(name="microbench_mod", cpp_sources=[SRC], functions=["launch_noop", "launch_copy"], verbose=False)
+mod = load_inline(name="microbench_mod", cpp_sources=[""], cuda_sources=[SRC],
+                  functions=["launch_noop", "launch_copy"], verbose=False)
 
 
 def time_loop(fn, iters=200, warmup=20, flush=True):

--- a//workspace/problems/05_topk_bitonic/scratch_overhead.py
+++ b//workspace/problems/05_topk_bitonic/scratch_overhead.py
@@ -11,6 +11,7 @@
 
 SRC = r"""
 #include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
 #include <cuda_runtime.h>
 
 __global__ void noop_kernel() {}
@@ -27,12 +28,14 @@
     int n = in.numel();
     int block = 256;
     int grid = (n + block - 1) / block;
-    copy_kernel<<<grid, block, 0, c10::cuda::getCurrentCUDAStream()>>>(
+    copy_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
         in.data_ptr<float>(), out.data_ptr<float>(), n);
 }
 """
 
-mod = load_inline(name="microbench_mod", cpp_sources=[""], cuda_sources=[SRC],
+DECLS = "void launch_noop(); void launch_copy(torch::Tensor in, torch::Tensor out);"
+
+mod = load_inline(name="microbench_mod", cpp_sources=[DECLS], cuda_sources=[SRC],
                   functions=["launch_noop", "launch_copy"], verbose=False)
 
 

--- /dev/null
+++ b//workspace/problems/05_topk_bitonic/scratch_overhead2.py
@@ -0,0 +1,125 @@
+"""Microbench v2: simulate the harness Model exactly, test __call__ bypass,
+and measure achievable bandwidth from a single real (copy) kernel.
+"""
+import torch
+import torch.nn as nn
+from torch.utils.cpp_extension import load_inline
+
+SRC = r"""
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda_runtime.h>
+
+__global__ void noop_kernel() {}
+__global__ void copy_kernel(const float* __restrict__ in, float* __restrict__ out, int n) {
+    int i = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i < n) out[i] = in[i];
+}
+void launch_noop() { noop_kernel<<<1, 32>>>(); }
+void launch_copy(torch::Tensor in, torch::Tensor out) {
+    int n = in.numel();
+    int block = 256;
+    int grid = (n + block - 1) / block;
+    copy_kernel<<<grid, block, 0, at::cuda::getCurrentCUDAStream()>>>(
+        in.data_ptr<float>(), out.data_ptr<float>(), n);
+}
+"""
+DECLS = "void launch_noop(); void launch_copy(torch::Tensor in, torch::Tensor out);"
+mod = load_inline(name="microbench_mod2", cpp_sources=[DECLS], cuda_sources=[SRC],
+                  functions=["launch_noop", "launch_copy"], verbose=False)
+
+
+def time_loop(fn, iters=300, warmup=30):
+    flush_buf = torch.empty(128 * 1024 * 1024 // 4, dtype=torch.float32, device="cuda")
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    times = []
+    for _ in range(iters):
+        flush_buf.zero_()
+        torch.cuda.synchronize()
+        s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True)
+        s.record(); fn(); e.record(); torch.cuda.synchronize()
+        times.append(s.elapsed_time(e))
+    times.sort()
+    return times[len(times)//2]
+
+
+# ---- Models that mimic the harness contract ----
+class GraphModel(nn.Module):
+    """forward replays a captured graph (the realistic best case)."""
+    def __init__(self, x):
+        super().__init__()
+        self.register_buffer("_dummy", torch.zeros(1))
+        self.out = torch.empty_like(x)
+        g = torch.cuda.CUDAGraph()
+        s = torch.cuda.Stream(); s.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(s):
+            for _ in range(3):
+                mod.launch_copy(x, self.out)
+        torch.cuda.current_stream().wait_stream(s)
+        with torch.cuda.graph(g):
+            mod.launch_copy(x, self.out)
+        self.g = g
+    def forward(self, x):
+        self.g.replay()
+        return self.out
+
+
+class GraphModelFastCall(GraphModel):
+    """Bypass nn.Module._call_impl hooks by overriding __call__."""
+    def __call__(self, x):
+        self.g.replay()
+        return self.out
+
+
+class DirectModel(nn.Module):
+    """forward calls the C++ launch directly (no graph)."""
+    def __init__(self, x):
+        super().__init__()
+        self.register_buffer("_dummy", torch.zeros(1))
+        self.out = torch.empty_like(x)
+    def forward(self, x):
+        mod.launch_copy(x, self.out)
+        return self.out
+
+
+class DirectModelFastCall(DirectModel):
+    def __call__(self, x):
+        mod.launch_copy(x, self.out)
+        return self.out
+
+
+def main():
+    x = torch.randn(64, 8192, device="cuda")  # 2MB, like shape 1
+    models = {
+        "DirectModel (nn.Module)":   DirectModel(x),
+        "DirectModel FastCall":      DirectModelFastCall(x),
+        "GraphModel (nn.Module)":    GraphModel(x),
+        "GraphModel FastCall":       GraphModelFastCall(x),
+    }
+    nbytes = x.numel() * 4
+    for name, m in models.items():
+        m.eval()
+        ms = time_loop(lambda m=m: m(x))
+        gbps = nbytes / (ms*1e-3) / 1e9
+        print(f"{name:28s}: {ms:.4f} ms  -> {gbps:6.1f} GB/s  peak_frac={gbps/1800:.3f}")
+
+    print("\n--- copy kernel raw bandwidth (per shape size, via graph) ---")
+    for desc, n in [("S0 0.5MB/131072", 131072), ("S1 2MB/524288", 524288),
+                    ("S3 0.7MB/192000", 192000), ("S4 2MB/524288", 524288)]:
+        xi = torch.randn(n, device="cuda"); xo = torch.empty(n, device="cuda")
+        g = torch.cuda.CUDAGraph()
+        s = torch.cuda.Stream(); s.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(s):
+            for _ in range(3): mod.launch_copy(xi, xo)
+        torch.cuda.current_stream().wait_stream(s)
+        with torch.cuda.graph(g): mod.launch_copy(xi, xo)
+        ms = time_loop(lambda: g.replay())
+        nbytes = n*4
+        gbps = nbytes/(ms*1e-3)/1e9
+        print(f"{desc:22s}: {ms:.4f} ms  -> {gbps:6.1f} GB/s  peak_frac={gbps/1800:.3f}")
+
+
+if __name__ == "__main__":
+    main()

--- /dev/null
+++ b//home/agent/.claude/projects/-workspace-problems-05-topk-bitonic/memory/topk-problem-facts.md
@@ -0,0 +1,18 @@
+---
+name: topk-problem-facts
+description: Core facts/constraints for the 05_topk_bitonic optimization task (shapes, scoring, environment, overhead floor)
+metadata:
+  type: project
+---
+
+05_topk_bitonic kernel optimization task at /workspace/problems/05_topk_bitonic.
+
+**Scoring**: geomean over 5 shapes of peak_fraction = gbps/1800 (DRAM BW of RTX PRO 6000). Memory regime. bytes_formula = batch*n*4 + batch*k*12. L2 is FLUSHED (128MB write) before each timed call in src/eval/timing.py, so it's a true HBM read. Must run `uv run python check.py` until PASS; `uv run python benchmark.py` for timing.
+
+**Shapes** (batch,n,k): (1,131072,64),(64,8192,8),(32,16384,32),(16,12000,16),(128,4096,1). All tiny: 0.5–2MB, all fit in 96MB L2 (but L2 flushed).
+
+**Correctness**: top-k returns values verbatim (no precision loss — trivially passes). Indices checked LENIENTLY via gather(x, sol_idx) vs ref_values. No NaN/inf stress cases, only scale (1e-4, 1e3). Forbidden ops (substring grep on solution.py source, even comments): torch.topk/kthvalue/sort/argsort + Tensor./aten. variants. torch.max IS allowed.
+
+**Env**: SM120 Blackwell, 188 SMs, 1536 thr/SM, 99KB optin shared. torch 2.11+cu130, nvcc 13.2 at /usr/local/cuda-host (CUDA_HOME set). Triton 3.6 available. MUST use `cuda_sources` (nvcc) not cpp_sources for CUDA code; declare wrapper fns in cpp_sources so load_inline auto-binding sees them; include ATen/cuda/CUDAContext.h for streams. Run python via `uv run`.
+
+**OVERHEAD FLOOR (measured)**: CUDA graph replay ~3.2us, nn.Module dispatch ~2.9us, raw kernel launch ~4.6us. Even a no-op copy kernel via graph caps ~2MB shapes at peak_frac ~0.23, ~0.5MB at ~0.06. So launch overhead is THE bottleneck, not compute/BW. torch.topk eager is ~48us/shape (geomean ~0.018). Strategy: single kernel + CUDA graph replay; override nn.Module.__call__ to skip hooks. See [[topk-progress]].

--- /dev/null
+++ b//home/agent/.claude/projects/-workspace-problems-05-topk-bitonic/memory/MEMORY.md
@@ -0,0 +1,4 @@
+# Project memory index
+
+- [topk-problem-facts](topk-problem-facts.md) — shapes, scoring, env, overhead floor for 05_topk_bitonic
+- [topk-progress](topk-progress.md) — current kernel design, what works, numbers, TODO

--- /dev/null
+++ b//workspace/problems/05_topk_bitonic/dev_kernel.py
@@ -0,0 +1,131 @@
+"""Dev: build the topk extension and validate the streaming+bitonic kernel."""
+import torch
+from torch.utils.cpp_extension import load_inline
+
+CUDA_SRC = r"""
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda_runtime.h>
+#include <c10/macros/Macros.h>
+
+constexpr int npow2(int x){ int p = 1; while (p < x) p <<= 1; return p; }
+
+// Bitonic sort M (power of 2) elements DESCENDING in shared mem, carrying idx.
+template <int BLK>
+__device__ __forceinline__ void bitonic_sort_desc(float* sval, int* sidx, int M) {
+    int tid = threadIdx.x;
+    for (int size = 2; size <= M; size <<= 1) {
+        for (int stride = size >> 1; stride > 0; stride >>= 1) {
+            __syncthreads();
+            for (int i = tid; i < M; i += BLK) {
+                int j = i ^ stride;
+                if (j > i) {
+                    float a = sval[i], b = sval[j];
+                    bool desc = ((i & size) == 0);
+                    bool swap = desc ? (a < b) : (a > b);
+                    if (swap) {
+                        sval[i] = b; sval[j] = a;
+                        int ti = sidx[i]; sidx[i] = sidx[j]; sidx[j] = ti;
+                    }
+                }
+            }
+        }
+    }
+    __syncthreads();
+}
+
+// One block per row. Streaming per-thread top-K (in shared) + bitonic merge.
+// Writes top-K (sorted desc) to out_val/out_idx (int64).
+template <int BLK, int K>
+__global__ void topk_select_kernel(
+    const float* __restrict__ xin,
+    float* __restrict__ out_val,
+    long long* __restrict__ out_idx,
+    int n)
+{
+    static_assert(BLK >= 32, "BLK must be >= 32");
+    constexpr int M = npow2(BLK * K);
+    int row = blockIdx.x;
+    int tid = threadIdx.x;
+    const long long base = (long long)row * n;
+
+    extern __shared__ char smem[];
+    float* sval = (float*)smem;
+    int* sidx = (int*)(smem + (long long)M * 4);
+
+    // init this thread's K slots to -inf
+    #pragma unroll
+    for (int s = 0; s < K; s++) { sval[tid*K + s] = -INFINITY; sidx[tid*K + s] = 0; }
+
+    float thr = -INFINITY;
+    int thr_pos = tid * K;
+
+    for (int e = tid; e < n; e += BLK) {
+        float val = xin[base + e];
+        if (val > thr) {
+            sval[thr_pos] = val;
+            sidx[thr_pos] = e;
+            // recompute min of my K slots
+            thr = sval[tid*K]; thr_pos = tid*K;
+            #pragma unroll
+            for (int s = 1; s < K; s++) {
+                float v = sval[tid*K + s];
+                if (v < thr) { thr = v; thr_pos = tid*K + s; }
+            }
+        }
+    }
+    __syncthreads();
+
+    bitonic_sort_desc<BLK>(sval, sidx, M);
+
+    if (tid < K) {
+        out_val[row*K + tid] = sval[tid];
+        out_idx[row*K + tid] = (long long)sidx[tid];
+    }
+}
+
+void launch_select_256_8(torch::Tensor x, torch::Tensor ov, torch::Tensor oi) {
+    int batch = x.size(0);
+    int n = x.size(1);
+    constexpr int BLK = 256, K = 8;
+    constexpr int M = npow2(BLK * K);
+    size_t smem = (size_t)M * 4 * 2;
+    cudaStream_t s = at::cuda::getCurrentCUDAStream();
+    topk_select_kernel<BLK, K><<<batch, BLK, smem, s>>>(
+        x.data_ptr<float>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), n);
+}
+"""
+
+DECLS = "void launch_select_256_8(torch::Tensor x, torch::Tensor ov, torch::Tensor oi);"
+
+mod = load_inline(name="topk_dev1", cpp_sources=[DECLS], cuda_sources=[CUDA_SRC],
+                  functions=["launch_select_256_8"],
+                  extra_cuda_cflags=["-O3", "--use_fast_math", "-gencode=arch=compute_120,code=sm_120"],
+                  verbose=False)
+
+
+def test(batch, n, k):
+    torch.manual_seed(0)
+    x = torch.randn(batch, n, device="cuda", dtype=torch.float32)
+    ov = torch.empty(batch, k, device="cuda", dtype=torch.float32)
+    oi = torch.empty(batch, k, device="cuda", dtype=torch.int64)
+    mod.launch_select_256_8(x, ov, oi)
+    torch.cuda.synchronize()
+    rv, ri = torch.topk(x, k, dim=-1, largest=True, sorted=True)
+    v_ok = torch.allclose(ov, rv, atol=1e-4, rtol=1e-4)
+    gathered = torch.gather(x, -1, oi.long())
+    i_ok = torch.allclose(gathered, rv, atol=1e-4, rtol=1e-4)
+    print(f"batch={batch} n={n} k={k}: values_ok={v_ok} indices_ok(gather)={i_ok}")
+    if not v_ok:
+        d = (ov - rv).abs().max().item()
+        print(f"   max val diff = {d}")
+        print("   my[0]:", ov[0][:10].tolist())
+        print("   rf[0]:", rv[0][:10].tolist())
+    if not i_ok:
+        print("   gathered[0]:", gathered[0][:10].tolist())
+        print("   ref[0]:     ", rv[0][:10].tolist())
+
+
+if __name__ == "__main__":
+    test(64, 8192, 8)
+    test(32, 16384, 8)   # reuse K=8 kernel just to test correctness on bigger n

--- /dev/null
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -0,0 +1,270 @@
+"""Build the full topk CUDA extension (dev module). Exposes `mod`."""
+import os
+import torch
+from torch.utils.cpp_extension import load_inline
+
+CUDA_SRC = r"""
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda_runtime.h>
+
+constexpr int npow2(int x){ int p = 1; while (p < x) p <<= 1; return p; }
+constexpr int ilog2(int x){ int r=0; while ((1<<r) < x) r++; return r; }
+
+// ---------- Bitonic sort M (pow2) elements DESCENDING, carry idx ----------
+template <int BLK>
+__device__ __forceinline__ void bitonic_sort_desc(float* sval, int* sidx, int M) {
+    int tid = threadIdx.x;
+    int logM = ilog2(M);
+    for (int s = 1; s <= logM; s++) {
+        int size = 1 << s;
+        for (int stride = size >> 1; stride > 0; stride >>= 1) {
+            __syncthreads();
+            for (int i = tid; i < M; i += BLK) {
+                int j = i ^ stride;
+                if (j > i) {
+                    float a = sval[i], b = sval[j];
+                    bool desc = ((i & size) == 0);
+                    bool swap = desc ? (a < b) : (a > b);
+                    if (swap) { sval[i] = b; sval[j] = a; int ti = sidx[i]; sidx[i] = sidx[j]; sidx[j] = ti; }
+                }
+            }
+        }
+    }
+    __syncthreads();
+}
+
+// ---------- One block per row: streaming per-thread top-K + bitonic merge ----------
+template <int BLK, int K>
+__global__ void topk_select_kernel(
+    const float* __restrict__ xin,
+    float* __restrict__ out_val,
+    long long* __restrict__ out_idx,
+    int n)
+{
+    constexpr int M = npow2(BLK * K);
+    int row = blockIdx.x;
+    int tid = threadIdx.x;
+    const long long base = (long long)row * n;
+
+    extern __shared__ char smem[];
+    float* sval = (float*)smem;
+    int* sidx = (int*)(smem + (size_t)M * 4);
+
+    #pragma unroll
+    for (int s = 0; s < K; s++) { sval[tid*K + s] = -INFINITY; sidx[tid*K + s] = 0; }
+
+    float thr = -INFINITY;
+    int thr_pos = tid * K;
+    for (int e = tid; e < n; e += BLK) {
+        float val = xin[base + e];
+        if (val > thr) {
+            sval[thr_pos] = val;
+            sidx[thr_pos] = e;
+            thr = sval[tid*K]; thr_pos = tid*K;
+            #pragma unroll
+            for (int s = 1; s < K; s++) {
+                float v = sval[tid*K + s];
+                if (v < thr) { thr = v; thr_pos = tid*K + s; }
+            }
+        }
+    }
+    __syncthreads();
+    bitonic_sort_desc<BLK>(sval, sidx, M);
+    if (tid < K) {
+        out_val[row*K + tid] = sval[tid];
+        out_idx[row*K + tid] = (long long)sidx[tid];
+    }
+}
+
+// ---------- Argmax (k=1): one block per row, warp-reduce ----------
+template <int BLK>
+__global__ void argmax_kernel(
+    const float* __restrict__ xin,
+    float* __restrict__ out_val,
+    long long* __restrict__ out_idx,
+    int n)
+{
+    constexpr int WARP = 32;
+    constexpr int NWARPS = BLK / WARP;
+    int row = blockIdx.x;
+    int tid = threadIdx.x;
+    int wid = tid / WARP, lane = tid % WARP;
+    const long long base = (long long)row * n;
+
+    __shared__ float sw[NWARPS];
+    __shared__ int si[NWARPS];
+
+    float myv = -INFINITY;
+    int myi = 0;
+    for (int e = tid; e < n; e += BLK) {
+        float val = xin[base + e];
+        if (val > myv) { myv = val; myi = e; }
+    }
+    // warp reduce
+    for (int off = 16; off > 0; off >>= 1) {
+        float ov = __shfl_xor_sync(0xffffffff, myv, off);
+        int oi = __shfl_xor_sync(0xffffffff, myi, off);
+        if (ov > myv) { myv = ov; myi = oi; }
+    }
+    if (lane == 0) { sw[wid] = myv; si[wid] = myi; }
+    __syncthreads();
+    if (wid == 0) {
+        myv = (lane < NWARPS) ? sw[lane] : -INFINITY;
+        myi = (lane < NWARPS) ? si[lane] : 0;
+        for (int off = 16; off > 0; off >>= 1) {
+            float ov = __shfl_xor_sync(0xffffffff, myv, off);
+            int oi = __shfl_xor_sync(0xffffffff, myi, off);
+            if (ov > myv) { myv = ov; myi = oi; }
+        }
+        if (lane == 0) { out_val[row] = myv; out_idx[row] = (long long)myi; }
+    }
+}
+
+// ---------- Multi-block stage 1: per-tile top-K into partial buffers ----------
+template <int BLK, int K>
+__global__ void topk_partial_kernel(
+    const float* __restrict__ xin,
+    float* __restrict__ tmp_val,   // (batch*tiles, K)
+    int* __restrict__ tmp_idx,     // (batch*tiles, K), row-relative indices
+    int n, int tile_len, int tiles)
+{
+    constexpr int M = npow2(BLK * K);
+    int seg = blockIdx.x;
+    int row = seg / tiles;
+    int tile = seg % tiles;
+    int tid = threadIdx.x;
+    long long base = (long long)row * n + (long long)tile * tile_len;
+    int my_len = min(tile_len, n - tile * tile_len);
+    int idx_base = tile * tile_len;
+
+    extern __shared__ char smem[];
+    float* sval = (float*)smem;
+    int* sidx = (int*)(smem + (size_t)M * 4);
+
+    #pragma unroll
+    for (int s = 0; s < K; s++) { sval[tid*K + s] = -INFINITY; sidx[tid*K + s] = 0; }
+    float thr = -INFINITY;
+    int thr_pos = tid * K;
+    for (int e = tid; e < my_len; e += BLK) {
+        float val = xin[base + e];
+        if (val > thr) {
+            sval[thr_pos] = val;
+            sidx[thr_pos] = idx_base + e;
+            thr = sval[tid*K]; thr_pos = tid*K;
+            #pragma unroll
+            for (int s = 1; s < K; s++) {
+                float v = sval[tid*K + s];
+                if (v < thr) { thr = v; thr_pos = tid*K + s; }
+            }
+        }
+    }
+    __syncthreads();
+    bitonic_sort_desc<BLK>(sval, sidx, M);
+    if (tid < K) {
+        tmp_val[seg*K + tid] = sval[tid];
+        tmp_idx[seg*K + tid] = sidx[tid];
+    }
+}
+
+// ---------- Multi-block stage 2: merge tiles*K candidates -> top-K ----------
+template <int BLK, int K, int M>   // M = pow2 >= tiles*K
+__global__ void topk_merge_kernel(
+    const float* __restrict__ tmp_val,
+    const int* __restri

... (3592 more chars)

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -1,5 +1,4 @@
-"""Build the full topk CUDA extension (dev module). Exposes `mod`."""
-import os
+"""Build the full topk CUDA extension (dev module). Exposes build()."""
 import torch
 from torch.utils.cpp_extension import load_inline
 
@@ -7,6 +6,7 @@
 #include <torch/extension.h>
 #include <ATen/cuda/CUDAContext.h>
 #include <cuda_runtime.h>
+#include <type_traits>
 
 constexpr int npow2(int x){ int p = 1; while (p < x) p <<= 1; return p; }
 constexpr int ilog2(int x){ int r=0; while ((1<<r) < x) r++; return r; }
@@ -77,7 +77,7 @@
     }
 }
 
-// ---------- Argmax (k=1): one block per row, warp-reduce ----------
+// ---------- Argmax (k=1): one block per row ----------
 template <int BLK>
 __global__ void argmax_kernel(
     const float* __restrict__ xin,
@@ -101,7 +101,6 @@
         float val = xin[base + e];
         if (val > myv) { myv = val; myi = e; }
     }
-    // warp reduce
     for (int off = 16; off > 0; off >>= 1) {
         float ov = __shfl_xor_sync(0xffffffff, myv, off);
         int oi = __shfl_xor_sync(0xffffffff, myi, off);
@@ -125,8 +124,8 @@
 template <int BLK, int K>
 __global__ void topk_partial_kernel(
     const float* __restrict__ xin,
-    float* __restrict__ tmp_val,   // (batch*tiles, K)
-    int* __restrict__ tmp_idx,     // (batch*tiles, K), row-relative indices
+    float* __restrict__ tmp_val,
+    int* __restrict__ tmp_idx,
     int n, int tile_len, int tiles)
 {
     constexpr int M = npow2(BLK * K);
@@ -167,14 +166,14 @@
     }
 }
 
-// ---------- Multi-block stage 2: merge tiles*K candidates -> top-K ----------
-template <int BLK, int K, int M>   // M = pow2 >= tiles*K
+// ---------- Multi-block stage 2: merge total=tiles*K candidates -> top-K ----------
+template <int BLK, int K, int M>
 __global__ void topk_merge_kernel(
     const float* __restrict__ tmp_val,
     const int* __restrict__ tmp_idx,
     float* __restrict__ out_val,
     long long* __restrict__ out_idx,
-    int total)   // tiles*K
+    int total)
 {
     int row = blockIdx.x;
     int tid = threadIdx.x;
@@ -208,29 +207,30 @@
         return;
     }
 
-    // Multi-block path for very large n per row (single/few rows): split into tiles.
-    const int MULTIBLOCK_THRESH = 32768;
-    if (n > MULTIBLOCK_THRESH) {
-        // choose tile_len so each tile ~fits good parallelism; tiles = n/tile_len
+    // Multi-block path: very large n (shape 0). Split into tiles, partial then merge.
+    if (n > 32768) {
+        constexpr int BLK = 256;
         int tile_len = 2048;
         int tiles = (n + tile_len - 1) / tile_len;
-        constexpr int BLK = 256;
-        // stage 1
         if (k == 64) {
             constexpr int K = 64;
-            constexpr int M = npow2(BLK * K);
-            topk_partial_kernel<BLK, K><<<batch*tiles, BLK, (size_t)M*8, s>>>(
+            constexpr int Mp = npow2(BLK * K);          // 8192
+            topk_partial_kernel<BLK, K><<<batch*tiles, BLK, (size_t)Mp*8, s>>>(
                 x.data_ptr<float>(), tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), n, tile_len, tiles);
             int total = tiles * K;
-            int M2 = npow2(total);
-            topk_merge_kernel<BLK, K, /*M templated*/ 0>; // placeholder; handled below
-            // call merge with M2 via dynamic
-        }
-        // (full impl below)
+            if (total <= 4096) {
+                topk_merge_kernel<BLK, K, 4096><<<batch, BLK, 4096*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
+            } else if (total <= 8192) {
+                topk_merge_kernel<BLK, K, 8192><<<batch, BLK, 8192*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
+            } else if (total <= 16384) {
+                topk_merge_kernel<BLK, K, 16384><<<batch, BLK, 16384*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
+            }
+            return;
+        }
     }
 
     // one-block-per-row path
-    auto launch = [&](auto k_const) {
+    auto launch_select = [&](auto k_const) {
         constexpr int K = decltype(k_const)::value;
         constexpr int BLK = 256;
         constexpr int M = npow2(BLK * K);
@@ -238,12 +238,12 @@
             x.data_ptr<float>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), n);
     };
     switch (k) {
-        case 2:  launch(std::integral_constant<int,2>{});  break;
-        case 4:  launch(std::integral_constant<int,4>{});  break;
-        case 8:  launch(std::integral_constant<int,8>{});  break;
-        case 16: launch(std::integral_constant<int,16>{}); break;
-        case 32: launch(std::integral_constant<int,32>{}); break;
-        case 64: launch(std::integral_constant<int,64>{}); break;
+        case 2:  launch_select(std::integral_constant<int,2>{});  break;
+        case 4:  launch_select(std::integral_constant<int,4>{});  break;
+        case 8:  launch_select(std::integral_constant<int,8>{});  break;
+        case 16: launch_select(std::integral_constant<int,16>{}); break;
+        case 32: launch_select(std::integral_constant<int,32>{}); break;
+        case 64: launch_select(std::integral_constant<int,64>{}); break;
         default: break;
     }
 }
@@ -253,6 +253,7 @@
 void topk_forward(torch::Tensor x, torch::Tensor ov, torch::Tensor oi,
                   torch::Tensor tmp_val, torch::Tensor tmp_idx, int64_t k);
 """
+
 
 def build():
     return load_inline(
@@ -265,6 +266,7 @@
         verbose=False,
     )
 
+
 if __name__ == "__main__":
-    m = build()
+    build()
     print("built OK")

--- /dev/null
+++ b//workspace/problems/05_topk_bitonic/dev_test.py
@@ -0,0 +1,111 @@
+"""Dev: correctness + timing of the full extension across all shapes."""
+import time
+import torch
+import topk_cuda
+
+mod = topk_cuda.build()
+
+SHAPES = [
+    (1, 131072, 64),
+    (64, 8192, 8),
+    (32, 16384, 32),
+    (16, 12000, 16),
+    (128, 4096, 1),
+]
+
+
+def tmp_size(n, k):
+    # max tiles needed (multiblock path) for tmp buffer
+    tiles = max(1, (n + 2047) // 2048)
+    return tiles
+
+
+def run(batch, n, k, x=None):
+    if x is None:
+        torch.manual_seed(0)
+        x = torch.randn(batch, n, device="cuda", dtype=torch.float32)
+    ov = torch.empty(batch, k, device="cuda", dtype=torch.float32)
+    oi = torch.empty(batch, k, device="cuda", dtype=torch.int64)
+    tiles = tmp_size(n, k)
+    tv = torch.empty(batch * tiles, k, device="cuda", dtype=torch.float32)
+    ti = torch.empty(batch * tiles, k, device="cuda", dtype=torch.int32)
+    mod.topk_forward(x, ov, oi, tv, ti, k)
+    torch.cuda.synchronize()
+    return ov, oi
+
+
+def check(batch, n, k):
+    for seed in (0, 1, 2):
+        torch.manual_seed(seed)
+        x = torch.randn(batch, n, device="cuda", dtype=torch.float32)
+        ov, oi = run(batch, n, k, x)
+        rv, ri = torch.topk(x, k, dim=-1, largest=True, sorted=True)
+        v_ok = torch.allclose(ov, rv, atol=1e-4, rtol=1e-4)
+        gathered = torch.gather(x, -1, oi.long())
+        i_ok = torch.allclose(gathered, rv, atol=1e-4, rtol=1e-4)
+        if not (v_ok and i_ok):
+            print(f"  FAIL batch={batch} n={n} k={k} seed={seed}: v={v_ok} i={i_ok}")
+            if not v_ok:
+                print(f"    maxdiff={(ov-rv).abs().max().item()}")
+                print(f"    my[0]={ov[0][:8].tolist()}")
+                print(f"    rf[0]={rv[0][:8].tolist()}")
+            return False
+    print(f"  OK  batch={batch} n={n} k={k}")
+    return True
+
+
+def timeit(batch, n, k):
+    torch.manual_seed(2026)
+    x = torch.randn(batch, n, device="cuda", dtype=torch.float32)
+    ov = torch.empty(batch, k, device="cuda", dtype=torch.float32)
+    oi = torch.empty(batch, k, device="cuda", dtype=torch.int64)
+    tiles = tmp_size(n, k)
+    tv = torch.empty(batch * tiles, k, device="cuda", dtype=torch.float32)
+    ti = torch.empty(batch * tiles, k, device="cuda", dtype=torch.int32)
+    # warmup
+    for _ in range(20):
+        mod.topk_forward(x, ov, oi, tv, ti, k)
+    torch.cuda.synchronize()
+    flush = torch.empty(128 * 1024 * 1024 // 4, device="cuda")
+    times = []
+    for _ in range(100):
+        flush.zero_(); torch.cuda.synchronize()
+        s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True)
+        s.record(); mod.topk_forward(x, ov, oi, tv, ti, k); e.record()
+        torch.cuda.synchronize()
+        times.append(s.elapsed_time(e))
+    times.sort()
+    ms = times[len(times)//2]
+    nbytes = batch*n*4 + batch*k*12
+    gbps = nbytes/(ms*1e-3)/1e9
+    sota_ms = time_sota(x, k)
+    sota_gbps = nbytes/(sota_ms*1e-3)/1e9 if sota_ms else float('nan')
+    print(f"  n={n:6d} k={k:3d}: sol {ms:.4f}ms {gbps:6.1f}GB/s frac={gbps/1800:.3f} | "
+          f"sota {sota_ms:.4f}ms {sota_gbps:6.1f}GB/s")
+
+
+def time_sota(x, k):
+    for _ in range(10):
+        torch.topk(x, k, dim=-1)
+    torch.cuda.synchronize()
+    flush = torch.empty(128 * 1024 * 1024 // 4, device="cuda")
+    times = []
+    for _ in range(50):
+        flush.zero_(); torch.cuda.synchronize()
+        s = torch.cuda.Event(enable_timing=True); e = torch.cuda.Event(enable_timing=True)
+        s.record(); torch.topk(x, k, dim=-1); e.record()
+        torch.cuda.synchronize()
+        times.append(s.elapsed_time(e))
+    times.sort()
+    return times[len(times)//2]
+
+
+if __name__ == "__main__":
+    print("=== correctness ===")
+    allok = True
+    for sh in SHAPES:
+        allok &= check(*sh)
+    print("ALL OK" if allok else "FAILURES")
+    print("\n=== timing (no graph) ===")
+    for sh in SHAPES:
+        timeit(*sh)

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -193,6 +193,23 @@
 }
 
 // ================= dispatch =================
+template <int BLK, int K>
+static inline void enable_big_smem_select() {
+    constexpr int M = npow2(BLK * K);
+    if ((size_t)M * 8 > 49152) {
+        cudaFuncSetAttribute(topk_select_kernel<BLK, K>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);
+    }
+}
+template <int BLK, int K>
+static inline void enable_big_smem_partial() {
+    constexpr int M = npow2(BLK * K);
+    if ((size_t)M * 8 > 49152) {
+        cudaFuncSetAttribute(topk_partial_kernel<BLK, K>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);
+    }
+}
+
 void topk_forward(torch::Tensor x, torch::Tensor ov, torch::Tensor oi,
                   torch::Tensor tmp_val, torch::Tensor tmp_idx, int64_t k)
 {

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -232,6 +232,7 @@
         if (k == 64) {
             constexpr int K = 64;
             constexpr int Mp = npow2(BLK * K);          // 8192
+            enable_big_smem_partial<BLK, K>();
             topk_partial_kernel<BLK, K><<<batch*tiles, BLK, (size_t)Mp*8, s>>>(
                 x.data_ptr<float>(), tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), n, tile_len, tiles);
             int total = tiles * K;
@@ -251,6 +252,7 @@
         constexpr int K = decltype(k_const)::value;
         constexpr int BLK = 256;
         constexpr int M = npow2(BLK * K);
+        enable_big_smem_select<BLK, K>();
         topk_select_kernel<BLK, K><<<batch, BLK, (size_t)M*8, s>>>(
             x.data_ptr<float>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), n);
     };

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -226,22 +226,22 @@
 
     // Multi-block path: very large n (shape 0). Split into tiles, partial then merge.
     if (n > 32768) {
-        constexpr int BLK = 256;
-        int tile_len = 2048;
-        int tiles = (n + tile_len - 1) / tile_len;
         if (k == 64) {
+            constexpr int BLK = 128;
             constexpr int K = 64;
             constexpr int Mp = npow2(BLK * K);          // 8192
+            int tile_len = 2048;
+            int tiles = (n + tile_len - 1) / tile_len;
             enable_big_smem_partial<BLK, K>();
             topk_partial_kernel<BLK, K><<<batch*tiles, BLK, (size_t)Mp*8, s>>>(
                 x.data_ptr<float>(), tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), n, tile_len, tiles);
             int total = tiles * K;
             if (total <= 4096) {
-                topk_merge_kernel<BLK, K, 4096><<<batch, BLK, 4096*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
+                topk_merge_kernel<256, K, 4096><<<batch, 256, 4096*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
             } else if (total <= 8192) {
-                topk_merge_kernel<BLK, K, 8192><<<batch, BLK, 8192*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
+                topk_merge_kernel<256, K, 8192><<<batch, 256, 8192*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
             } else if (total <= 16384) {
-                topk_merge_kernel<BLK, K, 16384><<<batch, BLK, 16384*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
+                topk_merge_kernel<256, K, 16384><<<batch, 256, 16384*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
             }
             return;
         }
@@ -250,7 +250,7 @@
     // one-block-per-row path
     auto launch_select = [&](auto k_const) {
         constexpr int K = decltype(k_const)::value;
-        constexpr int BLK = 256;
+        constexpr int BLK = (K <= 32) ? 256 : 128;
         constexpr int M = npow2(BLK * K);
         enable_big_smem_select<BLK, K>();
         topk_select_kernel<BLK, K><<<batch, BLK, (size_t)M*8, s>>>(

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -35,6 +35,8 @@
 }
 
 // ---------- One block per row: streaming per-thread top-K + bitonic merge ----------
+// Per-thread buffer laid out as [s][BLK] (index s*BLK+tid) -> bank-conflict-free
+// across threads at fixed s.
 template <int BLK, int K>
 __global__ void topk_select_kernel(
     const float* __restrict__ xin,
@@ -52,20 +54,20 @@
     int* sidx = (int*)(smem + (size_t)M * 4);
 
     #pragma unroll
-    for (int s = 0; s < K; s++) { sval[tid*K + s] = -INFINITY; sidx[tid*K + s] = 0; }
-
-    float thr = -INFINITY;
-    int thr_pos = tid * K;
+    for (int s = 0; s < K; s++) { sval[s*BLK + tid] = -INFINITY; sidx[s*BLK + tid] = 0; }
+
+    float thr = -INFINITY;   // = sval[0*BLK+tid]
+    int thr_s = 0;
     for (int e = tid; e < n; e += BLK) {
         float val = xin[base + e];
         if (val > thr) {
-            sval[thr_pos] = val;
-            sidx[thr_pos] = e;
-            thr = sval[tid*K]; thr_pos = tid*K;
+            sval[thr_s*BLK + tid] = val;
+            sidx[thr_s*BLK + tid] = e;
+            thr = sval[tid]; thr_s = 0;   // sval[tid] == sval[0*BLK+tid]
             #pragma unroll
             for (int s = 1; s < K; s++) {
-                float v = sval[tid*K + s];
-                if (v < thr) { thr = v; thr_pos = tid*K + s; }
+                float v = sval[s*BLK + tid];
+                if (v < thr) { thr = v; thr_s = s; }
             }
         }
     }

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -144,19 +144,19 @@
     int* sidx = (int*)(smem + (size_t)M * 4);
 
     #pragma unroll
-    for (int s = 0; s < K; s++) { sval[tid*K + s] = -INFINITY; sidx[tid*K + s] = 0; }
+    for (int s = 0; s < K; s++) { sval[s*BLK + tid] = -INFINITY; sidx[s*BLK + tid] = 0; }
     float thr = -INFINITY;
-    int thr_pos = tid * K;
+    int thr_s = 0;
     for (int e = tid; e < my_len; e += BLK) {
         float val = xin[base + e];
         if (val > thr) {
-            sval[thr_pos] = val;
-            sidx[thr_pos] = idx_base + e;
-            thr = sval[tid*K]; thr_pos = tid*K;
+            sval[thr_s*BLK + tid] = val;
+            sidx[thr_s*BLK + tid] = idx_base + e;
+            thr = sval[tid]; thr_s = 0;
             #pragma unroll
             for (int s = 1; s < K; s++) {
-                float v = sval[tid*K + s];
-                if (v < thr) { thr = v; thr_pos = tid*K + s; }
+                float v = sval[s*BLK + tid];
+                if (v < thr) { thr = v; thr_s = s; }
             }
         }
     }

--- /dev/null
+++ b//workspace/problems/05_topk_bitonic/dev_profile_target.py
@@ -0,0 +1,14 @@
+"""Minimal target for ncu profiling of the k=32 kernel."""
+import sys, torch, topk_cuda
+mod = topk_cuda.build()
+batch, n, k = 32, 16384, 32
+torch.manual_seed(0)
+x = torch.randn(batch, n, device="cuda")
+ov = torch.empty(batch, k, device="cuda")
+oi = torch.empty(batch, k, dtype=torch.int64, device="cuda")
+tiles = max(1, (n+2047)//2048)
+tv = torch.empty(batch*tiles, k, device="cuda")
+ti = torch.empty(batch*tiles, k, dtype=torch.int32, device="cuda")
+for _ in range(3):
+    mod.topk_forward(x, ov, oi, tv, ti, k)
+torch.cuda.synchronize()

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -169,13 +169,14 @@
 }
 
 // ---------- Multi-block stage 2: merge total=tiles*K candidates -> top-K ----------
-template <int BLK, int K, int M>
+// M (pow2 >= total) passed at runtime via launch shared-mem size.
+template <int BLK, int K>
 __global__ void topk_merge_kernel(
     const float* __restrict__ tmp_val,
     const int* __restrict__ tmp_idx,
     float* __restrict__ out_val,
     long long* __restrict__ out_idx,
-    int total)
+    int total, int M)
 {
     int row = blockIdx.x;
     int tid = threadIdx.x;
@@ -195,21 +196,32 @@
 }
 
 // ================= dispatch =================
-template <int BLK, int K>
-static inline void enable_big_smem_select() {
-    constexpr int M = npow2(BLK * K);
-    if ((size_t)M * 8 > 49152) {
-        cudaFuncSetAttribute(topk_select_kernel<BLK, K>,
-            cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);
-    }
-}
-template <int BLK, int K>
-static inline void enable_big_smem_partial() {
-    constexpr int M = npow2(BLK * K);
-    if ((size_t)M * 8 > 49152) {
+// Per-K block size: keeps per-block sort M = npow2(BLK*K) <= 4096 (<=32KB shared),
+// enabling many blocks/SM for high occupancy.
+template <int K> struct PartialBlk { static constexpr int VAL = 256; };
+template <> struct PartialBlk<16> { static constexpr int VAL = 128; };
+template <> struct PartialBlk<32> { static constexpr int VAL = 128; };
+template <> struct PartialBlk<64> { static constexpr int VAL = 64; };
+
+template <int K>
+static inline void launch_multiblock(int batch, int n, int tile_len, int tiles,
+        const float* xv, float* ov, long long* oi, float* tv, int* ti, cudaStream_t s)
+{
+    constexpr int BLK = PartialBlk<K>::VAL;
+    constexpr int Mp = npow2(BLK * K);
+    // big smem only if needed (these configs are <=32KB, so usually no)
+    if ((size_t)Mp * 8 > 49152)
         cudaFuncSetAttribute(topk_partial_kernel<BLK, K>,
             cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);
-    }
+    topk_partial_kernel<BLK, K><<<batch*tiles, BLK, (size_t)Mp*8, s>>>(
+        xv, tv, ti, n, tile_len, tiles);
+    int total = tiles * K;
+    int M = npow2(total);
+    constexpr int MB = 256;
+    if ((size_t)M * 8 > 49152)
+        cudaFuncSetAttribute(topk_merge_kernel<MB, K>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);
+    topk_merge_kernel<MB, K><<<batch, MB, (size_t)M*8, s>>>(tv, ti, ov, oi, total, M);
 }
 
 void topk_forward(torch::Tensor x, torch::Tensor ov, torch::Tensor oi,
@@ -218,53 +230,41 @@
     int batch = x.size(0);
     int n = x.size(1);
     cudaStream_t s = at::cuda::getCurrentCUDAStream();
+    const float* xv = x.data_ptr<float>();
+    float* ov_ = ov.data_ptr<float>();
+    long long* oi_ = (long long*)oi.data_ptr<int64_t>();
+    float* tv = tmp_val.data_ptr<float>();
+    int* ti = tmp_idx.data_ptr<int>();
 
     if (k == 1) {
         constexpr int BLK = 256;
-        argmax_kernel<BLK><<<batch, BLK, 0, s>>>(
-            x.data_ptr<float>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), n);
+        argmax_kernel<BLK><<<batch, BLK, 0, s>>>(xv, ov_, oi_, n);
         return;
     }
 
-    // Multi-block path: very large n (shape 0). Split into tiles, partial then merge.
-    if (n > 32768) {
-        if (k == 64) {
-            constexpr int BLK = 128;
-            constexpr int K = 64;
-            constexpr int Mp = npow2(BLK * K);          // 8192
-            int tile_len = 2048;
-            int tiles = (n + tile_len - 1) / tile_len;
-            enable_big_smem_partial<BLK, K>();
-            topk_partial_kernel<BLK, K><<<batch*tiles, BLK, (size_t)Mp*8, s>>>(
-                x.data_ptr<float>(), tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), n, tile_len, tiles);
-            int total = tiles * K;
-            if (total <= 4096) {
-                topk_merge_kernel<256, K, 4096><<<batch, 256, 4096*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
-            } else if (total <= 8192) {
-                topk_merge_kernel<256, K, 8192><<<batch, 256, 8192*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
-            } else if (total <= 16384) {
-                topk_merge_kernel<256, K, 16384><<<batch, 256, 16384*8, s>>>(tmp_val.data_ptr<float>(), tmp_idx.data_ptr<int>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), total);
-            }
-            return;
-        }
-    }
-
-    // one-block-per-row path
-    auto launch_select = [&](auto k_const) {
-        constexpr int K = decltype(k_const)::value;
-        constexpr int BLK = (K <= 32) ? 256 : 128;
-        constexpr int M = npow2(BLK * K);
-        enable_big_smem_select<BLK, K>();
-        topk_select_kernel<BLK, K><<<batch, BLK, (size_t)M*8, s>>>(
-            x.data_ptr<float>(), ov.data_ptr<float>(), (long long*)oi.data_ptr<int64_t>(), n);
-    };
+    // Pick tile_len so that batch*tiles ~>= ~256 blocks (fill the GPU) while
+    // keeping the merge (tiles*K) small. Larger batch -> fewer tiles per row.
+    int tile_len;
+    if (batch >= 64)      tile_len = 2048;
+    else if (batch >= 32) tile_len = 2048;
+    else if (batch >= 16) tile_len = 1024;
+    else                  tile_len = 2048;   // batch small (e.g. shape 0)
+    // shape-0-like (batch==1, huge n): cap tiles so merge fits in shared mem
+    if (batch == 1 && n > 32768) {
+        // total = tiles*K <= 8192 -> tiles <= 128
+        int max_tiles = (k > 0) ? (8192 / (int)k) : 8192;
+        tile_len = (n + max_tiles - 1) / max_tiles;
+        if (tile_len < 1024) tile_len = 1024;
+    }
+    int tiles = (n + tile_len - 1) / tile_len;
+
     switch (k) {
-        case 2:  launch_select(std::integral_constant<int,2>{});  break;
-        case 4:  launch_select(std::integr

... (959 more chars)

--- a//workspace/problems/05_topk_bitonic/dev_test.py
+++ b//workspace/problems/05_topk_bitonic/dev_test.py
@@ -15,8 +15,8 @@
 
 
 def tmp_size(n, k):
-    # max tiles needed (multiblock path) for tmp buffer
-    tiles = max(1, (n + 2047) // 2048)
+    # generous upper bound on tiles (tile_len >= 256 in dispatch)
+    tiles = max(1, (n + 255) // 256)
     return tiles
 
 

--- /dev/null
+++ b//workspace/problems/05_topk_bitonic/dev_sortonly.py
@@ -0,0 +1,69 @@
+"""Isolate bitonic sort cost: a kernel that reads 2048 floats, sorts, writes."""
+import torch
+from torch.utils.cpp_extension import load_inline
+
+SRC = r"""
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda_runtime.h>
+
+constexpr int npow2c(int x){ int p=1; while(p<x) p<<=1; return p; }
+constexpr int ilog2c(int x){ int r=0; while((1<<r)<x) r++; return r; }
+
+template <int BLK>
+__device__ __forceinline__ void bsort(float* sv, int* si, int M) {
+    int tid = threadIdx.x;
+    int logM = ilog2c(M);
+    for (int s=1;s<=logM;s++){ int size=1<<s;
+      for (int stride=size>>1;stride>0;stride>>=1){
+        __syncthreads();
+        for (int i=tid;i<M;i+=BLK){ int j=i^stride;
+          if(j>i){ float a=sv[i],b=sv[j]; bool d=((i&size)==0); bool sw=d?(a<b):(a>b);
+            if(sw){sv[i]=b;sv[j]=a;int t=si[i];si[i]=si[j];si[j]=t;} }
+        }
+      }
+    }
+    __syncthreads();
+}
+
+template <int BLK, int M>
+__global__ void sort_only(const float* __restrict__ in, float* __restrict__ out, int rowlen) {
+    __shared__ float sv[M];
+    __shared__ int si[M];
+    int tid=threadIdx.x;
+    int row=blockIdx.x;
+    long long base=(long long)row*rowlen;
+    for (int i=tid;i<M;i+=BLK){ sv[i] = (i<rowlen)? in[base+i] : -INFINITY; si[i]=i; }
+    __syncthreads();
+    bsort<BLK>(sv, si, M);
+    for (int i=tid;i<M;i+=BLK){ out[base+i]=sv[i]; }
+}
+
+void launch_sortonly(torch::Tensor in, torch::Tensor out, int blk, int m) {
+    cudaStream_t s = at::cuda::getCurrentCUDAStream();
+    int batch=in.size(0);
+    #define CASE(B,M) if(blk==B&&m==M) sort_only<B,M><<<batch,B,0,s>>>(in.data_ptr<float>(),out.data_ptr<float>(),in.size(1));
+    CASE(256,2048) CASE(256,4096) CASE(128,4096) CASE(128,2048) CASE(64,2048) CASE(64,4096)
+    #undef CASE
+}
+"""
+DECLS="void launch_sortonly(torch::Tensor in, torch::Tensor out, int blk, int m);"
+mod=load_inline(name="sortonly_v1",cpp_sources=[DECLS],cuda_sources=[SRC],functions=["launch_sortonly"],
+  extra_cuda_cflags=["-O3","--use_fast_math","-std=c++17","-gencode=arch=compute_120,code=sm_120"],verbose=False)
+
+def timeit(fn, iters=100, warmup=20):
+    flush=torch.empty(128*1024*1024//4,device="cuda")
+    for _ in range(warmup): fn()
+    torch.cuda.synchronize()
+    ts=[]
+    for _ in range(iters):
+        flush.zero_(); torch.cuda.synchronize()
+        s=torch.cuda.Event(enable_timing=True); e=torch.cuda.Event(enable_timing=True)
+        s.record(); fn(); e.record(); torch.cuda.synchronize()
+        ts.append(s.elapsed_time(e))
+    ts.sort(); return ts[len(ts)//2]
+
+for blk,m,nblocks,label in [(256,2048,256,"k8-ish"),(128,4096,256,"k32-ish"),(64,2048,512,"small"),(256,4096,256,"big")]:
+    x=torch.randn(nblocks,m,device="cuda"); o=torch.empty_like(x)
+    ms=timeit(lambda: mod.launch_sortonly(x,o,blk,m))
+    print(f"sort_only blk={blk} M={m} blocks={nblocks}: {ms:.4f} ms")

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -119,6 +119,96 @@
             if (ov > myv) { myv = ov; myi = oi; }
         }
         if (lane == 0) { out_val[row] = myv; out_idx[row] = (long long)myi; }
+    }
+}
+
+// ---------- Merge two ascending sorted K-arrays into top-K ascending ----------
+template <int K>
+__device__ __forceinline__ void merge_topk(
+    const float* A, const int* Ai,
+    const float* B, const int* Bi,
+    float* C, int* Ci)
+{
+    int ia = K - 1, ib = K - 1;
+    #pragma unroll
+    for (int c = K - 1; c >= 0; c--) {
+        float av = (ia >= 0) ? A[ia] : -INFINITY;
+        float bv = (ib >= 0) ? B[ib] : -INFINITY;
+        if (av >= bv) { C[c] = av; Ci[c] = Ai[ia]; ia--; }
+        else          { C[c] = bv; Ci[c] = Bi[ib]; ib--; }
+    }
+}
+
+// ---------- Multi-block stage 1 (merge-tree): per-tile top-K into partial ----------
+// Per-thread top-K kept sorted ascending in registers; then a shared-mem
+// pairwise merge tree reduces BLK buffers to 1 (log2(BLK) levels, O(BLK*K)).
+template <int BLK, int K>
+__global__ void topk_part_mt_kernel(
+    const float* __restrict__ xin,
+    float* __restrict__ tmp_val,
+    int* __restrict__ tmp_idx,
+    int n, int tile_len, int tiles)
+{
+    int seg = blockIdx.x;
+    int row = seg / tiles;
+    int tile = seg % tiles;
+    int tid = threadIdx.x;
+    long long base = (long long)row * n + (long long)tile * tile_len;
+    int my_len = min(tile_len, n - tile * tile_len);
+    int idx_base = tile * tile_len;
+
+    // phase 1: register sorted-ascending top-K, rv[0] = min
+    float rv[K]; int ri[K];
+    #pragma unroll
+    for (int s = 0; s < K; s++) { rv[s] = -INFINITY; ri[s] = 0; }
+    for (int e = tid; e < my_len; e += BLK) {
+        float val = xin[base + e];
+        if (val > rv[0]) {
+            rv[0] = val; ri[0] = idx_base + e;
+            #pragma unroll
+            for (int s = 1; s < K; s++) {
+                if (rv[s] < rv[s-1]) {
+                    float tf = rv[s]; rv[s] = rv[s-1]; rv[s-1] = tf;
+                    int ti = ri[s]; ri[s] = ri[s-1]; ri[s-1] = ti;
+                }
+            }
+        }
+    }
+
+    extern __shared__ char smem[];
+    float* sv = (float*)smem;                 // BLK*K
+    int* si = (int*)(sv + (size_t)BLK * K);   // BLK*K
+    float* tv = sv + (size_t)2 * BLK * K;     // K (temp for pair 0)
+    int* tw = si + (size_t)BLK * K;           // K (temp idx)
+
+    // write back buffers (buffer-major)
+    #pragma unroll
+    for (int s = 0; s < K; s++) { sv[tid*K + s] = rv[s]; si[tid*K + s] = ri[s]; }
+    __syncthreads();
+
+    // pairwise merge tree
+    int active = BLK;
+    while (active > 1) {
+        int pair = tid;
+        if (pair < active/2) {
+            if (pair == 0)
+                merge_topk<K>(sv, si, sv + K, si + K, tv, tw);
+            else
+                merge_topk<K>(sv + (2*pair)*K, si + (2*pair)*K,
+                              sv + (2*pair+1)*K, si + (2*pair+1)*K,
+                              sv + pair*K, si + pair*K);
+        }
+        __syncthreads();
+        if (tid < K) { sv[tid] = tv[tid]; si[tid] = tw[tid]; }   // copy pair0 result back
+        active >>= 1;
+        __syncthreads();
+    }
+
+    // buffer 0 now has top-K ascending; write descending, int64 idx
+    if (tid < K) {
+        int s = K - 1 - tid;
+        tmp_val[seg*K + tid] = sv[s];
+        tmp_idx[seg*K + tid] = si[s];
     }
 }
 

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -298,12 +298,12 @@
         const float* xv, float* ov, long long* oi, float* tv, int* ti, cudaStream_t s)
 {
     constexpr int BLK = PartialBlk<K>::VAL;
-    constexpr int Mp = npow2(BLK * K);
-    // big smem only if needed (these configs are <=32KB, so usually no)
-    if ((size_t)Mp * 8 > 49152)
-        cudaFuncSetAttribute(topk_partial_kernel<BLK, K>,
+    // shared = 2*BLK*K*4 (sv+si) + K*8 (temp)
+    size_t smem = (size_t)2 * BLK * K * 4 + (size_t)K * 8;
+    if (smem > 49152)
+        cudaFuncSetAttribute(topk_part_mt_kernel<BLK, K>,
             cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);
-    topk_partial_kernel<BLK, K><<<batch*tiles, BLK, (size_t)Mp*8, s>>>(
+    topk_part_mt_kernel<BLK, K><<<batch*tiles, BLK, smem, s>>>(
         xv, tv, ti, n, tile_len, tiles);
     int total = tiles * K;
     int M = npow2(total);

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -176,10 +176,10 @@
     }
 
     extern __shared__ char smem[];
-    float* sv = (float*)smem;                 // BLK*K
-    int* si = (int*)(sv + (size_t)BLK * K);   // BLK*K
-    float* tv = sv + (size_t)2 * BLK * K;     // K (temp for pair 0)
-    int* tw = si + (size_t)BLK * K;           // K (temp idx)
+    float* sv = (float*)smem;                       // BLK*K floats
+    int* si = (int*)(sv + (size_t)BLK * K);         // BLK*K ints
+    float* tv = (float*)(si + (size_t)BLK * K);     // K floats (temp for pair 0)
+    int* tw = (int*)(tv + K);                       // K ints (temp idx)
 
     // write back buffers (buffer-major)
     #pragma unroll

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -176,39 +176,37 @@
     }
 
     extern __shared__ char smem[];
-    float* sv = (float*)smem;                       // BLK*K floats
-    int* si = (int*)(sv + (size_t)BLK * K);         // BLK*K ints
-    float* tv = (float*)(si + (size_t)BLK * K);     // K floats (temp for pair 0)
-    int* tw = (int*)(tv + K);                       // K ints (temp idx)
-
-    // write back buffers (buffer-major)
+    // two ping-pong regions A,B (each BLK*K) for values and for indices
+    float* sv = (float*)smem;                            // 2*BLK*K floats
+    int* si = (int*)(sv + (size_t)2 * BLK * K);          // 2*BLK*K ints
+
+    // phase-1 buffers into region A
     #pragma unroll
     for (int s = 0; s < K; s++) { sv[tid*K + s] = rv[s]; si[tid*K + s] = ri[s]; }
     __syncthreads();
 
-    // pairwise merge tree
+    // pairwise merge tree, ping-pong A<->B (avoids in-place read/write races)
+    float* src_v = sv;            int* src_i = si;            // region A
+    float* dst_v = sv + BLK * K;  int* dst_i = si + BLK * K;  // region B
     int active = BLK;
     while (active > 1) {
         int pair = tid;
         if (pair < active/2) {
-            if (pair == 0)
-                merge_topk<K>(sv, si, sv + K, si + K, tv, tw);
-            else
-                merge_topk<K>(sv + (2*pair)*K, si + (2*pair)*K,
-                              sv + (2*pair+1)*K, si + (2*pair+1)*K,
-                              sv + pair*K, si + pair*K);
+            merge_topk<K>(src_v + (2*pair)*K,     src_i + (2*pair)*K,
+                          src_v + (2*pair+1)*K,   src_i + (2*pair+1)*K,
+                          dst_v + pair*K,         dst_i + pair*K);
         }
         __syncthreads();
-        if (tid < K) { sv[tid] = tv[tid]; si[tid] = tw[tid]; }   // copy pair0 result back
+        float* tv2 = src_v; src_v = dst_v; dst_v = tv2;
+        int* ti2 = src_i; src_i = dst_i; dst_i = ti2;
         active >>= 1;
-        __syncthreads();
-    }
-
-    // buffer 0 now has top-K ascending; write descending, int64 idx
+    }
+
+    // final result (top-K ascending) is in src region, buffer 0
     if (tid < K) {
         int s = K - 1 - tid;
-        tmp_val[seg*K + tid] = sv[s];
-        tmp_idx[seg*K + tid] = si[s];
+        tmp_val[seg*K + tid] = src_v[s];
+        tmp_idx[seg*K + tid] = src_i[s];
     }
 }
 

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -296,8 +296,8 @@
         const float* xv, float* ov, long long* oi, float* tv, int* ti, cudaStream_t s)
 {
     constexpr int BLK = PartialBlk<K>::VAL;
-    // shared = 2*BLK*K*4 (sv+si) + K*8 (temp)
-    size_t smem = (size_t)2 * BLK * K * 4 + (size_t)K * 8;
+    // shared = 2 ping-pong regions * (BLK*K floats + BLK*K ints) = 16*BLK*K bytes
+    size_t smem = (size_t)4 * BLK * K * 4;
     if (smem > 49152)
         cudaFuncSetAttribute(topk_part_mt_kernel<BLK, K>,
             cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -296,20 +296,42 @@
         const float* xv, float* ov, long long* oi, float* tv, int* ti, cudaStream_t s)
 {
     constexpr int BLK = PartialBlk<K>::VAL;
-    // shared = 2 ping-pong regions * (BLK*K floats + BLK*K ints) = 16*BLK*K bytes
-    size_t smem = (size_t)4 * BLK * K * 4;
-    if (smem > 49152)
-        cudaFuncSetAttribute(topk_part_mt_kernel<BLK, K>,
-            cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);
+    constexpr size_t smem = (size_t)4 * BLK * K * 4;   // 16*BLK*K bytes
     topk_part_mt_kernel<BLK, K><<<batch*tiles, BLK, smem, s>>>(
         xv, tv, ti, n, tile_len, tiles);
     int total = tiles * K;
     int M = npow2(total);
     constexpr int MB = 256;
-    if ((size_t)M * 8 > 49152)
-        cudaFuncSetAttribute(topk_merge_kernel<MB, K>,
+    topk_merge_kernel<MB, K><<<batch, MB, (size_t)M*8, s>>>(tv, ti, ov, oi, total, M);
+}
+
+// One-time init: opt in to >48KB dynamic shared mem where needed (call once,
+// NOT inside a graph capture).
+void topk_init() {
+    auto setp = [](auto kc) {
+        constexpr int K = decltype(kc)::value;
+        constexpr int BLK = PartialBlk<K>::VAL;
+        constexpr size_t smem = (size_t)4 * BLK * K * 4;
+        if (smem > 49152)
+            cudaFuncSetAttribute(topk_part_mt_kernel<BLK, K>,
+                cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);
+    };
+    setp(std::integral_constant<int,2>{});
+    setp(std::integral_constant<int,4>{});
+    setp(std::integral_constant<int,8>{});
+    setp(std::integral_constant<int,16>{});
+    setp(std::integral_constant<int,32>{});
+    setp(std::integral_constant<int,64>{});
+    // merge may sort up to 8192 elems (64KB) for shape 0
+    auto setm = [](auto kc) {
+        constexpr int K = decltype(kc)::value;
+        cudaFuncSetAttribute(topk_merge_kernel<256, K>,
             cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);
-    topk_merge_kernel<MB, K><<<batch, MB, (size_t)M*8, s>>>(tv, ti, ov, oi, total, M);
+    };
+    setm(std::integral_constant<int,8>{});
+    setm(std::integral_constant<int,16>{});
+    setm(std::integral_constant<int,32>{});
+    setm(std::integral_constant<int,64>{});
 }
 
 void topk_forward(torch::Tensor x, torch::Tensor ov, torch::Tensor oi,

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -1,12 +1,13 @@
 DECLS = r"""
 void topk_forward(torch::Tensor x, torch::Tensor ov, torch::Tensor oi,
                   torch::Tensor tmp_val, torch::Tensor tmp_idx, int64_t k);
+void topk_init();
 """
 
 
 def build():
     return load_inline(
-        name="topk_cuda_ext_v8",
+        name="topk_cuda_ext_v9",
         cpp_sources=[DECLS],
         cuda_sources=[CUDA_SRC],
-        functions=["topk_forward"],+        functions=["topk_forward", "topk_init"],

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,14 +1,34 @@
-"""Correct-but-slow baseline using only allowed ops (torch.max).
+"""Custom top-k kernel for RTX PRO 6000 (SM120 Blackwell).
 
-Placeholder to verify the check.py / benchmark.py plumbing and to measure the
-SOTA bar. Replaced later by a custom CUDA kernel.
+Strategy:
+  - One C++/CUDA extension with a family of kernels:
+      * argmax (k=1): warp-reduce per row.
+      * partial select (k>=2): per-tile register top-K + shared-mem pairwise
+        merge tree (log2(BLK) levels, O(BLK*K) -- far cheaper than full bitonic).
+      * merge: combine per-tile candidates into the final sorted top-K.
+  - Multi-block per row (tiled) so batch*tiles fills the 188 SMs -> high
+    occupancy (the full bitonic sort was sync-stall-bound at ~1 block/SM).
+  - CUDA graph capture keyed on the input data_ptr to remove Python/launch
+    overhead (these shapes are 0.5-2MB, dominated by per-call overhead).
 """
 import torch
 import torch.nn as nn
 
+import topk_cuda as _tc
+
 OP_TYPE = "topk"
 SUPPORTED_PRECISIONS = ["fp32"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
+
+_mod = None
+
+
+def _get_mod():
+    global _mod
+    if _mod is None:
+        _mod = _tc.build()
+        _mod.topk_init()
+    return _mod
 
 
 class Model(nn.Module):
@@ -16,18 +36,56 @@
         super().__init__()
         self.batch, self.n, self.k = batch, n, k
         self.register_buffer("_dummy", torch.zeros(1))
+        # lazily allocated on first forward (need a device)
+        self._ov = None
+        self._oi = None
+        self._tv = None
+        self._ti = None
+        self._graph = None
+        self._captured_ptr = None
+
+    def _ensure(self, device):
+        if self._ov is None:
+            mod = _get_mod()
+            self._ov = torch.empty(self.batch, self.k, dtype=torch.float32, device=device)
+            self._oi = torch.empty(self.batch, self.k, dtype=torch.int64, device=device)
+            tiles = max(1, (self.n + 255) // 256)
+            self._tv = torch.empty(self.batch * tiles, self.k, dtype=torch.float32, device=device)
+            self._ti = torch.empty(self.batch * tiles, self.k, dtype=torch.int32, device=device)
+
+    def _run(self, x):
+        _get_mod().topk_forward(x, self._ov, self._oi, self._tv, self._ti, self.k)
+
+    def _capture(self, x):
+        # warmup on a side stream (required before graph capture)
+        s = torch.cuda.Stream()
+        s.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(s):
+            for _ in range(3):
+                self._run(x)
+        torch.cuda.current_stream().wait_stream(s)
+        g = torch.cuda.CUDAGraph()
+        with torch.cuda.graph(g):
+            self._run(x)
+        self._graph = g
+        self._captured_ptr = x.data_ptr()
+
+    # Override __call__ to bypass nn.Module hook dispatch overhead.
+    def __call__(self, x: torch.Tensor):
+        self._ensure(x.device)
+        if x.data_ptr() != self._captured_ptr:
+            self._capture(x)
+        self._graph.replay()
+        return self._ov, self._oi
 
     def forward(self, x: torch.Tensor):
-        b, n, k = x.shape[0], x.shape[1], self.k
-        vals = torch.empty(b, k, dtype=torch.float32, device=x.device)
-        idxs = torch.empty(b, k, dtype=torch.int64, device=x.device)
-        work = x.clone()
-        for i in range(k):
-            v, idx = work.max(dim=-1)
-            vals[:, i] = v
-            idxs[:, i] = idx
-            work[torch.arange(b, device=x.device), idx] = float("-inf")
-        return vals, idxs
+        # kept for completeness / torch.compile; __call__ is the hot path
+        self._ensure(x.device)
+        if x.data_ptr() != self._captured_ptr:
+            self._capture(x)
+        else:
+            self._graph.replay()
+        return self._ov, self._oi
 
 
 # Module-level shims rebuilt by check.py / benchmark.py per shape.

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -202,11 +202,11 @@
         active >>= 1;
     }
 
-    // final result (top-K ascending) is in src region, buffer 0
+    // final result (top-K ascending) is in src region, buffer 0.
+    // Write ASCENDING (tmp[seg*K+0]=min) so the merge-tree stage can consume it.
     if (tid < K) {
-        int s = K - 1 - tid;
-        tmp_val[seg*K + tid] = src_v[s];
-        tmp_idx[seg*K + tid] = src_i[s];
+        tmp_val[seg*K + tid] = src_v[tid];
+        tmp_idx[seg*K + tid] = src_i[tid];
     }
 }
 

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -280,6 +280,52 @@
     if (tid < K) {
         out_val[row*K + tid] = sval[tid];
         out_idx[row*K + tid] = (long long)sidx[tid];
+    }
+}
+
+// ---------- Multi-block stage 2 (merge-tree): merge tiles sorted buffers ----------
+// tmp holds `tiles` ascending sorted K-buffers per row. Reduce via a pairwise
+// merge tree (log2(tiles) levels) -- far fewer syncs than bitonic for small tiles.
+template <int BLK, int K>
+__global__ void topk_merge_mt_kernel(
+    const float* __restrict__ tmp_val,
+    const int* __restrict__ tmp_idx,
+    float* __restrict__ out_val,
+    long long* __restrict__ out_idx,
+    int tiles, int total)
+{
+    int row = blockIdx.x;
+    int tid = threadIdx.x;
+    extern __shared__ char smem[];
+    float* sv = (float*)smem;                      // 2*total floats (ping-pong)
+    int* si = (int*)(sv + (size_t)2 * total);      // 2*total ints
+
+    for (int i = tid; i < total; i += BLK) {
+        sv[i] = tmp_val[(long long)row*total + i];
+        si[i] = tmp_idx[(long long)row*total + i];
+    }
+    __syncthreads();
+
+    float* src_v = sv;          int* src_i = si;
+    float* dst_v = sv + total;  int* dst_i = si + total;
+    int active = tiles;
+    while (active > 1) {
+        int pair = tid;
+        if (pair < active/2) {
+            merge_topk<K>(src_v + (2*pair)*K,     src_i + (2*pair)*K,
+                          src_v + (2*pair+1)*K,   src_i + (2*pair+1)*K,
+                          dst_v + pair*K,         dst_i + pair*K);
+        }
+        __syncthreads();
+        float* tv2 = src_v; src_v = dst_v; dst_v = tv2;
+        int* ti2 = src_i; src_i = dst_i; dst_i = ti2;
+        active >>= 1;
+    }
+    // result in src[0..K) ascending -> write descending
+    if (tid < K) {
+        int s = K - 1 - tid;
+        out_val[row*K + tid] = src_v[s];
+        out_idx[row*K + tid] = (long long)src_i[s];
     }
 }
 

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -346,9 +346,15 @@
     topk_part_mt_kernel<BLK, K><<<batch*tiles, BLK, smem, s>>>(
         xv, tv, ti, n, tile_len, tiles);
     int total = tiles * K;
-    int M = npow2(total);
+    // merge-tree merge: shared = 2*total*(4+4) = 16*total bytes
+    size_t m_smem = (size_t)2 * total * 4 + (size_t)2 * total * 4;
     constexpr int MB = 256;
-    topk_merge_kernel<MB, K><<<batch, MB, (size_t)M*8, s>>>(tv, ti, ov, oi, total, M);
+    if (m_smem <= 49152) {
+        topk_merge_mt_kernel<MB, K><<<batch, MB, m_smem, s>>>(tv, ti, ov, oi, tiles, total);
+    } else {
+        int M = npow2(total);
+        topk_merge_kernel<MB, K><<<batch, MB, (size_t)M*8, s>>>(tv, ti, ov, oi, total, M);
+    }
 }
 
 // One-time init: opt in to >48KB dynamic shared mem where needed (call once,

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -379,6 +379,8 @@
         constexpr int K = decltype(kc)::value;
         cudaFuncSetAttribute(topk_merge_kernel<256, K>,
             cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);
+        cudaFuncSetAttribute(topk_merge_mt_kernel<256, K>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);
     };
     setm(std::integral_constant<int,8>{});
     setm(std::integral_constant<int,16>{});

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -346,10 +346,10 @@
     topk_part_mt_kernel<BLK, K><<<batch*tiles, BLK, smem, s>>>(
         xv, tv, ti, n, tile_len, tiles);
     int total = tiles * K;
-    // merge-tree merge: shared = 2*total*(4+4) = 16*total bytes
+    // merge-tree merge: shared = 2*total*(4+4) = 16*total bytes (optin allows 99KB)
     size_t m_smem = (size_t)2 * total * 4 + (size_t)2 * total * 4;
     constexpr int MB = 256;
-    if (m_smem <= 49152) {
+    if (m_smem <= 99 * 1024) {
         topk_merge_mt_kernel<MB, K><<<batch, MB, m_smem, s>>>(tv, ti, ov, oi, tiles, total);
     } else {
         int M = npow2(total);

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -406,21 +406,15 @@
         return;
     }
 
-    // Pick tile_len so that batch*tiles ~>= ~256 blocks (fill the GPU) while
-    // keeping the merge (tiles*K) small. Larger batch -> fewer tiles per row.
-    int tile_len;
-    if (batch >= 64)      tile_len = 2048;
-    else if (batch >= 32) tile_len = 2048;
-    else if (batch >= 16) tile_len = 1024;
-    else                  tile_len = 2048;   // batch small (e.g. shape 0)
-    // shape-0-like (batch==1, huge n): cap tiles so merge fits in shared mem
-    if (batch == 1 && n > 32768) {
-        // total = tiles*K <= 8192 -> tiles <= 128
-        int max_tiles = (k > 0) ? (8192 / (int)k) : 8192;
-        tile_len = (n + max_tiles - 1) / max_tiles;
-        if (tile_len < 1024) tile_len = 1024;
-    }
-    int tiles = (n + tile_len - 1) / tile_len;
+    // Choose tiles so batch*tiles fills the GPU (~256+ blocks for BW) while
+    // keeping the merge-tree shared (16*tiles*K) <= ~99KB.
+    int max_tiles_merge = (99 * 1024) / (16 * (int)k);
+    int want_tiles = (256 + batch - 1) / batch;
+    int tiles = want_tiles;
+    if (tiles > max_tiles_merge) tiles = max_tiles_merge;
+    if (tiles < 1) tiles = 1;
+    int tile_len = (n + tiles - 1) / tiles;
+    tiles = (n + tile_len - 1) / tile_len;
 
     switch (k) {
         case 2:  launch_multiblock<2>(batch, n, tile_len, tiles, xv, ov_, oi_, tv, ti, s); break;

--- a//workspace/problems/05_topk_bitonic/dev_test.py
+++ b//workspace/problems/05_topk_bitonic/dev_test.py
@@ -4,6 +4,7 @@
 import topk_cuda
 
 mod = topk_cuda.build()
+mod.topk_init()
 
 SHAPES = [
     (1, 131072, 64),

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -406,12 +406,13 @@
         return;
     }
 
-    // Choose tiles so batch*tiles fills the GPU (~256+ blocks for BW) while
-    // keeping the merge-tree shared (16*tiles*K) <= ~99KB.
+    // Choose tiles (power of two) so batch*tiles fills the GPU while keeping the
+    // merge-tree shared (16*tiles*K) <= ~99KB. Pow2 so the pairwise tree halves
+    // cleanly with no orphaned buffer.
     int max_tiles_merge = (99 * 1024) / (16 * (int)k);
-    int want_tiles = (256 + batch - 1) / batch;
-    int tiles = want_tiles;
-    if (tiles > max_tiles_merge) tiles = max_tiles_merge;
+    int want = (256 + batch - 1) / batch;
+    int tiles = 1;
+    while (tiles * 2 <= want && tiles * 2 <= max_tiles_merge) tiles *= 2;
     if (tiles < 1) tiles = 1;
     int tile_len = (n + tiles - 1) / tiles;
     tiles = (n + tile_len - 1) / tile_len;

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -410,7 +410,7 @@
     // merge-tree shared (16*tiles*K) <= ~99KB. Pow2 so the pairwise tree halves
     // cleanly with no orphaned buffer.
     int max_tiles_merge = (99 * 1024) / (16 * (int)k);
-    int want = (256 + batch - 1) / batch;
+    int want = (768 + batch - 1) / batch;
     int tiles = 1;
     while (tiles * 2 <= want && tiles * 2 <= max_tiles_merge) tiles *= 2;
     if (tiles < 1) tiles = 1;

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -410,7 +410,7 @@
     // merge-tree shared (16*tiles*K) <= ~99KB. Pow2 so the pairwise tree halves
     // cleanly with no orphaned buffer.
     int max_tiles_merge = (99 * 1024) / (16 * (int)k);
-    int want = (768 + batch - 1) / batch;
+    int want = (256 + batch - 1) / batch;
     int tiles = 1;
     while (tiles * 2 <= want && tiles * 2 <= max_tiles_merge) tiles *= 2;
     if (tiles < 1) tiles = 1;

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -143,6 +143,7 @@
 // Per-thread top-K kept sorted ascending in registers; then a shared-mem
 // pairwise merge tree reduces BLK buffers to 1 (log2(BLK) levels, O(BLK*K)).
 template <int BLK, int K>
+__launch_bounds__(BLK, 512 / BLK)
 __global__ void topk_part_mt_kernel(
     const float* __restrict__ xin,
     float* __restrict__ tmp_val,

--- a//workspace/problems/05_topk_bitonic/topk_cuda.py
+++ b//workspace/problems/05_topk_bitonic/topk_cuda.py
@@ -143,7 +143,6 @@
 // Per-thread top-K kept sorted ascending in registers; then a shared-mem
 // pairwise merge tree reduces BLK buffers to 1 (log2(BLK) levels, O(BLK*K)).
 template <int BLK, int K>
-__launch_bounds__(BLK, 512 / BLK)
 __global__ void topk_part_mt_kernel(
     const float* __restrict__ xin,
     float* __restrict__ tmp_val,

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -1,32 +1,371 @@
-"""Custom top-k kernel for RTX PRO 6000 (SM120 Blackwell).
-
-Strategy:
-  - One C++/CUDA extension with a family of kernels:
-      * argmax (k=1): warp-reduce per row.
-      * partial select (k>=2): per-tile register top-K + shared-mem pairwise
-        merge tree (log2(BLK) levels, O(BLK*K) -- far cheaper than full bitonic).
-      * merge: combine per-tile candidates into the final sorted top-K.
-  - Multi-block per row (tiled) so batch*tiles fills the 188 SMs -> high
-    occupancy (the full bitonic sort was sync-stall-bound at ~1 block/SM).
-  - CUDA graph capture keyed on the input data_ptr to remove Python/launch
-    overhead (these shapes are 0.5-2MB, dominated by per-call overhead).
+"""Custom top-k kernel for RTX PRO 6000 (SM120 Blackwell, GDDR7).
+
+Same Model / get_inputs / get_init_inputs interface as reference.py.
+
+Algorithm (per row, last dim):
+  * k == 1 : warp-reduce argmax, one block per row.
+  * k >= 2 : tiled, two-stage.
+      stage 1 (per tile): each thread keeps a sorted-ascending top-K in registers,
+        then a shared-memory pairwise MERGE TREE (ping-pong, log2(BLK) levels,
+        O(BLK*K) compares) reduces BLK per-thread buffers to one top-K. Far fewer
+        syncs / compares than a full bitonic sort.
+      stage 2 (merge): a second merge-tree combines the per-tile sorted buffers
+        into the final sorted-descending top-K.
+    Tiles are chosen (power of two) so batch*tiles fills the GPU.
+  * CUDA graph capture keyed on the input data_ptr removes Python/launch
+    overhead (these 0.5-2MB shapes are dominated by per-call overhead, not by
+    the memory read). Falls back to direct execution when the input pointer
+    changes (e.g. the correctness runner uses many inputs).
+
+Implemented as a CUDA C++ extension via torch.utils.cpp_extension.load_inline.
 """
 import torch
 import torch.nn as nn
-
-import topk_cuda as _tc
+from torch.utils.cpp_extension import load_inline
 
 OP_TYPE = "topk"
 SUPPORTED_PRECISIONS = ["fp32"]
 HARDWARE_REQUIRED = ["RTX_PRO_6000", "H100", "B200"]
 
+_CUDA_SRC = r"""
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <cuda_runtime.h>
+#include <type_traits>
+
+constexpr int npow2(int x){ int p = 1; while (p < x) p <<= 1; return p; }
+constexpr int ilog2(int x){ int r=0; while ((1<<r) < x) r++; return r; }
+
+// ---------- Bitonic sort M (pow2) elements DESCENDING, carry idx ----------
+template <int BLK>
+__device__ __forceinline__ void bitonic_sort_desc(float* sval, int* sidx, int M) {
+    int tid = threadIdx.x;
+    int logM = ilog2(M);
+    for (int s = 1; s <= logM; s++) {
+        int size = 1 << s;
+        for (int stride = size >> 1; stride > 0; stride >>= 1) {
+            __syncthreads();
+            for (int i = tid; i < M; i += BLK) {
+                int j = i ^ stride;
+                if (j > i) {
+                    float a = sval[i], b = sval[j];
+                    bool desc = ((i & size) == 0);
+                    bool swap = desc ? (a < b) : (a > b);
+                    if (swap) { sval[i] = b; sval[j] = a; int ti = sidx[i]; sidx[i] = sidx[j]; sidx[j] = ti; }
+                }
+            }
+        }
+    }
+    __syncthreads();
+}
+
+// ---------- Argmax (k=1): one block per row ----------
+template <int BLK>
+__global__ void argmax_kernel(
+    const float* __restrict__ xin,
+    float* __restrict__ out_val,
+    long long* __restrict__ out_idx,
+    int n)
+{
+    constexpr int WARP = 32;
+    constexpr int NWARPS = BLK / WARP;
+    int row = blockIdx.x;
+    int tid = threadIdx.x;
+    int wid = tid / WARP, lane = tid % WARP;
+    const long long base = (long long)row * n;
+
+    __shared__ float sw[NWARPS];
+    __shared__ int si[NWARPS];
+
+    float myv = -INFINITY;
+    int myi = 0;
+    for (int e = tid; e < n; e += BLK) {
+        float val = xin[base + e];
+        if (val > myv) { myv = val; myi = e; }
+    }
+    for (int off = 16; off > 0; off >>= 1) {
+        float ov = __shfl_xor_sync(0xffffffff, myv, off);
+        int oi = __shfl_xor_sync(0xffffffff, myi, off);
+        if (ov > myv) { myv = ov; myi = oi; }
+    }
+    if (lane == 0) { sw[wid] = myv; si[wid] = myi; }
+    __syncthreads();
+    if (wid == 0) {
+        myv = (lane < NWARPS) ? sw[lane] : -INFINITY;
+        myi = (lane < NWARPS) ? si[lane] : 0;
+        for (int off = 16; off > 0; off >>= 1) {
+            float ov = __shfl_xor_sync(0xffffffff, myv, off);
+            int oi = __shfl_xor_sync(0xffffffff, myi, off);
+            if (ov > myv) { myv = ov; myi = oi; }
+        }
+        if (lane == 0) { out_val[row] = myv; out_idx[row] = (long long)myi; }
+    }
+}
+
+// ---------- Merge two ascending sorted K-arrays into top-K ascending ----------
+template <int K>
+__device__ __forceinline__ void merge_topk(
+    const float* A, const int* Ai,
+    const float* B, const int* Bi,
+    float* C, int* Ci)
+{
+    int ia = K - 1, ib = K - 1;
+    #pragma unroll
+    for (int c = K - 1; c >= 0; c--) {
+        float av = (ia >= 0) ? A[ia] : -INFINITY;
+        float bv = (ib >= 0) ? B[ib] : -INFINITY;
+        if (av >= bv) { C[c] = av; Ci[c] = Ai[ia]; ia--; }
+        else          { C[c] = bv; Ci[c] = Bi[ib]; ib--; }
+    }
+}
+
+// ---------- Stage 1 (merge-tree): per-tile top-K into partial buffers ----------
+// Per-thread top-K kept sorted ascending in registers; then a shared-mem
+// pairwise merge tree (ping-pong) reduces BLK buffers to 1.
+template <int BLK, int K>
+__global__ void topk_part_mt_kernel(
+    const float* __restrict__ xin,
+    float* __restrict__ tmp_val,
+    int* __restrict__ tmp_idx,
+    int n, int tile_len, int tiles)
+{
+    int seg = blockIdx.x;
+    int row = seg / tiles;
+    int tile = seg % tiles;
+    int tid = threadIdx.x;
+    long long base = (long long)row * n + (long long)tile * tile_len;
+    int my_len = min(tile

... (10293 more chars)

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -150,7 +150,7 @@
                 if (rv[s] < rv[s-1]) {
                     float tf = rv[s]; rv[s] = rv[s-1]; rv[s-1] = tf;
                     int ti = ri[s]; ri[s] = ri[s-1]; ri[s-1] = ti;
-                }
+                } else break;
             }
         }
     }

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -150,7 +150,7 @@
                 if (rv[s] < rv[s-1]) {
                     float tf = rv[s]; rv[s] = rv[s-1]; rv[s-1] = tf;
                     int ti = ri[s]; ri[s] = ri[s-1]; ri[s-1] = ti;
-                } else break;
+                }
             }
         }
     }

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -173,7 +173,9 @@
                           src_v + (2*pair+1)*K,   src_i + (2*pair+1)*K,
                           dst_v + pair*K,         dst_i + pair*K);
         }
-        __syncthreads();
+        // Late levels (active<=32) only touch the first warp -> syncwarp is enough
+        // and avoids the cross-warp barrier stall at low occupancy.
+        if (active <= 32) __syncwarp(); else __syncthreads();
         float* tv2 = src_v; src_v = dst_v; dst_v = tv2;
         int* ti2 = src_i; src_i = dst_i; dst_i = ti2;
         active >>= 1;

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -244,7 +244,7 @@
                           src_v + (2*pair+1)*K,   src_i + (2*pair+1)*K,
                           dst_v + pair*K,         dst_i + pair*K);
         }
-        __syncthreads();
+        if (active <= 32) __syncwarp(); else __syncthreads();
         float* tv2 = src_v; src_v = dst_v; dst_v = tv2;
         int* ti2 = src_i; src_i = dst_i; dst_i = ti2;
         active >>= 1;

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -156,22 +156,25 @@
     }
 
     extern __shared__ char smem[];
-    float* sv = (float*)smem;                            // 2*BLK*K floats
-    int* si = (int*)(sv + (size_t)2 * BLK * K);          // 2*BLK*K ints
+    // Pad buffer stride to KP=K+1 (odd -> coprime with 32 banks) to avoid the
+    // 32-way bank conflicts on the writeback sv[tid*KP+s] when K is a multiple of 32.
+    constexpr int KP = K + 1;
+    float* sv = (float*)smem;                            // 2*BLK*KP floats
+    int* si = (int*)(sv + (size_t)2 * BLK * KP);          // 2*BLK*KP ints
 
     #pragma unroll
-    for (int s = 0; s < K; s++) { sv[tid*K + s] = rv[s]; si[tid*K + s] = ri[s]; }
+    for (int s = 0; s < K; s++) { sv[tid*KP + s] = rv[s]; si[tid*KP + s] = ri[s]; }
     __syncthreads();
 
-    float* src_v = sv;            int* src_i = si;            // region A
-    float* dst_v = sv + BLK * K;  int* dst_i = si + BLK * K;  // region B
+    float* src_v = sv;             int* src_i = si;              // region A
+    float* dst_v = sv + BLK * KP;  int* dst_i = si + BLK * KP;   // region B
     int active = BLK;
     while (active > 1) {
         int pair = tid;
         if (pair < active/2) {
-            merge_topk<K>(src_v + (2*pair)*K,     src_i + (2*pair)*K,
-                          src_v + (2*pair+1)*K,   src_i + (2*pair+1)*K,
-                          dst_v + pair*K,         dst_i + pair*K);
+            merge_topk<K>(src_v + (2*pair)*KP,     src_i + (2*pair)*KP,
+                          src_v + (2*pair+1)*KP,   src_i + (2*pair+1)*KP,
+                          dst_v + pair*KP,         dst_i + pair*KP);
         }
         // Late levels (active<=32) only touch the first warp -> syncwarp is enough
         // and avoids the cross-warp barrier stall at low occupancy.

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -229,23 +229,26 @@
 {
     int row = blockIdx.x;
     int tid = threadIdx.x;
+    constexpr int KP = K + 1;   // padded buffer stride -> bank-conflict-free
     extern __shared__ char smem[];
-    float* sv = (float*)smem;                      // 2*total floats
-    int* si = (int*)(sv + (size_t)2 * total);      // 2*total ints
+    float* sv = (float*)smem;                          // 2*tiles*KP floats
+    int* si = (int*)(sv + (size_t)2 * tiles * KP);     // 2*tiles*KP ints
+    // scatter-load: element (b,s) -> sv[b*KP + s]
     for (int i = tid; i < total; i += BLK) {
-        sv[i] = tmp_val[(long long)row*total + i];
-        si[i] = tmp_idx[(long long)row*total + i];
+        int b = i / K, s = i - b * K;
+        sv[b*KP + s] = tmp_val[(long long)row*total + i];
+        si[b*KP + s] = tmp_idx[(long long)row*total + i];
     }
     __syncthreads();
-    float* src_v = sv;          int* src_i = si;
-    float* dst_v = sv + total;  int* dst_i = si + total;
+    float* src_v = sv;            int* src_i = si;
+    float* dst_v = sv + tiles*KP; int* dst_i = si + tiles*KP;
     int active = tiles;
     while (active > 1) {
         int pair = tid;
         if (pair < active/2) {
-            merge_topk<K>(src_v + (2*pair)*K,     src_i + (2*pair)*K,
-                          src_v + (2*pair+1)*K,   src_i + (2*pair+1)*K,
-                          dst_v + pair*K,         dst_i + pair*K);
+            merge_topk<K>(src_v + (2*pair)*KP,     src_i + (2*pair)*KP,
+                          src_v + (2*pair+1)*KP,   src_i + (2*pair+1)*KP,
+                          dst_v + pair*KP,         dst_i + pair*KP);
         }
         if (active <= 32) __syncwarp(); else __syncthreads();
         float* tv2 = src_v; src_v = dst_v; dst_v = tv2;

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -273,11 +273,12 @@
         const float* xv, float* ov, long long* oi, float* tv, int* ti, cudaStream_t s)
 {
     constexpr int BLK = PartialBlk<K>::VAL;
-    constexpr size_t smem = (size_t)4 * BLK * K * 4;
+    constexpr size_t smem = (size_t)4 * BLK * (K + 1) * 4;   // 2 regions * (val+idx), padded
     topk_part_mt_kernel<BLK, K><<<batch*tiles, BLK, smem, s>>>(
         xv, tv, ti, n, tile_len, tiles);
     int total = tiles * K;
-    size_t m_smem = (size_t)2 * total * 4 + (size_t)2 * total * 4;
+    // merge-tree merge: 2 regions * (tiles buffers of K+1 each) * (val+idx)
+    size_t m_smem = (size_t)4 * tiles * (K + 1) * 4;
     constexpr int MB = 256;
     if (m_smem <= 99 * 1024) {
         topk_merge_mt_kernel<MB, K><<<batch, MB, m_smem, s>>>(tv, ti, ov, oi, tiles, total);

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -293,7 +293,7 @@
     auto setp = [](auto kc) {
         constexpr int K = decltype(kc)::value;
         constexpr int BLK = PartialBlk<K>::VAL;
-        constexpr size_t smem = (size_t)4 * BLK * K * 4;
+        constexpr size_t smem = (size_t)4 * BLK * (K + 1) * 4;
         if (smem > 49152)
             cudaFuncSetAttribute(topk_part_mt_kernel<BLK, K>,
                 cudaFuncAttributeMaxDynamicSharedMemorySize, 99 * 1024);

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -176,9 +176,7 @@
                           src_v + (2*pair+1)*KP,   src_i + (2*pair+1)*KP,
                           dst_v + pair*KP,         dst_i + pair*KP);
         }
-        // Late levels (active<=32) only touch the first warp -> syncwarp is enough
-        // and avoids the cross-warp barrier stall at low occupancy.
-        if (active <= 32) __syncwarp(); else __syncthreads();
+        __syncthreads();
         float* tv2 = src_v; src_v = dst_v; dst_v = tv2;
         int* ti2 = src_i; src_i = dst_i; dst_i = ti2;
         active >>= 1;

--- a//workspace/problems/05_topk_bitonic/solution.py
+++ b//workspace/problems/05_topk_bitonic/solution.py
@@ -248,7 +248,7 @@
                           src_v + (2*pair+1)*KP,   src_i + (2*pair+1)*KP,
                           dst_v + pair*KP,         dst_i + pair*KP);
         }
-        if (active <= 32) __syncwarp(); else __syncthreads();
+        __syncthreads();
         float* tv2 = src_v; src_v = dst_v; dst_v = tv2;
         int* ti2 = src_i; src_i = dst_i; dst_i = ti2;
         active >>= 1;