--- /dev/null
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -0,0 +1,181 @@
+"""Top-k kernel via Triton — per-row processing with chunk-level selection."""
+from __future__ import annotations
+
+import math
+import torch
+import torch.nn as nn
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _topk_kernel(
+    x_ptr,
+    val_out_ptr,
+    idx_out_ptr,
+    N: tl.constexpr,
+    K: tl.constexpr,
+    stride_x_row,
+    stride_val_row,
+    stride_idx_row,
+    BLOCK_N: tl.constexpr,
+):
+    """One program per row.
+
+    Strategy:
+    1. Scan the row in chunks of BLOCK_N.
+    2. Per chunk: load into shared memory, each thread keeps its local top-K.
+    3. After scanning, sort the local top-K buffer and write output.
+    
+    For simplicity in v1: each thread maintains top-K in registers
+    from its strided portion, then we do a warp-level merge.
+    """
+    row_idx = tl.program_id(0)
+    
+    x_row = x_ptr + row_idx * stride_x_row
+    
+    # Each thread will maintain a small top-k buffer.
+    # We process elements in a strided loop.
+    # For K elements, we store them in registers.
+    # This kernel uses a simple approach: each thread scans its
+    # strided portion and keeps running top-k.
+    
+    pid = tl.arange(0, BLOCK_N)  # thread indices within block
+    mask = pid < N
+    
+    # For small N that fit in one block: load everything, sort, take top-k.
+    # For large N: we need multi-pass. 
+    
+    # For simplicity, handle N <= BLOCK_N directly:
+    if N <= BLOCK_N:
+        vals = tl.load(x_row + pid, mask=mask, other=-float('inf'))
+        vals_with_idx = vals
+        
+        # We need to sort descending and take top-k.
+        # Triton doesn't have a built-in sort, so we use a bitonic approach.
+        # But for large BLOCK_N, this is expensive in v1.
+        pass
+    
+    # Fallback general path: multi-pass scan.
+    # This first version does a simple approach with tl.max in a loop.
+    # It's not optimal but gets us started.
+
+
+@triton.jit
+def _select_topk_single_pass(
+    x_ptr,
+    val_out_ptr,
+    idx_out_ptr,
+    N: tl.constexpr,
+    K: tl.constexpr,
+    stride_x_row,
+    stride_val_row,
+    stride_idx_row,
+    BLOCK_N: tl.constexpr,
+):
+    """For rows where BLOCK_N >= N. Single block, load all, select top-k."""
+    row_idx = tl.program_id(0)
+    x_row = x_ptr + row_idx * stride_x_row
+    val_row = val_out_ptr + row_idx * stride_val_row
+    idx_row = idx_out_ptr + row_idx * stride_idx_row
+    
+    off = tl.arange(0, BLOCK_N)
+    mask = off < N
+    vals = tl.load(x_row + off, mask=mask, other=-float('inf'))
+    
+    # Build top-k using repeated tl.max extraction
+    for ki in range(K):
+        # Find max among remaining
+        best_val = tl.max(vals, axis=0)
+        # Find index of max
+        # Triton doesn't have argmax directly, need to construct it
+        best_mask = (vals == best_val)
+        best_idx_raw = tl.where(best_mask, off.to(tl.int64), N + ki).to(tl.int64)
+        best_idx = tl.min(best_idx_raw.to(tl.int32), axis=0)
+        
+        # Store at position ki
+        tl.store(val_row + ki, best_val)
+        tl.store(idx_row + ki, best_idx.to(tl.int64))
+        
+        # Remove from consideration (set to -inf at best_idx position)
+        # We can't directly modify at a position, so rebuild
+        vals = tl.where(off == best_idx, -float('inf'), vals)
+
+
+@triton.jit
+def _topk_scan_kernel(
+    x_ptr,
+    val_out_ptr,
+    idx_out_ptr,
+    N: tl.constexpr,
+    K: tl.constexpr,
+    stride_x_row,
+    stride_val_row,
+    stride_idx_row,
+):
+    """One program per row. Scan-based top-k selection.
+    
+    Each thread finds its local max in a strided loop,
+    then we reduce across threads.
+    Repeat K times for top-K.
+    
+    This is O(K * N) which is not optimal but gives a working baseline.
+    For K=64, N=131072 this would be ~8M ops - workable.
+    """
+    row_idx = tl.program_id(0)
+    tid = tl.program_id(1)
+    
+    # For this approach, we use the fact that we can 
+    # do a strided max-finding pass K times.
+    # Not the most efficient but simple and correct.
+    pass
+
+
+class Model(nn.Module):
+    """Top-k via Triton kernel."""
+    
+    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):
+        batch, n = x.shape[0], x.shape[1]
+        k = self.k
+        device = x.device
+        
+        val_out = torch.empty(batch, k, dtype=torch.float32, device=device)
+        idx_out = torch.empty(batch, k, dtype=torch.int64, device=device)
+        
+        # Choose block size: pad N to next power of 2
+        block_n = triton.next_power_of_2(n)
+        
+        # Use Triton launch with one program per row
+        grid = (batch,)
+        
+        _select_topk_single_pass[grid](
+            x, val_out, idx_out,
+            N=n, K=k,
+            stride_x_row=x.stride(0),
+            stride_val_row=val_out.stride(0),
+            stride_idx_row=idx_out.stride(0),
+            BLOCK_N=block_n,
+            num_ctas=1,
+        )
+        
+        return val_out, idx_out
+
+
+def get_inputs():
+    batch, n = 64, 8192
+    x = torch.randn(batch, n, dtype=torch.float32)
+    return [x]
+
+
+batch = 64
+n = 8192
+k = 8
+
+
+def get_init_inputs():
+    return [batch, n, k]

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -1,7 +1,6 @@
-"""Top-k kernel via Triton — per-row processing with chunk-level selection."""
+"""Top-k kernel via Triton — two-phase chunk select + merge."""
 from __future__ import annotations
 
-import math
 import torch
 import torch.nn as nn
 import triton
@@ -9,131 +8,100 @@
 
 
 @triton.jit
-def _topk_kernel(
+def _chunk_select_kernel(
     x_ptr,
-    val_out_ptr,
-    idx_out_ptr,
+    chunk_val_ptr,
+    chunk_idx_ptr,
     N: tl.constexpr,
     K: tl.constexpr,
     stride_x_row,
+    CHUNK_SIZE: tl.constexpr,
+):
+    """Phase 1: Each program handles one chunk of one row.
+    
+    For a chunk of CHUNK_SIZE elements, find the top-K and store.
+    Uses repeated max-extraction: K iterations, each finds the current max.
+    
+    grid = (batch, num_chunks)
+    """
+    row_idx = tl.program_id(0)
+    chunk_idx = tl.program_id(1)
+    
+    chunk_start = chunk_idx * CHUNK_SIZE
+    x_row = x_ptr + row_idx * stride_x_row
+    
+    out_row = chunk_idx * K
+    out_base = (row_idx * tl.cdiv(N, CHUNK_SIZE) + chunk_idx) * K
+    
+    offs = chunk_start + tl.arange(0, CHUNK_SIZE)
+    mask = offs < N
+    
+    vals = tl.load(x_row + offs, mask=mask, other=-float('inf'))
+    
+    for ki in tl.range(K):
+        m = tl.max(vals)
+        # Argmax: find the first position with value == m
+        m_idx = tl.where(vals == m, offs.to(tl.int32), 0x7FFFFFFF).to(tl.int32)
+        arg = tl.min(m_idx)
+        
+        tl.store(chunk_val_ptr + out_base + ki, m)
+        tl.store(chunk_idx_ptr + out_base + ki, arg.to(tl.int64))
+        
+        vals = tl.where(vals == m, -float('inf'), vals)
+
+
+@triton.jit
+def _merge_kernel(
+    chunk_val_ptr,
+    chunk_idx_ptr,
+    val_out_ptr,
+    idx_out_ptr,
+    num_chunks: tl.constexpr,
+    K_CHUNK: tl.constexpr,
+    K_OUT: tl.constexpr,
     stride_val_row,
     stride_idx_row,
-    BLOCK_N: tl.constexpr,
 ):
-    """One program per row.
-
-    Strategy:
-    1. Scan the row in chunks of BLOCK_N.
-    2. Per chunk: load into shared memory, each thread keeps its local top-K.
-    3. After scanning, sort the local top-K buffer and write output.
+    """Phase 2: Each program handles one row. Merge num_chunks * K_CHUNK into top-K_OUT.
     
-    For simplicity in v1: each thread maintains top-K in registers
-    from its strided portion, then we do a warp-level merge.
+    grid = (batch,)
+    Total candidates = num_chunks * K_CHUNK.
+    Extract top-K_OUT via repeated max.
     """
     row_idx = tl.program_id(0)
     
-    x_row = x_ptr + row_idx * stride_x_row
+    row_base = row_idx * num_chunks * K_CHUNK
+    n_candidates = num_chunks * K_CHUNK
     
-    # Each thread will maintain a small top-k buffer.
-    # We process elements in a strided loop.
-    # For K elements, we store them in registers.
-    # This kernel uses a simple approach: each thread scans its
-    # strided portion and keeps running top-k.
-    
-    pid = tl.arange(0, BLOCK_N)  # thread indices within block
-    mask = pid < N
-    
-    # For small N that fit in one block: load everything, sort, take top-k.
-    # For large N: we need multi-pass. 
-    
-    # For simplicity, handle N <= BLOCK_N directly:
-    if N <= BLOCK_N:
-        vals = tl.load(x_row + pid, mask=mask, other=-float('inf'))
-        vals_with_idx = vals
-        
-        # We need to sort descending and take top-k.
-        # Triton doesn't have a built-in sort, so we use a bitonic approach.
-        # But for large BLOCK_N, this is expensive in v1.
+    # Load all candidates
+    offs = tl.arange(0, tl.max_constexpr([n_candidates]))
+    if offs.max() < n_candidates:
         pass
     
-    # Fallback general path: multi-pass scan.
-    # This first version does a simple approach with tl.max in a loop.
-    # It's not optimal but gets us started.
+    # For variable size, load in chunks
+    # Since we know num_chunks and K_CHUNK at compile time, 
+    # total = num_chunks * K_CHUNK
+    
+    TOTAL = num_chunks * K_CHUNK
+    load_offs = tl.arange(0, TOTAL)
+    
+    cand_vals = tl.load(chunk_val_ptr + row_base + load_offs)
+    cand_idxs = tl.load(chunk_idx_ptr + row_base + load_offs)
+    
+    for ki in tl.range(K_OUT):
+        m = tl.max(cand_vals)
+        m_idx = tl.where(cand_vals == m, load_offs.to(tl.int32), 0x7FFFFFFF).to(tl.int32)
+        arg = tl.min(m_idx)
+        
+        orig_idx = tl.load(chunk_idx_ptr + row_base + arg)
+        
+        tl.store(val_out_ptr + row_idx * stride_val_row + ki, m)
+        tl.store(idx_out_ptr + row_idx * stride_idx_row + ki, orig_idx)
+        
+        cand_vals = tl.where(load_offs == arg, -float('inf'), cand_vals)
 
 
-@triton.jit
-def _select_topk_single_pass(
-    x_ptr,
-    val_out_ptr,
-    idx_out_ptr,
-    N: tl.constexpr,
-    K: tl.constexpr,
-    stride_x_row,
-    stride_val_row,
-    stride_idx_row,
-    BLOCK_N: tl.constexpr,
-):
-    """For rows where BLOCK_N >= N. Single block, load all, select top-k."""
-    row_idx = tl.program_id(0)
-    x_row = x_ptr + row_idx * stride_x_row
-    val_row = val_out_ptr + row_idx * stride_val_row
-    idx_row = idx_out_ptr + row_idx * stride_idx_row
-    
-    off = tl.arange(0, BLOCK_N)
-    mask = off < N
-    vals = tl.load(x_row + off, mask=mask, other=-float('inf'))
-    
-    # Build top-k using repeated tl.max extraction
-    for ki in range(K):
-        # Find max among remaining
-        best_val = tl.max(vals, axis=0)
-        # Find index of max
-        # Triton doesn't have argmax directly, need to construct it
-        best_mask = (vals == best_val)
-        best_idx_raw = tl.where(best_mask, off.to(tl.int64), N + ki).to(tl.int64)
-        best_idx = tl.min(best_idx_raw.to(tl.int32), axis=0)
-        
-        # Store at position ki
-        tl.store(val_row + ki, best_val)
-        tl.store(idx_row + ki, b

... (3093 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -7,13 +7,26 @@
 import triton.language as tl
 
 
+# Maximum chunk size for phase 1 (must fit in register file per thread)
+MAX_CHUNK = 1024
+
+
+def _get_chunk_size(n: int, k: int) -> int:
+    """Return chunk size for phase 1. Must be >= k and <= MAX_CHUNK, power of 2."""
+    chunk = min(MAX_CHUNK, triton.next_power_of_2(k))
+    # Ensure chunk >= k
+    if chunk < k:
+        chunk = triton.next_power_of_2(k)
+    return min(chunk, MAX_CHUNK)
+
+
 @triton.jit
 def _chunk_select_kernel(
     x_ptr,
     chunk_val_ptr,
     chunk_idx_ptr,
-    N: tl.constexpr,
-    K: tl.constexpr,
+    N,
+    K,
     stride_x_row,
     CHUNK_SIZE: tl.constexpr,
 ):
@@ -30,24 +43,24 @@
     chunk_start = chunk_idx * CHUNK_SIZE
     x_row = x_ptr + row_idx * stride_x_row
     
-    out_row = chunk_idx * K
-    out_base = (row_idx * tl.cdiv(N, CHUNK_SIZE) + chunk_idx) * K
+    # Number of chunks per row (computed at compile time since CHUNK_SIZE is constexpr)
+    num_chunks = tl.cdiv(N.to(tl.int32), CHUNK_SIZE.to(tl.int32))
+    out_base = (row_idx * num_chunks + chunk_idx) * K
     
     offs = chunk_start + tl.arange(0, CHUNK_SIZE)
     mask = offs < N
     
-    vals = tl.load(x_row + offs, mask=mask, other=-float('inf'))
+    vals = tl.load(x_row + offs, mask=mask, other=-1e30)
     
-    for ki in tl.range(K):
+    for ki in range(K):
         m = tl.max(vals)
-        # Argmax: find the first position with value == m
-        m_idx = tl.where(vals == m, offs.to(tl.int32), 0x7FFFFFFF).to(tl.int32)
+        m_idx = tl.where(vals == m, offs.to(tl.int32), 0x7FFFFFFF)
         arg = tl.min(m_idx)
         
         tl.store(chunk_val_ptr + out_base + ki, m)
         tl.store(chunk_idx_ptr + out_base + ki, arg.to(tl.int64))
         
-        vals = tl.where(vals == m, -float('inf'), vals)
+        vals = tl.where(offs == arg, -1e30, vals)
 
 
 @triton.jit
@@ -56,41 +69,31 @@
     chunk_idx_ptr,
     val_out_ptr,
     idx_out_ptr,
-    num_chunks: tl.constexpr,
-    K_CHUNK: tl.constexpr,
-    K_OUT: tl.constexpr,
+    num_chunks,
+    K_CHUNK,
+    K_OUT,
     stride_val_row,
     stride_idx_row,
+    TOTAL_CANDIDATES: tl.constexpr,
 ):
-    """Phase 2: Each program handles one row. Merge num_chunks * K_CHUNK into top-K_OUT.
+    """Phase 2: Each program handles one row. Merge num_chunks*K_CHUNK into top-K_OUT.
     
     grid = (batch,)
-    Total candidates = num_chunks * K_CHUNK.
-    Extract top-K_OUT via repeated max.
+    Total candidates = num_chunks * K_CHUNK, loaded in blocks.
     """
     row_idx = tl.program_id(0)
     
     row_base = row_idx * num_chunks * K_CHUNK
-    n_candidates = num_chunks * K_CHUNK
     
-    # Load all candidates
-    offs = tl.arange(0, tl.max_constexpr([n_candidates]))
-    if offs.max() < n_candidates:
-        pass
-    
-    # For variable size, load in chunks
-    # Since we know num_chunks and K_CHUNK at compile time, 
-    # total = num_chunks * K_CHUNK
-    
-    TOTAL = num_chunks * K_CHUNK
-    load_offs = tl.arange(0, TOTAL)
+    # Load all candidates at once using constexpr TOTAL_CANDIDATES
+    load_offs = tl.arange(0, TOTAL_CANDIDATES)
     
     cand_vals = tl.load(chunk_val_ptr + row_base + load_offs)
     cand_idxs = tl.load(chunk_idx_ptr + row_base + load_offs)
     
-    for ki in tl.range(K_OUT):
+    for ki in range(K_OUT):
         m = tl.max(cand_vals)
-        m_idx = tl.where(cand_vals == m, load_offs.to(tl.int32), 0x7FFFFFFF).to(tl.int32)
+        m_idx = tl.where(cand_vals == m, load_offs.to(tl.int32), 0x7FFFFFFF)
         arg = tl.min(m_idx)
         
         orig_idx = tl.load(chunk_idx_ptr + row_base + arg)
@@ -98,7 +101,7 @@
         tl.store(val_out_ptr + row_idx * stride_val_row + ki, m)
         tl.store(idx_out_ptr + row_idx * stride_idx_row + ki, orig_idx)
         
-        cand_vals = tl.where(load_offs == arg, -float('inf'), cand_vals)
+        cand_vals = tl.where(load_offs == arg, -1e30, cand_vals)
 
 
 class Model(nn.Module):    
@@ -113,22 +116,21 @@
         device = x.device
         
         # Phase 1: chunk-level selection
-        chunk_size = 1024  # power of 2, must be >= k
+        chunk_size = _get_chunk_size(n, k)
         num_chunks = triton.cdiv(n, chunk_size)
         
-        # If chunk_size < k, increase chunk_size
-        if chunk_size < k:
-            chunk_size = triton.next_power_of_2(k)
-            num_chunks = triton.cdiv(n, chunk_size)
+        # Phase 2 candidate count (must be constexpr for merge kernel)
+        # We pad total_candidates to next power of 2 for tl.arange
+        total_candidates = num_chunks * k
+        total_candidates_padded = triton.next_power_of_2(total_candidates)
         
-        # Temp storage for chunk-level top-k
-        chunk_vals = torch.empty(batch, num_chunks, k, dtype=torch.float32, device=device)
-        chunk_idxs = torch.empty(batch, num_chunks, k, dtype=torch.int64, device=device)
+        # Temp storage for chunk-level top-k (padded)
+        chunk_vals = torch.full((batch, total_candidates_padded,), -1e30, dtype=torch.float32, device=device)
+        chunk_idxs = torch.full((batch, total_candidates_padded,), -1, dtype=torch.int64, device=device)
         
         grid1 = (batch, num_chunks)
-        
         _chunk_select_kernel[grid1](
-            x, chunk_vals.view(-1), chunk_idxs.view(-1),
+            x, chunk_vals, chunk_idxs,
             N=n, K=k,
             stride_x_row=x.stride(0),
             CHUNK_SIZE=chunk_size,
@@ -140,11 +142,12 @@
         
         grid2 = (batch,)
         _merge_kernel[grid2](
-            chunk_vals.view(-1), chunk_idxs.view(-1),
+            chunk_vals, chunk_idxs,
             val_out, idx_out,
             num_chunks=num_chunks, K_CHUNK=k, K_OUT=k,
             stride_val_row=val_out.stride(0),
             stride_idx_row=idx_out.stride(0),
+ 

... (107 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -7,33 +7,20 @@
 import triton.language as tl
 
 
-# Maximum chunk size for phase 1 (must fit in register file per thread)
-MAX_CHUNK = 1024
-
-
-def _get_chunk_size(n: int, k: int) -> int:
-    """Return chunk size for phase 1. Must be >= k and <= MAX_CHUNK, power of 2."""
-    chunk = min(MAX_CHUNK, triton.next_power_of_2(k))
-    # Ensure chunk >= k
-    if chunk < k:
-        chunk = triton.next_power_of_2(k)
-    return min(chunk, MAX_CHUNK)
-
-
 @triton.jit
 def _chunk_select_kernel(
     x_ptr,
     chunk_val_ptr,
     chunk_idx_ptr,
     N,
-    K,
     stride_x_row,
+    num_chunks,
+    K: tl.constexpr,
     CHUNK_SIZE: tl.constexpr,
 ):
     """Phase 1: Each program handles one chunk of one row.
     
-    For a chunk of CHUNK_SIZE elements, find the top-K and store.
-    Uses repeated max-extraction: K iterations, each finds the current max.
+    For a chunk of CHUNK_SIZE elements, find the top-K via repeated max.
     
     grid = (batch, num_chunks)
     """
@@ -42,9 +29,6 @@
     
     chunk_start = chunk_idx * CHUNK_SIZE
     x_row = x_ptr + row_idx * stride_x_row
-    
-    # Number of chunks per row (computed at compile time since CHUNK_SIZE is constexpr)
-    num_chunks = tl.cdiv(N.to(tl.int32), CHUNK_SIZE.to(tl.int32))
     out_base = (row_idx * num_chunks + chunk_idx) * K
     
     offs = chunk_start + tl.arange(0, CHUNK_SIZE)
@@ -70,30 +54,26 @@
     val_out_ptr,
     idx_out_ptr,
     num_chunks,
-    K_CHUNK,
-    K_OUT,
     stride_val_row,
     stride_idx_row,
+    K_OUT: tl.constexpr,
     TOTAL_CANDIDATES: tl.constexpr,
 ):
-    """Phase 2: Each program handles one row. Merge num_chunks*K_CHUNK into top-K_OUT.
+    """Phase 2: Each program merges num_chunks*K candidates into top-K_OUT.
     
     grid = (batch,)
-    Total candidates = num_chunks * K_CHUNK, loaded in blocks.
     """
     row_idx = tl.program_id(0)
+    k_chunk = TOTAL_CANDIDATES // num_chunks
     
-    row_base = row_idx * num_chunks * K_CHUNK
-    
-    # Load all candidates at once using constexpr TOTAL_CANDIDATES
+    row_base = row_idx * num_chunks * k_chunk
     load_offs = tl.arange(0, TOTAL_CANDIDATES)
     
     cand_vals = tl.load(chunk_val_ptr + row_base + load_offs)
-    cand_idxs = tl.load(chunk_idx_ptr + row_base + load_offs)
     
     for ki in range(K_OUT):
         m = tl.max(cand_vals)
-        m_idx = tl.where(cand_vals == m, load_offs.to(tl.int32), 0x7FFFFFFF)
+        m_idx = tl.where(cand_vals == m, load_offs, 0x7FFFFFFF).to(tl.int32)
         arg = tl.min(m_idx)
         
         orig_idx = tl.load(chunk_idx_ptr + row_base + arg)
@@ -114,39 +94,42 @@
         batch, n = x.shape[0], x.shape[1]
         k = self.k
         device = x.device
+        dtype = x.dtype
         
-        # Phase 1: chunk-level selection
-        chunk_size = _get_chunk_size(n, k)
-        num_chunks = triton.cdiv(n, chunk_size)
+        # Phase 1: chunk-level top-k selection
+        CHUNK_SIZE = 1024
+        num_chunks = triton.cdiv(n, CHUNK_SIZE)
         
-        # Phase 2 candidate count (must be constexpr for merge kernel)
-        # We pad total_candidates to next power of 2 for tl.arange
+        # Pad total candidates to power of 2 for merge kernel
         total_candidates = num_chunks * k
         total_candidates_padded = triton.next_power_of_2(total_candidates)
         
-        # Temp storage for chunk-level top-k (padded)
-        chunk_vals = torch.full((batch, total_candidates_padded,), -1e30, dtype=torch.float32, device=device)
-        chunk_idxs = torch.full((batch, total_candidates_padded,), -1, dtype=torch.int64, device=device)
+        chunk_vals = torch.full(
+            (batch, total_candidates_padded,), -1e30, dtype=dtype, device=device
+        )
+        chunk_idxs = torch.full(
+            (batch, total_candidates_padded,), -1, dtype=torch.int64, device=device
+        )
         
         grid1 = (batch, num_chunks)
         _chunk_select_kernel[grid1](
             x, chunk_vals, chunk_idxs,
-            N=n, K=k,
-            stride_x_row=x.stride(0),
-            CHUNK_SIZE=chunk_size,
+            N=n, stride_x_row=x.stride(0),
+            num_chunks=num_chunks,
+            K=k, CHUNK_SIZE=CHUNK_SIZE,
         )
         
-        # Phase 2: merge all chunk results
-        val_out = torch.empty(batch, k, dtype=torch.float32, device=device)
+        # Phase 2: merge
+        val_out = torch.empty(batch, k, dtype=dtype, device=device)
         idx_out = torch.empty(batch, k, dtype=torch.int64, device=device)
         
         grid2 = (batch,)
         _merge_kernel[grid2](
-            chunk_vals, chunk_idxs,
-            val_out, idx_out,
-            num_chunks=num_chunks, K_CHUNK=k, K_OUT=k,
+            chunk_vals, chunk_idxs, val_out, idx_out,
+            num_chunks=num_chunks,
             stride_val_row=val_out.stride(0),
             stride_idx_row=idx_out.stride(0),
+            K_OUT=k,
             TOTAL_CANDIDATES=total_candidates_padded,
         )
         

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -57,19 +57,21 @@
     stride_val_row,
     stride_idx_row,
     K_OUT: tl.constexpr,
-    TOTAL_CANDIDATES: tl.constexpr,
+    K_CHUNK: tl.constexpr,
+    LOAD_SIZE: tl.constexpr,
 ):
-    """Phase 2: Each program merges num_chunks*K candidates into top-K_OUT.
+    """Phase 2: Each program merges num_chunks*K_CHUNK candidates into top-K_OUT.
     
     grid = (batch,)
+    LOAD_SIZE must be >= num_chunks * K_CHUNK and a power of 2.
     """
     row_idx = tl.program_id(0)
-    k_chunk = TOTAL_CANDIDATES // num_chunks
     
-    row_base = row_idx * num_chunks * k_chunk
-    load_offs = tl.arange(0, TOTAL_CANDIDATES)
+    row_base = row_idx * num_chunks * K_CHUNK
+    load_offs = tl.arange(0, LOAD_SIZE)
+    load_mask = load_offs < (num_chunks * K_CHUNK)
     
-    cand_vals = tl.load(chunk_val_ptr + row_base + load_offs)
+    cand_vals = tl.load(chunk_val_ptr + row_base + load_offs, mask=load_mask, other=-1e30)
     
     for ki in range(K_OUT):
         m = tl.max(cand_vals)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -84,6 +84,7 @@
         tl.store(idx_out_ptr + row_idx * stride_idx_row + ki, orig_idx)
         
         cand_vals = tl.where(load_offs == arg, -1e30, cand_vals)
+        cand_vals = tl.where(load_mask, cand_vals, -1e30)
 
 
 class Model(nn.Module):    

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -133,7 +133,8 @@
             stride_val_row=val_out.stride(0),
             stride_idx_row=idx_out.stride(0),
             K_OUT=k,
-            TOTAL_CANDIDATES=total_candidates_padded,
+            K_CHUNK=k,
+            LOAD_SIZE=total_candidates_padded,
         )
         
         return val_out, idx_out

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -1,143 +1,299 @@
-"""Top-k kernel via Triton — two-phase chunk select + merge."""
+"""Top-k kernel via raw CUDA C++ — shared-memory merge tree.
+
+Design:
+  1. One block per row (batch dimension).
+  2. Each thread scans its strided portion, maintaining a top-K buffer in
+     shared memory (insertion sort, rejecting most elements once the K-th
+     threshold stabilises).
+  3. Block-wide merge tree: log2(BLOCK_SIZE) rounds of pairwise merging two
+     sorted K-element buffers into one, keeping the top K.
+  4. Thread 0 writes the final K values + indices to global memory.
+
+Register / shared-memory pressure:
+  K=64 → 128 threads, 64 KiB per shared buffer (2 buffers = 128 KiB)
+  K=32 → 256 threads, 32 KiB per buffer
+  K=16 → 256 threads, 16 KiB per buffer
+  K= 8 → 256 threads,  8 KiB per buffer
+  K= 1 → 512 threads,  2 KiB per buffer   (special-cased to plain max)
+"""
 from __future__ import annotations
+
+import os
 
 import torch
 import torch.nn as nn
-import triton
-import triton.language as tl
-
-
-@triton.jit
-def _chunk_select_kernel(
-    x_ptr,
-    chunk_val_ptr,
-    chunk_idx_ptr,
-    N,
-    stride_x_row,
-    num_chunks,
-    K: tl.constexpr,
-    CHUNK_SIZE: tl.constexpr,
-):
-    """Phase 1: Each program handles one chunk of one row.
-    
-    For a chunk of CHUNK_SIZE elements, find the top-K via repeated max.
-    
-    grid = (batch, num_chunks)
-    """
-    row_idx = tl.program_id(0)
-    chunk_idx = tl.program_id(1)
-    
-    chunk_start = chunk_idx * CHUNK_SIZE
-    x_row = x_ptr + row_idx * stride_x_row
-    out_base = (row_idx * num_chunks + chunk_idx) * K
-    
-    offs = chunk_start + tl.arange(0, CHUNK_SIZE)
-    mask = offs < N
-    
-    vals = tl.load(x_row + offs, mask=mask, other=-1e30)
-    
-    for ki in range(K):
-        m = tl.max(vals)
-        m_idx = tl.where(vals == m, offs.to(tl.int32), 0x7FFFFFFF)
-        arg = tl.min(m_idx)
-        
-        tl.store(chunk_val_ptr + out_base + ki, m)
-        tl.store(chunk_idx_ptr + out_base + ki, arg.to(tl.int64))
-        
-        vals = tl.where(offs == arg, -1e30, vals)
-
-
-@triton.jit
-def _merge_kernel(
-    chunk_val_ptr,
-    chunk_idx_ptr,
-    val_out_ptr,
-    idx_out_ptr,
-    num_chunks,
-    stride_val_row,
-    stride_idx_row,
-    K_OUT: tl.constexpr,
-    K_CHUNK: tl.constexpr,
-    LOAD_SIZE: tl.constexpr,
-):
-    """Phase 2: Each program merges num_chunks*K_CHUNK candidates into top-K_OUT.
-    
-    grid = (batch,)
-    LOAD_SIZE must be >= num_chunks * K_CHUNK and a power of 2.
-    """
-    row_idx = tl.program_id(0)
-    
-    row_base = row_idx * num_chunks * K_CHUNK
-    load_offs = tl.arange(0, LOAD_SIZE)
-    load_mask = load_offs < (num_chunks * K_CHUNK)
-    
-    cand_vals = tl.load(chunk_val_ptr + row_base + load_offs, mask=load_mask, other=-1e30)
-    
-    for ki in range(K_OUT):
-        m = tl.max(cand_vals)
-        m_idx = tl.where(cand_vals == m, load_offs, 0x7FFFFFFF).to(tl.int32)
-        arg = tl.min(m_idx)
-        
-        orig_idx = tl.load(chunk_idx_ptr + row_base + arg)
-        
-        tl.store(val_out_ptr + row_idx * stride_val_row + ki, m)
-        tl.store(idx_out_ptr + row_idx * stride_idx_row + ki, orig_idx)
-        
-        cand_vals = tl.where(load_offs == arg, -1e30, cand_vals)
-        cand_vals = tl.where(load_mask, cand_vals, -1e30)
-
-
-class Model(nn.Module):    
+from torch.utils.cpp_extension import load_inline
+
+_CUDA_SOURCE = r"""
+#include <cuda_fp16.h>
+#include <algorithm>
+#include <cmath>
+
+static constexpr float NEG_INF = -1e30f;
+
+// --------------------------------------------------------------------------
+// Generic kernel: K >= 2
+// --------------------------------------------------------------------------
+template <int K>
+struct BlockSizeTraits;
+
+template <> struct BlockSizeTraits<64> { static constexpr int BLOCK = 128; };
+template <> struct BlockSizeTraits<32> { static constexpr int BLOCK = 256; };
+template <> struct BlockSizeTraits<16> { static constexpr int BLOCK = 256; };
+template <> struct BlockSizeTraits<8>  { static constexpr int BLOCK = 256; };
+
+// Bank-conflict padding: shared arrays are laid out as [BLOCK][K+PAD].
+// For K a multiple of 32, padding=1 ensures consecutive threads in a warp
+// hit consecutive banks.
+template <int K>
+static constexpr int SPAD = (K % 32 == 0) ? 1 : 0;
+
+template <int K, int BLOCK, int PAD>
+__global__ void topk_merge_tree_kernel(
+    const float* __restrict__ x,
+    float*       __restrict__ values,
+    int64_t*     __restrict__ indices,
+    int64_t row_stride,
+    int n
+)
+{
+    extern __shared__ char s_buf[];
+    float* s_val = reinterpret_cast<float*>(s_buf);
+    // Indices stored as int32 in shared mem; cast to int64 on output.
+    int*   s_idx = reinterpret_cast<int*>(s_buf + sizeof(float) * BLOCK * (K + PAD));
+
+    int tid = threadIdx.x;
+    int bid = blockIdx.x;
+
+    // -- Initialise local buffer to -inf / -1 ------------------------------------
+    int base = tid * (K + PAD);
+    #pragma unroll
+    for (int i = 0; i < K; ++i) {
+        s_val[base + i] = NEG_INF;
+        s_idx[base + i] = -1;
+    }
+    __syncthreads();
+
+    // -- Phase 1: strided scan with insertion-sort buffer --------------------------
+    const float* row = x + bid * row_stride;
+
+    for (int i = tid; i < n; i += BLOCK) {
+        float v = row[i];
+        float kth = s_val[base + K - 1];
+        if (v > kth) {
+            // Find insertion point by scanning from back.
+            // Unroll hint helps for moderate K.
+            int pos = K - 1;
+            while (pos > 0 && s_val[base + pos - 1] < v) {
+                s_val[base + pos] = s_val[base + pos - 1];
+                s_idx[base + pos] = s_idx[base + pos - 1];
+                --pos;
+            }
+            s_val[base 

... (8308 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -5,20 +5,11 @@
   2. Each thread scans its strided portion, maintaining a top-K buffer in
      shared memory (insertion sort, rejecting most elements once the K-th
      threshold stabilises).
-  3. Block-wide merge tree: log2(BLOCK_SIZE) rounds of pairwise merging two
+  3. Block-wide merge tree: log2(BLOCK) rounds of pairwise merging two
      sorted K-element buffers into one, keeping the top K.
   4. Thread 0 writes the final K values + indices to global memory.
-
-Register / shared-memory pressure:
-  K=64 → 128 threads, 64 KiB per shared buffer (2 buffers = 128 KiB)
-  K=32 → 256 threads, 32 KiB per buffer
-  K=16 → 256 threads, 16 KiB per buffer
-  K= 8 → 256 threads,  8 KiB per buffer
-  K= 1 → 512 threads,  2 KiB per buffer   (special-cased to plain max)
 """
 from __future__ import annotations
-
-import os
 
 import torch
 import torch.nn as nn
@@ -26,28 +17,17 @@
 
 _CUDA_SOURCE = r"""
 #include <cuda_fp16.h>
-#include <algorithm>
-#include <cmath>
+#include <cstdint>
 
 static constexpr float NEG_INF = -1e30f;
 
+// Bank-conflict padding for shared arrays.
+// For K = 64 = 2*32, pad=1 avoids bank conflicts; otherwise pad=0.
+template <int K> struct spad_v { static constexpr int value = (K % 32 == 0) ? 1 : 0; };
+
 // --------------------------------------------------------------------------
 // Generic kernel: K >= 2
 // --------------------------------------------------------------------------
-template <int K>
-struct BlockSizeTraits;
-
-template <> struct BlockSizeTraits<64> { static constexpr int BLOCK = 128; };
-template <> struct BlockSizeTraits<32> { static constexpr int BLOCK = 256; };
-template <> struct BlockSizeTraits<16> { static constexpr int BLOCK = 256; };
-template <> struct BlockSizeTraits<8>  { static constexpr int BLOCK = 256; };
-
-// Bank-conflict padding: shared arrays are laid out as [BLOCK][K+PAD].
-// For K a multiple of 32, padding=1 ensures consecutive threads in a warp
-// hit consecutive banks.
-template <int K>
-static constexpr int SPAD = (K % 32 == 0) ? 1 : 0;
-
 template <int K, int BLOCK, int PAD>
 __global__ void topk_merge_tree_kernel(
     const float* __restrict__ x,
@@ -59,13 +39,13 @@
 {
     extern __shared__ char s_buf[];
     float* s_val = reinterpret_cast<float*>(s_buf);
-    // Indices stored as int32 in shared mem; cast to int64 on output.
-    int*   s_idx = reinterpret_cast<int*>(s_buf + sizeof(float) * BLOCK * (K + PAD));
+    int*   s_idx = reinterpret_cast<int*>(s_buf
+                + sizeof(float) * BLOCK * (K + PAD));
 
     int tid = threadIdx.x;
     int bid = blockIdx.x;
 
-    // -- Initialise local buffer to -inf / -1 ------------------------------------
+    // ---- Initialise local buffer to -inf / -1 --------------------------------
     int base = tid * (K + PAD);
     #pragma unroll
     for (int i = 0; i < K; ++i) {
@@ -74,15 +54,13 @@
     }
     __syncthreads();
 
-    // -- Phase 1: strided scan with insertion-sort buffer --------------------------
+    // ---- Phase 1: strided scan with insertion-sort buffer --------------------
     const float* row = x + bid * row_stride;
 
     for (int i = tid; i < n; i += BLOCK) {
         float v = row[i];
         float kth = s_val[base + K - 1];
         if (v > kth) {
-            // Find insertion point by scanning from back.
-            // Unroll hint helps for moderate K.
             int pos = K - 1;
             while (pos > 0 && s_val[base + pos - 1] < v) {
                 s_val[base + pos] = s_val[base + pos - 1];
@@ -95,18 +73,18 @@
     }
     __syncthreads();
 
-    // -- Phase 2: merge-tree reduction --------------------------------------------
+    // ---- Phase 2: merge-tree reduction ---------------------------------------
     for (int stride = BLOCK / 2; stride > 0; stride /= 2) {
         if (tid < stride) {
-            int a_base = tid * (K + PAD);
-            int b_base = (tid + stride) * (K + PAD);
-
-            // Local merge into registers to avoid shared-mem races.
+            int a_base = tid        * (K + PAD);
+            int b_base = (tid+stride) * (K + PAD);
+
+            // Local merge into registers
             float mval[K];
             int   midx[K];
             int ai = 0, bi = 0, oi = 0;
             while (oi < K) {
-                if (ai < K && (bi >= K || s_val[a_base + ai] >= s_val[b_base + bi])) {
+                if (ai < K && (bi >= K || s_val[a_base+ai] >= s_val[b_base+bi])) {
                     mval[oi] = s_val[a_base + ai];
                     midx[oi] = s_idx[a_base + ai];
                     ++ai; ++oi;
@@ -125,10 +103,10 @@
         __syncthreads();
     }
 
-    // -- Phase 3: thread 0 writes result ------------------------------------------
+    // ---- Phase 3: thread 0 writes result -------------------------------------
     if (tid == 0) {
-        float*   out_val = values  + bid * K;
-        int64_t* out_idx = indices + bid * K;
+        float*   out_val = values  + (int64_t)bid * K;
+        int64_t* out_idx = indices + (int64_t)bid * K;
         #pragma unroll
         for (int i = 0; i < K; ++i) {
             out_val[i] = s_val[i];
@@ -138,8 +116,9 @@
 }
 
 // --------------------------------------------------------------------------
-// K == 1 : special-case argmax kernel (no shared-sort buffer needed)
-// --------------------------------------------------------------------------
+// K == 1 : argmax kernel (no shared sort buffer)
+// --------------------------------------------------------------------------
+template <int BLOCK>
 __global__ void topk_argmax_kernel(
     const float* __restrict__ x,
     float*       __restrict__ values,
@@ -152,26 +131,24 @@
     int bid = blockIdx.x;
     const float* row = x + bid * row_stride;
 
-    // Thread-local max
-    float best = NEG_INF;
+    float best     = NEG_INF;
     int   be

... (6724 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -257,7 +257,8 @@
 def _compile_module():
     return load_inline(
         name='topk_cuda_ext',
-        cpp_sources='',
+        cpp_sources='#include <torch/extension.h>\n'
+                     'std::tuple<torch::Tensor, torch::Tensor> run_topk(torch::Tensor x, int k);',
         cuda_sources=_CUDA_SOURCE,
         functions=['run_topk'],
         extra_cuda_cflags=[

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -248,9 +248,8 @@
     return {values, indices};
 }
 
-PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-    m.def("run_topk", &run_topk, "Top-K selection on last dim");
-}
+std::tuple<torch::Tensor, torch::Tensor> run_topk(torch::Tensor x, int k);
+"""
 """
 
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -247,9 +247,6 @@
 
     return {values, indices};
 }
-
-std::tuple<torch::Tensor, torch::Tensor> run_topk(torch::Tensor x, int k);
-"""
 """
 
 

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -37,10 +37,9 @@
     int n
 )
 {
-    extern __shared__ char s_buf[];
-    float* s_val = reinterpret_cast<float*>(s_buf);
-    int*   s_idx = reinterpret_cast<int*>(s_buf
-                + sizeof(float) * BLOCK * (K + PAD));
+    static constexpr int STRIDE = K + PAD;
+    __shared__ float s_val[BLOCK * STRIDE];
+    __shared__ int   s_idx[BLOCK * STRIDE];
 
     int tid = threadIdx.x;
     int bid = blockIdx.x;

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -45,7 +45,7 @@
     int bid = blockIdx.x;
 
     // ---- Initialise local buffer to -inf / -1 --------------------------------
-    int base = tid * (K + PAD);
+    int base = tid * STRIDE;
     #pragma unroll
     for (int i = 0; i < K; ++i) {
         s_val[base + i] = NEG_INF;

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -75,8 +75,8 @@
     // ---- Phase 2: merge-tree reduction ---------------------------------------
     for (int stride = BLOCK / 2; stride > 0; stride /= 2) {
         if (tid < stride) {
-            int a_base = tid        * (K + PAD);
-            int b_base = (tid+stride) * (K + PAD);
+            int a_base = tid        * STRIDE;
+            int b_base = (tid+stride) * STRIDE;
 
             // Local merge into registers
             float mval[K];

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -174,46 +174,37 @@
     }
     case 8: {
         constexpr int KK=8, BL=256, PD=spad_v<8>::value;
-        size_t smem = sizeof(float)*BL*(KK+PD) + sizeof(int)*BL*(KK+PD);
-        topk_merge_tree_kernel<KK,BL,PD><<<batch, BL, smem>>>(x,values,indices,n,n);
+        topk_merge_tree_kernel<KK,BL,PD><<<batch, BL>>>(x,values,indices,n,n);
         break;
     }
     case 16: {
         constexpr int KK=16, BL=256, PD=spad_v<16>::value;
-        size_t smem = sizeof(float)*BL*(KK+PD) + sizeof(int)*BL*(KK+PD);
-        topk_merge_tree_kernel<KK,BL,PD><<<batch, BL, smem>>>(x,values,indices,n,n);
+        topk_merge_tree_kernel<KK,BL,PD><<<batch, BL>>>(x,values,indices,n,n);
         break;
     }
     case 32: {
         constexpr int KK=32, BL=256, PD=spad_v<32>::value;
-        size_t smem = sizeof(float)*BL*(KK+PD) + sizeof(int)*BL*(KK+PD);
-        topk_merge_tree_kernel<KK,BL,PD><<<batch, BL, smem>>>(x,values,indices,n,n);
+        topk_merge_tree_kernel<KK,BL,PD><<<batch, BL>>>(x,values,indices,n,n);
         break;
     }
     case 64: {
         constexpr int KK=64, BL=128, PD=spad_v<64>::value;
-        size_t smem = sizeof(float)*BL*(KK+PD) + sizeof(int)*BL*(KK+PD);
-        topk_merge_tree_kernel<KK,BL,PD><<<batch, BL, smem>>>(x,values,indices,n,n);
+        topk_merge_tree_kernel<KK,BL,PD><<<batch, BL>>>(x,values,indices,n,n);
         break;
     }
     default: {
-        // Fallback: pick smallest template K0 >= k
         if (k <= 8) {
             constexpr int KK=8, BL=256, PD=spad_v<8>::value;
-            size_t smem = sizeof(float)*BL*(KK+PD) + sizeof(int)*BL*(KK+PD);
-            topk_merge_tree_kernel<KK,BL,PD><<<batch, BL, smem>>>(x,values,indices,n,n);
+            topk_merge_tree_kernel<KK,BL,PD><<<batch, BL>>>(x,values,indices,n,n);
         } else if (k <= 16) {
             constexpr int KK=16, BL=256, PD=spad_v<16>::value;
-            size_t smem = sizeof(float)*BL*(KK+PD) + sizeof(int)*BL*(KK+PD);
-            topk_merge_tree_kernel<KK,BL,PD><<<batch, BL, smem>>>(x,values,indices,n,n);
+            topk_merge_tree_kernel<KK,BL,PD><<<batch, BL>>>(x,values,indices,n,n);
         } else if (k <= 32) {
             constexpr int KK=32, BL=256, PD=spad_v<32>::value;
-            size_t smem = sizeof(float)*BL*(KK+PD) + sizeof(int)*BL*(KK+PD);
-            topk_merge_tree_kernel<KK,BL,PD><<<batch, BL, smem>>>(x,values,indices,n,n);
+            topk_merge_tree_kernel<KK,BL,PD><<<batch, BL>>>(x,values,indices,n,n);
         } else {
             constexpr int KK=64, BL=128, PD=spad_v<64>::value;
-            size_t smem = sizeof(float)*BL*(KK+PD) + sizeof(int)*BL*(KK+PD);
-            topk_merge_tree_kernel<KK,BL,PD><<<batch, BL, smem>>>(x,values,indices,n,n);
+            topk_merge_tree_kernel<KK,BL,PD><<<batch, BL>>>(x,values,indices,n,n);
         }
         break;
     }

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -1,13 +1,25 @@
-"""Top-k kernel via raw CUDA C++ — shared-memory merge tree.
+"""Top-k kernel via raw CUDA C++ — merge-tree with register scan.
+
+Key constraint: SM120 has a 48 KB per-block shared memory limit.
+For K=64, storing all per-thread buffers in shared memory would need
+256 * 65 * 8 = ~130 KB — impossible.
+
+Solution: each thread accumulates its top-K in **registers** during the scan.
+Only at the end do threads dump to shared memory for the merge tree.
+Shared memory for the merge then only needs to hold the *final* candidates:
+  K=64, BL=128 → 128*64*8 = 64 KB.  Still too big for 48 KB.
+  K=64, BL=64  →  64*65*8 = 33 KB.  Fits.
+  K=32, BL=128 → 128*33*8 = 33 KB.  Fits.
+  K=16, BL=256 → 256*16*8 = 32 KB.  Fits.
+  K= 8, BL=256 → 256* 8*8 = 16 KB.  Fits.
+  K= 1, BL=512 → 512* 1*8 =  4 KB.  Fits.
 
 Design:
-  1. One block per row (batch dimension).
-  2. Each thread scans its strided portion, maintaining a top-K buffer in
-     shared memory (insertion sort, rejecting most elements once the K-th
-     threshold stabilises).
-  3. Block-wide merge tree: log2(BLOCK) rounds of pairwise merging two
-     sorted K-element buffers into one, keeping the top K.
-  4. Thread 0 writes the final K values + indices to global memory.
+  1. One block per row.
+  2. Each thread scans strided elements, keeping top-K in registers.
+  3. Threads dump their register buffers to shared memory.
+  4. Block-wide merge tree to extract final top-K.
+  5. Thread 0 writes output.
 """
 from __future__ import annotations
 
@@ -15,21 +27,24 @@
 import torch.nn as nn
 from torch.utils.cpp_extension import load_inline
 
+_CPP_HEADER = (
+    '#include <torch/extension.h>\n'
+    'std::tuple<torch::Tensor, torch::Tensor> run_topk(torch::Tensor x, int k);'
+)
+
 _CUDA_SOURCE = r"""
 #include <cuda_fp16.h>
 #include <cstdint>
 
 static constexpr float NEG_INF = -1e30f;
 
-// Bank-conflict padding for shared arrays.
-// For K = 64 = 2*32, pad=1 avoids bank conflicts; otherwise pad=0.
 template <int K> struct spad_v { static constexpr int value = (K % 32 == 0) ? 1 : 0; };
 
 // --------------------------------------------------------------------------
-// Generic kernel: K >= 2
+// Generic kernel (K >= 2): register scan + shared-memory merge
 // --------------------------------------------------------------------------
 template <int K, int BLOCK, int PAD>
-__global__ void topk_merge_tree_kernel(
+__global__ void topk_kernel(
     const float* __restrict__ x,
     float*       __restrict__ values,
     int64_t*     __restrict__ indices,
@@ -38,71 +53,74 @@
 )
 {
     static constexpr int STRIDE = K + PAD;
-    __shared__ float s_val[BLOCK * STRIDE];
-    __shared__ int   s_idx[BLOCK * STRIDE];
+
+    // ---- Register-allocated per-thread top-K buffer --------------------------
+    float tv[K];
+    int   ti[K];
+    #pragma unroll
+    for (int i = 0; i < K; ++i) { tv[i] = NEG_INF; ti[i] = -1; }
 
     int tid = threadIdx.x;
     int bid = blockIdx.x;
-
-    // ---- Initialise local buffer to -inf / -1 --------------------------------
+    const float* row = x + bid * row_stride;
+
+    // Strided scan: each thread processes n/BLOCK elements
+    for (int i = tid; i < n; i += BLOCK) {
+        float v = row[i];
+        if (v > tv[K - 1]) {
+            // Insertion: shift from back
+            for (int j = K - 1; j > 0; --j) {
+                if (tv[j-1] < v) { tv[j] = tv[j-1]; ti[j] = ti[j-1]; }
+                else { break; }
+            }
+            // Find the actual insertion point
+            int pos;
+            for (pos = 0; pos < K && tv[pos] >= v; ++pos) {}
+            // Shift [pos..K-2] to [pos+1..K-1]
+            for (int j = K - 1; j > pos; --j) {
+                tv[j] = tv[j-1];
+                ti[j] = ti[j-1];
+            }
+            tv[pos] = v;
+            ti[pos] = i;
+        }
+    }
+
+    // ---- Dump to shared memory for merge tree --------------------------------
+    static constexpr int SMEM_SIZE = BLOCK * STRIDE;
+    __shared__ float s_val[SMEM_SIZE];
+    __shared__ int   s_idx[SMEM_SIZE];
+
     int base = tid * STRIDE;
     #pragma unroll
     for (int i = 0; i < K; ++i) {
-        s_val[base + i] = NEG_INF;
-        s_idx[base + i] = -1;
+        s_val[base + i] = tv[i];
+        s_idx[base + i] = ti[i];
     }
     __syncthreads();
 
-    // ---- Phase 1: strided scan with insertion-sort buffer --------------------
-    const float* row = x + bid * row_stride;
-
-    for (int i = tid; i < n; i += BLOCK) {
-        float v = row[i];
-        float kth = s_val[base + K - 1];
-        if (v > kth) {
-            int pos = K - 1;
-            while (pos > 0 && s_val[base + pos - 1] < v) {
-                s_val[base + pos] = s_val[base + pos - 1];
-                s_idx[base + pos] = s_idx[base + pos - 1];
-                --pos;
-            }
-            s_val[base + pos] = v;
-            s_idx[base + pos] = i;
-        }
-    }
-    __syncthreads();
-
-    // ---- Phase 2: merge-tree reduction ---------------------------------------
+    // ---- Merge-tree reduction ------------------------------------------------
     for (int stride = BLOCK / 2; stride > 0; stride /= 2) {
         if (tid < stride) {
-            int a_base = tid        * STRIDE;
-            int b_base = (tid+stride) * STRIDE;
-
-            // Local merge into registers
-            float mval[K];
-            int   midx[K];
+            int a = tid * STRIDE;
+            int b = (tid + stride) * STRIDE;
+            float mv[K];
+            int   mi[K];
             int ai = 0, bi = 0, oi = 0;
             while (oi < K) {
-                if (ai < K && (bi >= K || s_val[a_base+ai] >= s_val[b_base+bi])) {
-                    mval[oi] = s_val[a_base + ai];
-                    midx[oi] = s_idx[a_base + ai];
-         

... (4844 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -68,18 +68,11 @@
     for (int i = tid; i < n; i += BLOCK) {
         float v = row[i];
         if (v > tv[K - 1]) {
-            // Insertion: shift from back
-            for (int j = K - 1; j > 0; --j) {
-                if (tv[j-1] < v) { tv[j] = tv[j-1]; ti[j] = ti[j-1]; }
-                else { break; }
-            }
-            // Find the actual insertion point
-            int pos;
-            for (pos = 0; pos < K && tv[pos] >= v; ++pos) {}
-            // Shift [pos..K-2] to [pos+1..K-1]
-            for (int j = K - 1; j > pos; --j) {
-                tv[j] = tv[j-1];
-                ti[j] = ti[j-1];
+            int pos = K - 1;
+            while (pos > 0 && tv[pos - 1] < v) {
+                tv[pos] = tv[pos - 1];
+                ti[pos] = ti[pos - 1];
+                --pos;
             }
             tv[pos] = v;
             ti[pos] = i;

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -1,27 +1,19 @@
-"""Top-k kernel via raw CUDA C++ — merge-tree with register scan.
-
-Key constraint: SM120 has a 48 KB per-block shared memory limit.
-For K=64, storing all per-thread buffers in shared memory would need
-256 * 65 * 8 = ~130 KB — impossible.
-
-Solution: each thread accumulates its top-K in **registers** during the scan.
-Only at the end do threads dump to shared memory for the merge tree.
-Shared memory for the merge then only needs to hold the *final* candidates:
-  K=64, BL=128 → 128*64*8 = 64 KB.  Still too big for 48 KB.
-  K=64, BL=64  →  64*65*8 = 33 KB.  Fits.
-  K=32, BL=128 → 128*33*8 = 33 KB.  Fits.
-  K=16, BL=256 → 256*16*8 = 32 KB.  Fits.
-  K= 8, BL=256 → 256* 8*8 = 16 KB.  Fits.
-  K= 1, BL=512 → 512* 1*8 =  4 KB.  Fits.
-
-Design:
-  1. One block per row.
-  2. Each thread scans strided elements, keeping top-K in registers.
-  3. Threads dump their register buffers to shared memory.
-  4. Block-wide merge tree to extract final top-K.
-  5. Thread 0 writes output.
+"""Top-k via raw CUDA — two-phase segment-select + merge.
+
+Phase 1 (high parallelism):  grid = (batch, num_segments)
+  Each block handles a segment of a row, keeps per-thread top-M in registers,
+  then thread 0 selects the segment-level top-K sequentially from shared memory.
+
+Phase 2 (merge):  grid = (batch,)
+  Each block merges num_segments * K candidates into final top-K.
+  Candidates are loaded in chunks to stay within the 48 KB shared-memory cap.
+
+This design keeps good GPU utilisation even for batch=1 and n=131k:
+  131072 / 2048 = 64 segments  →  64 blocks of 256 threads = 16k threads.
 """
 from __future__ import annotations
+
+import math
 
 import torch
 import torch.nn as nn
@@ -35,214 +27,203 @@
 _CUDA_SOURCE = r"""
 #include <cuda_fp16.h>
 #include <cstdint>
+#include <algorithm>
 
 static constexpr float NEG_INF = -1e30f;
 
-template <int K> struct spad_v { static constexpr int value = (K % 32 == 0) ? 1 : 0; };
-
-// --------------------------------------------------------------------------
-// Generic kernel (K >= 2): register scan + shared-memory merge
-// --------------------------------------------------------------------------
-template <int K, int BLOCK, int PAD>
-__global__ void topk_kernel(
+// ========== Phase 1: segment-level top-K ====================================
+// Each block handles SEG_SIZE elements and outputs the top-K from that segment.
+//
+// Strategy:
+//   1. Each thread scans its strided portion, keeping top-M in registers
+//      where M = min(K, ceil(SEG_SIZE/BLOCK)).
+//   2. Threads dump their buffers to shared memory.
+//   3. Thread 0 linear-scans shared memory, picks the segment-wide top-K
+//      (K sequential iterations of max-finding).
+//
+// For SEG_SIZE=2048, BLOCK=256 → M = 8.  Shared = 2048*8 = 16 KB.  Fits.
+
+template <int K, int BLOCK, int M>
+__global__ void segment_select_kernel(
     const float* __restrict__ x,
+    float*       __restrict__ seg_vals,   // [batch, num_segs, K]
+    int*         __restrict__ seg_idxs,   // [batch, num_segs, K]
+    int64_t row_stride,
+    int n,
+    int seg_size
+)
+{
+    int bid_row  = blockIdx.x;            // which row
+    int bid_seg  = blockIdx.y;            // which segment in the row
+    int tid      = threadIdx.x;
+
+    int seg_start = bid_seg * seg_size;
+    int seg_end   = std::min(seg_start + seg_size, n);
+    int out_base  = ((int)bid_row * (int)gridDim.y + (int)bid_seg) * K;
+
+    // Per-thread register buffer (M elements, sorted descending)
+    float tv[M];
+    int   ti[M];
+    #pragma unroll
+    for (int i = 0; i < M; ++i) { tv[i] = NEG_INF; ti[i] = -1; }
+
+    // --- Strided scan -------------------------------------------------------
+    for (int i = seg_start + tid; i < seg_end; i += BLOCK) {
+        float v = x[(int64_t)bid_row * row_stride + i];
+        if (v > tv[M-1]) {
+            int pos = M - 1;
+            while (pos > 0 && tv[pos-1] < v) { tv[pos]=tv[pos-1]; ti[pos]=ti[pos-1]; --pos; }
+            tv[pos] = v;
+            ti[pos] = i;
+        }
+    }
+
+    // --- Dump to shared memory -----------------------------------------------
+    __shared__ float s_val[BLOCK * M];
+    __shared__ int   s_idx[BLOCK * M];
+    int base = tid * M;
+    #pragma unroll
+    for (int i = 0; i < M; ++i) { s_val[base+i] = tv[i]; s_idx[base+i] = ti[i]; }
+    __syncthreads();
+
+    // --- Thread 0 selects top-K from BLOCK*M candidates ---------------------
+    if (tid == 0) {
+        int total = BLOCK * M;
+        for (int ki = 0; ki < K; ++ki) {
+            float best = NEG_INF;
+            int best_pos = -1;
+            #pragma unroll
+            for (int j = 0; j < total; ++j) {
+                if (s_val[j] > best) { best = s_val[j]; best_pos = j; }
+            }
+            seg_vals[out_base + ki] = best;
+            seg_idxs[out_base + ki] = (best_pos >= 0) ? s_idx[best_pos] : -1;
+            if (best_pos >= 0) s_val[best_pos] = NEG_INF;  // mark consumed
+        }
+    }
+}
+
+// ========== Phase 2: merge across segments ===================================
+// Each block handles one row.  Reads (num_segs*K) candidates and outputs top-K.
+// Candidates are loaded in blocks of LOAD_BLK elements to respect shared-mem cap.
+
+template <int K, int BLOCK, int LOAD_BLK>
+__global__ void merge_kernel(
+    const float* __restrict__ seg_vals,
+    const int*   __restrict__ seg_idxs,
     float*       __restrict__ values,
     int64_t*     __restrict__ indices,
-    int64_t row_stride,
-    int n
+    int num_segs
 )
 {
-    static constexpr int STRIDE = K + PAD;
-
-    // ---- Register-allocated per-thread top-K buffer --------------------------
+    int bid = blockIdx.x;
+    int tid = threadIdx.x;
+    int row_base = bid * num_segs * K;
+    int n_cand   = num_segs * K;
+
+    // Per-

... (10123 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -213,17 +213,24 @@
     auto values  = torch::empty({batch, k_out}, x.options());
     auto indices = torch::empty({batch, k_out}, x.options().dtype(torch::kInt64));
 
-    #define PHASE2(K) merge_kernel<K, 256, 256>                          \
+    #define PHASE2(K, BL) merge_kernel<(K), 256, 256>                     \
         <<<batch, 256>>>(seg_vals.data_ptr<float>(), seg_idxs.data_ptr<int>(),\
             values.data_ptr<float>(), indices.data_ptr<int64_t>(), num_segs)
-
-    if      (k_out <= 8)  { PHASE2(8);  }
-    else if (k_out <= 16) { PHASE2(16); }
-    else if (k_out <= 32) { PHASE2(32); }
-    else if (k_out <= 64) { PHASE2(64); }
-    else                  { PHASE2(64); }
+    // For K>=16 use smaller blocks to fit smem cap:
+    // K=64 BL=64  → smem=32768   K=32 BL=128 → smem=32768
+    #define PHASE2_SMALL(K, BL) merge_kernel<(K), (BL), (BL)>             \
+        <<<batch, (BL)>>>(seg_vals.data_ptr<float>(), seg_idxs.data_ptr<int>(),\
+            values.data_ptr<float>(), indices.data_ptr<int64_t>(), num_segs)
+
+    if      (k_out <= 4)  { PHASE2(4,  256); }
+    else if (k_out <= 8)  { PHASE2(8,  256); }
+    else if (k_out <= 16) { PHASE2_SMALL(16, 256); }
+    else if (k_out <= 32) { PHASE2_SMALL(32, 128); }
+    else if (k_out <= 64) { PHASE2_SMALL(64, 64);  }
+    else                  { PHASE2_SMALL(64, 64);  }
 
     #undef PHASE2
+    #undef PHASE2_SMALL
 
     return {values, indices};
 }

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -196,41 +196,47 @@
 
     dim3 grid2d(batch, num_segs);
 
-    // Choose template based on k_out
-    #define PHASE1(K) segment_select_kernel<K, BLOCK1, (M < K ? M : K)>  \
-        <<<grid2d, BLOCK1>>>(x.data_ptr<float>(),                          \
-            seg_vals.data_ptr<float>(), seg_idxs.data_ptr<int>(),         \
+    // Select the template K (round up to next supported value)
+    int K;
+    if      (k_out <= 1)  K = 1;
+    else if (k_out <= 8)  K = 8;
+    else if (k_out <= 16) K = 16;
+    else if (k_out <= 32) K = 32;
+    else                  K = 64;
+
+    // Phase 1
+    #define PHASE1(KK) segment_select_kernel<KK, BLOCK1, (M < KK ? M : KK)>  \
+        <<<grid2d, BLOCK1>>>(x.data_ptr<float>(),                            \
+            seg_vals.data_ptr<float>(), seg_idxs.data_ptr<int>(),           \
             x.stride(0), n, SEG_SIZE)
 
-    if      (k_out <= 8)  { PHASE1(8);  }
-    else if (k_out <= 16) { PHASE1(16); }
-    else if (k_out <= 32) { PHASE1(32); }
-    else if (k_out <= 64) { PHASE1(64); }
-    else                  { PHASE1(64); }
+    switch (K) {
+        case 1:  PHASE1(1);  break;
+        case 8:  PHASE1(8);  break;
+        case 16: PHASE1(16); break;
+        case 32: PHASE1(32); break;
+        default: PHASE1(64); break;
+    }
 
     #undef PHASE1
 
     auto values  = torch::empty({batch, k_out}, x.options());
     auto indices = torch::empty({batch, k_out}, x.options().dtype(torch::kInt64));
 
-    #define PHASE2(K, BL) merge_kernel<(K), 256, 256>                     \
-        <<<batch, 256>>>(seg_vals.data_ptr<float>(), seg_idxs.data_ptr<int>(),\
-            values.data_ptr<float>(), indices.data_ptr<int64_t>(), num_segs)
-    // For K>=16 use smaller blocks to fit smem cap:
-    // K=64 BL=64  → smem=32768   K=32 BL=128 → smem=32768
-    #define PHASE2_SMALL(K, BL) merge_kernel<(K), (BL), (BL)>             \
+    // Phase 2 — smaller blocks for large K to fit 48 KB smem
+    #define PHASE2_S(KK, BL) merge_kernel<(KK), (BL), (BL)>             \
         <<<batch, (BL)>>>(seg_vals.data_ptr<float>(), seg_idxs.data_ptr<int>(),\
             values.data_ptr<float>(), indices.data_ptr<int64_t>(), num_segs)
 
-    if      (k_out <= 4)  { PHASE2(4,  256); }
-    else if (k_out <= 8)  { PHASE2(8,  256); }
-    else if (k_out <= 16) { PHASE2_SMALL(16, 256); }
-    else if (k_out <= 32) { PHASE2_SMALL(32, 128); }
-    else if (k_out <= 64) { PHASE2_SMALL(64, 64);  }
-    else                  { PHASE2_SMALL(64, 64);  }
-
-    #undef PHASE2
-    #undef PHASE2_SMALL
+    switch (K) {
+        case 1:  PHASE2_S(1,  512); break;
+        case 8:  PHASE2_S(8,  256); break;
+        case 16: PHASE2_S(16, 256); break;
+        case 32: PHASE2_S(32, 128); break;
+        default: PHASE2_S(64, 64);  break;
+    }
+
+    #undef PHASE2_S
 
     return {values, indices};
 }

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -1,19 +1,19 @@
-"""Top-k via raw CUDA — two-phase segment-select + merge.
-
-Phase 1 (high parallelism):  grid = (batch, num_segments)
-  Each block handles a segment of a row, keeps per-thread top-M in registers,
-  then thread 0 selects the segment-level top-K sequentially from shared memory.
+"""Top-k via raw CUDA — bitonic sort per segment, then K-way merge.
+
+Phase 1 (high parallelism):  grid = (batch, num_segs)
+  1. All threads cooperatively load a segment into shared memory (coalesced).
+  2. Parallel bitonic sort on shared memory (descending).
+  3. Output top-K to global memory.
 
 Phase 2 (merge):  grid = (batch,)
-  Each block merges num_segments * K candidates into final top-K.
-  Candidates are loaded in chunks to stay within the 48 KB shared-memory cap.
-
-This design keeps good GPU utilisation even for batch=1 and n=131k:
-  131072 / 2048 = 64 segments  →  64 blocks of 256 threads = 16k threads.
+  1. Load all segment top-K candidates into shared memory.
+  2. Bitonic sort again, output final top-K.
+
+Bitonic sort on shared memory uses BLOCK threads for BLOCK*ELEM_PER_THREAD
+elements.  Each sub-stage thread processes ELEM_PER_THREAD compare-swaps,
+giving near-perfect shared-memory throughput.
 """
 from __future__ import annotations
-
-import math
 
 import torch
 import torch.nn as nn
@@ -31,149 +31,191 @@
 
 static constexpr float NEG_INF = -1e30f;
 
-// ========== Phase 1: segment-level top-K ====================================
-// Each block handles SEG_SIZE elements and outputs the top-K from that segment.
-//
-// Strategy:
-//   1. Each thread scans its strided portion, keeping top-M in registers
-//      where M = min(K, ceil(SEG_SIZE/BLOCK)).
-//   2. Threads dump their buffers to shared memory.
-//   3. Thread 0 linear-scans shared memory, picks the segment-wide top-K
-//      (K sequential iterations of max-finding).
-//
-// For SEG_SIZE=2048, BLOCK=256 → M = 8.  Shared = 2048*8 = 16 KB.  Fits.
-
-template <int K, int BLOCK, int M>
-__global__ void segment_select_kernel(
+// ===== Bitonic compare-swap (descending, shared memory) =====================
+__device__ __forceinline__
+void cmp_swap(float* sv, int* si, int ixj_mask, int dir, int stage_mask) {
+    int ixj = __laneid() ^ ixj_mask;
+    int dir_x = (__laneid() & stage_mask) != 0;
+    bool descending = (dir == dir_x);
+    int i = __laneid();
+    if ((descending && sv[i] < sv[ixj]) || (!descending && sv[i] > sv[ixj])) {
+        float tmp_v = sv[i]; sv[i] = sv[ixj]; sv[ixj] = tmp_v;
+        int   tmp_i = si[i]; si[i] = si[ixj]; si[ixj] = tmp_i;
+    }
+}
+
+// =========================================================================
+// Phase 1:  bitonic sort per segment  →  segment-level top-K
+// =========================================================================
+// SEG_N  = number of elements per segment (power-of-2, fits in smem)
+// BLOCK  = number of threads per block
+// ELEM   = SEG_N / BLOCK     (elements per thread for loading)
+// K_OUT  = output K
+template <int SEG_N, int BLOCK, int ELEM, int K_OUT>
+__global__ void segment_bitonic_kernel(
     const float* __restrict__ x,
-    float*       __restrict__ seg_vals,   // [batch, num_segs, K]
-    int*         __restrict__ seg_idxs,   // [batch, num_segs, K]
+    float*       __restrict__ seg_vals,   // [batch, num_segs, K_OUT]
+    int*         __restrict__ seg_idxs,   // [batch, num_segs, K_OUT]
     int64_t row_stride,
-    int n,
-    int seg_size
+    int n
 )
 {
-    int bid_row  = blockIdx.x;            // which row
-    int bid_seg  = blockIdx.y;            // which segment in the row
+    int bid_row  = blockIdx.x;
+    int bid_seg  = blockIdx.y;
     int tid      = threadIdx.x;
-
-    int seg_start = bid_seg * seg_size;
-    int seg_end   = std::min(seg_start + seg_size, n);
-    int out_base  = ((int)bid_row * (int)gridDim.y + (int)bid_seg) * K;
-
-    // Per-thread register buffer (M elements, sorted descending)
-    float tv[M];
-    int   ti[M];
+    int seg_start = bid_seg * SEG_N;
+
+    // --- Load segment into shared memory (coalesced) ------------------------
+    __shared__ float sv[SEG_N];
+    __shared__ int   si[SEG_N];
+
     #pragma unroll
-    for (int i = 0; i < M; ++i) { tv[i] = NEG_INF; ti[i] = -1; }
-
-    // --- Strided scan -------------------------------------------------------
-    for (int i = seg_start + tid; i < seg_end; i += BLOCK) {
-        float v = x[(int64_t)bid_row * row_stride + i];
-        if (v > tv[M-1]) {
-            int pos = M - 1;
-            while (pos > 0 && tv[pos-1] < v) { tv[pos]=tv[pos-1]; ti[pos]=ti[pos-1]; --pos; }
-            tv[pos] = v;
-            ti[pos] = i;
-        }
-    }
-
-    // --- Dump to shared memory -----------------------------------------------
-    __shared__ float s_val[BLOCK * M];
-    __shared__ int   s_idx[BLOCK * M];
-    int base = tid * M;
+    for (int e = 0; e < ELEM; ++e) {
+        int gidx = seg_start + tid * ELEM + e;
+        int sidx = tid * ELEM + e;
+        if (gidx < n) {
+            sv[sidx] = x[(int64_t)bid_row * row_stride + gidx];
+            si[sidx] = gidx;
+        } else {
+            sv[sidx] = NEG_INF;
+            si[sidx] = -1;
+        }
+    }
+    __syncthreads();
+
+    // --- Bitonic sort (descending) on shared memory -------------------------
+    // SEG_N must be a power of 2.
+    for (int stage = 2; stage <= SEG_N; stage <<= 1) {
+        for (int d = stage >> 1; d > 0; d >>= 1) {
+            // Each thread handles ELEM pairs
+            #pragma unroll
+            for (int e = 0; e < ELEM; ++e) {
+                int my_i = tid * ELEM + e;
+                int jx = my_i ^ d;
+                if (jx > my_i && ((my_i & stage) == 0) != ((jx & stage) == 0)) {
+                    // Compare-swap for descending
+                    // direction:

... (12574 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -1,17 +1,7 @@
 """Top-k via raw CUDA — bitonic sort per segment, then K-way merge.
 
-Phase 1 (high parallelism):  grid = (batch, num_segs)
-  1. All threads cooperatively load a segment into shared memory (coalesced).
-  2. Parallel bitonic sort on shared memory (descending).
-  3. Output top-K to global memory.
-
-Phase 2 (merge):  grid = (batch,)
-  1. Load all segment top-K candidates into shared memory.
-  2. Bitonic sort again, output final top-K.
-
-Bitonic sort on shared memory uses BLOCK threads for BLOCK*ELEM_PER_THREAD
-elements.  Each sub-stage thread processes ELEM_PER_THREAD compare-swaps,
-giving near-perfect shared-memory throughput.
+Phase 1: grid = (batch, num_segs) — parallel bitonic sort per segment
+Phase 2: grid = (batch,) — parallel bitonic sort of merged candidates
 """
 from __future__ import annotations
 
@@ -31,82 +21,46 @@
 
 static constexpr float NEG_INF = -1e30f;
 
-// ===== Bitonic compare-swap (descending, shared memory) =====================
-__device__ __forceinline__
-void cmp_swap(float* sv, int* si, int ixj_mask, int dir, int stage_mask) {
-    int ixj = __laneid() ^ ixj_mask;
-    int dir_x = (__laneid() & stage_mask) != 0;
-    bool descending = (dir == dir_x);
-    int i = __laneid();
-    if ((descending && sv[i] < sv[ixj]) || (!descending && sv[i] > sv[ixj])) {
-        float tmp_v = sv[i]; sv[i] = sv[ixj]; sv[ixj] = tmp_v;
-        int   tmp_i = si[i]; si[i] = si[ixj]; si[ixj] = tmp_i;
-    }
-}
-
-// =========================================================================
-// Phase 1:  bitonic sort per segment  →  segment-level top-K
-// =========================================================================
-// SEG_N  = number of elements per segment (power-of-2, fits in smem)
-// BLOCK  = number of threads per block
-// ELEM   = SEG_N / BLOCK     (elements per thread for loading)
-// K_OUT  = output K
-template <int SEG_N, int BLOCK, int ELEM, int K_OUT>
-__global__ void segment_bitonic_kernel(
-    const float* __restrict__ x,
-    float*       __restrict__ seg_vals,   // [batch, num_segs, K_OUT]
-    int*         __restrict__ seg_idxs,   // [batch, num_segs, K_OUT]
-    int64_t row_stride,
-    int n
+// =========================================================================
+// Bitonic sort kernel:  sorts N elements (power-of-2) in shared memory
+// =========================================================================
+// Generic: N elements, B threads, E=ELEM=B/N elements per thread, K_OUT output.
+template <int N, int B, int ELEM, int K_OUT>
+__global__ void bitonic_topk(
+    const float* __restrict__ in_v,
+    const int*   __restrict__ in_i,
+    int64_t in_row_stride,
+    float*       __restrict__ out_v,
+    int*         __restrict__ out_i,
+    int row_idx
 )
 {
-    int bid_row  = blockIdx.x;
-    int bid_seg  = blockIdx.y;
-    int tid      = threadIdx.x;
-    int seg_start = bid_seg * SEG_N;
-
-    // --- Load segment into shared memory (coalesced) ------------------------
-    __shared__ float sv[SEG_N];
-    __shared__ int   si[SEG_N];
-
+    int tid = threadIdx.x;
+
+    __shared__ float sv[N];
+    __shared__ int   si[N];
+
+    // Load input
     #pragma unroll
     for (int e = 0; e < ELEM; ++e) {
-        int gidx = seg_start + tid * ELEM + e;
-        int sidx = tid * ELEM + e;
-        if (gidx < n) {
-            sv[sidx] = x[(int64_t)bid_row * row_stride + gidx];
-            si[sidx] = gidx;
-        } else {
-            sv[sidx] = NEG_INF;
-            si[sidx] = -1;
-        }
+        int gidx = tid * ELEM + e;
+        sv[gidx] = in_v[in_row_stride * row_idx + gidx];
+        si[gidx] = in_i ? in_i[in_row_stride * row_idx + gidx] : gidx;
     }
     __syncthreads();
 
-    // --- Bitonic sort (descending) on shared memory -------------------------
-    // SEG_N must be a power of 2.
-    for (int stage = 2; stage <= SEG_N; stage <<= 1) {
+    // Bitonic sort (descending)
+    for (int stage = 2; stage <= N; stage <<= 1) {
         for (int d = stage >> 1; d > 0; d >>= 1) {
-            // Each thread handles ELEM pairs
             #pragma unroll
             for (int e = 0; e < ELEM; ++e) {
-                int my_i = tid * ELEM + e;
-                int jx = my_i ^ d;
-                if (jx > my_i && ((my_i & stage) == 0) != ((jx & stage) == 0)) {
-                    // Compare-swap for descending
-                    // direction: ascending within each half of the block
-                    bool dir_up = (my_i & (stage >> 1)) == 0;
-                    // For descending sort: swap if wrong direction
-                    if (dir_up) {
-                        if (sv[my_i] < sv[jx]) {
-                            float tv = sv[my_i]; sv[my_i] = sv[jx]; sv[jx] = tv;
-                            int   ti2 = si[my_i]; si[my_i] = si[jx]; si[jx] = ti2;
-                        }
-                    } else {
-                        if (sv[my_i] > sv[jx]) {
-                            float tv = sv[my_i]; sv[my_i] = sv[jx]; sv[jx] = tv;
-                            int   ti2 = si[my_i]; si[my_i] = si[jx]; si[jx] = ti2;
-                        }
+                int mi = tid * ELEM + e;
+                int jx = mi ^ d;
+                if (jx > mi && ((mi & stage) != 0) != ((jx & stage) != 0)) {
+                    bool dir_up = (mi & (stage >> 1)) == 0;
+                    if ((dir_up && sv[mi] < sv[jx]) || (!dir_up && sv[mi] > sv[jx])) {
+                        float tv = sv[mi]; sv[mi] = sv[jx]; sv[jx] = tv;
+                        int   ti2= si[mi]; si[mi] = si[jx]; si[jx]   = ti2;
                     }
                 }
             }
@@ -114,36 +68,93 @@
         }
     }
 
-    // Now sv[0..K_OUT-1] are the largest K_OUT elements in descending order
-    // (assuming SEG_N >= K_OUT, which it is because SEG_N >= ELEM >= 1 and K

... (12471 more chars)

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -24,60 +24,6 @@
 // =========================================================================
 // Bitonic sort kernel:  sorts N elements (power-of-2) in shared memory
 // =========================================================================
-// Generic: N elements, B threads, E=ELEM=B/N elements per thread, K_OUT output.
-template <int N, int B, int ELEM, int K_OUT>
-__global__ void bitonic_topk(
-    const float* __restrict__ in_v,
-    const int*   __restrict__ in_i,
-    int64_t in_row_stride,
-    float*       __restrict__ out_v,
-    int*         __restrict__ out_i,
-    int row_idx
-)
-{
-    int tid = threadIdx.x;
-
-    __shared__ float sv[N];
-    __shared__ int   si[N];
-
-    // Load input
-    #pragma unroll
-    for (int e = 0; e < ELEM; ++e) {
-        int gidx = tid * ELEM + e;
-        sv[gidx] = in_v[in_row_stride * row_idx + gidx];
-        si[gidx] = in_i ? in_i[in_row_stride * row_idx + gidx] : gidx;
-    }
-    __syncthreads();
-
-    // Bitonic sort (descending)
-    for (int stage = 2; stage <= N; stage <<= 1) {
-        for (int d = stage >> 1; d > 0; d >>= 1) {
-            #pragma unroll
-            for (int e = 0; e < ELEM; ++e) {
-                int mi = tid * ELEM + e;
-                int jx = mi ^ d;
-                if (jx > mi && ((mi & stage) != 0) != ((jx & stage) != 0)) {
-                    bool dir_up = (mi & (stage >> 1)) == 0;
-                    if ((dir_up && sv[mi] < sv[jx]) || (!dir_up && sv[mi] > sv[jx])) {
-                        float tv = sv[mi]; sv[mi] = sv[jx]; sv[jx] = tv;
-                        int   ti2= si[mi]; si[mi] = si[jx]; si[jx]   = ti2;
-                    }
-                }
-            }
-            __syncthreads();
-        }
-    }
-
-    // Write top K_OUT
-    if (tid < K_OUT) {
-        out_v[out_base(row_idx) + tid] = sv[tid];
-        out_i[out_base(row_idx) + tid] = si[tid];
-    }
-}
-
-// I need a different approach since out_base depends on the kernel.
-// Let me just use direct parameters instead.
-
 // =========================================================================
 // Phase 1: segment-level bitonic sort
 // =========================================================================

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -1,17 +1,4 @@
 static constexpr float NEG_INF = -1e30f;
-
-// ===== Bitonic compare-swap (descending, shared memory) =====================
-__device__ __forceinline__
-void cmp_swap(float* sv, int* si, int ixj_mask, int dir, int stage_mask) {
-    int ixj = __laneid() ^ ixj_mask;
-    int dir_x = (__laneid() & stage_mask) != 0;
-    bool descending = (dir == dir_x);
-    int i = __laneid();
-    if ((descending && sv[i] < sv[ixj]) || (!descending && sv[i] > sv[ixj])) {
-        float tmp_v = sv[i]; sv[i] = sv[ixj]; sv[ixj] = tmp_v;
-        int   tmp_i = si[i]; si[i] = si[ixj]; si[ixj] = tmp_i;
-    }
-}
 
 // =========================================================================
 // Phase 1: segment-level bitonic sort

--- a//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
+++ b//home/infatoshi/cuda/KernelBench-Hard/problems/05_topk_bitonic/solution.py
@@ -228,8 +228,7 @@
     else if (nc2 <= 512)  P2_DISPATCH(512);
     else if (nc2 <= 1024) P2_DISPATCH(1024);
     else if (nc2 <= 2048) P2_DISPATCH(2048);
-    else if (nc2 <= 4096) P2_DISPATCH(4096);
-    else                  P2_DISPATCH(8192);
+    else                  P2_DISPATCH(4096);
 
     return {values, indices};
 }