Let me analyze this task carefully. I need to optimize a top-k kernel for the RTX PRO 6000 (SM120 Blackwell architecture). The key constraints are: 1. Must use custom kernel (CUDA C++, Triton, PTX, CUTLASS) - no PyTorch topk/sort/argsort 2. Must handle multiple shapes: batch x n x k with different values 3. Output must match torch.topk(largest=True, sorted=True) contract 4. Values tolerance: 1e-4 abs/rel, indices are lenient on ties 5. Must pass `python check.py` and benchmark with `python benchmark.py` Let me start by understanding the problem, reading the reference implementation, checking the existing files, and then designing an optimized kernel. First, I should read all the relevant files to understand what's expected.

Let me understand the problem: 1. I need to write a custom top-k kernel in `solution.py` 2. Must NOT use any forbidden ops (torch.topk, torch.sort, etc.) 3. Must handle shapes: (1, 131072, 64), (64, 8192, 8), (32, 16384, 32), (16, 12000, 16), (128, 4096, 1) 4. RTX PRO 6000 is Blackwell SM120 with GDDR7 1.8 TB/s bandwidth 5. The problem is memory-bound - need to maximize bandwidth utilization 6. Must pass `python check.py` and then optimize with `python benchmark.py` This is a classic top-k problem. For memory-bound top-k, the key is: - Maximize read bandwidth (coalesced loads) - Minimize register/shared memory usage - Use efficient selection/sorting algorithms For small k (like k=1, k=8, k=16, k=32, k=64), a common approach is: 1. Load data in tiles 2. Maintain a local sorted list of top-k elements in registers 3. Merge tiles Another approach is bitonic sort in shared memory, which is what the problem name hints at. For large n (131072) with moderate k (64), and many smaller cases, I think a tile-based local top-k with warp-level or block-level reduction would work well. Let me think about what approach to take. I could: 1. Use Triton - easier to write, but might have overhead 2. Use CUDA C++ via torch.utils.cpp_extension.load_inline - more control 3. Use CUTLASS - probably overkill for this Given the Blackwell architecture (SM120), Triton might be the fastest to iterate with. But for maximum performance on top-k, a carefully written CUDA kernel might be better. Let me look at what SOTA approaches exist. I know: - For small k, a register-based bubble/insertion sort maintaining top-k as you stream through data is very efficient - Warp-level primitives can help merge partial results - For k=1, it's just argmax and can be done very efficiently Actually, let me think about this more carefully. The largest shape is batch=1, n=131072, k=64. The total input is 512KB. At 1.8 TB/s, that's ~0.28 microseconds just to read. In practice, we'd be happy with a few microseconds. For many top-k implementations: - Tile-based: each thread processes a chunk, maintains top-k, then threads merge - Bitonic sort: sort everything then take top-k (wasteful for small k) A good approach for batched top-k with small k: - Each row processed by one block (or warp) - Each thread loads elements, maintains a local top-k heap/register array - Thread-level results merged within block via shared memory - Final top-k extracted For k up to 64, each thread could maintain 64 elements in registers. That's a lot of registers though (64 * (4+4) = 512 bytes per thread). With 1024 threads, that's 512KB registers, which might be too much. Better: Use a two-level approach: 1. Each thread processes a chunk and extracts top-k (using simple sorted array insertion) 2. All threads in block write to shared memory 3. Block-level reduction to get final top-k Or use a smaller k per thread. For example, each thread maintains top-2 or top-4, then merge. Actually, for k=64, a simpler approach might be: - Use a block of 256 threads - Each thread processes n/256 elements, keeping the top 1 (or top 2) - Then do a block-level bitonic sort on the 256 (or 512) candidates to extract top 64 Wait, that's not enough. If each thread keeps only top-1 from 512 elements, the top-64 of the whole might be distributed across many threads. We need each thread to keep at least enough that the union contains the true top-64. Actually, if each thread keeps top-k, then merging all threads' results and taking top-k gives the correct answer. So each thread needs to keep top-k from its chunk. For 256 threads and n=131072, each thread handles 512 elements, keeping top-64. Then we have 256*64 = 16384 candidates. A block-level sort of 16384 elements and taking top-64 gives the answer. 16384 elements in shared memory is 64KB (fp32) + 128KB (int64 indices) = 192KB. That's a lot of shared memory but might fit on Blackwell (max ~228KB per SM? Actually Blackwell increased this). Alternatively, we can do a tree reduction: - Each warp merges 32 top-k lists into one top-k list - Then warps merge For merging two sorted top-k lists, we can use a bitonic merge or simple merge. Actually, let me think about what's most efficient. The key bottleneck is reading memory. So we want: 1. Coalesced loads (each thread reads consecutive elements or strided) 2. Minimal computation per element 3. Minimal register pressure For reading, if each thread reads consecutive elements, we get good coalescing. For example, thread i reads elements i, i+blockDim.x, i+2*blockDim.x, etc. For maintaining top-k, a common trick is to use a sorted array in registers and insert each new element. For k=64, insertion is O(k) = 64 comparisons per element. For n=131072 and 256 threads, each thread processes 512 elements, so 512 * 64 = 32768 comparisons. That's fine. But 64 fp32 + 64 int32 indices per thread = 512 bytes in registers. With 256 threads = 131KB registers per block. SM120 has plenty of registers, so this is fine. Actually, let me think about a cleaner design: **Approach: Tile-based top-k with register queues** For each row (or batch of rows): - Launch one block per row - Block has `THREADS` threads (e.g., 256 or 512) - Each thread loads `n / THREADS` elements strided by `THREADS` - Each thread maintains a local sorted array of top-k (values + indices) - After all elements processed, write local top-k to shared memory - Perform block-level bitonic sort / merge to get global top-k - Write to output For the block-level merge: - We have `THREADS * k` candidates in shared memory - Need to extract top-k - Can use a partial bitonic sort (only need top-k, not full sort) For k=64 and THREADS=256, we have 16384 candidates. A full bitonic sort of 16384 is 16384 * log2(16384)^2 / 2 ≈ 16384 * 196 / 2 ≈ 1.6M comparisons. That's actually a lot. A better merge approach: - Use a tournament tree / heap to merge THREADS sorted lists - Or use a warp-level reduction repeatedly Actually, we can do this more cleverly. Instead of each thread keeping top-k, have each thread keep top-1, then do warp reductions to find the top element, then repeat k times. But that's O(k * n) with warp divergence and not great. Better: Each thread keeps top-k. Then for merging: - Use a parallel merge sort or bitonic sort on the candidates - But only need top-k, so we can use a "partial bitonic sort" that only sorts the first k elements Actually, a simple and efficient approach for the merge: - Write all candidates to shared memory - Use a bitonic sort network on chunks, but only propagate the top elements - Or use multiple rounds of parallel compare-and-swap Let me think about this differently. There's a known efficient approach for top-k on GPUs: **Partial bitonic sort in shared memory** 1. Load elements into shared memory (values + indices) 2. Run a bitonic sort network, but after each stage, only keep the "upper" half 3. Since we only need top-k, we can discard the bottom half progressively For n candidates and needing top-k: - The bitonic sort has log2(n) stages - After stage i, we can discard elements that are definitely not in top-k - Actually, bitonic sort builds sorted sequences of size 2, 4, 8, ..., n - After building sorted sequences of size m, if m > k, we only need the top-k from each m-sequence when merging Hmm, this gets complicated. Let me think of a simpler approach. **Radix Select** For fp32, we could do a radix-based selection. This is what PyTorch's cuTOPK uses. But implementing a full radix select in CUDA is complex. **Simple approach that might be fast enough:** Use Triton! Triton has good support for block-level operations and shared memory. I can write a kernel where each program instance handles one row: ```python @triton.jit def topk_kernel(x_ptr, out_val_ptr, out_idx_ptr, batch, n, k, BLOCK_N: tl.constexpr): row_idx = tl.program_id(0) # Each row processed by one block # Load tile, maintain top-k in registers # ... ``` For the top-k maintenance in Triton, I could: - Load BLOCK_N elements at a time - Merge with current top-k using a small sorting network Actually, Triton might have limitations with dynamic k. But k is known at kernel launch time (passed as parameter), though not at compile time. Wait, the k values are: 1, 8, 16, 32, 64. These are small and powers of 2 (except 1 which is trivial). I could write a specialized kernel for each k, or use a generic kernel with a loop. Let me look at what other people do for this. I should search for fast top-k CUDA implementations. Actually, I think the most practical approach is CUDA C++ with load_inline. This gives me full control over registers, shared memory, and warp primitives. Let me design a kernel: **Kernel Design: "Streaming Top-K with Warp Merge"** For each row, assign one warp (32 threads) or one block. If one block per row: - `THREADS = 256` threads per block - Each thread loads `n / THREADS` elements strided - Each thread maintains `LOCAL_K = k` best elements in a sorted array (value, idx) - After processing all elements, threads write to shared memory - Block does a bitonic sort / merge to extract top-k To optimize the merge, I'll use a technique where: - Each warp sorts its portion using warp shuffles - Warps merge hierarchically Actually, for simplicity and performance, let me use this: **Two-pass approach:** 1. **First pass**: Each thread processes a chunk and extracts top-k. Write all thread results to global memory. So we get `batch * num_threads * k` candidates. 2. **Second pass**: Each block reads `num_threads * k` candidates and extracts top-k. This is clean and easy to implement. The first pass is perfectly parallel and memory-bound. The second pass has much less data to process. For the second pass, `num_threads * k` is at most `256 * 64 = 16384` elements per row. This is small enough to fit in shared memory and sort efficiently. Let me calculate for each shape: 1. batch=1, n=131072, k=64: First pass reads 512KB. Second pass sorts 16384 elements. 2. batch=64, n=8192, k=8: First pass reads 64 * 8192 * 4 = 2MB. Second pass sorts 64 * 256 * 8 = 512 elements per row (but 64 rows in parallel). 3. batch=32, n=16384, k=32: First pass reads 32 * 16384 * 4 = 2MB. Second pass sorts 32 * 256 * 32 = 256K elements total. 4. batch=16, n=12000, k=16: First pass reads 16 * 12000 * 4 = 768KB. Second pass sorts 16 * 256 * 16 = 65536 elements total. 5. batch=128, n=4096, k=1: First pass reads 128 * 4096 * 4 = 2MB. Second pass sorts 128 * 256 * 1 = 32768 elements total. The first pass dominates for large n. For k=1, I can use a specialized path (argmax) which is much faster. Actually, I just realized: for k=1, I can use warp-level reductions very efficiently. Each thread finds local max, then warp reduces, then block reduces. Let me design the kernel more carefully. **First Pass Kernel:** ```cuda template<int K> __global__ void local_topk(const float* x, float* out_vals, int64_t* out_idxs, int n, int k) { int row = blockIdx.x; int tid = threadIdx.x; int nthreads = blockDim.x; const float* row_x = x + row * n; // Each thread maintains top-K elements float vals[K]; int64_t idxs[K]; // Initialize with -inf #pragma unroll for (int i = 0; i < K; i++) { vals[i] = -INFINITY; idxs[i] = -1; } // Strided loop over elements for (int i = tid; i < n; i += nthreads) { float v = row_x[i]; // Insert into sorted array if v > vals[K-1] if (v > vals[K-1]) { // Shift and insert int pos = K - 1; while (pos > 0 && v > vals[pos-1]) { vals[pos] = vals[pos-1]; idxs[pos] = idxs[pos-1]; pos--; } vals[pos] = v; idxs[pos] = i; } } // Write to global memory float* row_out_vals = out_vals + row * nthreads * k; int64_t* row_out_idxs = out_idxs + row * nthreads * k; #pragma unroll for (int i = 0; i < K; i++) { row_out_vals[tid * k + i] = vals[i]; row_out_idxs[tid * k + i] = idxs[i]; } } ``` **Second Pass Kernel:** ```cuda template<int K> __global__ void merge_topk(const float* in_vals, const int64_t* in_idxs, float* out_vals, int64_t* out_idxs, int num_candidates, int k) { int row = blockIdx.x; int tid = threadIdx.x; // Load all candidates into shared memory extern __shared__ char smem[]; float* s_vals = (float*)smem; int64_t* s_idxs = (int64_t*)(smem + num_candidates * sizeof(float)); const float* row_in_vals = in_vals + row * num_candidates; const int64_t* row_in_idxs = in_idxs + row * num_candidates; for (int i = tid; i < num_candidates; i += blockDim.x) { s_vals[i] = row_in_vals[i]; s_idxs[i] = row_in_idxs[i]; } __syncthreads(); // Bitonic sort in shared memory - only need top K // ... sort s_vals/s_idxs descending, then write top K } ``` The bitonic sort for the second pass needs to handle arbitrary `num_candidates`. I'll round up to the next power of 2, pad with -inf, and run a full bitonic sort, then take top K. Actually, full bitonic sort of 16384 elements with 256 threads is quite a bit of work but should be fast in shared memory. Let me think about whether there's a better merge approach. For the second pass, the candidates are already partially sorted (each thread's contribution is sorted descending). We could use a multi-way merge or a tournament tree, but on GPU a parallel bitonic sort is usually faster. Wait, I should think about this more. For k=64 and num_candidates=16384, a full bitonic sort does ~16384 * 14 * 14 / 2 = 1.6M comparisons. With 256 threads, that's ~6400 comparisons per thread. At 1 comparison per clock (optimistic), that's ~6400 cycles = ~4 microseconds at 1.5GHz. Plus memory loads. The first pass reads 512KB in ~0.3 microseconds at full bandwidth. So the second pass might dominate! Hmm, but we don't need to fully sort. We only need top-64. Can we do a partial bitonic sort? Yes! A partial bitonic sort for top-k works by: 1. Building bitonic sequences as usual 2. But when merging two sequences of size m, if the result is guaranteed to be outside top-k, we can stop tracking it Actually, a cleaner approach: use a selection algorithm. After the first pass, we have `num_candidates` elements. We can use a "quickselect" or "bitonic select" to find the top-k without full sort. For shared memory with many threads, the most efficient approach might be: 1. Load into shared memory 2. Use a parallel reduction to find the k-th largest element (or just repeatedly extract max) 3. Then filter and sort the top-k Repeatedly extracting max k times: - Each extraction requires a parallel reduction over num_candidates elements - k = 64, num_candidates = 16384 - Each reduction is ~16384 operations, 64 reductions = ~1M operations - This is comparable to sorting, but with more synchronization A better approach for the second pass: Use a small heap or sorted array in registers for each thread, and have each thread load a subset of candidates. Actually, wait. For the second pass, `num_candidates = nthreads * k`. If nthreads=256 and k=64, that's 16384. What if I use more threads? I could use 512 or 1024 threads per block. With 512 threads and n=131072, each thread processes 256 elements. That's still fine. num_candidates = 512 * 64 = 32768. More candidates, but each thread does less work in pass 1. With 1024 threads, each thread processes 128 elements. num_candidates = 1024 * 64 = 65536. Actually, fewer elements per thread means less register pressure and better latency hiding. Let me use 512 threads. For pass 2 with 32768 candidates: A bitonic sort of 32768 (padded to 32768, which is power of 2) is straightforward. 32768 * 15 * 15 / 2 = 3.7M comparisons. With 512 threads, ~7200 per thread. Still manageable. But wait - shared memory size! 32768 floats = 128KB, plus 32768 int64 = 256KB. Total 384KB. That's too much for most GPUs (max 164-228KB per block). So I need to limit shared memory. Let me use 256 threads for pass 1: - num_candidates = 256 * 64 = 16384 - Shared memory: 16384 * 4 + 16384 * 8 = 64KB + 128KB = 192KB - This might be too much for some GPUs (max 164KB on older ones, but Blackwell SM120 should have more) Actually, I just checked: Blackwell has up to 256KB shared memory per SM. But it's partitioned per block. I should check if 192KB is safe. To be safe, I could store indices as int32 instead of int64. The problem requires int64 output, but internally I can use int32 if n <= 2^31-1 (which it is, max n=131072). Then cast to int64 at the end. With int32 indices: 64KB + 64KB = 128KB. That's safe on all modern GPUs. Or I can use a two-level merge to keep shared memory smaller. Actually, let me think about a single-pass approach instead. **Single-pass block-level top-k:** - One block per row - Each thread loads elements and maintains a local top-k - Threads cooperatively merge using warp-level operations - Final top-k is in shared memory and written out For merging, here's a clever approach: - After each thread has its local top-k, we do a tree merge - Warp 0 merges threads 0-31's top-k into a top-k - Warp 1 merges threads 32-63's top-k - ... - Then warp 0 merges the warp results - etc. Each merge of two sorted top-k lists can be done in O(k) time with one thread, or in parallel with multiple threads. For merging two sorted arrays of size k into top-k: - We can use a parallel merge where thread i finds the i-th element - Or use a bitonic merge network for 2k elements and take top k A bitonic merge of two sorted k-element lists: 1. Reverse the second list 2. Concatenate: [list1, reversed(list2)] 3. Apply bitonic merge network 4. The first k elements are the top-k For k=64, a bitonic merge of 128 elements requires log2(128) = 7 stages, with 128/2 = 64 comparators per stage = 448 comparators total. With 64 threads (2 elements per thread), that's ~7 comparators per thread. Very fast! For merging 32 threads' results within a warp: - Merge pairs (16 merges of 2 lists -> 16 lists) - Merge pairs of results (8 merges -> 8 lists) - ... (4 -> 2 -> 1) Total merges: 16 + 8 + 4 + 2 + 1 = 31 merges of varying sizes. But wait, each merge produces a top-k list from two top-k lists. The size stays k. Actually, let me think about this. After the first round, we have 16 top-k lists. After the second, 8 top-k lists, etc. After 5 rounds, 1 top-k list. For k=64, each merge is a bitonic merge of 128 elements. With 64 threads per merge (since we need parallelism), that's 64 * 31 = 1984 thread-merges. But warps are 32 threads, so we'd use 2 warps per merge. Hmm, this is getting complex. Let me consider a simpler approach. **Simplest robust approach: Two-pass with bitonic sort in pass 2** I'll use: - Pass 1: Each block processes one row. Each thread maintains a sorted top-k array in registers. Write `blockDim.x * k` candidates to global memory. - Pass 2: Each block reads candidates for one row, sorts with bitonic sort in shared memory, writes top-k. I'll tune `blockDim.x` per shape or use a fixed value. Let me use 256 threads for simplicity. For pass 2 shared memory: I'll store values and int32 indices. After sorting, write int64 indices to output. For the bitonic sort, I'll write a generic kernel that works for any power-of-2 size. For non-power-of-2 candidates, I'll pad with -inf. Actually, I realize I can also do the entire thing in one pass with a different strategy: - Use a persistent kernel that keeps top-k in shared memory and updates it - Each thread loads a chunk, finds its top-k, then all threads merge into shared memory - But this has the same complexity Let me search for some existing top-k CUDA implementations to see what's actually fast. Wait, I should also consider using Triton. Let me check if there are any Triton top-k implementations I can reference. Actually, I recall that Triton is particularly good at fused reductions and has efficient shared memory access patterns. But implementing bitonic sort in Triton might be tricky due to the SIMT programming model. Actually, I know what - let me use CUDA C++. It gives me the most control. Let me draft the kernel step by step. Actually, for k=1 (argmax), I can write a specialized highly-optimized kernel. For k=1 with 4096 elements and 128 batches, a simple warp reduction is very fast. Let me also think about memory layout. The input is (batch, n) row-major. For coalesced access: - In pass 1, thread tid should read elements at tid, tid+blockDim.x, tid+2*blockDim.x, etc. within each row. - This means consecutive threads read consecutive elements = coalesced! Yes, because row elements are contiguous in memory. Thread 0 reads index 0, thread 1 reads index 1, etc. Perfect coalescing. For pass 2, each thread reads its portion of candidates. Candidates for a row are stored as [thread0_k0, thread0_k1, ..., thread0_kK-1, thread1_k0, ...]. This is also coalesced if thread tid reads starting at tid * k. OK let me write the code. I'll create a solution.py that compiles CUDA kernels inline. For the bitonic sort in pass 2, I need to handle the fact that different shapes have different numbers of candidates. The candidates = THREADS * k, where THREADS = 256. For each shape: 1. (1, 131072, 64): candidates = 16384 2. (64, 8192, 8): candidates = 2048 3. (32, 16384, 32): candidates = 8192 4. (16, 12000, 16): candidates = 4096 5. (128, 4096, 1): candidates = 256 All are powers of 2 except shape 4 has candidates = 4096 (power of 2) and shape 1 has 16384 (power of 2). Wait, all are powers of 2 with THREADS=256! 256 * 1 = 256 256 * 8 = 2048 256 * 16 = 4096 256 * 32 = 8192 256 * 64 = 16384 Great! All powers of 2. That makes bitonic sort much easier. But wait, n=12000 for shape 4. With 256 threads, each thread processes 12000/256 = 46.875 elements. I'll just have each thread handle 47 elements (some threads will go up to 12000, others stop at 11999). That's fine. Now, for the bitonic sort, I need to write a kernel that can handle different sizes. I can use a template parameter for the candidate count, or write a generic one. Actually, I think a cleaner approach is to write a single kernel that takes the candidate count as a runtime parameter and uses a loop-based bitonic sort. The performance difference is minimal for shared memory sorts. Loop-based bitonic sort: ```cuda for (int size = 2; size <= N; size *= 2) { for (int stride = size / 2; stride > 0; stride /= 2) { // Compare-and-swap elements for (int i = tid; i < N/2; i += blockDim.x) { // calculate indices to compare int idx = i * 2 + ((i / stride) % 2) * stride; // Actually, standard bitonic sort indexing... } __syncthreads(); } } ``` Wait, standard bitonic sort indexing for descending order: For each pair of elements to compare: ```cuda int ixj = i ^ j; if (ixj > i) { if ((i & k2) == 0) { // ascending in this phase... actually for descending: if (s_vals[i] < s_vals[ixj]) swap; } else { if (s_vals[i] > s_vals[ixj]) swap; } } ``` Let me write a cleaner bitonic sort. The standard algorithm for descending: ```cuda // N must be power of 2 for (int k = 2; k <= N; k *= 2) { for (int j = k / 2; j > 0; j /= 2) { for (int i = tid; i < N; i += blockDim.x) { int ixj = i ^ j; if (ixj > i) { bool ascending = (i & k) == 0; // For descending overall, we want: // In first half of k-block: larger values go to lower indices // In second half: larger values go to higher indices... wait } } __syncthreads(); } } ``` Actually, let me be careful. For a descending sort (largest first), the standard bitonic sort works as follows: In the outer loop over k = 2, 4, 8, ..., N: - For each k, we create bitonic sequences of length k - The first k/2 elements should be in ascending order, the second k/2 in descending (or vice versa, depending on position) For a fully descending sort: - At each stage k, elements are divided into blocks of size k - Within each block, we want the first half sorted ascending and second half descending (forming a bitonic sequence) - Then we merge Actually, let me just use the well-known indexing. For descending order: ```cuda for (int k = 2; k <= N; k = k * 2) { for (int j = k / 2; j > 0; j = j / 2) { for (int i = threadIdx.x; i < N; i += blockDim.x) { int ixj = i ^ j; if (ixj > i) { if ((i & k) == 0) { if (s_vals[i] < s_vals[ixj]) { swap(s_vals[i], s_vals[ixj]); swap(s_idxs[i], s_idxs[ixj]); } } else { if (s_vals[i] > s_vals[ixj]) { swap(s_vals[i], s_vals[ixj]); swap(s_idxs[i], s_idxs[ixj]); } } } } __syncthreads(); } } ``` Wait, I need to double-check this. For k=2, j=1: - i=0: ixj=1, (0 & 2)==0, so if s[0] < s[1], swap -> s[0] > s[1] (descending pair) - i=1: ixj=0, but ixj > i is false, so skip - i=2: ixj=3, (2 & 2)!=0, so if s[2] > s[3], swap -> s[2] < s[3] (ascending pair) - i=3: ixj=2, skip After k=2: pairs are [desc, asc, desc, asc, ...] - this is correct for bitonic sequence building. For k=4, j=2: - i=0: ixj=2, (0&4)==0, if s[0] < s[2], swap -> larger to 0 - i=1: ixj=3, (1&4)==0, if s[1] < s[3], swap -> larger to 1 - i=2: ixj=0, skip - i=3: ixj=1, skip - i=4: ixj=6, (4&4)!=0, if s[4] > s[6], swap -> smaller to 4 - i=5: ixj=7, (5&4)!=0, if s[5] > s[7], swap -> smaller to 5 After j=2: within each 4-block, first half has larger elements from first&second quarters, second half has smaller. Then j=1: - i=0: ixj=1, (0&4)==0, if s[0] < s[1], swap - i=2: ixj=3, (2&4)==0, if s[2] < s[3], swap - i=4: ixj=5, (4&4)!=0, if s[4] > s[5], swap - i=6: ixj=7, (6&4)!=0, if s[6] > s[7], swap After k=4: each 4-block is fully descending. Great! So the code above produces a fully descending sort. Good. Now, each thread accesses elements i and ixj. When i and ixj are far apart, this can cause bank conflicts if they're in the same bank. But for float4 or float, adjacent threads access adjacent elements, so there might be some conflicts. Actually, with 32 banks and stride j, if j is a multiple of 32, all threads in a warp access the same bank -> major conflict. To avoid bank conflicts, we can use a different indexing scheme or use shared memory with padding. But for small sorts in shared memory, bank conflicts might not be terrible. Actually, for j >= 32, all threads in a warp read the same bank but different addresses. Wait, no - for j=32, thread i reads bank (i % 32) and bank ((i^32) % 32) = bank (i % 32). So yes, both reads are from the same bank! That's a 2-way conflict for all threads. One fix is to pad the shared memory array so that index i maps to bank (i * PADDING % 32). For example, with PADDING=33, adjacent elements are in different banks. But this makes indexing more complex. Another fix is to have each thread process two adjacent elements that are compared. But in bitonic sort, the comparison distance varies. For the sizes we're dealing with (up to 16384), the number of stages with j >= 32 is small (log2(16384) - 5 = 9 stages). Bank conflicts in these stages might not be a huge issue. Alternatively, I can use a "shuffled" indexing to avoid conflicts. But let's keep it simple first and optimize if needed. Actually, there's a known trick: store as `struct { float val; int idx; }` and align properly. Or interleave the storage. Let me just write a working version first, then profile and optimize. Now let me think about the overall structure of solution.py. ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline # CUDA source code cuda_source = """ ... kernels ... """ # Compile module = load_inline( name="topk_cuda", cpp_sources="", cuda_sources=cuda_source, functions=["topk_forward"], extra_cuda_cflags=["-O3", "--use_fast_math"], ) class Model(nn.Module): def __init__(self, batch, n, k): super().__init__() self.batch = batch self.n = n self.k = k self.register_buffer("_dummy", torch.zeros(1)) def forward(self, x): return module.topk_forward(x, self.k) ``` Wait, I need to handle the case where k varies. But for a given model instance, k is fixed. For the kernel dispatch, I'll have: - A specialized path for k=1 (argmax) - A general two-pass path for k > 1 Actually, for k=1, the two-pass approach works too (candidates = 256, sort 256 elements to get top 1). But an argmax reduction is much faster. Let me implement both. For k=1: ```cuda __global__ void topk1_kernel(const float* x, float* out_vals, int64_t* out_idxs, int n) { int row = blockIdx.x; int tid = threadIdx.x; int lane = tid % 32; int warp = tid / 32; const float* row_x = x + row * n; float max_val = -INFINITY; int max_idx = -1; for (int i = tid; i < n; i += blockDim.x) { float v = row_x[i]; if (v > max_val) { max_val = v; max_idx = i; } } // Warp reduce #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { float other_val = __shfl_down_sync(0xFFFFFFFF, max_val, offset); int other_idx = __shfl_down_sync(0xFFFFFFFF, max_idx, offset); if (other_val > max_val) { max_val = other_val; max_idx = other_idx; } } // Write warp results to shared memory __shared__ float s_vals[32]; __shared__ int s_idxs[32]; if (lane == 0) { s_vals[warp] = max_val; s_idxs[warp] = max_idx; } __syncthreads(); // Final reduction by warp 0 if (tid < 32) { float v = (tid < (blockDim.x + 31) / 32) ? s_vals[tid] : -INFINITY; int idx = (tid < (blockDim.x + 31) / 32) ? s_idxs[tid] : -1; #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { float other_val = __shfl_down_sync(0xFFFFFFFF, v, offset); int other_idx = __shfl_down_sync(0xFFFFFFFF, idx, offset); if (other_val > v) { v = other_val; idx = other_idx; } } if (tid == 0) { out_vals[row] = v; out_idxs[row] = idx; } } } ``` For k > 1, the two-pass approach. Let me write a templated version for the first pass. Actually, since k can be 8, 16, 32, 64, I can template on k for better performance. But CUDA kernels with template parameters need to be instantiated. I can use `switch(k)` to dispatch to template instantiations. Let me draft the full CUDA code. For the first pass template: ```cuda template <int K> __global__ void local_topk_kernel(const float* __restrict__ x, float* __restrict__ tmp_vals, int* __restrict__ tmp_idxs, int n, int k) { int row = blockIdx.x; int tid = threadIdx.x; int nthreads = blockDim.x; const float* row_x = x + row * n; float vals[K]; int idxs[K]; #pragma unroll for (int i = 0; i < K; i++) { vals[i] = -INFINITY; idxs[i] = -1; } for (int i = tid; i < n; i += nthreads) { float v = row_x[i]; if (v > vals[K-1]) { int pos = K - 1; while (pos > 0 && v > vals[pos-1]) { vals[pos] = vals[pos-1]; idxs[pos] = idxs[pos-1]; pos--; } vals[pos] = v; idxs[pos] = i; } } float* row_tmp_vals = tmp_vals + row * nthreads * k; int* row_tmp_idxs = tmp_idxs + row * nthreads * k; #pragma unroll for (int i = 0; i < K; i++) { row_tmp_vals[tid * k + i] = vals[i]; row_tmp_idxs[tid * k + i] = idxs[i]; } } ``` For the second pass, bitonic sort: ```cuda template <int K> __global__ void merge_topk_kernel(const float* __restrict__ tmp_vals, const int* __restrict__ tmp_idxs, float* __restrict__ out_vals, int64_t* __restrict__ out_idxs, int n_candidates, int k) { int row = blockIdx.x; int tid = threadIdx.x; extern __shared__ char smem[]; float* s_vals = (float*)smem; int* s_idxs = (int*)(smem + n_candidates * sizeof(float)); const float* row_tmp_vals = tmp_vals + row * n_candidates; const int* row_tmp_idxs = tmp_idxs + row * n_candidates; for (int i = tid; i < n_candidates; i += blockDim.x) { s_vals[i] = row_tmp_vals[i]; s_idxs[i] = row_tmp_idxs[i]; } __syncthreads(); // Bitonic sort descending for (int kk = 2; kk <= n_candidates; kk *= 2) { for (int j = kk / 2; j > 0; j /= 2) { for (int i = tid; i < n_candidates; i += blockDim.x) { int ixj = i ^ j; if (ixj > i) { if ((i & kk) == 0) { if (s_vals[i] < s_vals[ixj]) { float tmp_v = s_vals[i]; s_vals[i] = s_vals[ixj]; s_vals[ixj] = tmp_v; int tmp_i = s_idxs[i]; s_idxs[i] = s_idxs[ixj]; s_idxs[ixj] = tmp_i; } } else { if (s_vals[i] > s_vals[ixj]) { float tmp_v = s_vals[i]; s_vals[i] = s_vals[ixj]; s_vals[ixj] = tmp_v; int tmp_i = s_idxs[i]; s_idxs[i] = s_idxs[ixj]; s_idxs[ixj] = tmp_i; } } } } __syncthreads(); } } // Write top K if (tid < k) { out_vals[row * k + tid] = s_vals[tid]; out_idxs[row * k + tid] = (int64_t)s_idxs[tid]; } } ``` Wait, there's an issue with shared memory size. The extern shared memory is allocated per block, and I need to tell the kernel how much to allocate. In the PyTorch load_inline API, I can pass `shared_mem` parameter when launching, or I can use `extern __shared__` and specify at kernel launch. In PyTorch, when using `torch.cuda.ExternalResource` or custom ops, I'd typically use `cpp_extension` and write a C++ wrapper. But with `load_inline`, I need to be careful. Actually, `torch.utils.cpp_extension.load_inline` loads a module with the functions I declare. But I need a C++ function that calls the kernels with the right parameters. Let me write the C++ wrapper: ```cpp #include <torch/extension.h> #include <cuda.h> #include <cuda_runtime.h> // ... kernel declarations ... torch::Tensor topk_forward(torch::Tensor x, int k) { int batch = x.size(0); int n = x.size(1); auto vals = torch::empty({batch, k}, x.options()); auto idxs = torch::empty({batch, k}, torch::dtype(torch::kInt64).device(x.device())); const float* x_ptr = x.data_ptr<float>(); float* vals_ptr = vals.data_ptr<float>(); int64_t* idxs_ptr = idxs.data_ptr<int64_t>(); int threads = 256; if (k == 1) { topk1_kernel<<<batch, threads>>>(x_ptr, vals_ptr, idxs_ptr, n); } else { int n_candidates = threads * k; auto tmp_vals = torch::empty({batch, n_candidates}, x.options()); auto tmp_idxs = torch::empty({batch, n_candidates}, torch::dtype(torch::kInt32).device(x.device())); float* tmp_vals_ptr = tmp_vals.data_ptr<float>(); int* tmp_idxs_ptr = tmp_idxs.data_ptr<int>(); switch (k) { case 8: local_topk_kernel<8><<<batch, threads>>>(x_ptr, tmp_vals_ptr, tmp_idxs_ptr, n, k); merge_topk_kernel<8><<<batch, threads, n_candidates * (sizeof(float) + sizeof(int))>>>(tmp_vals_ptr, tmp_idxs_ptr, vals_ptr, idxs_ptr, n_candidates, k); break; case 16: local_topk_kernel<16><<<batch, threads>>>(x_ptr, tmp_vals_ptr, tmp_idxs_ptr, n, k); merge_topk_kernel<16><<<batch, threads, n_candidates * (sizeof(float) + sizeof(int))>>>(tmp_vals_ptr, tmp_idxs_ptr, vals_ptr, idxs_ptr, n_candidates, k); break; case 32: local_topk_kernel<32><<<batch, threads>>>(x_ptr, tmp_vals_ptr, tmp_idxs_ptr, n, k); merge_topk_kernel<32><<<batch, threads, n_candidates * (sizeof(float) + sizeof(int))>>>(tmp_vals_ptr, tmp_idxs_ptr, vals_ptr, idxs_ptr, n_candidates, k); break; case 64: local_topk_kernel<64><<<batch, threads>>>(x_ptr, tmp_vals_ptr, tmp_idxs_ptr, n, k); merge_topk_kernel<64><<<batch, threads, n_candidates * (sizeof(float) + sizeof(int))>>>(tmp_vals_ptr, tmp_idxs_ptr, vals_ptr, idxs_ptr, n_candidates, k); break; default: // Generic fallback - not needed for our shapes break; } } return {vals, idxs}; } ``` Wait, actually I should be returning a tuple from the C++ function. In PyTorch C++ extensions: ```cpp std::vector<torch::Tensor> topk_forward(torch::Tensor x, int k) { // ... return {vals, idxs}; } ``` And in Python: ```python module = load_inline( name="topk_cuda", cpp_sources=cpp_source, cuda_sources=cuda_source, functions=["topk_forward"], extra_cuda_cflags=["-O3", "--use_fast_math"], with_cuda=True, ) ``` Now let me think about optimizing the first pass. The insertion sort in registers might not be optimal. For small K, unrolling the insertion is important. For K=8, each insertion requires up to 8 comparisons and shifts. For 512 elements per thread (n=131072, threads=256), that's 512 * 8 = 4096 operations per thread. At 1.5GHz, that's ~2.7 microseconds. With 256 threads, the block takes ~2.7 microseconds (they run in parallel). Total for batch=1: ~2.7 microseconds + memory reads. Actually, the memory read of 512KB at 1.8 TB/s is ~0.28 microseconds. But in practice, with latency and coalescing, it's probably 1-2 microseconds. So the compute might be comparable. Wait, 256 threads can saturate memory bandwidth if they're all reading. For 512KB, with perfect coalescing, each warp reads 128 bytes per instruction. That's 512KB / 128B = 4096 warp-instructions. At 1 warp/SM/clock (roughly), with 128 warps total (4 blocks of 32 warps on a large GPU), this might take a few microseconds. Hmm, let me think about whether I can optimize further. Actually, for the first pass, instead of maintaining a sorted array, I could use a different approach: 1. Process elements in groups of K 2. Sort each group of K using a sorting network 3. Merge with the current top-K using a merge network But for K=64, this might be more complex than simple insertion. Another idea: Use vectorized loads. Each thread loads 4 floats at a time (float4). This reduces instruction count and improves memory throughput. For n=131072 and 256 threads, each thread loads 128 float4s. That's efficient. But we need to handle non-power-of-2 n (like 12000). 12000/256 = 46.875, so some threads load 47 elements, some 46. With float4, we'd need to handle remainders carefully. Actually, let me add float4 loads. It's a common optimization: ```cuda float4* x4 = (float4*)row_x; int n4 = n / 4; for (int i = tid; i < n4; i += nthreads) { float4 v4 = x4[i]; // process v4.x, v4.y, v4.z, v4.w } // Handle remaining elements for (int i = tid + n4 * 4; i < n; i += nthreads) { float v = row_x[i]; // process v } ``` This is good! Let me incorporate this. Now, for the second pass, the bitonic sort in shared memory. The issue is that with many threads accessing shared memory with varying strides, bank conflicts can occur. For j = 1: threads access i and i^1, which are adjacent -> no bank conflict (different banks) For j = 2: access i and i^2, stride 2 -> might have some conflicts depending on bank size Actually, for 32-byte banks (as in modern GPUs), float is 4 bytes, so 8 floats per bank. Stride j=1 means accessing elements 1 apart, which is 4 bytes apart. With 32 threads in a warp, thread i accesses banks (i % 8) and ((i^1) % 8). For even i, (i^1) = i+1, so banks are i%8 and (i+1)%8 -> different banks. For odd i, same but reversed. No conflicts. For j=2: i and i^2. For i=0: 0 and 2, banks 0 and 2. i=1: 1 and 3, banks 1 and 3. i=2: 2 and 0, same as i=0. Wait, but threads 0 and 2 both access banks 0 and 2. Are there conflicts? Thread 0 accesses shared[0] and shared[2]. Thread 1 accesses shared[1] and shared[3]. Thread 2 accesses shared[2] and shared[0]. Thread 3 accesses shared[3] and shared[1]. In a warp, threads 0,1,2,3 are active. Bank 0 is accessed by thread 0 (shared[0]) and thread 2 (shared[0] via shared[2]'s partner). So yes, bank conflict! For j=4: even worse. Many threads access the same banks. This is a well-known issue with bitonic sort. One solution is to pad the shared memory array so that index i maps to a different bank than index i^j for all j. A common padding is: ```cuda // Instead of s[i], use s[i + (i >> 5)] // Or use a struct with padding ``` Actually, a simpler fix for small N (like up to 16384) is to not worry too much about it. The number of stages with significant conflicts is small, and shared memory is fast anyway. But we can optimize: store values and indices in a struct and add padding: ```cuda struct Item { float val; int idx; }; __shared__ Item s_items[N + N/32]; // pad every 32 elements ``` With padding, index i maps to `i + (i >> 5)` or similar. Let me use `i + (i / 32)` as the padded index. Actually, a cleaner way: `s[i + (i >> 5)]` gives one extra element every 32 elements, shifting elements 32+ apart by 1, which changes the bank. Hmm, this is getting complicated. Let me just use the simple version first and optimize if profiling shows it's a bottleneck. Actually, wait. Let me check: what GPU do we have and what's its shared memory bandwidth? Blackwell SM120 has much higher shared memory bandwidth than previous generations. Bank conflicts in shared memory might be less of an issue. Let me proceed with the simple implementation and then profile. One more thing: for the merge pass with k=1 and candidates=256, the bitonic sort is overkill. I already have a specialized k=1 kernel. Let me also think about whether I should process multiple rows per block for small shapes. For example, batch=128, n=4096, k=1. If I launch 128 blocks of 256 threads each, that's 32768 threads. With ~100+ SMs, each SM gets ~300 threads, which is fine. But the overhead of 128 kernel launches might be noticeable. Actually, in CUDA, launching one block per row is standard. The GPU scheduler handles it well. For batch=128, n=4096, k=1, each block reads 4096 floats = 16KB. With 128 blocks, total read = 2MB. This should be fast. Let me think about another optimization: for the merge pass, I can reduce the number of threads to just what's needed. For candidates=2048 (k=8), I only need 1024 threads to cover all elements. But with 256 threads, each thread processes 8 elements in the bitonic sort loops. That's fine. Actually, using fewer threads in the merge pass means fewer thread loops but more work per thread. With 256 threads and 2048 candidates, each thread handles 8 elements per stage. With 1024 threads, each thread handles 2 elements. 256 threads might be better due to less synchronization overhead. Let me think about whether I can fuse the two passes. In some cases, especially for small n, a single pass might be faster. For n=4096, k=1: single pass argmax is optimal. For n=8192, k=8: single pass where each thread keeps top-8 and then block reduces. Actually, for n=8192 with 256 threads, each thread processes 32 elements. Keeping top-8 in registers and then doing a block-level merge of 256 * 8 = 2048 candidates. This is the two-pass approach, but the "passes" are both in registers/shared memory within the same kernel launch? No, pass 1 writes to global memory, pass 2 reads from global memory. Can I avoid the global memory write/read? Yes, by doing everything in one kernel: 1. Each thread processes its chunk and keeps top-k in registers 2. Write to shared memory 3. Bitonic sort in shared memory 4. Write top-k to output This eliminates the global memory temporary! That's a big win. For n=131072, k=64, with 256 threads: - Each thread processes 512 elements, keeps top-64 - Write 256 * 64 = 16384 candidates to shared memory (192KB with int64, 128KB with int32) - Bitonic sort 16384 in shared memory - Write top-64 to output 192KB shared memory might be too much. But 128KB (with int32 indices) might fit on Blackwell. Actually, I don't know the exact shared memory limit per block on SM120. Ampere/Hopper had 164KB configurable, but Blackwell increased it. Let me check. Actually, for safety, I can use a two-level approach within one kernel: 1. Each warp (32 threads) processes its portion and keeps top-k 2. Warp-level shuffle to find top-k per warp 3. Then block-level merge of warp results For 256 threads = 8 warps. Each warp keeps top-k. Then we merge 8 top-k lists. For k=64, merging 8 lists of 64 into top-64: - Load all 512 candidates into shared memory - Bitonic sort 512 elements (padded to 512) -> 512 * 9 * 9 / 2 = ~20K comparisons - Take top-64 Shared memory: 512 * 4 + 512 * 4 = 4KB (with int32 indices). Very small! Wait, that's brilliant! Here's the single-kernel approach: **Warp-local top-k + block-level merge** 1. Divide block into warps (8 warps of 32 threads for 256 threads) 2. Each warp processes its strided elements: - Warp w processes elements starting at w*32 + lane, then +256, +512, etc. - Each thread in warp keeps a local top-k in registers - Within warp, merge all 32 top-k lists into one warp top-k list using warp shuffles 3. After all warps have their top-k, write all warp top-k lists to shared memory 4. One warp (or all threads) bitonic sorts the candidates and extracts top-k Let me detail step 2 - warp-level merge: Each thread in the warp has top-k in registers. To merge 32 top-k lists into one: - Write all 32*k candidates to shared memory (within warp) - Bitonic sort 32*k elements (padded to next power of 2) - Each thread reads back top-k For k=64, 32*k = 2048. Shared memory per warp: 2048 * 4 + 2048 * 4 = 16KB. For 8 warps, if they use shared memory sequentially (not simultaneously), total is still 16KB. Or we can partition shared memory into 8 regions of 16KB each = 128KB total. Actually, warps in a block can share shared memory. We can have 8 * 2048 = 16384 elements total = 128KB. That's the same as before. But the key insight: instead of having each thread keep top-k and then sorting 16384 elements at the block level, we first merge within warps to reduce 256 top-k lists to 8 top-k lists. Then we sort 512 elements at the block level instead of 16384. That's 32x fewer elements to sort at the block level! Wait, let me recalculate: - Original: 256 threads * 64 = 16384 candidates -> sort 16384 - With warp merge: 8 warps * 64 = 512 candidates -> sort 512 For k=8: 256 * 8 = 2048 -> sort 2048 With warp merge: 8 * 8 = 64 -> sort 64 For k=16: 256 * 16 = 4096 -> sort 4096 With warp merge: 8 * 16 = 128 -> sort 128 For k=32: 256 * 32 = 8192 -> sort 8192 With warp merge: 8 * 32 = 256 -> sort 256 This is a HUGE reduction in sorting work! Let me do this. **Detailed warp-level merge:** Within a warp, each of 32 threads has a sorted top-k array `vals[K]`, `idxs[K]`. To merge these into a single top-k array: 1. Write all 32*K candidates to shared memory (warp-private region) 2. Bitonic sort the 32*K candidates 3. Thread 0 reads top-k and broadcasts to all threads in warp... or each thread reads its portion of top-k Actually, after sorting, we only need the top-k. We can have thread lane read `s_vals[lane]` for lane < k. Then each thread in the warp has one element of the top-k. But for the subsequent processing, it might be better if each thread has the full top-k (for k=64, that's 64 registers per thread). Or we can have thread 0 keep the full top-k and the rest just continue processing. Wait, in this design, the warp merge happens after all elements are processed. Each warp first processes all its assigned elements (strided by blockDim.x), building thread-local top-k. Then merges within the warp. Then writes warp top-k to block shared memory. Then block-level bitonic sort on warp results. For step 2 (warp merge), I can do it like this: ```cuda // Within warp, merge 32 top-k lists __shared__ float s_warp_vals[32][128]; // 8 warps * 2048 max candidates, but let's use dynamic __shared__ int s_warp_idxs[32][128]; ``` Actually, let's use a simpler approach. I'll allocate shared memory for all warps: ```cuda // Max candidates per warp = 32 * 64 = 2048 // 8 warps -> 16384 total elements // But we process warp by warp __shared__ float s_vals[WARP_MAX_CANDIDATES]; __shared__ int s_idxs[WARP_MAX_CANDIDATES]; ``` For merging within warp: ```cuda // Each thread writes its top-k for (int i = 0; i < k; i++) { s_vals[lane * k + i] = vals[i]; s_idxs[lane * k + i] = idxs[i]; } __syncwarp(); // Bitonic sort 32*k candidates int warp_candidates = 32 * k; for (int kk = 2; kk <= warp_candidates; kk *= 2) { for (int j = kk / 2; j > 0; j /= 2) { for (int i = lane; i < warp_candidates; i += 32) { int ixj = i ^ j; if (ixj > i) { if ((i & kk) == 0) { if (s_vals[i] < s_vals[ixj]) swap; } else { if (s_vals[i] > s_vals[ixj]) swap; } } } __syncwarp(); } } // Each thread reads back top-k portion for (int i = 0; i < k; i++) { vals[i] = s_vals[lane * k + i]; // Wait, top-k are at indices 0..k-1 idxs[i] = s_idxs[lane * k + i]; } ``` Wait, after sorting, the top-k are at indices 0 to k-1. I want all threads in the warp to have the top-k. But with 32 threads and k=64, each thread can hold 2 elements. Or I can have thread 0 hold all 64 in its registers and then... that's too many. Actually, I don't need each thread to hold all top-k after the warp merge. I just need the warp top-k to be written to shared memory, and then the block-level sort will pick it up. Here's a revised plan: 1. Each thread processes its chunk, keeps top-k in registers 2. Warp merge: all threads write to warp shared memory, sort, result is top-k in shared memory indices 0..k-1 3. Warp 0 (or a designated thread) copies warp top-k from shared memory to block shared memory at warp_offset 4. After all warps done, block-level bitonic sort on `num_warps * k` candidates 5. Write top-k to output For step 3, I need to coordinate. Let me think... Actually, a cleaner design: - After warp-level sort, the top-k are at s_vals[0..k-1] in the warp's shared memory region - Thread 0 of each warp copies these k elements to the block-level shared memory - All warps do this simultaneously (each to its own region) - Then __syncthreads() - Block-level bitonic sort on block_candidates = num_warps * k Shared memory layout: - Partition 1: warp working area (max 2048 floats + 2048 ints) = 16KB - Partition 2: block candidate area (8 * 64 = 512 floats + 512 ints) = 4KB Total: 20KB. Very manageable! Actually, we can reuse partition 1 as the block working area for the final sort. The block candidates (512 elements) can be stored in a separate small array. Or we can sort in place. Let me think about the control flow more carefully. ```cuda template <int K> __global__ void topk_kernel(const float* __restrict__ x, float* __restrict__ out_vals, int64_t* __restrict__ out_idxs, int n, int k) { int row = blockIdx.x; int tid = threadIdx.x; int lane = tid & 31; int warp_id = tid >> 5; int num_warps = blockDim.x >> 5; const float* row_x = x + row * n; // Thread-local top-K float vals[K]; int idxs[K]; #pragma unroll for (int i = 0; i < K; i++) { vals[i] = -INFINITY; idxs[i] = -1; } // Process elements for (int i = tid; i < n; i += blockDim.x) { float v = row_x[i]; if (v > vals[K-1]) { int pos = K - 1; while (pos > 0 && v > vals[pos-1]) { vals[pos] = vals[pos-1]; idxs[pos] = idxs[pos-1]; pos--; } vals[pos] = v; idxs[pos] = i; } } // Warp-level shared memory // Layout: [warp0_area][warp1_area]... but actually each warp uses it sequentially // Better: all warps use the same area but at different times? No, they run concurrently. // So partition by warp_id. extern __shared__ char smem[]; // Warp work area: each warp gets 32*K elements float* s_warp_vals = (float*)smem; int* s_warp_idxs = (int*)(smem + num_warps * 32 * K * sizeof(float)); float* my_warp_vals = s_warp_vals + warp_id * 32 * K; int* my_warp_idxs = s_warp_idxs + warp_id * 32 * K; // Write thread top-k to warp shared memory #pragma unroll for (int i = 0; i < K; i++) { my_warp_vals[lane * K + i] = vals[i]; my_warp_idxs[lane * K + i] = idxs[i]; } __syncwarp(); // Bitonic sort within warp int warp_cand = 32 * K; for (int kk = 2; kk <= warp_cand; kk *= 2) { for (int j = kk / 2; j > 0; j /= 2) { for (int i = lane; i < warp_cand; i += 32) { int ixj = i ^ j; if (ixj > i) { bool dir = (i & kk) == 0; // true for ascending in bitonic, but we want... // For descending sort: // In first half of k-block: larger at lower index // In second half: larger at higher index if (dir) { // i in first half if (my_warp_vals[i] < my_warp_vals[ixj]) { swap(my_warp_vals[i], my_warp_vals[ixj]); swap(my_warp_idxs[i], my_warp_idxs[ixj]); } } else { if (my_warp_vals[i] > my_warp_vals[ixj]) { swap(my_warp_vals[i], my_warp_vals[ixj]); swap(my_warp_idxs[i], my_warp_idxs[ixj]); } } } } __syncwarp(); } } // Now top-K of warp are at my_warp_vals[0..K-1] // Copy warp top-K to block candidate area // Block candidate area is after warp work areas float* s_block_vals = (float*)(smem + num_warps * 32 * K * (sizeof(float) + sizeof(int))); int* s_block_idxs = (int*)(s_block_vals + num_warps * K); if (lane < K) { s_block_vals[warp_id * K + lane] = my_warp_vals[lane]; s_block_idxs[warp_id * K + lane] = my_warp_idxs[lane]; } __syncthreads(); // Block-level bitonic sort on num_warps * K candidates int block_cand = num_warps * K; if (tid < block_cand) { // Only need block_cand threads, but all can participate for (int kk = 2; kk <= block_cand; kk *= 2) { for (int j = kk / 2; j > 0; j /= 2) { for (int i = tid; i < block_cand; i += blockDim.x) { int ixj = i ^ j; if (ixj > i) { bool dir = (i & kk) == 0; if (dir) { if (s_block_vals[i] < s_block_vals[ixj]) { swap(s_block_vals[i], s_block_vals[ixj]); swap(s_block_idxs[i], s_block_idxs[ixj]); } } else { if (s_block_vals[i] > s_block_vals[ixj]) { swap(s_block_vals[i], s_block_vals[ixj]); swap(s_block_idxs[i], s_block_idxs[ixj]); } } } } __syncthreads(); } } } // Write output if (tid < k) { out_vals[row * k + tid] = s_block_vals[tid]; out_idxs[row * k + tid] = (int64_t)s_block_idxs[tid]; } } ``` Wait, there's a subtle issue. In the block-level sort, I'm using `tid < block_cand` as a condition inside the loop, but the `__syncthreads()` requires all threads in the block to reach it. So I can't have an `if (tid < block_cand)` around the `__syncthreads()`. I need all threads to participate in the block-level sort. For small block_cand (like 64), some threads will have `i >= block_cand` and do nothing, but they'll still hit `__syncthreads()`. That's fine. Actually, let me restructure: ```cuda for (int kk = 2; kk <= block_cand; kk *= 2) { for (int j = kk / 2; j > 0; j /= 2) { for (int i = tid; i < block_cand; i += blockDim.x) { // compare and maybe swap } __syncthreads(); } } ``` This works because all threads participate in the loop (even if their `i` is out of bounds, the loop just doesn't execute for them), and all hit `__syncthreads()`. Now let me think about shared memory size: - For K=64, num_warps=8: - Warp work area: 8 * 32 * 64 * 4 = 64KB (vals) + 64KB (idxs) = 128KB - Block candidate area: 8 * 64 * 4 = 2KB (vals) + 2KB (idxs) = 4KB - Total: 132KB This is on the high side but should fit on Blackwell (which has larger shared memory). For K=32: 64KB + 32KB + 2KB + 2KB = 100KB. For K=16: 64KB + 16KB + 1KB + 1KB = 82KB. For K=8: 64KB + 8KB + 0.5KB + 0.5KB = 73KB. Hmm, 132KB might exceed the limit on some GPUs. Ampere has 164KB total but only 100KB is guaranteed for a single block without special config. Hopper has 228KB. Blackwell SM120... I think it has 256KB. To be safe, I could reduce threads per block for K=64. With 128 threads (4 warps): - Warp work area: 4 * 32 * 64 * 4 = 32KB + 32KB = 64KB - Block candidates: 4 * 64 * 4 = 1KB + 1KB = 2KB - Total: 66KB Each thread then processes 2x more elements (1024 instead of 512 for n=131072). That's fine - insertion sort of 1024 elements into top-64 is still reasonable. Alternatively, I can keep 256 threads but reduce the warp work area by having the warp merge happen in stages. For example, first merge pairs of threads (2*k), then pairs of pairs (4*k), etc. This uses less shared memory at any one time. Actually, the pairwise merge approach: - Step 1: 16 pairs of threads merge their top-k -> 16 lists of 2k candidates, sort, keep top-k. Shared: 2k per pair = 32k total. - Step 2: 8 pairs of pairs merge -> 8 lists. Shared: 2k per merge = 16k? No, each merge takes 2k from previous step and produces k. We need 2k workspace. This is essentially a tree. At each level, the workspace needed is `2 * k * (num_threads_at_level / 2)`. Level 0 (thread pairs): 2 * k * 16 = 32k elements Level 1 (4 threads): 2 * k * 8 = 16k elements ... Actually, we can reuse the same workspace. Max needed at any level is 32k elements = 128KB for k=64. That's still a lot. Let me just use fewer threads for larger k. I'll dispatch: - k <= 16: 256 threads - k == 32: 256 threads (100KB shared mem) - k == 64: 128 threads (66KB shared mem) Or better yet, I can make the block size a runtime parameter. Actually, let me also consider another approach: instead of doing warp-level merge via shared memory bitonic sort, I can use warp shuffles to do a merge of two sorted lists. Merge of two sorted k-element lists using warp shuffles: - This is complex. There are papers on parallel merge, but it's not trivial. Let me stick with the shared memory approach and just use 128 threads for k=64. For n=131072 and 128 threads, each thread processes 1024 elements. With top-64 insertion sort, that's 1024 * 64 = 65536 comparisons per thread. At 1.5GHz, ~44 microseconds. Wait, that's a lot! Actually, 65536 comparisons at 1 clock each = 65536 clocks = 43 microseconds per thread. But we're running 128 threads in parallel. The issue is that this is per block, and for batch=1, we have 1 block. So 43 microseconds just for comparisons. Plus memory reads. Hmm, that's slow. With 256 threads, each thread does 512 * 64 = 32768 comparisons = 22 microseconds. Still not great. Can I make the insertion faster? Yes! For maintaining top-k, instead of shifting elements one by one, I can use a different data structure. But in registers, we're limited. Actually, for k=64, if each thread processes 512 elements, we do 512 insertions into a size-64 sorted array. Each insertion is up to 64 shifts. What if we use a tournament/tree approach within the thread? 1. Load 64 elements 2. Sort them using a bitonic sort network (fast, no branches) 3. Load next 64 elements, sort them 4. Merge the two sorted 64-element lists, keep top 64 5. Repeat For n=512, we do 8 loads of 64 elements. Each 64-element sort: bitonic sort of 64 = 64 * 6 * 6 / 2 = 1152 comparisons. Each merge of two 64-element lists into top-64: we can use a bitonic merge of 128 elements and take top 64. Bitonic merge of 128 = 128 * 7 / 2 = 448 comparisons? Actually, a bitonic merge is just log2(128) = 7 stages with 64 comparators each = 448 comparisons. Total: 8 sorts + 7 merges = 8 * 1152 + 7 * 448 = 9216 + 3136 = 12352 comparisons. vs. original: 512 * 64 = 32768 comparisons. That's 2.6x fewer comparisons! And bitonic sort networks have no branches (or minimal branches), so they execute faster. But implementing bitonic sort for 64 elements in registers is complex. We'd need to do compare-and-swap of register pairs. For 64 elements per thread, we need 64 registers for values and 64 for indices. That's 128 registers. Plus temp registers for swapping. A GPU has 256 registers per thread (typically), so this fits but leaves little room for other variables. Actually, let me count: - 64 float registers - 64 int registers - Loop counters, pointers: ~10 registers Total: ~138 registers. Should fit in 255. The bitonic sort of 64 elements in registers: ```cuda // 64-element bitonic sort (descending) // Using CUDA's ability to swap registers #define CMP_SWAP(i, j) \ if (vals[i] < vals[j]) { \ float tmp_v = vals[i]; vals[i] = vals[j]; vals[j] = tmp_v; \ int tmp_i = idxs[i]; idxs[i] = idxs[j]; idxs[j] = tmp_i; \ } // Stage 1: pairs CMP_SWAP(0, 1); CMP_SWAP(2, 3); ... CMP_SWAP(62, 63); // Stage 2: 4-element sequences CMP_SWAP(0, 2); CMP_SWAP(1, 3); ... // ... ``` This is a lot of code but can be fully unrolled. For 64 elements, log2(64)=6 stages, with 32*6 = 192 comparators per stage on average? Actually, total comparators for bitonic sort of N is (N * log2(N) * (log2(N)+1)) / 4 = 64 * 6 * 7 / 4 = 672 comparators. Each comparator is a compare and possible swap. With registers, this is very fast. But writing all this out is tedious. We can use a loop-based approach too: ```cuda for (int k = 2; k <= 64; k *= 2) { for (int j = k / 2; j > 0; j /= 2) { for (int i = 0; i < 64; i++) { int ixj = i ^ j; if (ixj > i) { bool dir = (i & k) == 0; if (dir) { if (vals[i] < vals[ixj]) swap; } else { if (vals[i] > vals[ixj]) swap; } } } } } ``` This is much cleaner and the compiler will unroll it since 64 is a compile-time constant (if K is a template parameter and we use a constexpr). Wait, K is a template parameter, but the loop bounds use `K` which is known at compile time. The compiler should unroll this nicely. Actually, for the inner loop `i = 0; i < K; i++`, the compiler knows K at compile time and can unroll. But there are branches inside. For register operations, branches might not be terrible but they do cause some divergence within a thread (not within a warp, since all threads in the warp do the same thing). Wait, within a thread, the branch is deterministic based on `i` and `j`. The compiler might even optimize it away since `dir` is known at compile time for each unrolled iteration. Let me check: For k=2, j=1: - i=0: ixj=1, dir=(0&2)==0=true. Compare vals[0] < vals[1] - i=1: ixj=0, ixj>i false, skip - i=2: ixj=3, dir=(2&2)==0=false. Compare vals[2] > vals[3] - etc. The compiler can see all of this at compile time and generate straight-line code. So the loop-based version should compile to essentially the same code as fully unrolled CMP_SWAP macros. Great! Let me use the loop-based bitonic sort for register arrays. **Revised thread-level processing:** Instead of streaming insertion sort, I'll do: 1. Process elements in chunks of K 2. For each chunk, copy to a temporary buffer, bitonic sort it 3. Merge with current top-K using a bitonic merge 4. Keep only top-K Wait, merging two sorted K-element lists to get top-K: - Concatenate: [current_topK, new_chunk_sorted] - This is a bitonic sequence? No. [descending, descending] is not bitonic. - But [descending, ascending_reversed] is bitonic. - If I reverse the new chunk, I get a bitonic sequence of size 2K. - Then I can bitonic merge it (log2(2K) stages) and take top K. So: 1. Load K elements into `chunk_vals[K]`, `chunk_idxs[K]` 2. Bitonic sort `chunk` descending 3. Reverse `chunk` (or load in reverse order) 4. Concatenate with `topK`: [topK_vals, reversed_chunk_vals] 5. Bitonic merge of size 2K 6. First K elements are the new topK Steps 3-5 can be optimized. Actually, [descending, descending] can be turned into a bitonic sequence by reversing the second half. Then bitonic merge. For K=64, merging two sorted lists: - Create bitonic sequence: topK[0..63] + chunk[63..0] (reversed) - Bitonic merge of 128 elements: 7 stages - Result: first 64 are the merged top-64 This is efficient! But wait, we need to handle the fact that chunk might have fewer than K elements (last chunk). We can pad with -inf. Also, for loading elements, we can process them in groups of K. For n=1024 and K=64, that's 16 chunks. Each chunk is sorted (672 comparators), then merged (448 comparators). Total: 16 * 672 + 15 * 448 = 10752 + 6720 = 17472 comparators. vs. insertion sort: 1024 * 64 = 65536 comparisons. So the chunk+merge approach is ~3.7x fewer comparisons. And it's branchless (or mostly branchless). This should be significantly faster! Let me implement this. ```cuda template <int K> __device__ void bitonic_sort(float* vals, int* idxs, int n) { for (int k = 2; k <= n; k *= 2) { for (int j = k / 2; j > 0; j /= 2) { #pragma unroll for (int i = 0; i < n; i++) { int ixj = i ^ j; if (ixj > i) { bool dir = (i & k) == 0; if (dir) { if (vals[i] < vals[ixj]) { float tmp_v = vals[i]; vals[i] = vals[ixj]; vals[ixj] = tmp_v; int tmp_i = idxs[i]; idxs[i] = idxs[ixj]; idxs[j] = tmp_i; } } else { if (vals[i] > vals[ixj]) { float tmp_v = vals[i]; vals[i] = vals[ixj]; vals[ixj] = tmp_v; int tmp_i = idxs[i]; idxs[i] = idxs[ixj]; idxs[j] = tmp_i; } } } } } } } template <int K> __device__ void bitonic_merge_topk(float* topk_vals, int* topk_idxs, float* chunk_vals, int* chunk_idxs) { // chunk is already sorted descending // Reverse chunk to create bitonic sequence: [topk_desc, chunk_asc] float seq_vals[2*K]; int seq_idxs[2*K]; #pragma unroll for (int i = 0; i < K; i++) { seq_vals[i] = topk_vals[i]; seq_idxs[i] = topk_idxs[i]; seq_vals[K + i] = chunk_vals[K - 1 - i]; seq_idxs[K + i] = chunk_idxs[K - 1 - i]; } // Bitonic merge of size 2K for (int j = K; j > 0; j /= 2) { #pragma unroll for (int i = 0; i < 2*K; i++) { int ixj = i ^ j; if (ixj > i) { // For descending merge: first half of each 2j block keeps larger bool is_first_half = (i & (2*j)) == 0; if (is_first_half) { if (seq_vals[i] < seq_vals[ixj]) { swap } } else { if (seq_vals[i] > seq_vals[ixj]) { swap } } } } } // Copy top K back #pragma unroll for (int i = 0; i < K; i++) { topk_vals[i] = seq_vals[i]; topk_idxs[i] = seq_idxs[i]; } } ``` Wait, for the bitonic merge, I'm not sure about the direction logic. Let me think carefully. A bitonic sequence is one that first increases then decreases (or vice versa). [descending, ascending] is bitonic. The bitonic merge algorithm for a bitonic sequence of size N: 1. Compare elements i and i + N/2 for all i < N/2 2. If the sequence is [descending, ascending] and we want descending: - In the first half, keep the larger of each pair - In the second half, keep the smaller of each pair 3. Recursively merge each half Actually, the standard bitonic merge for descending: ```cuda for (int j = N/2; j > 0; j /= 2) { for (int i = 0; i < N; i++) { int ixj = i ^ j; if (ixj > i) { if (i < (i | j)) { // or some condition // keep larger in lower half } } } } ``` Hmm, let me think about this more carefully. In a bitonic merge: - We compare i and i + j where j starts at N/2 and halves each iteration - For a descending sort, we want the larger elements in the first half after each comparison stage Actually, for the merge of a bitonic sequence into descending order: ```cuda for (int j = N/2; j > 0; j /= 2) { for (int i = 0; i < N; i++) { int ixj = i ^ j; if (ixj > i) { // Determine direction based on position in the current sub-problem bool up = ((i / (2*j)) % 2) == 0; // first half of 2j block if (up) { if (vals[i] < vals[ixj]) swap; // larger goes to first half } else { if (vals[i] > vals[ixj]) swap; // smaller goes to second half } } } } ``` Wait, in the first iteration j = N/2. The 2j block is the whole N. So `up = ((i / N) % 2) == 0` is always true for i < N. So we always do `if (vals[i] < vals[ixj]) swap`. This means after the first stage, the larger of each pair (i, i+N/2) goes to position i. After this stage, we have two bitonic sequences of size N/2, each with larger elements in the first half. We recursively merge each half. For the second iteration j = N/4: - For i in [0, N/2): 2j block is [0, N/2) or [N/2, N). - `up = ((i / (N/2)) % 2) == 0`. For i in [0, N/2), this is true. For i in [N/2, N), this is false. - So for first half: keep larger. For second half: keep smaller. This is correct! So the merge code is: ```cuda for (int j = K; j > 0; j /= 2) { // K is N/2 where N=2K for (int i = 0; i < 2*K; i++) { int ixj = i ^ j; if (ixj > i) { bool up = ((i / (2*j)) % 2) == 0; if (up) { if (seq_vals[i] < seq_vals[ixj]) swap; } else { if (seq_vals[i] > seq_vals[ixj]) swap; } } } } ``` Wait, for j=K, 2*j = 2*K = N. For all i in [0, N), i/N = 0, so up is always true. Good. For j=K/2, 2*j = K. For i in [0, K), i/K = 0, up=true. For i in [K, 2K), i/K = 1, up=false. Good. This looks correct! Now, for the chunk processing, I need to load elements in chunks of K, sort each chunk, then merge with topk. But what about the last chunk which might be smaller than K? I can just load the remaining elements and fill the rest with -inf. Actually, an even simpler approach: process elements one at a time but use a more efficient insertion. Hmm, I think the chunk+merge approach is good but might use too many registers (2K for the merge sequence). For K=64, that's 128 float registers + 128 int registers = 256 registers. That's a lot and might spill to local memory. Let me check: for K=64, each thread needs: - topk_vals[64], topk_idxs[64] = 128 registers - chunk_vals[64], chunk_idxs[64] = 128 registers - seq_vals[128], seq_idxs[128] = 256 registers Total: 512 registers. That's way too much. GPU limit is 255 per thread. This will spill badly. I need to be more register-efficient. Alternative: instead of keeping chunk and seq in separate arrays, I can merge directly: 1. chunk_vals and chunk_idxs are loaded from global memory into registers 2. Sort chunk in place (using chunk_vals, chunk_idxs) 3. Merge with topk by comparing elements directly without a full seq array A merge of two sorted descending arrays A[0..K-1] and B[0..K-1] into C[0..K-1]: - We can do this with a parallel merge. Thread i finds the i-th largest element. - But with one thread, a sequential merge is O(K). Sequential merge: ```cuda int a = 0, b = 0; for (int i = 0; i < K; i++) { if (a < K && (b >= K || A[a] >= B[b])) { C[i] = A[a]; c_idx[i] = A_idx[a]; a++; } else { C[i] = B[b]; c_idx[i] = B_idx[b]; b++; } } ``` This is K comparisons and K assignments. For K=64, that's 64 operations. And it uses only A, B, C arrays (3 * 64 * 2 = 384 values, but they can overlap). Wait, we can do this in place or with a temporary: ```cuda float new_vals[K]; int new_idxs[K]; int a = 0, b = 0; for (int i = 0; i < K; i++) { if (a < K && (b >= K || topk_vals[a] >= chunk_vals[b])) { new_vals[i] = topk_vals[a]; new_idxs[i] = topk_idxs[a]; a++; } else { new_vals[i] = chunk_vals[b]; new_idxs[i] = chunk_idxs[b]; b++; } } // Copy back for (int i = 0; i < K; i++) { topk_vals[i] = new_vals[i]; topk_idxs[i] = new_idxs[i]; } ``` This uses 2K extra registers (new_vals, new_idxs). Total: 4K = 256 registers for K=64. Still a lot but might fit. Actually, we can optimize: instead of a separate new array, we can write directly to topk if we're careful. But since we read from topk and write to topk, we'd overwrite elements before reading them. One trick: since both arrays are sorted descending, and we want the top K, we can use a double-buffered approach with just 2 arrays that swap roles. Actually, let me just use the sequential merge. It's simple and only uses one extra array of K. Total registers: 3K values + K values temp = 4K values. For K=64: 256 values. Each value is a float or int, so ~256 registers. This is tight but might fit. For K=32: 128 registers. Very comfortable. For K=16: 64 registers. For K=8: 32 registers. So the sequential merge approach works well for smaller K and is acceptable for K=64. Let me also think about the chunk size. For sequential merge, the cost per chunk is: - Sort chunk of size C: bitonic sort = C * log2(C)^2 / 2 comparators - Merge with topk: K comparisons If C = K, sort cost = K * log2(K)^2 / 2, merge cost = K. Total per chunk: K * (log2(K)^2 / 2 + 1). For n elements, number of chunks = n / K. Total cost: n * (log2(K)^2 / 2 + 1). For K=64: log2(64) = 6, so cost per element = 6^2/2 + 1 = 19 comparisons. For n=1024: 1024 * 19 = 19456 comparisons. vs. simple insertion: 1024 * 64 = 65536 comparisons. Improvement: 3.4x. For n=512: 512 * 19 = 9728 comparisons. At 1.5GHz, ~6.5 microseconds. Much better! And register usage: chunk_vals[64], chunk_idxs[64], topk_vals[64], topk_idxs[64], new_vals[64], new_idxs[64] = 384 values. 384 registers is too many. But some of these can be optimized: - chunk is loaded from memory, sorted, then merged. We can reuse chunk storage as temp during merge. - Actually, during merge we need topk_vals, chunk_vals, and new_vals. That's 3 * 64 = 192 float registers + 192 int registers = 384 registers. That's still a lot. But we don't need to keep chunk_vals after merge. So: 1. Load 64 elements into chunk_vals, chunk_idxs 2. Sort chunk_vals, chunk_idxs in place (using chunk arrays) 3. Merge topk_vals and chunk_vals into new_vals 4. Copy new_vals to topk_vals After step 2, we still need chunk_vals for step 3. After step 3, we can discard chunk_vals. But registers are allocated statically, so the compiler will allocate space for all of them. Actually, the compiler might optimize this. If we write it carefully, the compiler might reuse registers. But it's hard to guarantee. Alternative: use a smaller chunk size, like C=32. - Sort chunk of 32: 32 * 5 * 5 / 2 = 400 comparators - Merge with topk (64): 64 comparisons - Per chunk: 464 operations - For n=1024: 32 chunks * 464 = 14848 operations. - Registers: chunk[32], topk[64], new[64] = 160 values. Much better! For K=64, chunk size 32: - Number of chunks = ceil(1024 / 32) = 32 - Total ops = 32 * (400 + 64) = 14848 - Still much better than insertion (65536) For K=32, chunk size 32: - Sort chunk: 32 * 5 * 5 / 2 = 400 - Merge: 32 comparisons - Per chunk: 432 - For n=512: 16 * 432 = 6912 This is great! Let me use chunk size = min(K, 32) or something. Actually, let me just use a fixed chunk size of 32 for all K >= 8. For K=8, chunk size 8. Wait, if K=8, chunk size should be 8. Sorting 8 elements is 8 * 3 * 3 / 2 = 36 comparators. Merging with topk of 8: 8 comparisons. For n=1024: 128 chunks * 44 = 5632 operations. vs insertion: 1024 * 8 = 8192. Still better, but not as dramatic. For K=1, we don't need any of this. OK, I think using chunk-based processing with bitonic sort of chunks and sequential merge is a good approach. Let me formalize it. Actually, let me step back. Maybe I'm overcomplicating this. Let me first try a simple insertion sort approach and see how it performs. If it's fast enough, I don't need the complexity. For the largest shape (1, 131072, 64) with 256 threads: - Each thread: 512 insertions into top-64 - 512 * 64 = 32768 comparisons - At ~1 comparison per clock (optimistic), 32768 / 1.5e9 = 22 microseconds - Memory read: 512KB at 1.8 TB/s = 0.28 microseconds (theoretical), probably 2-5 microseconds in practice - So compute dominates at ~22 microseconds Total: ~25-30 microseconds. Is that good? Let me check what PyTorch's topk does. For similar sizes, torch.topk might take 10-20 microseconds on H100. So 25-30 microseconds might be ~50-70% of torch.topk. We want to beat or match torch.topk. If I optimize to 3x fewer comparisons with chunk+merge, I could get down to ~10-15 microseconds. That would be excellent. Let me try the chunk+merge approach with chunk size = K (no extra register overhead for chunk arrays beyond what's needed for sorting). Actually, I just realized: for bitonic sort of K elements in registers, I can do it in-place on the topk array if I'm processing the first chunk. Then for subsequent chunks, I need a temporary chunk array. But there's another approach: instead of maintaining a sorted topk array and doing insertion/merge, I can: 1. Load all n/thread elements into a local buffer 2. Sort the entire buffer using a fast sorting network or quicksort 3. Take top-k For n/thread = 512, sorting 512 elements is 512 * 9 * 9 / 2 = 20736 comparisons. Then take top-64. This is actually more comparisons than insertion sort (32768 vs 20736)! Wait, 20736 is LESS than 32768. But sorting 512 elements in registers requires 512 float registers + 512 int registers = 1024 registers. That's way too many. So we can't sort the entire local buffer. What if we use a partial sort? We only need top-64 out of 512. A partial heap or tournament tree would be O(512 * log(64)) = 512 * 6 = 3072 operations. That's much better! A tournament tree / heap of size 64: - Insert 512 elements one by one - Each insertion: compare with root (smallest in heap), if larger, replace root and sift down - Sift down: at most log2(64) = 6 comparisons Total: 512 * 6 = 3072 comparisons. vs. 32768 for insertion sort. That's 10x better! But a heap in registers is tricky. The sift-down requires accessing elements at indices that depend on the tree structure. In a binary heap stored in an array: - Root at 0 - Children of i at 2i+1, 2i+2 - Parent of i at (i-1)/2 For sift-down: ```cuda int pos = 0; while (2*pos+1 < K) { int child = 2*pos+1; if (child+1 < K && vals[child+1] > vals[child]) child++; if (vals[pos] >= vals[child]) break; swap(vals[pos], vals[child]); swap(idxs[pos], idxs[child]); pos = child; } ``` This is up to log2(K) = 6 swaps, but each swap is 3 assignments, and there are branches. For K=64 and 512 elements: 512 * 6 * (some constant) = ~3000 operations. But wait, for heap-based top-k, we actually want a MIN-heap of the current top-k. The root is the smallest of the top-k. When we see a new element: - If new element > root, replace root and sift-down - If new element <= root, ignore After processing all elements, the heap contains the top-k but not sorted. We need to extract them in sorted order. We can do this by repeatedly extracting the max... but a min-heap doesn't support efficient max extraction. Alternatively, after building the min-heap of top-k, we can sort the heap array using a sort (64 * 6 * 6 / 2 = 1152 comparisons) to get them in descending order. Total: 512 * 6 + 1152 = 4224 operations for heap + sort. vs. insertion: 32768. That's 7.8x better! And register usage: just the heap array of 64 + temp variables. Very efficient! Let me implement the heap approach. ```cuda template <int K> __device__ void heap_sift_down(float* vals, int* idxs, int pos) { while (true) { int left = 2 * pos + 1; int right = 2 * pos + 2; int smallest = pos; if (left < K && vals[left] < vals[smallest]) smallest = left; if (right < K && vals[right] < vals[smallest]) smallest = right; if (smallest == pos) break; float tmp_v = vals[pos]; vals[pos] = vals[smallest]; vals[smallest] = tmp_v; int tmp_i = idxs[pos]; idxs[pos] = idxs[smallest]; idxs[smallest] = tmp_i; pos = smallest; } } template <int K> __device__ void build_min_heap(float* vals, int* idxs) { for (int i = K / 2 - 1; i >= 0; i--) { heap_sift_down<K>(vals, idxs, i); } } template <int K> __device__ void heap_sort_descending(float* vals, int* idxs) { // Convert min-heap to sorted descending for (int i = K - 1; i > 0; i--) { // Move max (actually, for min-heap, the max is somewhere in the leaves) // Hmm, heap sort from min-heap gives ascending order, not descending. // We need max-heap for descending order. } } ``` Wait, I want the final output sorted descending. If I use a min-heap, at the end I have the top-k elements but unsorted. To get descending order: - Build a max-heap instead? No, for streaming top-k we need a min-heap to quickly check if a new element belongs. - After processing, we can just sort the K elements using bitonic sort. So: 1. Initialize heap with -inf 2. For each element: if > root, replace root and sift-down 3. After all elements, sort the K elements descending For step 2, the heap is a min-heap (root is smallest of top-k). For step 3, bitonic sort of K elements descending. This is clean! Let me write it. ```cuda template <int K> __device__ void heap_replace_root(float* vals, int* idxs, float new_val, int new_idx) { vals[0] = new_val; idxs[0] = new_idx; int pos = 0; while (true) { int left = 2 * pos + 1; int right = 2 * pos + 2; int smallest = pos; if (left < K && vals[left] < vals[smallest]) smallest = left; if (right < K && vals[right] < vals[smallest]) smallest = right; if (smallest == pos) break; float tmp_v = vals[pos]; vals[pos] = vals[smallest]; vals[smallest] = tmp_v; int tmp_i = idxs[pos]; idxs[pos] = idxs[smallest]; idxs[smallest] = tmp_i; pos = smallest; } } template <int K> __device__ void bitonic_sort_desc(float* vals, int* idxs, int n) { for (int k = 2; k <= n; k *= 2) { for (int j = k / 2; j > 0; j /= 2) { for (int i = 0; i < n; i++) { int ixj = i ^ j; if (ixj > i) { bool dir = (i & k) == 0; if (dir) { if (vals[i] < vals[ixj]) { float tv = vals[i]; vals[i] = vals[ixj]; vals[ixj] = tv; int ti = idxs[i]; idxs[i] = idxs[ixj]; idxs[ixj] = ti; } } else { if (vals[i] > vals[ixj]) { float tv = vals[i]; vals[i] = vals[ixj]; vals[ixj] = tv; int ti = idxs[i]; idxs[i] = idxs[ixj]; idxs[ixj] = ti; } } } } } } } ``` Now the thread-level processing: ```cuda float vals[K]; int idxs[K]; #pragma unroll for (int i = 0; i < K; i++) { vals[i] = -INFINITY; idxs[i] = -1; } // Optional: build initial heap with first K elements // But since we stream, we can just process all elements for (int i = tid; i < n; i += blockDim.x) { float v = row_x[i]; if (v > vals[0]) { heap_replace_root<K>(vals, idxs, v, i); } } // Sort descending bitonic_sort_desc<K>(vals, idxs, K); ``` For K=64, 512 elements per thread: - Heap operations: 512 sift-downs, each up to 6 levels - Each level: 2 comparisons + 3 swaps (9 ops) - Total: 512 * 6 * 9 = 27648 operations - Plus sort: 1152 comparisons - Total: ~28800 operations Hmm, that's actually similar to insertion sort (32768). The heap has less comparison but more memory movement (swaps). In registers, swaps are fast but still take instructions. Wait, I can optimize the heap by not doing full swaps. Instead, track the position and do a single assignment at the end: ```cuda template <int K> __device__ void heap_replace_root_fast(float* vals, int* idxs, float new_val, int new_idx) { int pos = 0; while (true) { int left = 2 * pos + 1; int right = 2 * pos + 2; int smallest = pos; if (left < K && vals[left] < vals[smallest]) smallest = left; if (right < K && vals[right] < vals[smallest]) smallest = right; if (smallest == pos) break; pos = smallest; } // Shift elements along the path and insert at pos // Wait, we need to shift from root to pos, but in a heap the parent is always <= children. // Actually, if we just put new_val at root and sift down, we compare with children and move // the smaller child up. But in our case, we want to maintain min-heap property. // Standard sift down: // Start at root with new_val // While current node > smaller child: // Move smaller child up // Go to that child position // Insert new_val at final position } ``` Ah, the standard sift-down is: ```cuda float v = new_val; int idx = new_idx; int pos = 0; while (true) { int left = 2 * pos + 1; int right = 2 * pos + 2; int smallest = pos; if (left < K && vals[left] < v) smallest = left; if (right < K && vals[right] < v) smallest = right; if (smallest == pos) { vals[pos] = v; idxs[pos] = idx; break; } vals[pos] = vals[smallest]; idxs[pos] = idxs[smallest]; pos = smallest; } ``` This does at most log2(K) iterations, each with 2 comparisons and 2 assignments (moving child up). Then 2 assignments at the end. Total: 6 * (2+2) + 2 = 26 operations per sift-down. For 512 elements: 512 * 26 = 13312 operations. Plus sort: 1152. Total: ~14500. vs. insertion sort: 512 * 64 = 32768 comparisons + 32768 shifts = ~65536 operations. Heap approach: ~14500 operations. 4.5x better! And register usage: just vals[K] and idxs[K]. No extra temp arrays needed for sift-down (just a few scalar variables). This is excellent! Let me use the heap approach. Wait, there's a subtle issue. The heap requires `vals[0]` to be the smallest of the top-k. But initially, all vals are -inf. After replacing root with the first K elements, the heap property might not hold. Actually, if we start with all -inf, and we replace root with element 0, then root is still -inf (if element 0 is -inf, which it isn't... wait, element 0 is some value > -inf). Let me trace: initial vals = [-inf, -inf, ..., -inf]. Process element 0 (v=1.0): v > vals[0] (-inf), replace root with 1.0, sift down. Since children are -inf, 1.0 stays at root. vals = [1.0, -inf, ..., -inf]. Process element 1 (v=2.0): v > vals[0] (1.0), replace root with 2.0, sift down. vals = [2.0, -inf, ..., -inf]. ... Process element K-1 (v=K): vals = [K, -inf, ..., -inf]. At this point, the heap property is violated! The root is K, but many children are -inf. A min-heap requires parent <= children. Here parent > children. But wait, the sift-down algorithm I wrote compares the new value with children and moves smaller children up. If children are -inf, they'll move up, and the new value will go down. Let me trace more carefully: - vals = [1.0, -inf, -inf, ...] - Replace root with 2.0, sift down: - pos=0, left=1 (-inf), right=2 (-inf). smallest=1 (or 2). Move vals[1] (-inf) to pos 0. pos=1. - pos=1, left=3 (-inf), right=4 (-inf). Move vals[3] to pos 1. pos=3. - ... continue down to a leaf. - Insert 2.0 at leaf. After this, the heap is almost all -inf except one leaf has 2.0. That's not right either. Hmm, I see the problem. In a min-heap, when we replace the root (minimum), we should put the new element at the root and sift it down to maintain the heap property. But in our case, the new element is LARGER than the root (that's why we're replacing). So it will sift down until it finds a position where it's <= its children. Since many children might be -inf, it might go all the way to a leaf. This means after processing the first few elements, the large elements end up at the leaves, and -inf values bubble up to the root. Then subsequent elements are compared with -inf at the root, so almost all of them replace the root and sift down. This is actually correct behavior! After processing all elements, the heap will contain the K largest elements, with the smallest of those K at the root. The large elements will be scattered throughout the heap. Wait, let me verify. After processing many elements: - The root should be the K-th largest (smallest among top-K) - All elements in the heap should be >= root - This is exactly the min-heap property! Yes, this works. When we replace root with a new element that's larger, it sifts down to its correct position. When we replace with an element that's smaller, it might not sift down at all and stay at root. Actually, in our algorithm, we only replace if new element > root. So the new element is always > current root. It will sift down to some position. The old root is discarded. After processing all elements, the heap contains the K largest elements seen so far, with the smallest of them at the root. Perfect! Then we sort the heap in descending order using bitonic sort. One more thing: the sift-down algorithm I wrote compares `v` with children, not `vals[pos]` with children. This is the standard sift-down where we start with a "hole" at the root and move it down: ```cuda float v = new_val; int idx = new_idx; int pos = 0; while (true) { int left = 2 * pos + 1; int right = 2 * pos + 2; int smallest = pos; if (left < K && vals[left] < v) smallest = left; if (right < K && vals[right] < v) smallest = right; if (smallest == pos) break; vals[pos] = vals[smallest]; idxs[pos] = idxs[smallest]; pos = smallest; } vals[pos] = v; idxs[pos] = idx; ``` Yes! This is correct and efficient. Now, what about when there are duplicate values or ties? The heap handles this fine. The final sort is stable enough for our purposes (we don't care about tie-breaking for indices). Great, I think the heap approach is the way to go. Let me now design the full kernel with warp-level and block-level merging. Actually, wait. I was going to do warp merge + block merge. But with the heap approach, each thread processes its chunk and produces a sorted top-k. Then we still need to merge. For warp merge, I can use the same heap approach within the warp! Each thread has sorted top-k. We want to merge 32 sorted top-k lists into one top-k list. One approach: 1. Write all 32*k candidates to shared memory 2. One thread builds a heap of size k from all candidates 3. But that's O(32*k) operations for one thread, which is slow. Better: use a tournament tree or parallel heap. Actually, for warp merge, a bitonic sort of 32*k elements in shared memory is probably the most parallel and efficient. Let me stick with that. For block-level merge of warp results, also use bitonic sort. OK here's the final kernel design: **Kernel: `topk_kernel<K, THREADS>`** Parameters: - `K`: compile-time k - `THREADS`: threads per block (128 or 256) Shared memory layout: - Warp work area: `num_warps * 32 * K * (sizeof(float) + sizeof(int))` - Block candidate area: `num_warps * K * (sizeof(float) + sizeof(int))` Thread processing: 1. Initialize heap with -inf 2. Strided loop over elements, heap replace if larger 3. Bitonic sort thread's top-k in registers 4. Write to warp shared memory 5. Warp-level bitonic sort of 32*K candidates 6. Warp's thread 0 writes top-k to block candidate area 7. __syncthreads() 8. Block-level bitonic sort of num_warps*K candidates 9. Write top-k to output Wait, in step 6, I said "warp's thread 0 writes top-k". But we can have all threads in the warp participate. After the warp-level sort, top-k are at indices 0..K-1 in warp shared memory. I can have all threads in the block participate in writing to the block candidate area. Actually, since each warp sorts independently in its own shared memory region, after sort the top-k are at the start of each warp's region. Then I need to copy them to the block candidate area. Let me use a dedicated step: - After warp sort, if (lane < K), write `my_warp_vals[lane]` to `s_block_vals[warp_id * K + lane]` This way, all warps write simultaneously and no thread is idle. Then `__syncthreads()` and block sort. For the block sort, all threads participate. After sort, threads 0..K-1 write to output. This is clean! Now let me think about the bitonic sort in shared memory for the warp level. With 32 threads in a warp, each thread can process multiple elements. For warp_cand = 32 * K: - Number of threads = 32 - Elements per thread = K In each stage, thread `lane` processes elements `lane, lane+32, lane+64, ...` up to `warp_cand`. For K=64, warp_cand = 2048. Each thread processes 64 elements. In each `__syncwarp()`, the threads cooperate on the sort. For j=1: thread lane compares `lane*64 + i` and `lane*64 + i ^ 1` for various i. Wait, the indexing is `i` and `i ^ j`, not `lane*64 + i`. The standard loop is: ```cuda for (int i = lane; i < warp_cand; i += 32) { int ixj = i ^ j; if (ixj > i) { // compare and swap } } ``` For j=1: i=lane, lane+32, lane+64, ... ixj = i ^ 1. For even i, ixj = i+1. For odd i, ixj = i-1. So thread lane compares pairs (i, i^1) where i = lane, lane+32, etc. For even lane: compares (lane, lane+1), (lane+32, lane+33), etc. For odd lane: compares (lane-1, lane), (lane+31, lane+32), etc. But the `ixj > i` check means odd lanes skip. So only even lanes do the compare-and-swap. That's fine - half the threads are active. For j=2: i = lane, lane+32, ... ixj = i ^ 2. For lane=0: compares (0,2), (32,34), etc. For lane=1: (1,3), (33,35), etc. For lane=2: ixj=0, but 0 < 2, so skip. For lane=3: ixj=1, 1 < 3, skip. So lanes 0 and 1 are active, lanes 2 and 3 skip. This pattern continues. In each stage, half the threads are active. This is actually quite efficient. But there's still the bank conflict issue. For large j (like j=1024), thread lane accesses i=lane and ixj=lane^1024. Both map to the same bank if 1024 is a multiple of the number of banks. With 32-byte banks (8 floats per bank), 1024 / 8 = 128, so i and i^1024 are in the same bank. But different threads access different pairs, so there might not be conflicts within a warp. Wait, for j=1024 and lane=0: accesses 0 and 1024. Bank(0) = 0, Bank(1024) = 1024 % 8 = 0. Same bank! For lane=1: accesses 1 and 1025. Bank(1) = 1, Bank(1025) = 1. Same bank! For lane=2: accesses 2 and 1026. Bank(2) = 2, Bank(1026) = 2. Same bank! So ALL threads in the warp access two addresses in the SAME bank. That's a 32-way bank conflict! But wait, modern GPUs (since Volta) have memory access multicast/broadcast. If all threads in a warp access the same bank but different addresses, it's serialized. If they access the same address, it's broadcast. For j=1024, each thread accesses two DIFFERENT addresses in the same bank. So the warp needs 2 * 32 = 64 accesses to that bank, serialized into 64 transactions. That's terrible! However, the bank width is 4 bytes on older GPUs and 8 bytes on newer ones. Let me check Blackwell... I think it's 8 bytes (2 floats per bank) or maybe 32 bytes. Actually, for shared memory banks: - Pre-Volta: 4 bytes per bank, 32 banks - Volta+: 4 bytes per bank, 32 banks (but independent thread scheduling) - Some sources say Ampere/Hopper have 32-byte banks for certain access patterns Actually, I think for NVIDIA GPUs, shared memory has 32 banks, each 4 bytes wide. So consecutive 4-byte words are in consecutive banks. Two addresses in the same bank if they differ by a multiple of 128 bytes (32 * 4). For j=1024, addresses are i*4 and (i^1024)*4 bytes. The difference is 1024*4 = 4096 bytes, which is a multiple of 128. So yes, same bank. For j=512: diff = 512*4 = 2048, multiple of 128. Same bank. For j=256: diff = 1024, multiple of 128. Same bank. For j=128: diff = 512, multiple of 128. Same bank. For j=64: diff = 256, not multiple of 128. Different banks! For j=32: diff = 128, multiple of 128. Same bank. Wait, 32*4 = 128. Yes, j=32 also causes same bank. So for j = 32, 64, 128, 256, 512, 1024, most of these cause same-bank accesses! That's many stages with bank conflicts. To fix this, we can use padded indexing. Instead of `s[i]`, use `s[i + (i >> 5)]` or similar. But then `i^j` needs to map correctly too. A common padding is: ```cuda #define IDX(i) ((i) + ((i) >> 5)) ``` This adds one extra element every 32 elements, shifting elements 32+ apart by 1. So `IDX(0)=0, IDX(1)=1, ..., IDX(31)=31, IDX(32)=33, IDX(33)=34, ...`. Now, for j=32: i and i^32. - i=0: IDX(0)=0, IDX(32)=33. Banks: 0 % 32 = 0, 33 % 32 = 1. Different! - i=1: IDX(1)=1, IDX(33)=34. Banks: 1, 2. Different! For j=64: i=0: IDX(0)=0, IDX(64)=66. Banks: 0, 2. Different! i=32: IDX(32)=33, IDX(96)=99. Banks: 1, 3. Different! For j=1024: i=0: IDX(0)=0, IDX(1024)=1056. Banks: 0, 1056 % 32 = 0. Still same bank! Hmm, j=1024 corresponds to offset of 1024 elements = 4096 bytes. Even with padding, 4096 / 128 = 32 padding slots inserted, so the offset in memory is 4096 + 32 * something. Actually, for i=0: IDX(0)=0. For i=1024: there are 1024/32 = 32 padding slots inserted before index 1024. So IDX(1024) = 1024 + 32 = 1056. The byte offset is 1056 * 4 = 4224. 4224 / 128 = 33. So bank = 4224 % 128 / 4 = ... wait, bank = (byte_offset / 4) % 32 = 1056 % 32 = 0. Same as IDX(0) = 0 % 32 = 0. So j=1024 still conflicts. But j=1024 is large, and there are only log2(N) - 10 stages with j >= 1024 (for N=2048, only j=1024). For N=2048, the stages with j=1024 is just one stage. The conflict is 2-way (each thread accesses 2 addresses in the same bank), so 2 transactions per thread. With 32 threads, that's 64 transactions total. It's bad but only for one stage. For j=512: i=0, IDX(0)=0, IDX(512)=512+16=528. Bank: 528 % 32 = 16. Different from 0! Great. For j=256: i=0, IDX(0)=0, IDX(256)=256+8=264. Bank: 264 % 32 = 8. Different! For j=128: i=0, IDX(0)=0, IDX(128)=128+4=132. Bank: 132 % 32 = 4. Different! For j=64: i=0, IDX(0)=0, IDX(64)=64+2=66. Bank: 66 % 32 = 2. Different! For j=32: i=0, IDX(0)=0, IDX(32)=32+1=33. Bank: 33 % 32 = 1. Different! So with `IDX(i) = i + (i >> 5)`, the only remaining conflict is j=1024 (and j=2048, etc., powers of 2 >= 1024). For N <= 2048, only j=1024 conflicts. For N=16384, j=8192, 16384 would also conflict. This is a huge improvement! Let me use this padding. The total shared memory size with padding: for N elements, we need N + N/32 elements. For 2048: 2048 + 64 = 2112. For 16384: 16384 + 512 = 16896. Wait, for the warp-level sort, warp_cand = 32 * K. For K=64, that's 2048. With padding: 2112 floats + 2112 ints = 16.5KB + 16.5KB = 33KB per warp. For 8 warps: 264KB. That's too much! Actually, I miscalculated. Each warp uses its own shared memory region. The regions don't overlap in time (all warps run concurrently). So total shared memory = max per warp * num_warps = 33KB * 8 = 264KB if they're all active simultaneously. But wait, shared memory is per block, and all warps in the block share it. So I need 264KB. That's way too much. I can't use this much shared memory. Let me reconsider. Options: 1. Don't use padding, accept bank conflicts 2. Use fewer warps (fewer threads) 3. Process warp merges sequentially instead of in parallel For option 3: have one warp do all the warp merges one at a time. But then we're not using the full block parallelism. For option 2: use 128 threads (4 warps). Shared memory: 4 * 33KB = 132KB. Still a lot but maybe fits on Blackwell. For option 1: no padding. Let me see how bad the conflicts really are. With j=1024, each thread accesses two addresses in the same bank. On modern GPUs, shared memory can handle 2-way conflicts at half bandwidth. For 32 threads accessing the same bank, it's serialized 64-way. Wait, each thread accesses 2 addresses in the SAME bank, but different addresses. For 32 threads, that's 64 different addresses in one bank. This requires 64 separate transactions. At 1 transaction per clock, that's 64 clocks = ~43 nanoseconds. For one stage. There are log2(N) stages total, so this is a small fraction. Actually, for a warp-level sort of 2048 elements, there are log2(2048) * (log2(2048)+1) / 2 = 11 * 12 / 2 = 66 stages. Each stage has some threads active. The total time is dominated by the many small-j stages where all threads are active and there are no bank conflicts. The large-j stages with conflicts are few in number. So maybe bank conflicts aren't that bad. Let me just use un-padded shared memory and see. Actually, let me think about this from a different angle. For the warp-level sort, with 32 threads and 2048 elements: - j=1: 1024 comparators, 32 threads active -> 32 iterations, no conflicts - j=2: 1024 comparators, 16 threads active -> 16 iterations, some conflicts - ... - j=1024: 1024 comparators, 1 thread active -> 1 iteration, conflict The total "work" in terms of comparator executions is 2048 * 11 * 12 / 4 = 67584 comparators. At 1 comparator per thread per clock, with 32 threads, that's ~2112 clocks = 1.4 microseconds. Even with some slowdown from bank conflicts, it's still very fast. Similarly for block-level sort of 512 elements with 128 threads: - 512 * 9 * 10 / 4 = 11520 comparators - With 128 threads, ~90 clocks = 60 nanoseconds. Very fast! So the sorting time is negligible compared to the heap processing time (~10+ microseconds). Bank conflicts in shared memory sort are not the bottleneck. Let me not worry about padding and just use simple indexing. Now, let me also consider vectorized loads for the element processing. Each thread can load float4 (4 floats at a time) from global memory. This is especially beneficial for large n. For n=131072 and 128 threads, each thread reads 1024 elements = 256 float4 loads. For n=8192 and 256 threads, each thread reads 32 elements = 8 float4 loads. For n=4096 and 256 threads, each thread reads 16 elements = 4 float4 loads. Vectorized loads can significantly improve memory bandwidth. Let me add them. ```cuda // Process 4 elements at a time using float4 int n4 = n / 4; float4* x4 = (float4*)(row_x); for (int i = tid; i < n4; i += blockDim.x) { float4 v4 = x4[i]; if (v4.x > vals[0]) heap_replace<K>(vals, idxs, v4.x, i * 4); if (v4.y > vals[0]) heap_replace<K>(vals, idxs, v4.y, i * 4 + 1); if (v4.z > vals[0]) heap_replace<K>(vals, idxs, v4.z, i * 4 + 2); if (v4.w > vals[0]) heap_replace<K>(vals, idxs, v4.w, i * 4 + 3); } // Process remaining elements for (int i = tid + n4 * 4; i < n; i += blockDim.x) { float v = row_x[i]; if (v > vals[0]) heap_replace<K>(vals, idxs, v, i); } ``` Wait, there's an issue: the float4 load requires 16-byte alignment. For a row starting at `row * n`, if `n` is not a multiple of 4... actually, `row * n` is a multiple of 4 bytes (since each float is 4 bytes), so `row_x` is always 4-byte aligned. But float4 requires 16-byte alignment! For n=12000, row_x + 0 is aligned to 4 bytes but not necessarily 16 bytes (it is aligned to 4 bytes because row * 12000 * 4 = row * 48000, which is divisible by 4 but not necessarily by 16). Actually, for row-major layout, the start of each row is at `batch * n` offset. If `n` is not a multiple of 4, then rows after row 0 are not 16-byte aligned. For n=12000, 12000 * 4 = 48000, which is divisible by 16? 48000 / 16 = 3000. Yes, it is! Wait, 12000 = 16 * 750. So 12000 * 4 = 48000 = 16 * 3000. Yes, divisible by 16. For n=131072: 131072 * 4 = 524288 = 16 * 32768. Divisible by 16. For n=8192: 8192 * 4 = 32768 = 16 * 2048. Divisible by 16. For n=16384: 16384 * 4 = 65536 = 16 * 4096. Divisible by 16. For n=4096: 4096 * 4 = 16384 = 16 * 1024. Divisible by 16. All n values are multiples of 4, and since the tensor starts at a 16-byte aligned address (PyTorch allocates aligned memory), all rows are 16-byte aligned. So float4 loads are safe! Actually, let me double-check. PyTorch tensors are typically allocated with at least 16-byte alignment for CUDA. And all our n values are multiples of 4. So yes, float4 is safe. For n=12000, the last row starts at 15 * 12000 * 4 = 720000, which is 16 * 45000. Still aligned. Great, I'll use float4 loads. Now let me also think about the k=1 case. For k=1, the heap approach degenerates to a simple max reduction. I already wrote a specialized argmax kernel. But I can also just use the heap approach with K=1: - vals[0] starts at -inf - For each element, if > vals[0], replace - After all elements, sort 1 element (trivial) - Then warp-level and block-level sorts are also trivial For K=1, warp_cand = 32 * 1 = 32. Block_cand = 8 * 1 = 8. The sorts are tiny. But the argmax reduction using warp shuffles is much faster. I'll keep the specialized k=1 kernel. Actually, I can even use warp shuffles for the warp-level merge of K=1. But let me just keep the specialized kernel - it's simple and fast. Let me now think about warp divergence in the heap replace. The condition `if (v > vals[0])` causes divergence within the warp. Some threads might replace, others might not. The `heap_replace` function has loops with dynamic iteration counts, which causes more divergence. However, for random data, roughly half the elements will be > root (initially) and then fewer as the heap fills. Actually, for random Gaussian data, the probability that a new element is > the K-th largest of previously seen elements decreases. But with K=64 and n=131072, after seeing enough elements, only about K/n = 0.05% of new elements will be > root. Wait, that's for the final stages. But initially, almost all elements are > -inf. So for the first 64 elements, all threads replace. Then the heap fills and replacements become rare. The divergence is not terrible because: 1. Threads in a warp see similar data patterns (they're processing nearby elements) 2. When replacements are rare, most threads skip quickly Also, the `if (v > vals[0])` is very cheap. The expensive part is `heap_replace` which only runs for a small fraction of elements after the heap is full. For the first K elements per thread, we can avoid the heap entirely and just fill the array. But since each thread processes strided elements (tid, tid+THREADS, tid+2*THREADS, ...), the "first K elements" are spread out. It's simpler to just use the heap from the start. One optimization: initialize vals with the first K elements (or as many as available), then build the heap. This avoids the initial -inf replacements. ```cuda // Load first K elements (or as many as this thread sees) int count = 0; for (int i = tid; i < n && count < K; i += blockDim.x) { vals[count] = row_x[i]; idxs[count] = i; count++; } // If we have fewer than K elements, fill rest with -inf for (int i = count; i < K; i++) { vals[i] = -INFINITY; idxs[i] = -1; } // Build min-heap for (int i = K / 2 - 1; i >= 0; i--) { heap_sift_down<K>(vals, idxs, i); } // Process remaining elements for (int i = tid + count * blockDim.x; i < n; i += blockDim.x) { float v = row_x[i]; if (v > vals[0]) { heap_replace_root<K>(vals, idxs, v, i); } } ``` But wait, `count` is usually 1 or 2 for the first loop (since blockDim.x is 128 or 256 and n is large). So this doesn't save much. For n=4096 and 256 threads, each thread sees exactly 16 elements. For K=1, count=1. For K=8, count might be 1 if 256 threads and n=4096. Actually, the first loop `i = tid; i < n && count < K; i += blockDim.x` gives each thread at most 1 element (since blockDim.x >> K for most cases). So this is basically just loading the first element and then building a heap of mostly -infs. Not very helpful. I'll just stick with initializing to -inf. Actually, there's another optimization for the case where each thread sees only a few elements. For n=4096 and 256 threads, each thread sees 16 elements. For K=64, 16 < 64, so each thread's local top-k is just its 16 elements (sorted) plus 48 -infs. In this case, the heap is overkill - we can just load all elements into an array and sort it. I can add a branch: ```cuda if (elements_per_thread <= K) { // Load all elements, sort, done } else { // Use heap } ``` For n=4096 and 256 threads: elements_per_thread = 16. For K=1, 8, 16, 32, 64, only K=1 and K=8 have elements_per_thread > K. For n=8192 and 256 threads: elements_per_thread = 32. K=1, 8, 16 have elements_per_thread > K. For n=12000 and 256 threads: elements_per_thread = 47. K=1, 8, 16, 32 have elements_per_thread > K. For n=16384 and 256 threads: elements_per_thread = 64. K=1, 8, 16, 32 have elements_per_thread > K. K=64 has elements_per_thread == K. For n=131072 and 128 threads: elements_per_thread = 1024. All K have elements_per_thread > K. This branch could help for small element counts. But adding branches per thread might complicate the code. Let me keep it simple and just use the heap for all cases. The overhead of processing 16 elements with a heap of size 64 is minimal. Actually, let me do a quick calculation. For 16 elements and K=64: - Initialize 64 elements to -inf: trivial - Process 16 elements: each does `if (v > vals[0])` (always true initially), then heap_replace - After first element, vals[0] is no longer -inf. After a few elements, sift-down happens. - Total operations: ~16 * 26 = 416 operations. - Then bitonic sort 64 elements: 1152 comparisons. - Total: ~1568 operations. vs. loading 16 elements, filling rest with -inf, sorting 64: still 1152 comparisons. So the heap adds ~416 operations, which is small. OK, I'll just use the heap universally. It's simpler. Now let me think about another thing: the `__syncwarp()` in the warp-level bitonic sort. Modern GPUs (Volta+) have independent thread scheduling, so `__syncwarp()` is needed to ensure all threads in the warp have written to shared memory before others read. Similarly, `__syncthreads()` in the block-level sort. Now let me plan the actual code structure. I'll write a single CUDA source file with: 1. `topk1_kernel` for k=1 2. `topk_kernel<K>` template for k > 1 3. C++ wrapper function `topk_forward` that dispatches based on k And in Python, I'll compile it with `load_inline`. Let me draft the full code. I'll start with a simpler version and iterate. Actually, I realize I should check if there are any existing files in the directory that might help, like sota.py. Let me also check what PyTorch version is available and the CUDA architecture. Let me also read the problem prompt to see if there are any hints.

Good, I have the environment info: - PyTorch 2.11.0+cu130 - CUDA 13.0 - RTX PRO 6000 Blackwell, SM 12.0 Now let me write a first version of the kernel and test it. I'll start with a CUDA C++ kernel using the heap + bitonic sort approach. Let me draft the code carefully. I'll create solution.py with the inline CUDA extension. Key decisions: 1. For k=1: specialized argmax kernel with warp shuffle reduction 2. For k > 1: single kernel with heap-based local top-k, warp-level bitonic sort, block-level bitonic sort 3. Block size: 256 threads for k <= 32, 128 threads for k >= 64 (to manage shared memory) 4. Use float4 loads where possible 5. Use templates for k values (8, 16, 32, 64) Wait, I need to handle k=16 for n=12000 (non-power-of-2). The bitonic sort requires power-of-2 candidates. For warp_cand = 32 * 16 = 512 (power of 2). For block_cand = num_warps * 16. With 256 threads = 8 warps, block_cand = 128 (power of 2). With 128 threads = 4 warps, block_cand = 64 (power of 2). All good! For k=8: warp_cand = 256, block_cand = 64 (8 warps) or 32 (4 warps). All powers of 2. For k=32: warp_cand = 1024, block_cand = 256 or 128. All powers of 2. For k=64: warp_cand = 2048, block_cand = 128 (4 warps). Power of 2. Great, all candidate counts are powers of 2. This makes bitonic sort straightforward. Let me write the code. I need to be careful with: 1. The heap sift-down implementation 2. The bitonic sort indexing 3. Shared memory allocation 4. Float4 alignment and bounds checking Actually, for n=12000 with float4 loads: - n4 = 12000 / 4 = 3000 - The remaining elements: 12000 - 3000*4 = 0. So no remainder! For n=131072: 131072 / 4 = 32768. No remainder. For n=8192: 2048. No remainder. For n=16384: 4096. No remainder. For n=4096: 1024. No remainder. All n values are multiples of 4! Great, no remainder handling needed for float4. Wait, but the strided access with float4 means thread tid loads at offset tid. If blockDim.x is 256, then thread 0 loads at 0, thread 1 at 4, thread 2 at 8, etc. But then the next iteration loads at 256*4 = 1024 offset. Thread 0 loads elements 0, 1024, 2048, etc. Thread 1 loads 4, 1028, 2052, etc. This is NOT coalesced! Consecutive threads should read consecutive memory locations for coalescing. With float4, thread 0 reads bytes 0-15, thread 1 reads bytes 16-31, etc. But in my strided loop, thread 0 reads at 0, then 1024, etc. Thread 1 reads at 4, then 1028, etc. Wait, is this coalesced? Let's see: - Warp reads: thread 0 at addr 0, thread 1 at addr 16, thread 2 at addr 32, ..., thread 31 at addr 496. - These are consecutive 16-byte chunks. The warp can read this in one or two 128-byte transactions. Yes, this is coalesced! Then next iteration: thread 0 at addr 1024*4 = 4096, thread 1 at 1032*4 = 4128? No! Wait, I said `i = tid; i < n4; i += blockDim.x`. For tid=1, i starts at 1, which corresponds to element index 4 (since x4[i] loads 4 floats starting at i*16). For tid=2, i=2, element index 8. So consecutive threads load consecutive float4s. In the next iteration, i += 256. Thread 0 loads x4[256] = element 1024. Thread 1 loads x4[257] = element 1028. These are also consecutive. So yes, the access pattern IS coalesced! Each warp reads 32 consecutive float4s = 512 floats = 2048 bytes. Perfect. Great, float4 loads work well here. Now let me write the code. I'll create a first draft, test correctness, then optimize. One more consideration: for the bitonic sort in shared memory, I need to make sure the loop bounds and conditions are correct. Let me verify with a small example. For descending sort of 8 elements: - k=2, j=1: compare (0,1), (2,3), (4,5), (6,7). For even i, if s[i] < s[i^1], swap. This makes each pair descending. - k=4, j=2: compare (0,2), (1,3), (4,6), (5,7). For i in [0,4), if s[i] < s[i^1], swap. For i in [4,8), if s[i] > s[i^1], swap. After this, first 4 are the larger elements from pairs (0,2) and (1,3). Second 4 are the smaller from (4,6) and (5,7). - k=4, j=1: compare (0,1), (2,3), (4,5), (6,7). For i in [0,4), keep larger. For i in [4,8), keep smaller. Now first 4 are descending, last 4 are ascending. - k=8, j=4: compare (0,4), (1,5), (2,6), (3,7). For i in [0,8), keep larger (since i & 8 == 0 for all i < 8). After this, larger elements from first and second halves are in positions 0-3. - k=8, j=2: compare (0,2), (1,3), (4,6), (5,7). For i in [0,4), keep larger. For i in [4,8), keep smaller. After this, positions 0-3 have the largest 4, sorted descending. Positions 4-7 have the smallest 4, sorted ascending. - k=8, j=1: compare (0,1), (2,3), (4,5), (6,7). For i in [0,4), keep larger. For i in [4,8), keep smaller. Now all 8 are sorted descending. Yes! The algorithm is correct. Now, in the loop: ```cuda for (int kk = 2; kk <= N; kk *= 2) { for (int j = kk / 2; j > 0; j /= 2) { for (int i = tid; i < N; i += blockDim.x) { int ixj = i ^ j; if (ixj > i) { if ((i & kk) == 0) { if (s[i] < s[ixj]) swap; } else { if (s[i] > s[ixj]) swap; } } } __syncthreads(); // or __syncwarp() } } ``` For the warp-level sort, `N = warp_cand = 32 * K`. `tid` is `lane` (0-31). `blockDim.x` should be replaced with `32` (warp size). And `__syncwarp()`. For the block-level sort, `N = block_cand = num_warps * K`. `tid` is `threadIdx.x`. `blockDim.x` is the block size. And `__syncthreads()`. One issue: in the block-level sort, `blockDim.x` might be larger than `N`. For example, with 256 threads and block_cand = 64 (k=8). Threads 64-255 will have `i >= N` and do nothing. But they still participate in `__syncthreads()`. That's fine. Now let me also handle the case where `block_cand < blockDim.x` in the output writing. Only threads 0..k-1 should write output. OK, I think I have enough to write the code. Let me create solution.py. Actually, one more optimization: for the heap, I can unroll the sift-down loop since K is a template parameter and the max depth is log2(K). For K=64, max depth is 6. I can write the sift-down with a `#pragma unroll` on the loop. ```cuda template <int K> __device__ inline void heap_replace(float* vals, int* idxs, float new_val, int new_idx) { float v = new_val; int idx = new_idx; int pos = 0; #pragma unroll for (int d = 0; d < 6; d++) { // log2(64) = 6 int left = 2 * pos + 1; int right = 2 * pos + 2; int smallest = pos; if (left < K && vals[left] < v) smallest = left; if (right < K && vals[right] < v) smallest = right; if (smallest == pos) { vals[pos] = v; idxs[pos] = idx; return; } vals[pos] = vals[smallest]; idxs[pos] = idxs[smallest]; pos = smallest; } vals[pos] = v; idxs[pos] = idx; } ``` But the compiler might not like the early return inside a loop that we want unrolled. Let me restructure: ```cuda template <int K> __device__ inline void heap_replace(float* vals, int* idxs, float new_val, int new_idx) { float v = new_val; int idx = new_idx; int pos = 0; bool done = false; #pragma unroll for (int d = 0; d < 6; d++) { if (!done) { int left = 2 * pos + 1; int right = 2 * pos + 2; int smallest = pos; if (left < K && vals[left] < v) smallest = left; if (right < K && vals[right] < v) smallest = right; if (smallest == pos) { vals[pos] = v; idxs[pos] = idx; done = true; } else { vals[pos] = vals[smallest]; idxs[pos] = idxs[smallest]; pos = smallest; } } } if (!done) { vals[pos] = v; idxs[pos] = idx; } } ``` Hmm, the `done` flag might add overhead. For small K, the loop is short anyway. Let me just use a regular while loop and trust the compiler. Or better, since K is known at compile time, the compiler can unroll and optimize. Actually, for a cleaner implementation, let me just write the while loop. For K=64, it iterates at most 6 times. The compiler should inline and unroll it nicely. ```cuda template <int K> __device__ inline void heap_replace(float* vals, int* idxs, float new_val, int new_idx) { float v = new_val; int idx = new_idx; int pos = 0; while (true) { int left = 2 * pos + 1; int right = 2 * pos + 2; int smallest = pos; if (left < K && vals[left] < v) smallest = left; if (right < K && vals[right] < v) smallest = right; if (smallest == pos) break; vals[pos] = vals[smallest]; idxs[pos] = idxs[smallest]; pos = smallest; } vals[pos] = v; idxs[pos] = idx; } ``` This is clean and the compiler should handle it well. Similarly for the bitonic sort in registers: ```cuda template <int K> __device__ inline void bitonic_sort_regs(float* vals, int* idxs) { #pragma unroll for (int k = 2; k <= K; k *= 2) { #pragma unroll for (int j = k / 2; j > 0; j /= 2) { #pragma unroll for (int i = 0; i < K; i++) { int ixj = i ^ j; if (ixj > i) { bool dir = (i & k) == 0; if (dir) { if (vals[i] < vals[ixj]) { float tv = vals[i]; vals[i] = vals[ixj]; vals[ixj] = tv; int ti = idxs[i]; idxs[i] = idxs[ixj]; idxs[ixj] = ti; } } else { if (vals[i] > vals[ixj]) { float tv = vals[i]; vals[i] = vals[ixj]; vals[ixj] = tv; int ti = idxs[i]; idxs[i] = idxs[ixj]; idxs[ixj] = ti; } } } } } } } ``` Since K is a template parameter, all loops should be fully unrolled by the compiler. For K=64, this generates 672 compare-swap operations. Each is a branch, but the branch is predictable (based on compile-time known i, j, k). The compiler might even convert some to predicated instructions. Wait, actually for the inner loop over i, the branch `if (ixj > i)` is compile-time evaluable for unrolled iterations. For example, when i=0, j=1: ixj=1 > 0, so always true. When i=1, j=1: ixj=0 > 1 is false, so the body is eliminated. Great! Similarly, `dir = (i & k) == 0` is compile-time known. And the comparisons `vals[i] < vals[ixj]` might be optimized. This should generate very efficient code. OK, let me write the full solution.py now. Wait, I should also consider using `__forceinline__` for device functions to ensure inlining. Let me draft the complete CUDA source code. One more thing: I need to handle the template dispatch for k values. I'll instantiate for k=8, 16, 32, 64. Also, I should think about whether to use 128 or 256 threads. Let me use: - k=1: 256 threads - k=8, 16, 32: 256 threads - k=64: 128 threads (to reduce shared memory) For k=64 with 128 threads: - warp_cand = 32 * 64 = 2048 per warp - block_cand = 4 * 64 = 256 - Shared memory: 4 * 2048 * 8 + 256 * 8 = 65536 + 2048 = 67584 bytes per block = 66KB - This is comfortable. For k=32 with 256 threads: - warp_cand = 32 * 32 = 1024 per warp - block_cand = 8 * 32 = 256 - Shared memory: 8 * 1024 * 8 + 256 * 8 = 65536 + 2048 = 67584 bytes = 66KB - Also comfortable. For k=16 with 256 threads: - warp_cand = 32 * 16 = 512 per warp - block_cand = 8 * 16 = 128 - Shared memory: 8 * 512 * 8 + 128 * 8 = 32768 + 1024 = 33792 bytes = 33KB For k=8 with 256 threads: - warp_cand = 32 * 8 = 256 per warp - block_cand = 8 * 8 = 64 - Shared memory: 8 * 256 * 8 + 64 * 8 = 16384 + 512 = 16896 bytes = 16.5KB All fit comfortably in shared memory. Good. Now, let me think about the `topk1_kernel`. I can use 256 threads with float4 loads. ```cuda __global__ void topk1_kernel(const float* __restrict__ x, float* __restrict__ out_vals, int64_t* __restrict__ out_idxs, int n) { int row = blockIdx.x; int tid = threadIdx.x; const float* row_x = x + row * n; float max_val = -INFINITY; int max_idx = -1; // Vectorized load int n4 = n >> 2; const float4* x4 = (const float4*)row_x; for (int i = tid; i < n4; i += blockDim.x) { float4 v4 = x4[i]; if (v4.x > max_val) { max_val = v4.x; max_idx = (i << 2); } if (v4.y > max_val) { max_val = v4.y; max_idx = (i << 2) + 1; } if (v4.z > max_val) { max_val = v4.z; max_idx = (i << 2) + 2; } if (v4.w > max_val) { max_val = v4.w; max_idx = (i << 2) + 3; } } // Handle remaining (none in our case, but for completeness) for (int i = (tid + (n4 << 2)); i < n; i += blockDim.x) { float v = row_x[i]; if (v > max_val) { max_val = v; max_idx = i; } } // Warp reduction unsigned mask = 0xFFFFFFFF; #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { float other_val = __shfl_down_sync(mask, max_val, offset); int other_idx = __shfl_down_sync(mask, max_idx, offset); if (other_val > max_val) { max_val = other_val; max_idx = other_idx; } } __shared__ float s_vals[32]; __shared__ int s_idxs[32]; int lane = tid & 31; int warp_id = tid >> 5; if (lane == 0) { s_vals[warp_id] = max_val; s_idxs[warp_id] = max_idx; } __syncthreads(); if (tid < 32) { float v = (tid < ((blockDim.x + 31) >> 5)) ? s_vals[tid] : -INFINITY; int idx = (tid < ((blockDim.x + 31) >> 5)) ? s_idxs[tid] : -1; #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { float other_val = __shfl_down_sync(mask, v, offset); int other_idx = __shfl_down_sync(mask, idx, offset); if (other_val > v) { v = other_val; idx = other_idx; } } if (tid == 0) { out_vals[row] = v; out_idxs[row] = idx; } } } ``` This looks good. Now the main kernel template. ```cuda template <int K> __global__ void topk_kernel(const float* __restrict__ x, float* __restrict__ out_vals, int64_t* __restrict__ out_idxs, int n) { int row = blockIdx.x; int tid = threadIdx.x; int lane = tid & 31; int warp_id = tid >> 5; int num_warps = blockDim.x >> 5; const float* row_x = x + row * n; // Thread-local top-K using min-heap float vals[K]; int idxs[K]; #pragma unroll for (int i = 0; i < K; i++) { vals[i] = -INFINITY; idxs[i] = -1; } // Vectorized processing int n4 = n >> 2; const float4* x4 = (const float4*)row_x; for (int i = tid; i < n4; i += blockDim.x) { float4 v4 = x4[i]; if (v4.x > vals[0]) heap_replace<K>(vals, idxs, v4.x, i << 2); if (v4.y > vals[0]) heap_replace<K>(vals, idxs, v4.y, (i << 2) + 1); if (v4.z > vals[0]) heap_replace<K>(vals, idxs, v4.z, (i << 2) + 2); if (v4.w > vals[0]) heap_replace<K>(vals, idxs, v4.w, (i << 2) + 3); } for (int i = (n4 << 2) + tid; i < n; i += blockDim.x) { float v = row_x[i]; if (v > vals[0]) heap_replace<K>(vals, idxs, v, i); } // Sort thread's top-K descending bitonic_sort_regs<K>(vals, idxs); // Shared memory layout // Warp work areas + block candidate area extern __shared__ char smem[]; int warp_cand = 32 * K; int block_cand = num_warps * K; float* s_warp_vals = (float*)smem; int* s_warp_idxs = (int*)(smem + num_warps * warp_cand * sizeof(float)); float* my_warp_vals = s_warp_vals + warp_id * warp_cand; int* my_warp_idxs = s_warp_idxs + warp_id * warp_cand; // Write to warp shared memory #pragma unroll for (int i = 0; i < K; i++) { my_warp_vals[lane * K + i] = vals[i]; my_warp_idxs[lane * K + i] = idxs[i]; } __syncwarp(); // Bitonic sort within warp for (int kk = 2; kk <= warp_cand; kk <<= 1) { for (int j = kk >> 1; j > 0; j >>= 1) { for (int i = lane; i < warp_cand; i += 32) { int ixj = i ^ j; if (ixj > i) { if ((i & kk) == 0) { if (my_warp_vals[i] < my_warp_vals[ixj]) { float tv = my_warp_vals[i]; my_warp_vals[i] = my_warp_vals[ixj]; my_warp_vals[ixj] = tv; int ti = my_warp_idxs[i]; my_warp_idxs[i] = my_warp_idxs[ixj]; my_warp_idxs[ixj] = ti; } } else { if (my_warp_vals[i] > my_warp_vals[ixj]) { float tv = my_warp_vals[i]; my_warp_vals[i] = my_warp_vals[ixj]; my_warp_vals[ixj] = tv; int ti = my_warp_idxs[i]; my_warp_idxs[i] = my_warp_idxs[ixj]; my_warp_idxs[ixj] = ti; } } } } __syncwarp(); } } // Copy warp top-K to block candidate area float* s_block_vals = (float*)(smem + num_warps * warp_cand * (sizeof(float) + sizeof(int))); int* s_block_idxs = (int*)(s_block_vals + block_cand); if (lane < K) { s_block_vals[warp_id * K + lane] = my_warp_vals[lane]; s_block_idxs[warp_id * K + lane] = my_warp_idxs[lane]; } __syncthreads(); // Block-level bitonic sort for (int kk = 2; kk <= block_cand; kk <<= 1) { for (int j = kk >> 1; j > 0; j >>= 1) { for (int i = tid; i < block_cand; i += blockDim.x) { int ixj = i ^ j; if (ixj > i) { if ((i & kk) == 0) { if (s_block_vals[i] < s_block_vals[ixj]) { float tv = s_block_vals[i]; s_block_vals[i] = s_block_vals[ixj]; s_block_vals[ixj] = tv; int ti = s_block_idxs[i]; s_block_idxs[i] = s_block_idxs[ixj]; s_block_idxs[ixj] = ti; } } else { if (s_block_vals[i] > s_block_vals[ixj]) { float tv = s_block_vals[i]; s_block_vals[i] = s_block_vals[ixj]; s_block_vals[ixj] = tv; int ti = s_block_idxs[i]; s_block_idxs[i] = s_block_idxs[ixj]; s_block_idxs[ixj] = ti; } } } } __syncthreads(); } } // Write output if (tid < K) { out_vals[row * K + tid] = s_block_vals[tid]; out_idxs[row * K + tid] = (int64_t)s_block_idxs[tid]; } } ``` Wait, there's a potential issue with shared memory layout. The `extern __shared__ char smem[]` means I need to specify the shared memory size at kernel launch. The size is: `num_warps * warp_cand * sizeof(float) + num_warps * warp_cand * sizeof(int) + block_cand * sizeof(float) + block_cand * sizeof(int)` For K=64, num_warps=4, warp_cand=2048, block_cand=256: - Warp vals: 4 * 2048 * 4 = 32768 - Warp idxs: 4 * 2048 * 4 = 32768 - Block vals: 256 * 4 = 1024 - Block idxs: 256 * 4 = 1024 - Total: 67584 bytes For K=32, num_warps=8, warp_cand=1024, block_cand=256: - Warp vals: 8 * 1024 * 4 = 32768 - Warp idxs: 8 * 1024 * 4 = 32768 - Block vals: 256 * 4 = 1024 - Block idxs: 256 * 4 = 1024 - Total: 67584 bytes Interesting, both use the same total shared memory! For K=16, num_warps=8, warp_cand=512, block_cand=128: - Warp vals: 8 * 512 * 4 = 16384 - Warp idxs: 8 * 512 * 4 = 16384 - Block vals: 128 * 4 = 512 - Block idxs: 128 * 4 = 512 - Total: 33792 bytes For K=8, num_warps=8, warp_cand=256, block_cand=64: - Warp vals: 8 * 256 * 4 = 8192 - Warp idxs: 8 * 256 * 4 = 8192 - Block vals: 64 * 4 = 256 - Block idxs: 64 * 4 = 256 - Total: 16896 bytes All well within limits. Now for the C++ wrapper: ```cpp #include <torch/extension.h> #include <cuda_runtime.h> #include <float.h> // ... device functions and kernels ... torch::Tensor topk_forward_cuda(torch::Tensor x, int k) { int batch = x.size(0); int n = x.size(1); auto out_vals = torch::empty({batch, k}, x.options()); auto out_idxs = torch::empty({batch, k}, torch::dtype(torch::kInt64).device(x.device())); const float* x_ptr = x.data_ptr<float>(); float* vals_ptr = out_vals.data_ptr<float>(); int64_t* idxs_ptr = out_idxs.data_ptr<int64_t>(); if (k == 1) { topk1_kernel<<<batch, 256>>>(x_ptr, vals_ptr, idxs_ptr, n); } else { int threads = (k >= 64) ? 128 : 256; int num_warps = threads / 32; int warp_cand = 32 * k; int block_cand = num_warps * k; size_t smem_size = num_warps * warp_cand * sizeof(float) + num_warps * warp_cand * sizeof(int) + block_cand * sizeof(float) + block_cand * sizeof(int); switch (k) { case 8: topk_kernel<8><<<batch, threads, smem_size>>>(x_ptr, vals_ptr, idxs_ptr, n); break; case 16: topk_kernel<16><<<batch, threads, smem_size>>>(x_ptr, vals_ptr, idxs_ptr, n); break; case 32: topk_kernel<32><<<batch, threads, smem_size>>>(x_ptr, vals_ptr, idxs_ptr, n); break; case 64: topk_kernel<64><<<batch, threads, smem_size>>>(x_ptr, vals_ptr, idxs_ptr, n); break; } } return torch::stack({out_vals, out_idxs}, 0); // Wait, we need to return a tuple } ``` Actually, returning a tuple from C++ to Python: ```cpp std::vector<torch::Tensor> topk_forward_cuda(torch::Tensor x, int k) { // ... return {out_vals, out_idxs}; } ``` And in Python: ```python def topk_forward(x, k): return module.topk_forward_cuda(x, k) ``` Now, there's an issue: the Model's forward takes x and returns (values, indices). But the C++ function takes x and k. Since k is fixed per Model instance, I can pass it. Wait, actually I need to think about how to integrate this with the PyTorch nn.Module. The Model has `self.k` and forward takes just `x`. So: ```python class Model(nn.Module): def __init__(self, batch, n, k): super().__init__() self.batch = batch self.n = n self.k = k self.register_buffer("_dummy", torch.zeros(1)) def forward(self, x): return topk_forward(x, self.k) ``` Where `topk_forward` is the compiled CUDA function. Let me also think about `torch::stack` - no, I don't want to stack. I want to return a tuple. In the C++ function: ```cpp std::vector<torch::Tensor> topk_forward_cuda(torch::Tensor x, int k) { // ... compute ... return {out_vals, out_idxs}; } ``` In Python, `module.topk_forward_cuda(x, k)` returns a tuple `(out_vals, out_idxs)`. One issue: `load_inline` with C++ sources that include CUDA code. I need to make sure the C++ source and CUDA source are properly combined. Actually, `load_inline` takes `cpp_sources` and `cuda_sources`. The CUDA sources contain `__global__` functions and device code. The C++ sources contain the wrapper function that calls the kernels. But with `load_inline`, I can put everything in `cuda_sources` and leave `cpp_sources` empty, as long as the C++ wrapper is also in the CUDA source file (nvcc can compile mixed C++/CUDA). Let me verify: `torch.utils.cpp_extension.load_inline` with `cuda_sources` containing both device and host code should work. Actually, looking at PyTorch docs, `cuda_sources` is compiled with nvcc and `cpp_sources` with the C++ compiler. nvcc can compile host code, so putting everything in `cuda_sources` is fine. Let me write the complete solution.py. One more thing: I need to make sure the CUDA code compiles with `-arch=sm_120` or similar for Blackwell. PyTorch's `load_inline` should handle architecture detection, but I can add extra flags. Actually, PyTorch 2.11 with CUDA 13.0 should support sm_120. But let me be safe and not specify architecture explicitly - PyTorch should figure it out. Wait, CUDA 13.0 is very new and supports Blackwell. But the default architecture flags in PyTorch might not include sm_120. Let me check. Actually, I can add `"--gpu-architecture=sm_120"` to `extra_cuda_cflags`. But if PyTorch doesn't support it, it might fail. Let me check what architectures are supported. Let me just try without explicit arch first. PyTorch usually detects the GPU architecture and compiles for it. Let me also add `"-O3"`, `"--use_fast_math"` to the flags. OK, let me write solution.py now. Actually, one more concern: the bitonic sort in registers uses loops over K, and for K=64, it generates a lot of code. The compiler might struggle with register allocation or instruction cache. Let me see if there's a way to reduce code size. Alternatively, I can use a smaller bitonic sort for the thread-local sort. For example, since we only need the top-K from a min-heap, and the min-heap is already somewhat structured, maybe a simpler sort would suffice. Actually, a heap is NOT sorted. After extracting elements from a min-heap, they come out in ascending order. But we want descending order. So we need to sort the K elements. For K=64, bitonic sort is 672 comparators. That's a lot of code but still manageable. One optimization: instead of fully sorting, we can use a partial sort that only ensures the first K elements are in descending order. But since we're sorting exactly K elements, a full sort is fine. Another idea: use a heap sort instead of bitonic sort. Heap sort of 64 elements: 1. Build max-heap (but we have a min-heap...) 2. Actually, to get descending order from a min-heap, we can just extract all elements and reverse. Extracting from a min-heap gives ascending order. So if we extract K elements, we get ascending order. Then reverse to get descending. Heap extraction of K elements: K * log2(K) = 64 * 6 = 384 operations. Plus we need a temp array. But wait, our heap is a min-heap in an array. We can convert it to a sorted ascending array by repeated extraction: ```cuda float sorted_vals[K]; int sorted_idxs[K]; for (int i = 0; i < K; i++) { sorted_vals[i] = vals[0]; sorted_idxs[i] = idxs[0]; // Replace root with last element and sift down vals[0] = vals[K - 1 - i]; idxs[0] = idxs[K - 1 - i]; // sift down... } ``` Then reverse `sorted_vals` and `sorted_idxs` to get descending. This is 384 sift-down operations vs 672 comparators for bitonic sort. But each sift-down has branches and more complex control flow. Bitonic sort is branch-predicated and regular. On GPU, bitonic sort in registers might be faster despite more operations. I'll stick with bitonic sort. Now, there's one more optimization I want to make. For the warp-level bitonic sort, instead of each thread handling elements `lane, lane+32, lane+64, ...`, I can have each thread handle a contiguous chunk. This might improve shared memory locality. Actually, the standard strided approach is: ```cuda for (int i = lane; i < warp_cand; i += 32) { ... } ``` With 32 threads and warp_cand=2048, each thread handles 64 elements at stride 32. This means thread 0 handles 0, 32, 64, ... and thread 1 handles 1, 33, 65, ... In bitonic sort, when j is small (1, 2, 4, etc.), elements i and i^j are close together. With the strided layout, thread 0 handles both i and i^j for many cases, which means no inter-thread communication is needed! That's great! For example, j=1: i and i^1 are adjacent. Thread lane handles i = lane, lane+32, lane+64, ... For i=lane (even if lane is even), i^1 = lane+1, which is handled by thread lane+1. So inter-thread communication IS needed. Wait, but in the loop, each thread only processes its own i values. It doesn't process ixj. So thread 0 reads s[0] and s[1], compares, and swaps if needed. Thread 1 doesn't process i=1 because `ixj = 0 < 1`, so it skips. So thread 0 does both the read and write for the pair (0, 1). For i=32, thread 0 processes it. ixj=33, handled by thread 1. Thread 0 reads s[32] and s[33], compares, swaps. Thread 1 does nothing for i=33. So each pair (i, i^j) is handled by the thread with the smaller i. This means threads with smaller lane handle more work? No, because for each j, roughly half the i values have ixj > i. These are distributed across all threads. Actually, for j=1: i values with ixj > i are the even i's. Thread lane handles i = lane, lane+32, lane+64, ... If lane is even, all these i's are even, so all are processed. If lane is odd, all are odd, so none are processed (ixj = i^1 is even < odd i). So for j=1, only even lanes do work! Lanes 0, 2, 4, ... are active. 16 threads do work, 16 are idle. For j=2: i values with ixj > i are those where the second bit is 0. i = ...00, ...01. Wait: - i=0: 00, ixj=2: 10. 2>0, process. - i=1: 01, ixj=3: 11. 3>1, process. - i=2: 10, ixj=0: 00. 0<2, skip. - i=3: 11, ixj=1: 01. 1<3, skip. - i=4: 100, ixj=6: 110. 6>4, process. - i=5: 101, ixj=7: 111. 7>5, process. So for j=2, i % 4 in {0, 1} are processed. With i = lane + 32*m: - lane=0: i=0, 32, 64, ... all have i%32=0, so i%4=0. All processed. - lane=1: i%4=1. All processed. - lane=2: i%4=2. None processed. - lane=3: i%4=3. None processed. So lanes 0,1 are active. Lanes 2,3 idle. 8 pairs of active/idle as m increases. This pattern means that for each j, 1/2 of the threads are active. This is actually fine - GPU warps can handle this with some threads masked off. The advantage of this approach is that each active thread reads two elements that are close together in memory (for small j). For j=1, they're adjacent. For j=2, they're 2 apart. This is cache-friendly and has minimal bank conflicts (for j < 32, they're in different banks!). For large j, like j=1024, elements are far apart. But as we discussed, large j stages are few in number. I think this implementation is solid. Let me write it. Actually, I want to reconsider the warp-level sort shared memory layout. Currently: ```cuda my_warp_vals[lane * K + i] = vals[i]; ``` This means thread 0 stores at 0..K-1, thread 1 at K..2K-1, etc. So the array is organized as [thread0_vals, thread1_vals, ...]. In the bitonic sort, element i is at offset i. Thread lane accesses i = lane, lane+32, lane+64, etc. With the current layout: - i=0: thread 0's first element - i=1: thread 0's second element? No! Wait, my layout is `lane * K + i`. For K=64: - Thread 0: indices 0..63 - Thread 1: indices 64..127 - ... But in the bitonic sort, i and i^j refer to global indices 0..warp_cand-1. For j=1, i=0 and i^1=1. These are both in thread 0's region (0..63). For i=63 and i^1=62, also in thread 0's region. For j=64, i=0 and i^j=64. These are in thread 0 (0) and thread 1 (64). This layout means that for small j (< K), both elements are in the same thread's region, so the same thread handles the comparison. This is good! For larger j, elements span across threads. The thread with smaller i handles the pair. This layout seems good. Let me keep it. One issue: for j=K (e.g., j=64), i and i^64 might be in adjacent threads' regions. Thread 0 handles i=0, reads s[0] and s[64]. These are 64 floats apart = 256 bytes. With 32-byte banks, that's 8 banks apart. No conflict! Actually, let me check bank conflicts more carefully. Bank = (byte_offset / 4) % 32 = index % 32. For j=64: i=0, ixj=64. Banks: 0 and 0. SAME BANK! Oh no! For j that are multiples of 32, i and i^j map to the same bank. This is a conflict. For j=32: i=0, ixj=32. Banks: 0 and 0. Same bank. For j=64: same. For j=128: same. This is the bank conflict issue I identified earlier. But now with our layout, for j < K, both elements are in the same thread's region, so no inter-thread conflict (only one thread accesses both). For j >= K, elements are in different threads' regions, and multiple threads might access the same bank. For j=K=64: thread 0 accesses banks 0 and 0 (s[0] and s[64]). Thread 1 accesses banks 1 and 1 (s[1] and s[65])? Wait: - Thread 0: i=0, ixj=64. Both in bank 0. Thread 0 reads both. - Thread 1: i=1, ixj=65. Bank 1 and bank 1. Thread 1 reads both. - ... - Each thread reads two addresses in the same bank. No inter-thread conflict! Each thread's two accesses are serialized by the bank, but threads run in parallel. Actually, is this better or worse than having threads access different banks? With a 2-way conflict per thread, the warp still executes in one go (all threads access their first address, then their second). The bank is busy for 2 cycles instead of 1. That's a 2x slowdown for that stage. For j=128: similar, 2-way conflict per thread. For j=256: same. For j=512: same. For j=1024: same. The number of such stages is small. So it's acceptable. For j=32: i=0, ixj=32. Both in thread 0's region (0..63). Thread 0 reads both. No inter-thread issue. For j=16: i=0, ixj=16. Both in thread 0's region. Good. So the only "inter-thread" conflicts are for j values where i and i^j are in different threads' regions AND multiple threads access the same bank. Let me check when that happens. For j < K: both in same thread. No inter-thread conflict. For j >= K: elements are in different threads. - j=K=64: thread 0 handles i=0 (in thread 0) and ixj=64 (in thread 1). Thread 0 reads both. No other thread reads bank 0 at the same time. Thread 1 reads i=1 and ixj=65. No conflict with thread 0. Wait, I need to check: does thread 1 also handle i=64? No, because for i=64, ixj=0, and 0 < 64, so thread that handles i=0 handles this pair. Thread 1 handles i=1, 33, 65, etc. So thread 1 never handles i=64. So for j=64, each pair (i, i^64) is handled by the thread with smaller i. That thread reads both elements. No other thread touches that pair. So no inter-thread bank conflict! The only bank conflict is that a single thread reads two addresses in the same bank. That's a 2-way conflict, handled by the bank's dual-port or serialization. For j=128: i=0 handles (0, 128). i=1 handles (1, 129). etc. Thread 0 handles i=0, 32, 64, 96. For i=0: pair (0, 128). For i=32: pair (32, 160). For i=64: pair (64, 192). For i=96: pair (96, 224). All handled by thread 0. No inter-thread conflict. For j=1024: i=0 handles (0, 1024). Thread 0 reads both. Thread 1 handles (1, 1025). Etc. No inter-thread conflict. So actually, with this layout and the "smaller i handles the pair" convention, there are NO inter-thread bank conflicts! The only conflicts are intra-thread 2-way conflicts for j that are multiples of 32. This is excellent! Wait, let me double-check for j values that are NOT multiples of 32 but still >= K. j=K+1=65: i=0 handles (0, 65). Bank 0 and bank 1. Different banks. Good! j=K+2=66: i=0 handles (0, 66). Bank 0 and bank 2. Good! ... j=2*K-1=127: i=0 handles (0, 127). Bank 0 and bank 31. Good! So the only 2-way intra-thread conflicts are for j = 32, 64, 96, 128, ... i.e., multiples of 32. And for j < K, there's no conflict because both elements are in the same thread and the bank difference is j % 32. For j=32 and K=64: j < K. Both elements in same thread. Bank difference = 0. Conflict. For j=64: j = K. Elements in adjacent threads. Handled by thread with smaller i. Bank difference = 0. Conflict. For j=96: j > K. Elements in threads 0 and 1 (for i=0, ixj=96, which is in thread 1 since 64 <= 96 < 128). Handled by thread 0. Banks: 0 and 0. Conflict. Wait, for j=96, i=0: bank 0. ixj=96: bank 96 % 32 = 0. Same bank. Conflict. i=1: bank 1. ixj=97: bank 97 % 32 = 1. Same bank. Conflict. Yes, multiples of 32 cause conflicts. How many such j values are there? For warp_cand = 2048: - j = 32, 64, 96, 128, 160, 192, 224, 256, 320, 384, 448, 512, 640, 768, 1024 - Actually, j values in bitonic sort are powers of 2 and their halves: 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024. - Multiples of 32 among these: 32, 64, 128, 256, 512, 1024. That's 6 stages out of 11. For each of these 6 stages, 2-way intra-thread conflict. 2x slowdown for those stages. The total slowdown is small since the other 5 stages run at full speed. This is totally acceptable! No need for padding. Great, I'm confident in the design. Let me write the code. One more thing: I should use `constexpr` where possible to help the compiler. Also, for the bitonic sort in registers, since K is a template parameter, I can use `constexpr int N = K;` but it's not necessary. Let me now write solution.py. After writing, I'll run `python check.py` to verify correctness, then `python benchmark.py` to see performance. Let me also create a test script to quickly verify the kernel works before running the full check.py. Actually, check.py is the authority. Let me just run it after creating solution.py. Let me draft the code now.

Let me write the first draft of solution.py. I need to be very careful with the CUDA code to make sure it's correct. Let me think about one more thing: the C++ wrapper needs to handle the fact that the Model's forward takes x and returns (values, indices). The C++ function should take x and k. Also, I need to make sure the compilation works. Let me write a clean version. Wait, I realize there might be an issue with the heap_replace function. When K is a template parameter, the compiler knows K at compile time. But the loop `while (true)` with `left = 2 * pos + 1` might not be unrolled optimally. Let me check if I should manually unroll. For K=64, max depth is 6. For K=32, max depth is 5. For K=16, max depth is 4. For K=8, max depth is 3. I can write: ```cuda template <int K> __device__ inline void heap_replace(float* vals, int* idxs, float new_val, int new_idx) { float v = new_val; int idx = new_idx; int pos = 0; // Unrolled sift-down for known K #pragma unroll for (int d = 0; d < 6; d++) { // 6 is log2(64), enough for all our K values if (d >= 6) break; // Actually, we need dynamic break } } ``` Hmm, manual unrolling with early exit is tricky. Let me just trust the compiler with the while loop. Actually, I just realized: for K=8, 16, 32, the heap_replace with a while loop that iterates at most log2(K) times should be very fast. The compiler will likely unroll it since the loop bound is small and K is a compile-time constant... well, K is a template parameter, so the compiler knows it. But the while loop exit depends on runtime data (the heap structure), so the compiler might not unroll it fully. I can use `#pragma unroll` on the while loop if I convert it to a for loop with a fixed upper bound: ```cuda template <int K> __device__ inline void heap_replace(float* vals, int* idxs, float new_val, int new_idx) { float v = new_val; int idx = new_idx; int pos = 0; bool done = false; #pragma unroll for (int d = 0; d < 6; d++) { if (!done) { int left = 2 * pos + 1; int right = 2 * pos + 2; int smallest = pos; if (left < K && vals[left] < v) smallest = left; if (right < K && vals[right] < v) smallest = right; if (smallest == pos) { vals[pos] = v; idxs[pos] = idx; done = true; } else { vals[pos] = vals[smallest]; idxs[pos] = idxs[smallest]; pos = smallest; } } } if (!done) { vals[pos] = v; idxs[pos] = idx; } } ``` For K=8, the loop runs 6 times but the first 3 iterations do real work and the rest are masked by `done`. The compiler might optimize away the masked iterations, or it might not. Actually, with `if (!done)` inside a loop, the compiler can see that once `done` is true, the rest of the iterations do nothing. But whether it eliminates them depends on optimization level. Let me just use the while loop with `#pragma unroll`. NVCC supports `#pragma unroll` on while loops too, but you need to specify the count. ```cuda #pragma unroll 6 while (true) { // ... if (smallest == pos) break; // ... } ``` Wait, I don't think `#pragma unroll` works well with `while(true)` that has a `break`. Let me use the for loop with `done` flag. It's safer. Actually, another option is to manually unroll for each K. But that's tedious. Let me just write a clean while loop. The compiler on `-O3` should inline and optimize it well. If profiling shows it's a bottleneck, I'll optimize further. OK, let me write the full solution.py now.

I wrote the first draft. Let me run check.py to verify correctness.

The compilation failed for two reasons: 1. `C10_CUDA_KERNEL_LAUNCH_CHECK()` is undefined - I should just remove it or use `AT_CUDA_CHECK(cudaGetLastError())` or simply omit it 2. The pybind11 module definition is being generated automatically by PyTorch, and it's trying to find `topk_forward_cuda` in the C++ scope. But since I put everything in `cuda_sources`, the C++ compiler doesn't see the CUDA code when compiling `main.cpp`. Wait, looking at the error more carefully: - `main.cpp` is auto-generated by PyTorch's load_inline and it contains the pybind11 wrapper. It references `topk_forward_cuda` but can't find it because `topk_forward_cuda` is in the CUDA source file, not visible to the C++ compiler. The solution is to put the declaration of `topk_forward_cuda` in `cpp_sources` and the implementation in `cuda_sources`. Or better, put everything in `cpp_sources` and let PyTorch handle it... but that won't work for CUDA code. Actually, the standard approach with `load_inline` is: - `cpp_sources`: contains the function declarations and pybind11 module definition - `cuda_sources`: contains the CUDA kernel implementations Wait, but `load_inline` already generates the pybind11 module automatically. The issue is that the auto-generated `main.cpp` includes the function declarations from `cpp_sources` but not from `cuda_sources`. Let me check how `load_inline` works. The `cpp_sources` are compiled with the C++ compiler, and `cuda_sources` with nvcc. The pybind11 binding is generated in `main.cpp` based on the `functions` parameter. For `load_inline` to work, the functions listed in `functions` must be declared in `cpp_sources` (which gets included in the generated main.cpp), and defined in either `cpp_sources` or `cuda_sources`. So I need to add a header/declaration of `topk_forward_cuda` in `cpp_sources`. Let me restructure: - `cpp_sources`: forward declaration of `topk_forward_cuda` - `cuda_sources`: all the CUDA code including the implementation Actually, looking at PyTorch docs more carefully, when you use `load_inline`, the `cpp_sources` get compiled as C++ and `cuda_sources` as CUDA. The generated binding code includes `cpp_sources` but not `cuda_sources`. So any function you want to bind must be declared in `cpp_sources`. So: ```python cpp_source = """ std::vector<torch::Tensor> topk_forward_cuda(torch::Tensor x, int k); """ cuda_source = """ #include <torch/extension.h> #include <cuda_runtime.h> ... all kernels ... std::vector<torch::Tensor> topk_forward_cuda(torch::Tensor x, int k) { ... } """ ``` Let me fix this.

Let me run check.py again.

The CUDA error is "invalid argument". This usually means one of the kernel launch parameters is wrong. Let me debug by running with CUDA_LAUNCH_BLOCKING=1 and also adding some print statements. Actually, looking at the stack trace, the error happens during `torch.isnan(solution_out).any()`. This means the kernel returned invalid data or the output tensor is corrupted. Common causes: 1. Shared memory size too large for the block 2. Wrong number of threads/blocks 3. Memory access out of bounds in the kernel Let me check if the shared memory size is valid. For k=64 with 128 threads: - num_warps = 4 - warp_cand = 32 * 64 = 2048 - block_cand = 4 * 64 = 256 - smem = 4 * 2048 * 4 + 4 * 2048 * 4 + 256 * 4 + 256 * 4 = 32768 + 32768 + 1024 + 1024 = 67584 bytes For Blackwell (SM120), what's the max shared memory per block? I think it's at least 164KB, probably 228KB or 256KB. 66KB should be fine. Wait, but the problem might be with how I'm accessing shared memory. Let me check if there are any out-of-bounds accesses. In the warp-level sort, `my_warp_vals` points to `s_warp_vals + warp_id * warp_cand`. With warp_id up to num_warps-1 = 3, and warp_cand = 2048, the warp area is 4 * 2048 = 8192 floats = 32768 bytes. In the block-level sort, `s_block_vals` points after the warp areas. The offset is `num_warps * warp_cand * (sizeof(float) + sizeof(int)) = 4 * 2048 * 8 = 65536 bytes`. And `s_block_idxs = s_block_vals + block_cand = 65536 + 256*4 = 65536 + 1024 = 66560 bytes`. Wait, `s_block_vals` is a float pointer. `s_block_vals + block_cand` advances by `block_cand * sizeof(float)` bytes, not `block_cand` bytes. So `s_block_idxs` is at `s_block_vals + block_cand` which is 65536 + 1024 = 66560 bytes from smem start. Total shared memory: 66560 + 1024 = 67584 bytes. Matches our calculation. Now, in the block-level sort, threads access `s_block_vals[i]` and `s_block_vals[i^j]` where i goes up to block_cand-1 = 255. And j goes up to block_cand/2 = 128. So i^j can be at most 255. This is within bounds. In the warp-level sort, i goes up to warp_cand-1 = 2047. j up to 1024. i^j up to 2047. Within bounds. Wait, what about the extern shared memory declaration? In CUDA, when you declare `extern __shared__ char smem[]`, the actual size is specified at kernel launch. If the size is wrong, you get undefined behavior. But I'm calculating smem correctly. Let me double-check: ```cpp size_t smem = num_warps * warp_cand * sizeof(float) + num_warps * warp_cand * sizeof(int) + block_cand * sizeof(float) + block_cand * sizeof(int); ``` For k=64: 4*2048*4 + 4*2048*4 + 256*4 + 256*4 = 32768 + 32768 + 1024 + 1024 = 67584. Yes. For k=32: 8*1024*4 + 8*1024*4 + 256*4 + 256*4 = 32768 + 32768 + 1024 + 1024 = 67584. Yes. For k=16: 8*512*4 + 8*512*4 + 128*4 + 128*4 = 16384 + 16384 + 512 + 512 = 33792. Yes. For k=8: 8*256*4 + 8*256*4 + 64*4 + 64*4 = 8192 + 8192 + 256 + 256 = 16896. Yes. All should fit. Hmm, let me think about another potential issue. The `float4` loads. For `n=12000`, `n4 = 3000`. Thread tid loads `x4[i]` where `i = tid, tid+256, tid+512, ...`. For tid=255, i=255, which is within [0, 3000). For tid=0, i=256, also within bounds. The max i is 255 + k*256 for some k. We need i < 3000. 255 + 11*256 = 3071 > 3000. So the loop stops at i = 255 + 10*256 = 2815, and next would be 3071 which is > 3000. So all accesses are in bounds. Wait, but `x4` is a `float4*`, so `x4[i]` accesses bytes `i*16` to `i*16+15`. For i=2999, that's bytes 47984 to 47999. The row has 12000 floats = 48000 bytes. So the last float4 covers the last 4 floats. Good. For n=131072, n4=32768. i goes up to 32767. x4[32767] covers bytes 524272 to 524287. Row has 131072*4=524288 bytes. So the last float4 is exactly at the end. Good. Let me run a simple test to see if the kernel works at all.

Interesting! The error happens for k=64 and k=32 but not for k=8, k=16, and k=1. This strongly suggests the issue is with the shared memory size or block configuration for larger k. For k=64: threads=128, num_warps=4, warp_cand=2048, block_cand=256, smem=67584 bytes. For k=32: threads=256, num_warps=8, warp_cand=1024, block_cand=256, smem=67584 bytes. For k=16: threads=256, num_warps=8, warp_cand=512, block_cand=128, smem=33792 bytes. For k=8: threads=256, num_warps=8, warp_cand=256, block_cand=64, smem=16896 bytes. Both failing cases have smem=67584 bytes. The passing cases have smaller smem. This suggests 67584 bytes might exceed the shared memory limit! But wait, I thought Blackwell had more shared memory. Let me check what the actual limit is on this GPU. Actually, the limit might be 48KB by default unless explicitly configured. In CUDA, the default shared memory per block is 48KB. To use more, you need to call `cudaFuncSetAttribute` with `cudaFuncAttributeMaxDynamicSharedMemorySize`. Ah! That's the issue. PyTorch's CUDA launch doesn't automatically set the dynamic shared memory size beyond the default. The `smem` parameter in `<<<...>>>` only specifies how much dynamic shared memory to allocate, but if it exceeds the kernel's maximum, you get "invalid argument". Actually, wait. For kernels launched with `<<<batch, threads, smem>>>`, CUDA allocates `smem` bytes of dynamic shared memory. But there's a limit per block (e.g., 48KB, 96KB, or more depending on the architecture and configuration). On Ampere, the max is 164KB but you need to explicitly request more than 48KB. On Hopper, it's 228KB. On Blackwell, it might be higher. But the issue is that the kernel attribute `maxDynamicSharedSizeBytes` needs to be set if you want more than the default (usually 48KB). Let me check: for smem=67584 = 66KB. This exceeds 48KB. So I need to set the kernel attribute. In PyTorch's inline CUDA, I can set the attribute using `cudaFuncSetAttribute`: ```cpp cudaFuncSetAttribute(topk_kernel<64>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); ``` But I need to make sure smem doesn't exceed the device's limit. Actually, let me check what the limit is on this device. ```bash python -c "import torch; print(torch.cuda.get_device_properties(0))" ``` Let me check the properties. Actually, let me just fix the code by adding `cudaFuncSetAttribute` calls. Or better, I can reduce shared memory usage for k=64 and k=32. Wait, let me first check the device properties.

Interesting! `shared_memory_per_multiprocessor` is 102400 bytes = 100KB. That's much less than I expected for Blackwell. And `shared_memory_per_block` is probably 48KB or similar. Wait, 102400 bytes per SM? That seems low. Let me check the correct attribute name.

So the default shared memory per block is 49152 bytes (48KB), and the opt-in maximum is 101376 bytes (~99KB). My kernel uses 67584 bytes (~66KB) for k=64 and k=32, which exceeds 48KB. I need to use `cudaFuncSetAttribute` to request more shared memory. I need to add something like: ```cpp cudaFuncSetAttribute(topk_kernel<64>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); ``` But I need to do this for each kernel instantiation that uses more than 48KB. And I need to make sure smem <= 101376. For k=64: smem = 67584 bytes. This is < 101376, so it's allowed with opt-in. For k=32: smem = 67584 bytes. Same. For k=16: smem = 33792 bytes. This is < 49152, so no opt-in needed. For k=8: smem = 16896 bytes. No opt-in needed. Let me update the C++ wrapper to set the attribute for k=32 and k=64. Wait, actually I can set it unconditionally for all k >= 32. Or I can just set it for all kernels to be safe. Let me also think about whether I can reduce shared memory usage instead. For k=64 with 128 threads: - num_warps = 4 - warp_cand = 2048 per warp - Total warp area = 4 * 2048 = 8192 elements = 32768 bytes for vals + 32768 bytes for idxs = 65536 bytes - Block area = 256 * 4 + 256 * 4 = 2048 bytes - Total = 67584 bytes I could reduce the warp work area by not having each thread keep K elements, but instead having each warp process elements sequentially. For example, have the warp maintain a shared top-k and update it as threads process elements. But that's complex. Alternatively, I can process warp merges in two stages to reduce shared memory. But that adds synchronization complexity. The simplest fix is to use `cudaFuncSetAttribute`. Let me do that. Actually, wait - the warp work area is `num_warps * warp_cand`. What if I don't pre-allocate per-warp areas, but instead have all warps share one area? Since warps execute the warp-level sort concurrently, they can't share the same area. But what if I reduce the warp-level sort to use registers instead of shared memory? For warp_cand = 2048 and 32 threads, each thread could keep 64 elements in registers and do a register-based bitonic sort... but then I'd need inter-thread communication for the compare-and-swap, which is exactly what shared memory is for. Alternatively, I could use warp shuffles for the warp-level merge. But comparing and swapping elements across threads with shuffles is possible but complex. Actually, let me try a different approach: instead of a full bitonic sort within the warp, I can use a simpler merge algorithm. Each thread has a sorted top-k. I want to merge 32 sorted top-k lists into one top-k list. I can use a tournament/tree approach: 1. Load all 32*k candidates into shared memory 2. Use a parallel reduction to repeatedly find the maximum For finding the maximum of N elements with 32 threads: - Tree reduction: 32 threads load 32 elements, warp shuffle to find max, repeat. To get top-k, I'd need to do this k times, each time removing the found max. That's k * log2(N) operations with warp shuffles. For k=64 and N=2048: 64 * 11 = 704 warp shuffle operations. Each shuffle is ~1 clock. So ~704 clocks = 0.5 microseconds. Plus memory loads. Wait, but to "remove" the found max, I need to mark it and ignore it in subsequent iterations. I can do this by setting it to -inf after finding it. So the algorithm would be: 1. Load 32*k candidates into shared memory 2. For iter = 0 to k-1: a. Each thread finds max of its portion b. Warp shuffle reduction to find global max c. Thread that owns the max writes it to output[iter] d. Set that element to -inf in shared memory This is O(k * log2(32*k)) operations. For k=64: 64 * 11 = 704 operations. Very fast! And shared memory is just 32*k elements = 2048 floats + 2048 ints = 16KB per warp. For 4 warps, if they share the same area (used sequentially), it's still 16KB. But they run concurrently... Actually, we can do the tournament in a single warp's shared memory area. After all threads have loaded their elements, we do the tournament. But all 32 threads in the warp need to participate. Since warps execute in lockstep, this is fine. The shared memory is per-block, so all warps share it. If each warp uses a different region, we need num_warps * region_size. But we can have all warps use smaller regions. For 4 warps, 4 * 16KB = 64KB. Plus block area 2KB = 66KB. Still the same. Actually, the tournament tree approach doesn't need the full warp_cand * num_warps storage. Each warp needs only its 32*k candidates. But 4 warps still need 4 * 16KB = 64KB. Hmm. What if I do the tournament tree in two stages: 1. First, merge pairs of threads within the warp using shuffles 2. Then merge pairs of warps using shared memory For merging two sorted k-element lists using shuffles: - Thread A has sorted list A[0..k-1] - Thread B has sorted list B[0..k-1] - We want top-k of A union B This can be done with a parallel merge or bitonic merge using shuffles. For k=64, a bitonic merge of 128 elements can be done with warp shuffles if we have enough registers. Actually, I think the cleanest solution is just to use `cudaFuncSetAttribute`. Let me do that. Wait, but if I use `cudaFuncSetAttribute`, I need to call it on the host side before launching the kernel. And I need to make sure the attribute is set for each template instantiation. Let me update the C++ code: ```cpp std::vector<torch::Tensor> topk_forward_cuda(torch::Tensor x, int k) { // ... if (k == 1) { topk1_kernel<<<batch, 256>>>(x_ptr, v_ptr, i_ptr, n); } else { // ... calculate threads, smem ... if (smem > 49152) { switch (k) { case 8: cudaFuncSetAttribute(topk_kernel<8>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); break; case 16: cudaFuncSetAttribute(topk_kernel<16>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); break; case 32: cudaFuncSetAttribute(topk_kernel<32>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); break; case 64: cudaFuncSetAttribute(topk_kernel<64>, cudaFuncAttributeMaxDynamicSharedMemorySize, smem); break; } } switch (k) { // launch kernels } } // ... } ``` Let me implement this.

I need to recompile and test. Since torch extensions are cached, I might need to clear the cache or change the extension name to force recompilation. Let me change the name to "topk_cuda_v2" to force recompilation.

Let me test again.

The kernel now runs without errors for all shapes. Let me run the full check.py to verify correctness.

Great, correctness passes! Now let me run benchmark.py to see the initial performance.

The performance is terrible for most shapes. Let me analyze: 1. shape=0 (1, 131072, 64): solution takes 1.709ms vs eager 0.055ms. That's 31x slower! 2. shape=1 (64, 8192, 8): solution takes 0.059ms vs eager 0.049ms. OK but not great. 3. shape=2 (32, 16384, 32): solution takes 0.278ms vs eager 0.051ms. 5.4x slower. 4. shape=3 (16, 12000, 16): solution takes 0.138ms vs eager 0.047ms. 2.9x slower. 5. shape=4 (128, 4096, 1): solution takes 0.009ms vs eager 0.022ms. Actually faster! The k=1 case is already faster than eager. Good. For the others, especially k=64, the performance is abysmal. The issue is likely: 1. Too much computation per element (heap operations + bitonic sorts) 2. Register pressure causing spills 3. Warp divergence in the heap For shape=0 (1, 131072, 64) with 128 threads: - Each thread processes 1024 elements - For each element, heap_replace with up to 6 iterations - Then bitonic sort 64 elements in registers - Then warp bitonic sort 2048 elements in shared memory - Then block bitonic sort 256 elements in shared memory The heap is the bottleneck. 1024 * 6 * (2 comparisons + 2 assignments) = ~24000 operations per thread. But with branches and divergence, it's probably much worse in practice. Let me profile with ncu or torch.profiler to see what's actually happening. But first, let me think about fundamental optimizations. For the large shape (1, 131072, 64), the key insight is that this is memory-bound. torch.topk achieves ~10 GB/s on this shape, which is about 0.5% of peak bandwidth (1.8 TB/s). Wait, 10 GB/s is terrible! The input is 512KB. At 10 GB/s, that's 0.05ms. At 1.8 TB/s, it should be 0.00028ms. So torch.topk is achieving about 2% of peak bandwidth. Actually, let me recalculate: 512KB at 10 GB/s = 51.2 microseconds = 0.051ms. Yes. At 1800 GB/s, it would be 0.28 microseconds. So torch.topk is at about 0.6% of theoretical peak. The benchmark shows 9.6 GB/s for eager. Hmm, 512KB * 1000 / 9.6GB/s = 53 microseconds. That seems very slow for just reading memory. Wait, the benchmark says gbps=9.632 for shape=0. But 512KB in 0.055ms = 512KB / 0.055ms = 9.3 GB/s. Yes. So even torch.topk is only getting ~10 GB/s. Why so slow? Because top-k is not a simple read - it involves comparison and selection. The memory access pattern is not perfectly sequential (though it should be). My solution gets 0.307 GB/s = 1.7ms for 512KB. That's much worse. The problem is clearly the heap-based approach. For k=64, each element requires a heap sift-down with up to 6 levels, which is ~20-30 operations. With 1024 elements per thread, that's 20,000-30,000 operations. At 1.5GHz, that's 13-20 microseconds per thread. But with 128 threads, the block still takes that long (they run in parallel but the block has only 128 threads, so one warp at a time on an SM). Actually, 128 threads = 4 warps. An SM can run many warps concurrently. But if the kernel is compute-bound, the time is determined by the operations per thread. Wait, but 4 warps on an SM with 1.5GHz should finish 30,000 operations in about 20 microseconds. My actual time is 1700 microseconds. That's 85x slower! Something is very wrong. Let me think... could it be that the bitonic sort in registers is generating massive code that causes instruction cache thrashing? Or register spilling? For K=64, bitonic_sort_regs generates 672 compare-swap operations, each with branches. The compiler might generate a lot of instructions. But 672 instructions is tiny. Or maybe the issue is that the compiler is NOT unrolling the loops properly, and instead is doing dynamic indexing into the register arrays, which causes register spills to local memory? Ah! That's a very important point. In CUDA, arrays in registers (`float vals[64]`) are typically mapped to registers if the indexing is compile-time constant. But if the compiler can't determine the index at compile time, it will spill to local memory. In `heap_replace`, the index `pos` is dynamic. So `vals[pos]` and `idxs[pos]` access local memory, not registers! Local memory is cached but much slower than registers. Similarly, in `bitonic_sort_regs`, the loops over `i`, `j`, `k` are unrolled, so `vals[i]` and `vals[ixj]` should be compile-time constants. But the compiler might not fully unroll the inner loop if it's too large. Wait, `bitonic_sort_regs` has: ```cuda for (int i = 0; i < K; ++i) { int ixj = i ^ j; if (ixj > i) { ... vals[i] ... vals[ixj] ... } } ``` For K=64, this loop has 64 iterations. The `#pragma unroll` should fully unroll it. For each iteration, `i` and `ixj` are compile-time constants (since `j` and `k` are also compile-time known after unrolling). So `vals[i]` should map to specific registers. But `vals` is declared as `float vals[K];` which is a local array. The compiler can either allocate registers for each element or spill to local memory. With 64 floats and 64 ints = 128 values per thread, plus other variables, the total register usage might be high but should fit. For K=64 with 128 threads per block, if each thread uses ~150 registers, the block uses 128 * 150 = 19200 registers. An SM has 65536 registers, so this fits easily. But wait, the warp-level bitonic sort uses shared memory. The block-level sort also uses shared memory. These don't use registers. Let me check if the compiler is spilling registers by looking at the ptxas info. But with `load_inline`, I don't have direct access to ptxas output. Actually, I can add `"-Xptxas=-v"` to the nvcc flags to see register usage and spill information. Let me try that. Actually, a more practical approach: let me first run ncu to profile the kernel and see what's taking time. Or I can use `torch.profiler`. But before profiling, let me think about a much simpler and likely faster approach. **Alternative approach: Two-pass with much simpler first pass** Instead of maintaining a heap per thread, just have each thread find its top-1 or top-2 element. Then at the warp/block level, merge all thread results to get top-k. Wait, that doesn't work. If each thread only keeps top-1, and the true top-64 are all in one thread's chunk, we'd miss 63 of them. But what if we use a larger number of threads and have each thread keep fewer elements? For example, with 1024 threads and k=64, each thread processes 128 elements and keeps top-1. Then we have 1024 candidates, from which we need top-64. A bitonic sort of 1024 elements is fast. But 1024 threads per block is a lot. Actually, what about using 512 threads for k=64? - Each thread processes 256 elements, keeps top-1. - 512 candidates, sort to get top-64. - But the problem is that the top-64 might be spread across many threads. If each thread keeps only top-1, we'd only get 1 true top element if all top-64 are in one thread's chunk. So each thread MUST keep at least k elements to guarantee correctness. Hmm, but what if we use a two-pass approach where the first pass uses many threads and each keeps top-k, and the second pass merges? That's what I was doing originally. The issue is the first pass is too slow. Wait, what if we use a block-level heap in shared memory instead of per-thread heaps? All threads in a block cooperatively insert elements into a shared top-k heap. But a shared memory heap would have massive contention. Each insertion requires atomic operations or serialized access. What about a different data structure? Instead of a heap, use a sorted array and binary search for insertion position. For K=64, binary search is 6 comparisons, then shift up to 64 elements. That's 6 + 64 = 70 operations per insertion. Worse than heap. Or use a tournament tree / selection tree in registers? Complex. Actually, let me reconsider the performance numbers. 1.7ms for shape (1, 131072, 64) means I'm doing something very wrong. Let me check if maybe the issue is compilation with debug info or low optimization. Actually, the `-O3` flag is there. And `--use_fast_math`. Could it be that `cudaFuncSetAttribute` is causing some overhead? No, that's a one-time host-side call. Wait, I just realized something. For k=64 with 128 threads: - block_cand = 4 * 64 = 256 - In the block-level bitonic sort, the loop does: ```cuda for (int kk = 2; kk <= block_cand; kk <<= 1) { for (int j = kk >> 1; j > 0; j >>= 1) { for (int i = tid; i < block_cand; i += blockDim.x) { ``` With blockDim.x = 128 and block_cand = 256: - For kk=2, j=1: i = tid, tid+128. 2 iterations per thread. - For kk=4, j=2,1: 2+2 = 4 iterations per thread. - ... - For kk=256, j=128,64,32,16,8,4,2,1: 8 stages, each with 2 iterations per thread = 16 iterations. Total: about 2+4+6+8+10+12+14+16 = 72 iterations per thread. Each iteration does a compare-and-swap. So ~72 operations per thread. Very fast. The warp-level sort with warp_cand=2048 and 32 threads: - Total stages: log2(2048) * (log2(2048)+1) / 2 = 11 * 12 / 2 = 66 stages. - Each stage: 2048 / 32 = 64 iterations per thread. - Total: 66 * 64 = 4224 iterations per thread. Each with compare-and-swap. - At 1 operation per iteration, ~4224 operations. Fast. The bitonic sort in registers for K=64: - 672 compare-swap operations. Fast. The heap_replace for 1024 elements: - Each call: up to 6 iterations, each with 2 comparisons, 2 assignments. - Say ~20 operations per call. - 1024 calls = ~20480 operations. Total operations per thread: ~20480 + 672 + 4224 + 72 = ~25500 operations. At 1.5GHz, that's ~17 microseconds. But actual time is 1700 microseconds = 1.7ms. That's a 100x gap! Something is seriously wrong. Wait... could it be that the compiler is not unrolling the heap_replace loop, and instead of using registers, it's storing `vals` and `idxs` in local memory? Let me think about this. `vals` is declared as `float vals[K];` in the kernel. K=64. With 128 threads, each thread has a 64-element float array. The compiler tries to map this to registers. But `heap_replace` does dynamic indexing: `vals[pos]` where `pos` is computed at runtime. If the compiler can't prove that `pos` is bounded and known, it might spill the entire array to local memory. For `bitonic_sort_regs`, all indices `i` and `ixj` are compile-time constants after unrolling. So the compiler should definitely use registers there. But `vals` is the SAME array as in `heap_replace`. If the compiler decides to spill it for `heap_replace`, it might spill it everywhere. Actually, in PTX/local memory terminology, a local array can be partially in registers and partially in local memory. But typically, if ANY access is dynamic, the compiler might put the whole array in local memory. If `vals[64]` and `idxs[64]` are in local memory, that's 512 bytes per thread. For 128 threads = 64KB local memory per block. Local memory is backed by global memory (L1/L2 cached). Each heap_replace operation would require multiple local memory loads/stores, which are very slow. This would explain the 100x slowdown! To fix this, I need to ensure the compiler keeps `vals` and `idxs` in registers. One trick is to use individual scalar variables instead of arrays: ```cuda float v0, v1, v2, ..., v63; int i0, i1, i2, ..., i63; ``` But this is extremely tedious to write and doesn't scale. Another trick: use a struct array and hope the compiler optimizes better. Not likely. Actually, for small K, the compiler usually does keep arrays in registers if the total size is small. For K=8, it definitely does. For K=64, it might spill. Let me check by adding `-Xptxas=-v` to see the register count and spill info. Actually, I just realized: with 128 threads per block and each thread using many registers, the block might not fit on an SM if register count is too high. But that would just reduce occupancy, not cause a 100x slowdown. Let me run with profiling flags to see register usage.

Interesting! Looking at the ptxas output: For `topk_kernel<64>`: - 512 bytes stack frame - 0 bytes spill stores, 0 bytes spill loads - 79 registers - 512 bytes cumulative stack size For `topk_kernel<32>`: - 256 bytes stack frame - 48 registers - 256 bytes cumulative stack size For `topk_kernel<16>`: - 128 bytes stack frame - 48 registers For `topk_kernel<8>`: - 64 bytes stack frame - 39 registers The stack frame size equals `K * (sizeof(float) + sizeof(int))` = K * 8 bytes. This means the compiler is storing `vals` and `idxs` in local memory (stack frame), NOT in registers! Even though there are 0 spill stores/loads, the stack frame indicates the arrays are in local memory. Local memory accesses go through the L1 cache, but they're still much slower than registers. This explains the terrible performance! For K=64, every access to `vals[pos]` or `idxs[pos]` in `heap_replace` goes to local memory. With 1024 elements and 6 heap levels per element, that's ~6000 local memory accesses per thread. At L1 latency of ~20-30 cycles, that's ~180,000 cycles = 120 microseconds per thread. But with 128 threads, and cache misses... it could easily be milliseconds. Wait, but ptxas says 0 spill stores/loads. How is that possible if the array is in local memory? Actually, in CUDA terminology, "stack frame" refers to local memory that is explicitly allocated for local arrays, while "spill stores/loads" refer to register spills. The stack frame is where the compiler puts arrays that it decides not to map to registers. It's not counted as "spills" because it's a deliberate allocation. So the compiler is intentionally putting `vals[K]` and `idxs[K]` in local memory because K=64 is too large for register allocation, or because dynamic indexing prevents register mapping. For K=8, the stack frame is 64 bytes. That's `8 * (4 + 4) = 64`. Same pattern. How do I force the compiler to use registers? One approach: use individual variables instead of arrays. But for K=64, that's 128 variables. Very tedious. Another approach: use smaller K per thread. For example, each thread maintains top-8 instead of top-64. But then we need more threads or more passes. Wait, I have an idea. Instead of each thread maintaining a heap of K elements, what if each thread maintains only a few elements (like top-2 or top-4), and then we do multiple levels of merging? For example: - Each thread processes n/blockDim.x elements and keeps top-2. - Warp merge: 32 threads * 2 = 64 candidates. Bitonic sort to get top-k. - For k <= 64, if each thread keeps top-2, the warp has 64 candidates. A bitonic sort of 64 gives us top-k directly! Wait, that's brilliant! For k=64, if each thread keeps top-2 and there are 32 threads in a warp, we get 64 candidates. Bitonic sort 64 to get top-64. But is this correct? NO! If each thread only keeps top-2 from its chunk, the true top-64 might include elements that are the 3rd or 4th largest in a thread's chunk. For example, if one thread's chunk has 100 very large elements, its top-2 might be #1 and #2, but the global top-64 might include its #3, #4, etc. If we only keep top-2 per thread, we'd miss those. So each thread must keep AT LEAST K elements. There's no way around this for correctness. Hmm, unless we process elements in a way that guarantees the top-k can't be concentrated in one thread's chunk. With strided access, each thread sees a representative sample of the row. For random data, the probability that one thread's chunk contains many top-k elements follows a hypergeometric distribution. But we can't rely on probabilities for correctness. Actually, we CAN guarantee correctness with a two-pass approach: 1. First pass: each thread processes a chunk and keeps top-k. But if threads keep only top-m where m < k, we can't guarantee correctness in one pass. 2. But we can do MULTIPLE passes over the data. For example: - Pass 1: each thread keeps top-1. Write all top-1s to global memory. - Pass 2: each block reads the top-1s and keeps top-k. But with 256 threads, pass 1 produces 256 candidates. For k=64, 256 candidates might not contain all top-64 if the true top-64 are all in one thread's chunk. Wait, each thread's top-1 is its maximum. If the global top-64 are all in one thread's chunk, then only 1 of them is the thread's top-1. The other 63 are NOT in the candidate set. So this fails. The fundamental issue: each thread MUST keep at least k elements to guarantee that the union of all thread results contains the global top-k. Proof: Consider a row where the largest k elements are all in one thread's chunk. That thread must keep all k of them. Any other thread keeping fewer than k might drop some of them. So per-thread storage of k elements is necessary. But does it need to be in registers? No, it can be in shared memory. But then multiple threads compete for shared memory. Wait, what if I use a completely different approach: a block-level algorithm where all threads in a block cooperatively process one row? **Block-level streaming top-k:** 1. Divide shared memory into two buffers: current top-k and incoming chunk. 2. Each block loads a chunk of elements (e.g., 1024 elements) into shared memory. 3. Sort the chunk and merge with current top-k. 4. Repeat until all elements processed. For n=131072 and block size 1024, we need 128 iterations. Each iteration: - Load 1024 floats = 4KB - Sort 1024 elements in shared memory: 1024 * 10 * 10 / 4 = 25600 comparisons - Merge with top-64: we can use a bitonic merge of 1088 elements (1024 + 64) But sorting 1024 elements in shared memory with 256 threads is fast. And we only need to do it 128 times. Actually, there's an even better approach: **partial bitonic sort**. Instead of fully sorting each chunk, we can use a bitonic network that only keeps the top-k elements. For a chunk of size M and keeping top-k: 1. Load M elements into shared memory 2. Apply a partial bitonic sort that sorts only the first k elements 3. The rest can be discarded There's a known algorithm for this: after each merge stage, only propagate the upper half if it's part of the top-k region. Actually, a simpler approach for small k relative to chunk size: 1. Load chunk into shared memory 2. Find the k-th largest element using a selection algorithm (e.g., quickselect or bitonic select) 3. Filter elements >= k-th largest 4. Sort the top-k elements But finding the k-th largest in shared memory with many threads is non-trivial. Another idea: use a **register-based bitonic sort for chunks**. Since we need to keep k elements per thread, and k=64, what if we process data in chunks of 64, sort each chunk in registers, and keep the best 64 using a merge? Wait, I already thought about this. The problem was register pressure. But what if we process chunks of 32 or 16 instead? For chunk size 32 and k=64: 1. Load 32 elements into registers 2. Sort them (bitonic sort of 32 = 400 comparators) 3. Merge with top-64 using sequential merge (64 comparisons) 4. Repeat Total registers needed: chunk_vals[32], chunk_idxs[32], topk_vals[64], topk_idxs[64], temp for merge[64]. That's 256 values. Still a lot, but maybe the compiler handles it better than a dynamic heap. Actually, if I write the merge carefully with compile-time known indices, the compiler might keep everything in registers. Let me try a different first-pass strategy: **First pass: chunk-based bitonic sort + sequential merge** For each thread: - Maintain `topk_vals[K]` and `topk_idxs[K]` in registers - Process elements in chunks of CHUNK_SIZE (e.g., 16 or 32) - For each chunk: - Load CHUNK_SIZE elements into `chunk_vals`, `chunk_idxs` - Bitonic sort `chunk_vals` descending - Merge `topk_vals` (sorted descending, size K) with `chunk_vals` (sorted descending, size CHUNK_SIZE) into a new sorted array of size K The merge can be done sequentially: ```cuda float merged_vals[K]; int merged_idxs[K]; int a = 0, b = 0; for (int i = 0; i < K; i++) { if (a < K && (b >= CHUNK_SIZE || topk_vals[a] >= chunk_vals[b])) { merged_vals[i] = topk_vals[a]; merged_idxs[i] = topk_idxs[a]; a++; } else { merged_vals[i] = chunk_vals[b]; merged_idxs[i] = chunk_idxs[b]; b++; } } // Copy merged back to topk ``` This uses a temporary array of K. Total arrays: topk[K], chunk[CHUNK_SIZE], merged[K] = 2K + CHUNK_SIZE values. For K=64, CHUNK_SIZE=32: 160 values. Much less than 256 for heap + bitonic sort. But the issue is still dynamic indexing in the merge. `topk_vals[a]` where `a` is dynamic. The compiler might put these arrays in local memory. However, for the bitonic sort of chunk, all indices are compile-time known (if CHUNK_SIZE is a template parameter and we fully unroll). So `chunk_vals` might be in registers. But `topk_vals` and `merged_vals` would still be in local memory due to dynamic indexing. Hmm, is there any way to do top-k with K=64 entirely with compile-time known indices? Not really, unless we unroll all possible code paths, which is impossible. Wait, what about using shared memory for the per-thread top-k? If we allocate shared memory for all threads' top-k arrays, then dynamic indexing is fine (shared memory supports dynamic indexing). For 256 threads and K=64: - Shared memory for vals: 256 * 64 * 4 = 64KB - Shared memory for idxs: 256 * 64 * 4 = 64KB (using int32) - Total: 128KB This fits within the 48KB default limit? No, 128KB > 48KB. We'd need opt-in shared memory. But wait, if we use 128 threads: - 128 * 64 * 8 = 64KB. Still > 48KB. If we use 64 threads: - 64 * 64 * 8 = 32KB. Fits in 48KB! - Each thread processes 131072 / 64 = 2048 elements. - 64 threads = 2 warps. But with only 64 threads, we have less parallelism. However, each thread still runs on the same SM, so the total throughput might be similar if memory-bound. Actually, let me think about this differently. If we put the per-thread top-k in shared memory: 1. Each thread loads elements, compares with its top-k in shared memory 2. Insertion/heap operations use shared memory (fast, no local memory) 3. After all elements, sort each thread's top-k in shared memory or registers 4. Then merge across threads For 64 threads, shared memory = 64 * 64 * 8 = 32KB. Plus we need workspace for merging. Total might be ~40KB, which fits in 48KB! And each thread processes 2048 elements. Heap operations in shared memory: each element requires reading vals[0] (root), then if larger, doing a sift-down with reads/writes to shared memory. Shared memory latency is ~20-30 cycles, much faster than local memory which goes to L1 and might miss. But we still have warp divergence in the heap. And shared memory bank conflicts. Actually, wait. If each thread has its own 64-element region in shared memory, and threads access different regions, there are no bank conflicts (assuming proper alignment). For example, thread tid accesses `s_vals[tid * 64 + i]`. The bank is `(tid * 64 + i) % 32`. Since 64 is a multiple of 32, `tid * 64 % 32 = 0`. So thread 0 accesses banks 0..31, thread 1 accesses banks 0..31, etc. ALL threads access the SAME banks! Major conflicts! To avoid this, we need padding: ```cuda s_vals[tid * (64 + 1) + i] // pad by 1 ``` Bank = `(tid * 65 + i) % 32`. For tid=0: banks 0..31. tid=1: banks 1..32%32=0..30,31,0,1. So adjacent threads access shifted banks. Much better! With pad=1, total shared memory = 64 * 65 * 8 = 33.6KB. Still under 48KB. For the heap sift-down in shared memory: ```cuda float v = new_val; int idx = new_idx; int pos = 0; while (true) { int left = 2 * pos + 1; int right = 2 * pos + 2; float left_val = s_warp_vals[tid * 65 + left]; float right_val = s_warp_vals[tid * 65 + right]; int smallest = pos; if (left < K && left_val < v) smallest = left; if (right < K && right_val < v) smallest = right; if (smallest == pos) { s_warp_vals[tid * 65 + pos] = v; s_warp_idxs[tid * 65 + pos] = idx; break; } s_warp_vals[tid * 65 + pos] = s_warp_vals[tid * 65 + smallest]; s_warp_idxs[tid * 65 + pos] = s_warp_idxs[tid * 65 + smallest]; pos = smallest; } ``` This accesses shared memory with dynamic indices. But since each thread accesses its own padded region, there are no inter-thread bank conflicts. Intra-thread accesses might have conflicts if `pos` and `smallest` map to the same bank. But `smallest = 2*pos+1` or `2*pos+2`, so the bank difference is `(2*pos+1 - pos) % 32 = (pos+1) % 32`. For most pos values, this is a different bank. So minimal conflicts. But wait, all threads in a warp execute the heap operations in lockstep ONLY if they all take the same path. But `if (v > vals[0])` causes divergence. Some threads do the heap replace, others skip. When threads diverge, they execute sequentially. For random data, after the heap is full (after K elements), about K/n of new elements are larger than the root. For n=2048 and K=64, that's 64/2048 = 3.125%. So only 3% of elements trigger heap_replace. Most of the time, threads just do the comparison `v > vals[0]` and skip. With 64 threads in a block, each warp has 32 threads. The probability that any thread in a warp needs to do heap_replace is low. When no threads need it, the warp skips quickly. When some threads need it, those threads execute the heap_replace while others are masked off. This seems promising! But the overhead of shared memory access is still higher than registers. Actually, let me reconsider the fundamental issue. The problem is that the compiler puts arrays in local memory. But what if I help the compiler by making the arrays smaller and using a different data structure? For K=64, what if I use a **pair of 32-element arrays** and a merge tree? - Each thread maintains two heaps of size 32. - Process elements, insert into appropriate heap. - After all elements, merge the two heaps into one sorted array of 64. - Then sort 64 elements using bitonic sort. The heaps of size 32 might fit in registers better. But I'm not sure if the compiler would map them to registers either. Actually, looking at ptxas output again: - topk_kernel<8>: 39 registers, 64 bytes stack frame - topk_kernel<16>: 48 registers, 128 bytes stack frame - topk_kernel<32>: 48 registers, 256 bytes stack frame - topk_kernel<64>: 79 registers, 512 bytes stack frame The stack frame grows linearly with K, confirming it's the vals/idxs arrays. But notice that for K=32, it uses only 48 registers. For K=64, 79 registers. The arrays are definitely in local memory. But wait - 0 spill stores/loads means the compiler is NOT spilling registers. It's just allocating local memory for the arrays. The registers are used for loop variables, pointers, etc. Now, the key question: why is local memory so slow? Local memory is cached in L1. For sm_120 (Blackwell), L1 is 128KB per SM. With 128 threads and 512 bytes each = 64KB local memory. This fits in L1! So local memory accesses should hit L1 cache and be relatively fast (a few cycles). But the actual performance is 100x slower than expected. Why? Ah! I think I know. The issue might not be local memory at all. Let me reconsider the algorithm. For shape (1, 131072, 64) with batch=1 and 128 threads: - Only 1 block is launched! - 128 threads = 4 warps - The GPU has 188 SMs, each can run many warps. But we only have 4 warps total! - So only 1 SM is used (or a fraction of 1 SM), and the rest are idle. This is a critical insight! For batch=1, we launch 1 block. With 128 threads, we have very little parallelism. The kernel is compute-bound on a single SM. torch.topk, on the other hand, probably uses many more threads per block or launches multiple blocks per row. It might process the row with hundreds or thousands of threads. For memory-bound operations with batch=1, we need to launch enough threads to saturate memory bandwidth. A single block with 128 threads cannot saturate 1.8 TB/s. So the fundamental issue is lack of parallelism for small batch sizes! Let me think about how to fix this. **Solution: Use many more threads per row** Instead of 1 block per row with 128-256 threads, use multiple blocks per row, or much larger blocks. Option A: One block per row with 1024 threads. - For k=64, each thread keeps top-64 in local memory. - 1024 threads * 512 bytes = 512KB local memory per block. - This might be too much local memory, causing L1 thrashing or spilling to L2. Option B: Multiple blocks per row. - Block 0 processes elements 0..N/2, block 1 processes N/2..N. - Each block keeps top-k in shared memory or global memory. - Write block results to global memory. - A final block merges all block results. This is the two-pass approach I considered earlier! Let me revisit it. **Two-pass approach with many blocks:** Pass 1: Local top-k per block - Launch `batch * num_blocks_per_row` blocks - Each block has `THREADS` threads (e.g., 256) - Each thread processes a subset of elements and keeps top-k in registers/local memory - Block writes `THREADS * k` candidates to global memory Pass 2: Merge candidates - Launch `batch` blocks - Each block reads `num_blocks_per_row * THREADS * k` candidates - Bitonic sort in shared memory - Write top-k to output For batch=1, n=131072, k=64: - If we use num_blocks_per_row = 8, each block has 256 threads. - Each thread processes 131072 / (8 * 256) = 64 elements. - Each thread keeps top-64 from 64 elements. Since 64 == 64, each thread just sorts its 64 elements. - Block writes 256 * 64 = 16384 candidates to global memory. - Pass 2 reads 8 * 16384 = 131072 candidates. That's a lot! Bitonic sort of 131072 in shared memory would need 131072 * 8 = 1MB of shared memory. Too much. Wait, that's too many candidates. We need fewer blocks or fewer threads. Alternative: num_blocks_per_row = 4, threads = 256. - Each thread processes 131072 / (4 * 256) = 128 elements. - Each thread keeps top-64 from 128 elements using heap. - Block writes 256 * 64 = 16384 candidates. - Pass 2 reads 4 * 16384 = 65536 candidates. Still too many for shared memory. Hmm. The issue is that with many blocks, the number of candidates in pass 2 grows. What if we use a tree reduction? - Pass 1: 32 blocks per row, each with 256 threads. Each thread keeps top-2 (not top-64). - Total candidates = 32 * 256 * 2 = 16384. - Pass 2: bitonic sort 16384 in shared memory. Need 128KB shared memory. Fits with opt-in! - But is top-2 per thread sufficient? No! As I proved earlier, each thread must keep at least K elements. Wait, what if we use more threads but each thread still keeps K elements? The total candidates = num_blocks * threads * K. We want this to be manageable in pass 2. For pass 2, the maximum we can sort in shared memory is limited by shared memory size. With 101KB opt-in shared memory, we can store ~100KB / 8 bytes per element = ~12800 elements (values + int32 indices). Actually, 100KB / 8 = 12500. For power of 2, maybe 8192 or 16384. If we can sort 8192 candidates in pass 2: - num_blocks * threads * K <= 8192 - For K=64: num_blocks * threads <= 128 - If threads=256: num_blocks <= 0.5. So only 1 block. This brings us back to 1 block per row for k=64. Which means for batch=1, only 1 block total. Very little parallelism. But wait, what if pass 2 doesn't fully sort all candidates? What if it uses a selection algorithm (like quickselect) to find the top-k without full sort? For 8192 candidates and k=64, we can: 1. Load into shared memory 2. Use a parallel quickselect or bitonic select to find the 64th largest element 3. Then filter and sort the top 64 Parallel quickselect on GPU is tricky but possible. Alternatively, we can use a "partial bitonic sort" where we only run enough stages to get the top-k elements sorted. For bitonic sort, after the stages where kk <= k, the first k elements are not necessarily sorted. But after kk = 2*k, they might be... Actually, there's a known result: in a bitonic sort network, after completing all stages for kk = 2, 4, 8, ..., 2*k, the first k elements contain the k smallest (or largest) elements, but not necessarily sorted. Then a final sort of the first k elements gives the answer. Wait, I'm not sure about this. Let me think. In bitonic sort, the array is divided into sorted sequences that double in size. After stage kk, sequences of size kk are sorted. So after kk = k, sequences of size k are sorted. But these might be ascending or descending depending on position. If we want the top-k of an array of size N, after running bitonic sort stages up to kk = k, do the first k elements contain the global top-k? I don't think so. Consider N=8, k=2. After kk=2, we have 4 descending pairs. The first 2 elements are the max of each pair, but the global top-2 might both be in the same pair. So partial bitonic sort doesn't guarantee the top-k are in the first k positions. What about a tournament tree / heap-based merge in pass 2? For pass 2, we have M candidates (e.g., 8192) and need top-k (64). We can: 1. Load all candidates into shared memory 2. Build a max-heap of size M? No, M is too large. 3. Use a block-level parallel reduction to repeatedly find the max, k times. Block-level parallel reduction to find max of M elements: - Tree reduction over M elements with B threads - For M=8192 and B=256, each thread finds max of 32 elements, then warp shuffle reduction, then block reduction. - Total: ~M operations + B log(B) operations. Doing this k=64 times: 64 * (8192 + 256*8) = 64 * 10240 = 655360 operations. With 256 threads at 1.5GHz: ~655360 / 256 / 1.5e9 * 1e6 = 1.7 microseconds. Wait, that's fast! But we need to "remove" each found max so we don't find it again. We can set it to -inf after finding it. So pass 2 algorithm: ```cuda // Load M candidates into shared memory for (int iter = 0; iter < k; iter++) { // Find max of all M elements using parallel reduction float max_val = -INF; int max_idx = -1; for (int i = tid; i < M; i += blockDim.x) { if (s_vals[i] > max_val) { max_val = s_vals[i]; max_idx = i; } } // Block reduction to find global max // ... warp shuffle + shared memory reduction ... if (tid == 0) { out_vals[iter] = global_max_val; out_idxs[iter] = global_max_idx; s_vals[global_max_idx] = -INF; // Remove for next iteration } __syncthreads(); } ``` For M=8192, B=256: - Local max per thread: 32 comparisons - Warp reduction: 5 shuffle stages - Block reduction: 3 stages (8 warps -> 1) - Total per iteration: ~40 operations per thread - 64 iterations: 2560 operations per thread - At 1.5GHz: ~1.7 microseconds This is extremely fast! And it only needs M * 8 bytes of shared memory (for vals + idxs). For M=8192: 64KB. Well under 48KB limit! Wait, M=8192, each element is float + int = 8 bytes. 8192 * 8 = 65536 bytes = 64KB. This is ABOVE the 48KB default limit. We'd need opt-in shared memory for 64KB. The opt-in limit is 101KB, so 64KB is fine. Actually, if we store indices as int16 (for n <= 32768) or we don't store indices in shared memory and compute them separately... no, we need indices. For M=8192 with int32 indices: 8192 * 8 = 64KB. With opt-in, this works. But if we want to avoid opt-in, we can use M=4096: 4096 * 8 = 32KB. This fits in 48KB. For M=4096 and K=64: num_blocks * threads <= 4096 / 64 = 64. If threads=256: num_blocks <= 0.25. So 1 block. If threads=128: num_blocks <= 0.5. So 1 block. If threads=64: num_blocks <= 1. So 1 block. Still 1 block per row for k=64. For k=32: M=4096, num_blocks * threads <= 4096 / 32 = 128. threads=256: num_blocks <= 0.5. 1 block. threads=128: num_blocks <= 1. 1 block. threads=64: num_blocks <= 2. 2 blocks! For k=16: M=4096, num_blocks * threads <= 4096 / 16 = 256. threads=256: num_blocks <= 1. 1 block. threads=128: num_blocks <= 2. 2 blocks. For k=8: M=4096, num_blocks * threads <= 4096 / 8 = 512. threads=256: num_blocks <= 2. 2 blocks. So even with M=4096, for most cases we still have only 1 block per row. Hmm. What about using M=8192 (64KB shared memory, needs opt-in)? For k=64: num_blocks * threads <= 8192 / 64 = 128. threads=256: num_blocks <= 0.5. 1 block. threads=128: num_blocks <= 1. 1 block. threads=64: num_blocks <= 2. 2 blocks. Still 1-2 blocks for k=64. For k=32: num_blocks * threads <= 8192 / 32 = 256. threads=256: num_blocks <= 1. 1 block. threads=128: num_blocks <= 2. 2 blocks. For k=16: num_blocks * threads <= 8192 / 16 = 512. threads=256: num_blocks <= 2. 2 blocks. For k=8: num_blocks * threads <= 8192 / 8 = 1024. threads=256: num_blocks <= 4. 4 blocks. So for k=64, we're stuck with 1-2 blocks per row. For batch=1, that's 1-2 blocks total. Not enough to saturate the GPU. Wait, what if we use M=16384 (128KB shared memory)? That would need opt-in and might not fit (limit is 101KB). So no. What about a different approach: instead of having pass 2 sort candidates in shared memory, have pass 2 use a **persistent block** that processes multiple rows? Or use a **single global memory merge**? Actually, for pass 2 with M candidates, we can store them in global memory and use a simple parallel merge tree across multiple blocks. But then pass 2 also becomes a multi-block kernel. Two-pass with global memory: - Pass 1: `batch * num_blocks_per_row` blocks, each writes `threads * k` candidates to global memory. - Total candidates in global memory: `batch * num_blocks_per_row * threads * k`. - Pass 2: For each row, launch `merge_blocks` blocks that read all candidates and produce top-k. For pass 2, if we have many candidates, we can use a tree reduction: - Block 0 reads first half, extracts top-k. - Block 1 reads second half, extracts top-k. - Then a final block merges the two top-k lists. But this requires multiple passes. We can generalize to a log-depth tree. Actually, for k=64, each merge of two top-k lists produces a top-k list. So we can do a binary tree reduction: - Level 0: num_blocks leaves, each has top-k - Level 1: num_blocks/2 nodes, each merges 2 top-k lists - ... - Root: 1 node, final top-k For num_blocks = 32: - Level 0: 32 blocks - Level 1: 16 blocks - Level 2: 8 blocks - Level 3: 4 blocks - Level 4: 2 blocks - Level 5: 1 block Total kernel launches: 6. But each launch is small. Wait, but for batch=1, launching 6 kernels in sequence with only a few blocks each has significant kernel launch overhead. Alternatively, we can do the reduction within a single kernel using multiple thread blocks and global memory atomics/grid synchronization. But that's complex. Actually, let me think about this from a completely different angle. **What if we use many threads per row but have them all operate on a shared top-k in global memory?** No, atomic contention would be terrible. **What if we use a warp-based approach where each warp processes a chunk of elements and maintains its top-k in registers, then all warps merge using a tree in shared memory?** With 1024 threads = 32 warps: - Each warp (32 threads) processes n/1024 elements per thread. - Each thread keeps top-k in local memory. - Warp merge: 32 threads merge their top-k into warp top-k using shared memory. - Block merge: 32 warps' top-k are merged using a tree in shared memory. For n=131072 and 1024 threads: - Each thread processes 128 elements. - Local top-k from 128 elements using heap in local memory. - Warp merge: bitonic sort 32*k elements in shared memory. Wait, I already had this design but with 128 threads. The issue was that local memory was slow. But if I use 1024 threads, each thread processes only 128 elements instead of 1024. The heap work is 8x less per thread. And with 1024 threads per block, I have 32 warps. The GPU can run multiple blocks concurrently, but for batch=1, only 1 block. Still, 1024 threads is much more parallelism than 128 threads. But can 1024 threads fit with K=64? - Shared memory for warp merge: 32 warps * 32 * 64 * 8 bytes = 32 * 2048 * 8 = 524288 bytes = 512KB. WAY too much. So I can't use 1024 threads with K=64 because the warp merge needs too much shared memory. What if I use 512 threads? - Shared memory: 16 warps * 32 * 64 * 8 = 16 * 2048 * 8 = 262144 bytes = 256KB. Still too much. 256 threads: 8 warps * 32 * 64 * 8 = 8 * 2048 * 8 = 131072 bytes = 128KB. Still too much for default 48KB, and exceeds opt-in 101KB. 128 threads: 4 warps * 32 * 64 * 8 = 4 * 2048 * 8 = 65536 bytes = 64KB. With opt-in, this fits! And we already used this configuration. But 128 threads = 4 warps is still not enough parallelism for batch=1. So the single-block-per-row approach fundamentally cannot use enough threads for batch=1 without exceeding shared memory. **Solution: Multi-block per row with global memory reduction tree** For batch=1, launch many blocks per row. Each block processes a chunk and writes its top-k to global memory. Then use a reduction kernel to merge the block results. Let me design this: Pass 1: `batch * num_blocks` blocks, each with `THREADS` threads. - Each thread processes chunk_size = n / (num_blocks * THREADS) elements. - Each thread keeps top-k in local memory. - Block-level warp merge + block merge to produce final block top-k. - Write block top-k to global memory buffer. Wait, but with many blocks, each block still has the same shared memory requirements for warp/block merge. But we can simplify: instead of doing block-level merge, each thread just writes its top-k to global memory. Then pass 2 merges all thread results. Pass 1 simplified: - `batch * num_blocks` blocks, `THREADS` threads each. - Each thread processes chunk_size elements, keeps top-k. - Writes top-k directly to global memory at offset `(batch_idx * num_blocks + block_idx) * THREADS * k + tid * k`. No shared memory needed in pass 1! Just local memory for top-k. Pass 2: For each row, read `num_blocks * THREADS * k` candidates and extract top-k. For n=131072, k=64, THREADS=256: - If num_blocks=4: chunk_size = 131072 / (4*256) = 128 elements per thread. - Candidates = 4 * 256 * 64 = 65536. - Pass 2 needs to find top-64 from 65536 candidates. For pass 2 with 65536 candidates per row: - We can launch 1 block per row with 256 threads. - Load candidates in tiles into shared memory. - Use the tournament tree / repeated max reduction to extract top-64. Wait, loading 65536 candidates is 65536 * 8 = 512KB. That's too much for shared memory. We'd need to process in tiles. Tile-based pass 2: 1. Maintain a top-k array in shared memory (or registers) 2. Load a tile of candidates from global memory into shared memory 3. Merge tile with current top-k 4. Repeat For merging a tile of 1024 candidates with current top-k (64): - Load tile into shared memory (1024 * 8 = 8KB) - Each thread finds its local top-k from the tile - Block-level merge to update global top-k This is getting complex. Let me think of a simpler pass 2. Actually, for pass 2, if we have M candidates and need top-k, and M is large (e.g., 65536), we can: 1. Use a block-level heap in shared memory of size k. 2. Stream through all M candidates in tiles. 3. For each candidate, if it's larger than heap root, replace and sift-down. This is the same heap approach but in shared memory. And with many threads (e.g., 256), each thread can process multiple candidates. For M=65536 and 256 threads: - Each thread processes 256 candidates. - For each candidate, do `if (cand > root) replace_root(heap, cand)`. - But `root` is in shared memory, and `replace_root` requires synchronizing all threads because the heap is shared! We can't have 256 threads independently updating the same heap. That requires atomics or serialization. What if each thread maintains its own local top-k from its subset of candidates, and then we merge all thread results? That's back to the same problem! We'd have 256 * 64 = 16384 candidates to merge. OK, let me try a completely different approach. What if pass 2 uses a **bitonic sort on chunks**? For M=65536 candidates stored in global memory: - We can't sort all 65536 in one go. - But we can sort chunks of 1024 in shared memory, then merge chunks. This is a standard merge sort: 1. Sort chunks of 1024 in shared memory 2. Merge sorted chunks pairwise, then quadruple, etc. For M=65536, chunk=1024: - Level 0: 64 sorted chunks - Level 1: 32 chunks of 2048 - Level 2: 16 chunks of 4096 - Level 3: 8 chunks of 8192 - Level 4: 4 chunks of 16384 - Level 5: 2 chunks of 32768 - Level 6: 1 chunk of 65536 But we only need the top-64! After each merge, we can truncate to top-64. For merging two sorted lists of size 1024 into top-64: - Use a parallel merge: each thread finds its position in the merged list - Or use a bitonic merge network Actually, merging two sorted lists A and B into top-k is simple: - Compare A[0] with B[0], take larger. Then compare A[1] with B[0] (or A[0] with B[1]), etc. - This is a sequential merge, O(k) time. - With 64 threads, each thread can find one element of the top-64 in O(1) time if we use binary search. Wait, thread i can find the i-th largest element by finding the split point between A and B. For descending lists, the i-th largest is either A[a] or B[b] where a + b = i. We can binary search for the split. But this is complex to implement efficiently. Let me step back. What do the fastest GPU top-k implementations do? I know that NVIDIA's CUB library has a `DeviceRadixSort` and `DeviceSelect`. PyTorch's topk uses a radix select for large k and a bitonic sort for small k. For our case (small k: 1, 8, 16, 32, 64), the fastest approach on GPU is typically: 1. Each thread loads a chunk 2. Each thread maintains a small sorted array of top-k in registers 3. Threads merge within warps using shuffles or shared memory 4. Warps merge within blocks 5. Blocks write to global memory 6. A second kernel merges block results The key insight from CUB and other libraries is that they use **register-level sorting networks** for small arrays and **warp-level shuffles** for merging. Actually, I just realized something important. In my current implementation, the arrays are in local memory. But what if I use a **struct array** and access it with compile-time known indices? Wait, the compiler puts arrays in local memory when there's dynamic indexing. But in `heap_replace`, the index `pos` depends on the data. The compiler cannot know `pos` at compile time. However, for K=64, the heap has only 6 levels. I can manually unroll the entire heap operation without a loop: ```cuda // Level 0 int p0 = 0; int l0 = 1, r0 = 2; int s0 = p0; if (l0 < K && vals[l0] < v) s0 = l0; if (r0 < K && vals[r0] < v) s0 = r0; if (s0 == p0) { vals[p0] = v; idxs[p0] = idx; return; } vals[p0] = vals[s0]; idxs[p0] = idxs[s0]; // Level 1 int p1 = s0; int l1 = 2*p1+1, r1 = 2*p1+2; int s1 = p1; if (l1 < K && vals[l1] < v) s1 = l1; if (r1 < K && vals[r1] < v) s1 = r1; if (s1 == p1) { vals[p1] = v; idxs[p1] = idx; return; } vals[p1] = vals[s1]; idxs[p1] = idxs[s1]; // ... repeat for 6 levels ``` But `vals[l0]`, `vals[r0]`, etc. still use dynamic indices `l0`, `r0`, which depend on `s0`, which depends on data. The compiler still can't map these to specific registers. Unless K is very small (like 8), the compiler will still use local memory. Actually, wait. For K=8, ptxas shows 39 registers and 64 bytes stack frame. Even K=8 uses local memory! So manual unrolling won't help for K > 4 or so. The compiler simply refuses to map arrays to registers when there's any dynamic indexing, regardless of size. So for K >= 8, the arrays WILL be in local memory. Period. Given this constraint, how can I make local memory faster? 1. **Reduce the number of local memory accesses** 2. **Increase occupancy so local memory accesses are interleaved and L1 cache is well-utilized** For point 1: instead of a heap with up to 6 levels per element, use a simpler data structure. What if we just load all elements into local memory, sort them, and take top-k? For n/thread = 128 and K=64: load 128 elements, sort 128 elements using bitonic sort, take top-64. Sorting 128 elements with bitonic sort: 128 * 7 * 7 / 4 = 1568 comparators. vs. heap: 128 * 6 * 2 = 1536 comparisons (roughly). Comparable. But bitonic sort has NO dynamic indexing! All indices are compile-time known. Wait, but we need to load 128 elements into a local array first. That's 128 floats + 128 ints = 1024 bytes of local memory. Then bitonic sort on this local array... but the array is in local memory, so each comparator does 2 loads and 2 stores to local memory. Hmm, same problem. What if we load elements directly into registers using scalar variables? For chunk_size = 8: ```cuda float v0, v1, v2, v3, v4, v5, v6, v7; int i0, i1, i2, i3, i4, i5, i6, i7; // Load v0 = row_x[tid]; i0 = tid; v1 = row_x[tid+blockDim.x]; i1 = tid+blockDim.x; // ... // Bitonic sort of 8 // All indices are compile-time known, so all variables stay in registers! ``` Then after sorting 8 elements in registers, merge with top-64 (in local memory). The merge requires dynamic indexing. But we can process 8 elements at a time: load 8, sort 8, then merge the sorted 8 with the sorted 64. The merge can be done with a sequential loop over the 8 elements, inserting each into the 64-array. For each of the 8 elements: - Binary search its position in the 64-array: 6 comparisons - Shift elements to make room: up to 64 moves - Insert: 1 move - Total: ~71 operations per element, all on local memory. 8 elements * 71 = 568 operations. Plus bitonic sort of 8 = 36 comparators (in registers). Total ~604 operations per chunk. For 128 elements = 16 chunks: 16 * 604 = 9664 operations. vs. heap: 128 * 6 * 4 = 3072 operations. The heap is still fewer operations. But local memory operations are cached. Let me think about which is faster in practice. Actually, I realize the real issue might not be the heap algorithm itself, but the fact that with 128 threads and batch=1, we have terrible GPU utilization. Let me verify by checking what happens with batch=64, n=8192, k=8: - 64 blocks of 256 threads = 16384 threads total. - My kernel takes 0.059ms vs eager 0.049ms. Only 20% slower! - For this shape, local memory might still be a factor, but with many blocks, the GPU is well-utilized and the L1 cache is effective. For batch=1, n=131072, k=64: - 1 block of 128 threads = 128 threads total. - My kernel takes 1.709ms vs eager 0.055ms. 31x slower! - With only 1 block, the GPU is mostly idle. The single block has 4 warps, which can run on one SM. But the local memory accesses within those warps might be the bottleneck. So the primary issue for batch=1 is lack of parallelism, and the secondary issue is local memory latency. To fix batch=1, I MUST use multiple blocks per row. Let me design a multi-block approach. **Multi-block design for pass 1:** Launch `num_blocks` blocks per row. Each block has `THREADS` threads. - Block `b` processes elements `[b * chunk_size, (b+1) * chunk_size)`. - chunk_size = n / num_blocks. - Each thread within the block processes `chunk_size / THREADS` elements. - Each thread keeps top-k in local memory. - Block writes `THREADS * k` candidates to global memory. Actually, instead of having each thread write its top-k, we can do a block-level merge within each block to reduce candidates. This reduces pass 2 work. For block-level merge with `THREADS` threads, we merge `THREADS * k` candidates into `k` candidates using shared memory bitonic sort. This is what I already do in my single-block kernel. The shared memory for this is `THREADS * k * 8` bytes. For THREADS=256, k=64: 256 * 64 * 8 = 131072 bytes = 128KB. Too much for default, too much for opt-in (101KB). For THREADS=128, k=64: 128 * 64 * 8 = 65536 bytes = 64KB. Needs opt-in. For THREADS=64, k=64: 64 * 64 * 8 = 32768 bytes = 32KB. Fits in default! But with 64 threads per block, each thread processes more elements. For n=131072 and num_blocks=8: - chunk_size = 131072 / 8 = 16384 - Each thread processes 16384 / 64 = 256 elements - Heap work: 256 * 6 * 4 = ~6000 operations And we have 8 blocks! So 8 * 64 = 512 threads total. Much better parallelism. For pass 2: 8 blocks * 64 threads * 64 k = 32768 candidates per row. Pass 2 needs to find top-64 from 32768 candidates. 32768 candidates = 32768 * 8 = 256KB. Too much for shared memory. So pass 2 needs to stream through candidates in tiles. Pass 2 design: - Launch 1 block per row with 256 threads. - Shared memory: block top-k heap (64 * 8 = 512 bytes) + tile buffer. - Tile size: 1024 candidates = 8KB. - Stream through all 32768 candidates in tiles of 1024. - For each tile, each thread processes 4 candidates, checks if > heap root. - But the heap is shared! We need synchronization. Better pass 2 design: - Each thread maintains a local top-k from its portion of candidates. - Then block-level merge to final top-k. With 256 threads and 32768 candidates, each thread processes 128 candidates. - Each thread loads 128 candidates, keeps top-64 in local memory. - Block-level merge: 256 * 64 = 16384 candidates in shared memory. 128KB. Too much. Alternative: two-level merge in pass 2. - Warp-level merge first: each warp merges 32 * 64 = 2048 candidates into 64. - 8 warps produce 512 candidates. - Block-level merge: 512 candidates into 64. - 512 * 8 = 4KB shared memory. Easy! So pass 2 can be: 1. Each thread loads 128 candidates and keeps top-64 in local memory. 2. Warp-level bitonic sort of 32*64=2048 candidates in shared memory. 3. Each warp writes top-64 to block shared memory. 4. Block-level bitonic sort of 8*64=512 candidates. 5. Write top-64 to output. This is exactly the same as my original single-block design! The only difference is that pass 2 processes candidates instead of raw elements. And pass 2 has 256 threads (8 warps). Shared memory for warp sort: 8 * 2048 * 8 = 131072 bytes = 128KB. Too much! Wait, same shared memory issue. For 8 warps with warp_cand=2048 each, we need 8 * 2048 * 8 = 128KB for the warp work areas. This exceeds 101KB opt-in limit. Can we reduce the warp work area? Yes! Instead of having each warp sort 2048 candidates in its own shared memory area, we can sort them sequentially: - Warp 0 sorts its 2048 candidates in a shared memory area. - After warp 0 finishes, warp 1 sorts its 2048 candidates in the SAME area. - And so on. But warps in a block execute concurrently, not sequentially. We can't have them reuse the same area unless we synchronize between warps. Actually, we CAN synchronize warps using `__syncthreads()` between warp sorts. But this serializes the warps, reducing parallelism. For 8 warps serialized: 8x slowdown of the warp sort. But the warp sort is fast (~1-2 microseconds). 8x = ~16 microseconds. Still fast compared to the element processing. Wait, but with 256 threads, do we really need 8 separate warp sorts? What if we just use 1 warp to do all the merging? If we have 256 threads each with top-64 in local memory, we need to merge 256 * 64 = 16384 candidates. We can't sort all 16384 in shared memory at once (needs 128KB). But we can do it in stages: 1. Sort 1024 candidates at a time in shared memory, keep top-64. 2. Repeat 16 times. For each batch of 1024 candidates: - Load 1024 candidates into shared memory (8KB). - Bitonic sort 1024 in shared memory with 256 threads. - Take top-64. - Merge with current top-64. How to merge current top-64 with new top-64? - Both are sorted descending. - Sequential merge: 64 + 64 = 128 comparisons, produce new top-64. - But the current top-64 is in local memory (or we can keep it in registers if we use scalar variables... no, K=64 is too big for registers). Actually, we can keep the running top-64 in shared memory too! Shared memory for running top-64: 64 * 8 = 512 bytes. So pass 2 algorithm: 1. Allocate shared memory: s_cand_vals[1024], s_cand_idxs[1024], s_top_vals[64], s_top_idxs[64]. 2. Initialize s_top to -inf. 3. For each batch of 1024 candidates (from global memory): a. Load 1024 candidates into s_cand. b. Bitonic sort s_cand descending. c. Merge s_top (64) with s_cand top-64 (64) into new s_top. 4. Write s_top to output. Total shared memory: 1024*4 + 1024*4 + 64*4 + 64*4 = 8192 + 256 + 256 = 8704 bytes = 8.5KB. Very small! For M=32768 candidates, we have 32 batches of 1024. Each batch: - Sort 1024: 1024 * 10 * 10 / 4 = 25600 comparisons. With 256 threads, ~100 operations per thread. - Merge 64+64: 128 comparisons. With 256 threads, we can parallelize but it's small. - Total per batch: ~100 operations per thread. - 32 batches: ~3200 operations per thread. - At 1.5GHz: ~2 microseconds. This is fast! And it uses very little shared memory. Now, for the merge of two sorted 64-element lists into one 64-element list, we can do: - Thread i (0 <= i < 64) finds the i-th largest element. - To find the i-th largest in the union of A[0..63] and B[0..63]: - We can binary search in A: for A[j], there are j elements in A larger than A[j], and up to 64 elements in B larger than A[j]. - The rank of A[j] in the union is j + (# of B elements > A[j]). - We want elements with rank < 64. Actually, a simpler parallel merge: - For each element A[j], count how many B elements are >= A[j]. - The position of A[j] in the merged array is j + count. But this requires binary search in B for each A[j]. For 64 elements and 64 threads, each thread does one binary search (6 comparisons). Total: 64 * 6 = 384 comparisons. Very fast! But implementing this in shared memory with synchronization is a bit complex. Let me use a simpler approach: just do a sequential merge with one thread, since it's only 128 comparisons. ```cuda if (tid == 0) { float a[64], b[64]; // copy s_top and s_cand_top to a and b int ai = 0, bi = 0; for (int i = 0; i < 64; i++) { if (ai < 64 && (bi >= 64 || a[ai] >= b[bi])) { s_top_vals[i] = a[ai]; s_top_idxs[i] = a_idx[ai]; ai++; } else { s_top_vals[i] = b[bi]; s_top_idxs[i] = b_idx[bi]; bi++; } } } __syncthreads(); ``` This is 128 operations for one thread. At 1.5GHz, 0.1 microseconds. Negligible. So pass 2 is dominated by sorting 1024 elements, 32 times. 32 * 25600 comparisons = ~819200 comparisons total. With 256 threads: 3200 comparisons per thread. At 1 comparison per clock: 2 microseconds. Fast! Now let me design the full two-pass kernel. **Pass 1: Multi-block local top-k** - Launch `batch * num_blocks` blocks. - Each block has `THREADS` threads (e.g., 128 or 256). - Block `b` of row `r` processes elements `[b * n/num_blocks, (b+1) * n/num_blocks)`. - Each thread keeps top-k in local memory (heap). - After processing, each thread's top-k is written to global memory at: `tmp_vals[((r * num_blocks + b) * THREADS + tid) * k + i]` No block-level merge in pass 1. Just thread-level top-k written directly. **Pass 2: Global merge** - Launch `batch` blocks. - Each block has `THREADS` threads (e.g., 256). - Shared memory: s_cand[1024] vals + s_cand[1024] idxs + s_top[64] vals + s_top[64] idxs = ~8704 bytes. - Initialize s_top to -inf. - Stream through `num_blocks * THREADS * k` candidates in tiles of 1024. - For each tile: 1. Load into s_cand 2. Bitonic sort s_cand descending 3. Sequential merge s_top with s_cand[0..63] - Write s_top to output. Wait, but we have `num_blocks * THREADS * k` candidates. For batch=1, n=131072, k=64, num_blocks=8, THREADS=128: - Candidates = 8 * 128 * 64 = 65536. - Tiles of 1024: 64 tiles. - Pass 2 time: 64 * sort_1024 + 64 * merge_64. - sort_1024 = 25600 comparators. With 256 threads, ~100 comparators per thread. - 64 * 100 = 6400 comparators per thread. - At 1.5GHz: ~4 microseconds. Fast! And pass 1: - 8 blocks, 128 threads each. - Each thread processes 131072 / (8 * 128) = 128 elements. - Heap: 128 * 6 * 4 = ~3000 operations per thread. - With 8 blocks running on 8 SMs (or more if occupancy allows), total time ~2 microseconds. Total: ~6 microseconds. Compared to torch.topk's ~55 microseconds. We'd be 9x faster! Wait, that seems too optimistic. Let me recalculate. For pass 1 with local memory heap: - Each heap_replace does ~6 iterations. - Each iteration accesses local memory: read vals[left], vals[right], write vals[pos], read idxs[left]... - Local memory is cached in L1. Each access is ~10-20 cycles. - 128 elements * 6 iterations * 4 accesses * 15 cycles = ~46080 cycles per thread. - At 1.5GHz: ~30 microseconds per thread. - With 8 blocks on 8 SMs: ~30 microseconds total. For pass 2: - 64 tiles * sort_1024. - Each sort_1024 comparison does 2 shared memory reads, 2 shared memory writes. - Shared memory is ~20 cycles. - 64 tiles * 25600 comps * 4 accesses * 20 cycles = ~131M cycles... wait, that's total across all threads. - With 256 threads: 131M / 256 = 512K cycles per thread. - At 1.5GHz: ~340 microseconds per thread. - But only 1 block! So 340 microseconds. Hmm, 340 microseconds for pass 2 is not great. But it's still better than 1700 microseconds. Actually, let me recalculate more carefully. A bitonic sort of 1024 with 256 threads: - Number of stages: log2(1024) * (log2(1024)+1) / 2 = 10 * 11 / 2 = 55 stages. - Each stage: 1024 / 2 = 512 comparators. - With 256 threads, each thread does 512 / 256 = 2 comparators per stage. - Total: 55 * 2 = 110 comparators per thread. - Each comparator: 2 reads, 1 comparison, 2 writes (if swap). Let's say 5 operations. - 110 * 5 = 550 operations per thread. - 64 tiles: 64 * 550 = 35200 operations per thread. - At 1.5GHz: ~23 microseconds. That's much more reasonable! And with shared memory, each operation is fast. So pass 2: ~23 microseconds. Pass 1: ~30 microseconds. Total: ~53 microseconds. Compare to torch.topk: ~55 microseconds. So we'd be roughly equal, not 9x faster. But wait, torch.topk is only getting 10 GB/s. If we can get closer to peak bandwidth (1800 GB/s), we'd be much faster. For 512KB input, at 1800 GB/s = 0.28 microseconds. But the comparison and selection overhead makes this impossible. Actually, for pass 1, if we use many more blocks, we can get better memory bandwidth utilization. With num_blocks=32 instead of 8: - 32 blocks, 128 threads each. - Each thread processes 32 elements. - Heap: 32 * 6 * 4 * 15 cycles = ~11520 cycles = 7.7 microseconds. - 32 blocks running in parallel on 32 SMs: ~7.7 microseconds. - Candidates = 32 * 128 * 64 = 262144. - Pass 2 tiles = 262144 / 1024 = 256 tiles. - Pass 2 time = 256 * 550 ops * 5 cycles / 256 threads = ~2750 cycles per thread = 1.8 microseconds... wait no. - 256 tiles * 110 comparators per thread = 28160 comparators per thread. - 28160 * 5 cycles = 140800 cycles = 94 microseconds. Hmm, with more blocks in pass 1, pass 1 gets faster but pass 2 gets slower. There's a tradeoff. For num_blocks=8: total ~53 microseconds. For num_blocks=16: pass 1 ~15 us, pass 2 ~47 us, total ~62 us. For num_blocks=32: total ~102 us. So num_blocks=8 seems optimal for this shape. But ~53 microseconds is still roughly equal to torch.topk. We want to beat it. How can we make pass 1 faster? The bottleneck is local memory accesses in the heap. What if we use a different algorithm for pass 1 that doesn't need a heap? For small chunk_size (like 32 elements), each thread can just load all 32 elements, sort them, and keep top-64 (which is just the 32 elements + padding). No heap needed! For chunk_size = 32 and K = 64: 1. Load 32 elements into `chunk_vals[32]`, `chunk_idxs[32]` (local memory). 2. Fill remaining 32 with -inf. 3. Bitonic sort 64 in local memory. Wait, but `chunk_vals[32]` and the sort are still in local memory. Dynamic indexing in bitonic sort! The compiler will put `chunk_vals` in local memory. Same problem. But what if chunk_size is small enough that the compiler maps it to registers? For 32 floats + 32 ints = 256 bytes. The compiler might still put it in local memory if there's dynamic indexing. Actually, ptxas showed that even K=8 (64 bytes) uses local memory. So any array with dynamic indexing goes to local memory, regardless of size. Given this constraint, the only way to avoid local memory is to avoid dynamic indexing entirely. This means: 1. All array indices must be compile-time constants. 2. No loops with dynamic iteration counts that access arrays. For top-k with K=64, this means we cannot use arrays for the top-k storage. We must use scalar variables. But with 64 scalar variables, writing a generic kernel is extremely tedious. However, we can use macros or template metaprogramming to generate the code. Wait, there's another option: use Triton. Triton handles registers automatically and can do block-level operations efficiently. But implementing top-k in Triton is also complex. Actually, I just had another idea. What if I use **warp-level parallel sorting**? With 32 threads in a warp, we can sort up to 32 elements entirely in registers using warp shuffles. Each thread holds one element. The bitonic sort network uses warp shuffles to exchange elements between threads. For sorting 32 elements in a warp: - Bitonic sort of 32 requires log2(32) * (log2(32)+1) / 2 = 15 stages. - Each stage: some threads exchange elements with other threads using `__shfl_sync`. - No shared memory or local memory needed for the elements themselves! But we also need indices. Each thread can hold one (value, index) pair. Exchanges use `__shfl_sync` for both value and index. This is brilliant! For K=32, we can sort exactly one warp's worth of elements entirely in registers using shuffles. For K=64, we'd need 2 warps. But we need to MERGE the sorted warp results with the running top-k. And the running top-k needs to be stored somewhere. What if we maintain the running top-k across multiple warps in shared memory, but the per-chunk sorting is done in registers with shuffles? Actually, let me think about a different architecture entirely. **Warp-level streaming top-k with shuffle-based merge** Each warp (32 threads) processes a stream of elements: 1. Each thread loads one element. 2. The warp sorts the 32 elements using shuffle-based bitonic sort. 3. The sorted 32 elements are merged with the warp's running top-k in shared memory. 4. Repeat. For K=64, the running top-k is 64 elements. We can store it in shared memory: 64 * 8 = 512 bytes per warp. For 8 warps: 4KB. Very small! The merge of sorted 32 with sorted 64 into sorted 64: - We can use a bitonic merge of 96 elements (or just sequential merge). - With 32 threads, each thread can participate in the merge. But the merge still needs dynamic indexing or complex shuffle patterns. Actually, for merging sorted A (size 32) and sorted B (size 64) into top-64, we can: 1. Concatenate A and B. 2. Use a partial bitonic sort that only keeps the top 64. But implementing partial bitonic sort with shuffles is complex. Let me look at this from yet another angle. What if I use **CUDA cooperative groups** or just simple warp shuffles to do a full sort of all candidates? For a block with 256 threads and K=64: - Total candidates after all threads process their chunks: 256 * 64 = 16384. - We can't sort 16384 with 256 threads in shared memory (needs 128KB). But what if we use **multiple warps to sort in a tree fashion**? - Each warp sorts its 64 candidates using a shuffle-based sort. - Then warps merge pairwise. For warp sort of 64 with 32 threads: - Each thread holds 2 elements (since 64 / 32 = 2). - Bitonic sort of 64 with each thread holding 2 elements. - Compare-and-swap within a thread (no shuffle) or across threads (shuffle). This is possible! Each thread stores 2 values in registers. For a compare-swap between thread i and thread j, we shuffle the values. For bitonic sort of 64 with 32 threads (2 elements per thread): - There are log2(64) stages. - Some comparators are within a thread (between its 2 elements). - Some comparators are between threads (using shuffles). - The pattern is known and can be hardcoded. After the warp sort, each warp has 64 sorted elements. Then we need to merge 8 warps' results (512 elements) into top-64. We can do this with shared memory: load all 512 into shared memory, bitonic sort 512, take top-64. Shared memory for 512 elements: 512 * 8 = 4KB. Very small! And for the per-thread element processing, each thread still needs a local top-k of 64. But instead of a heap, each thread can: 1. Load elements into a local buffer. 2. Sort the buffer using a local bitonic sort. 3. The buffer is in local memory, so it's slow. Wait, but if each thread only processes a small number of elements (like 32), it can load them into scalar registers and sort them without arrays! For chunk_size = 16: ```cuda float v0 = row_x[base + 0]; float v1 = row_x[base + 1]; // ... v15 int i0 = base + 0; // ... i15 // Bitonic sort of 16 using scalar variables // All operations are on scalars in registers! CMP_SWAP(v0, v1); CMP_SWAP(v2, v3); ... // stage 1 CMP_SWAP(v0, v2); CMP_SWAP(v1, v3); ... // stage 2 // etc. ``` Then merge the sorted 16 with the running top-64. But the running top-64 is in local memory. Hmm. What if each thread processes 16 elements, sorts them in registers, then writes them to shared memory. Then the block merges all threads' sorted chunks? For 256 threads * 16 elements = 4096 elements per block. Block-level merge of 4096 sorted 16-element lists into top-64. This is complex but possible. We can do it in stages: 1. Each warp merges 32 sorted 16-element lists into one sorted 512-element list? No, we only need top-64. 2. A warp can merge 32 sorted 16-element lists into top-64 by loading all 512 candidates and doing a selection. With 32 threads in a warp, each loads 16 candidates = 512 total. Warp-level selection of top-64 from 512: - We can use a parallel tournament tree or repeated max reduction. - 64 iterations of finding the max among 512 elements. - Each iteration: tree reduction over 512 elements with 32 threads. - 512 / 32 = 16 local comparisons per thread. - Warp shuffle reduction: 5 stages. - Total per iteration: ~21 operations. - 64 iterations: ~1344 operations per thread. - Plus setting the found max to -inf. With 8 warps doing this in parallel: each warp processes 512 candidates and produces top-64. Then we have 8 * 64 = 512 candidates. Block-level merge of 512 candidates into top-64 using shared memory bitonic sort: 4KB shared memory, fast. Total operations per block: - Element loading: 4096 loads - Warp-level selection: 8 warps * 1344 ops = ~10752 ops - Block-level sort: 512-element bitonic sort = ~2304 ops - Total: ~13000 ops for the block - With 256 threads: ~50 ops per thread. Wait, that can't be right. 10752 ops / 256 threads = 42 ops per thread. Plus 2304 / 256 = 9 ops. Plus 4096 loads / 256 = 16 loads. Total ~67 ops per thread. At 1.5GHz, 67 ops = 45 nanoseconds! That's incredibly fast. But the warp-level selection uses `__shfl_sync` for reductions. Each shuffle is ~1-4 cycles. And we need 64 iterations * 5 shuffles = 320 shuffles per thread. At 4 cycles each: 1280 cycles = 0.85 microseconds. Plus local comparisons: 64 * 16 = 1024 comparisons. Plus 64 writes to shared memory to mark -inf. Total per thread: ~2000-3000 cycles = 1.3-2 microseconds. With 8 warps in parallel, the block takes ~2 microseconds. And for batch=1, n=131072: - We need 131072 / 4096 = 32 blocks per row. - 32 blocks can run on 32 SMs in parallel! - Total time: ~2 microseconds (if memory bandwidth is saturated). Memory bandwidth: 32 blocks * 4096 floats * 4 bytes = 512KB total. At 1.8 TB/s, reading 512KB takes 0.28 microseconds. So memory is not the bottleneck; compute is. This design is very promising! Let me formalize it. **New kernel design: "Warp-Select Top-K"** For each block: 1. Divide block into warps (e.g., 8 warps of 32 threads for 256 threads). 2. Each thread loads `ELEMENTS_PER_THREAD` (e.g., 16) elements into scalar registers. 3. Sort the thread's elements using scalar bitonic sort (all in registers). 4. Write sorted elements to warp shared memory. 5. Each warp performs warp-level selection of top-k from all threads in the warp: - Load all candidates into registers/shuffles. - Repeatedly find max using warp shuffle reduction, k times. - After finding each max, mark it as -inf. 6. Write warp top-k to block shared memory. 7. Block-level bitonic sort of all warp top-k arrays to produce final top-k. Wait, steps 3-5 can be optimized. Instead of each thread sorting its elements first, what if we just have each thread load its elements and then the warp collectively finds the top-k? Actually, let me simplify: **Warp-level selection from a shared memory array** 1. All threads in a warp load their elements into a warp-shared memory array. - Thread i writes at positions i * EPT + 0, 1, ..., EPT-1. - Total warp array size: 32 * EPT elements. 2. The warp performs k iterations of max-finding: - Each thread finds local max of its portion of the array. - Warp shuffle reduction to find global max. - Thread 0 records the global max. - The thread that owns the global max sets it to -inf in shared memory. For EPT=16, warp array size=512 elements. k iterations=64. Each iteration: local max of 512/32=16 elements (16 comparisons), then 5 shuffle reductions. 64 * (16 + 5) = 1344 operations per thread. After k iterations, we have top-k values in a register array (or shared memory). Then sort the top-k using bitonic sort. Wait, the top-k are found in descending order (largest first). So they're already sorted! No need for a final sort. Actually, the max-finding approach finds them one by one in descending order. So after k iterations, we have k sorted values. We just need to record their indices too. For indices: when we find the global max, we also find its index. We can store the (value, index) pairs. Now, for the block-level merge: - Each warp produces top-k. - We have `num_warps` top-k arrays. - We need to merge them into one top-k. For `num_warps=8` and `k=64`: 512 candidates. We can load all 512 into block shared memory and do a bitonic sort. Or we can do another round of warp-level selection. Block-level merge using bitonic sort of 512: - 512 * 9 * 9 / 4 = 10368 comparators. - With 256 threads: ~40 comparators per thread. - Fast. Total shared memory: - Warp arrays: 8 warps * 512 elements * 8 bytes = 32KB. This is per-warp area. Wait, all warps run concurrently and each needs its own area. So 8 * 512 * 8 = 32KB. - Block merge area: 512 * 8 = 4KB. - Total: 36KB. Fits in 48KB! Wait, 8 warps * 512 elements * 8 bytes: 512 elements = 512 * 8 = 4096 bytes per warp. 8 warps = 32768 bytes = 32KB. Block merge area: 512 * 8 = 4096 bytes. Total: 36KB + some overhead. Fits! This is excellent! Let me design the kernel in detail. Actually, wait. The warp-level selection requires each warp to have its own shared memory area. But warps in a block execute concurrently. So we need to partition shared memory: ```cuda __shared__ float s_warp_vals[8][512]; __shared__ int s_warp_idxs[8][512]; ``` Or dynamically: ```cuda float* warp_vals = s_vals + warp_id * 512; int* warp_idxs = s_idxs + warp_id * 512; ``` Total: 8 * 512 * 4 + 8 * 512 * 4 = 16384 + 16384 = 32768 bytes. For block merge: ```cuda float* block_vals = s_block_vals; // 512 floats = 2048 bytes int* block_idxs = s_block_idxs; // 512 ints = 2048 bytes ``` Total: 32768 + 4096 = 36864 bytes = 36KB. Under 48KB! For k=64 and 256 threads: works. For k=64 and 128 threads (4 warps): warp area = 4 * 512 * 8 = 16KB. Block merge = 4 * 64 * 8 = 2KB. Total = 18KB. Even better. But with 128 threads, each thread processes more elements. For n=131072 and 4 warps * 32 threads = 128 threads per block. If we use num_blocks=1: each thread processes 1024 elements. EPT=32. Warp array size = 32 * 32 = 1024 elements. Warp area = 4 * 1024 * 8 = 32KB. Block merge = 4 * 64 * 8 = 2KB. Total = 34KB. Still under 48KB. But with 1 block, parallelism is limited. We need multiple blocks. If we use num_blocks=8 with 128 threads each: - Each block processes 131072 / 8 = 16384 elements. - 128 threads * EPT = 16384. EPT = 128. - Warp array size = 32 * 128 = 4096 elements. - Warp area = 4 * 4096 * 8 = 131072 bytes = 128KB. EXCEEDS 48KB! So with multiple blocks, EPT must be smaller to fit shared memory. For num_blocks=8, THREADS=128: - Elements per block = 16384. - If EPT=16: warp array = 32 * 16 = 512 elements. Warp area = 4 * 512 * 8 = 16KB. Block merge = 2KB. Total = 18KB. OK. - Each thread processes 16384 / 128 = 128 elements. But EPT=16 means each thread loads 16 elements at a time, and we do 128/16 = 8 rounds. So the kernel loops 8 times: 1. Load 16 elements into registers. 2. Write to warp shared memory. 3. Warp selects top-64 from 512 candidates. 4. After 8 rounds, each warp has 8 * 64 = 512 candidates in shared memory (or we keep a running top-64). Wait, we can keep a running top-64! After each round: - Warp has current top-64. - New 512 candidates come in. - Warp selects top-64 from 512 + 64 = 576 candidates. How to select top-64 from 576 candidates with a warp? - Load all 576 into shared memory. - 64 iterations of max-finding. - Each iteration: local max of 576/32 = 18 elements, then warp reduction. 64 * (18 + 5) = 1472 operations per thread per round. 8 rounds: ~11776 operations per thread. At 1.5GHz: ~7.8 microseconds. With 8 blocks on 8 SMs: ~7.8 microseconds total for pass 1. Then pass 2: merge 8 blocks * 4 warps * 64 = 2048 candidates. We can use a single block with bitonic sort of 2048 in shared memory. 2048 * 11 * 11 / 4 = 61952 comparators. With 256 threads: ~242 comparators per thread. Fast. Or even simpler: use another warp-select kernel for pass 2. Actually, for pass 2 with 2048 candidates, we can just use the same warp-select approach with 1 block: - 256 threads. - Load candidates in tiles of 512. - Each warp selects top-64 from its tile. - Then block merges 8 warp results (512 candidates) into top-64. Pass 2 time: loading 2048 candidates + 4 rounds of warp-select + block merge. ~2-3 microseconds. Total: ~10 microseconds. Compare to torch.topk's ~55 microseconds. That's 5.5x faster! This is very promising. But I need to implement the warp-level selection carefully. **Warp-level selection algorithm:** Input: `warp_vals[M]` and `warp_idxs[M]` in shared memory (M = 32 * EPT). Output: top-k values and indices, stored in shared memory or registers. Method: k iterations of parallel max-finding with masking. ```cuda // Each thread maintains local max float local_max = -FLT_MAX; int local_idx = -1; // k iterations for (int iter = 0; iter < k; iter++) { // Each thread finds max of its portion local_max = -FLT_MAX; for (int i = lane; i < M; i += 32) { if (s_vals[i] > local_max) { local_max = s_vals[i]; local_idx = s_vals[i]; // actually index } } // Warp reduction to find global max #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { float other_max = __shfl_down_sync(mask, local_max, offset); int other_idx = __shfl_down_sync(mask, local_idx, offset); if (other_max > local_max) { local_max = other_max; local_idx = other_idx; } } // Broadcast global max and its index to all threads float global_max = __shfl_sync(mask, local_max, 0); int global_idx = __shfl_sync(mask, local_idx, 0); // Record in output if (lane == 0) { out_vals[iter] = global_max; out_idxs[iter] = global_idx; } // Mark the found element as -inf // We need to know which thread owns it int owner = __shfl_sync(mask, 0, 0); // Wait, we need the lane that found it // Actually, after reduction, lane 0 has the global max. // But we need to know which element in shared memory it was. // We can have lane 0 also track the shared memory index. } ``` Wait, there's a subtle issue. After warp reduction, lane 0 has the global max and its `local_idx`. But `local_idx` is the index within the thread's portion, not the shared memory index. We need the shared memory index to set it to -inf. We can track the shared memory index instead: ```cuda int local_sm_idx = -1; for (int i = lane; i < M; i += 32) { if (s_vals[i] > local_max) { local_max = s_vals[i]; local_sm_idx = i; } } ``` Then after warp reduction, lane 0 has `local_sm_idx` which is the shared memory index. We broadcast this to all threads, and the thread with `lane == (local_sm_idx % 32)` sets `s_vals[local_sm_idx] = -FLT_MAX`. Wait, but `local_sm_idx` might not be in lane 0's portion. After reduction, lane 0 has the winning `local_sm_idx`. We can broadcast it: ```cuda int winner_sm_idx = __shfl_sync(mask, local_sm_idx, 0); if (lane == (winner_sm_idx & 31)) { s_vals[winner_sm_idx] = -FLT_MAX; } ``` This works! Only one thread writes to shared memory. However, there's a problem: what if two elements have the same max value? The reduction picks one arbitrarily, which is fine. Now, for the inner loop `for (int i = lane; i < M; i += 32)`, if M=512, each thread does 16 iterations. For k=64, the outer loop runs 64 times. Total: 64 * 16 = 1024 comparisons per thread. Plus 64 * 5 = 320 shuffles. Plus 64 broadcasts. Total: ~1400 operations per thread. For M=512 and k=64: ~1400 ops per thread. For 8 rounds: 8 * 1400 = 11200 ops per thread. At 1.5GHz: ~7.5 microseconds. This is very reasonable! But wait, we need to run this for each warp. With 4 warps per block, and 8 blocks total: - Total warp-executions: 4 * 8 = 32. - These run across 32 SMs (since 8 blocks * 4 warps = 32 warps, and each SM can run multiple warps). - An SM has 4 warp schedulers, so each SM can run ~16-32 warps concurrently. - 32 warps total can run on 2-4 SMs concurrently. - So the actual parallelism is limited. Wait, 8 blocks, each with 4 warps = 32 warps total. Each SM can run up to 32 warps (limited by registers/shared memory). So all 32 warps might fit on 1 SM! But we have 188 SMs. So we're only using 1 SM out of 188. That's terrible! Actually, no. CUDA blocks are distributed across SMs. 8 blocks can run on 8 different SMs simultaneously. Each SM runs 4 warps (1 block). So 8 SMs are used. The remaining 180 SMs are idle. To use more SMs, we need more blocks. For batch=1, we can only launch blocks for that one row. So we need `num_blocks >= num_SMs` to saturate the GPU. That's 188 blocks! But if we use 188 blocks, each block processes 131072 / 188 = 697 elements. With 128 threads, each thread processes 5.4 elements. EPT would be small. Let's calculate for num_blocks=188, threads=128: - Elements per block = 131072 / 188 = 697.2 - Elements per thread = 697 / 128 = 5.45 - We can round up to 6 elements per thread (EPT=6). - Warp array size = 32 * 6 = 192 elements. - Warp area per block = 4 * 192 * 8 = 6144 bytes = 6KB. - Block merge = 4 * 64 * 8 = 2KB. - Total shared memory = 8KB. Fits easily. Pass 1: - 188 blocks. - Each thread processes 6 elements (load into registers, write to shared memory). - Each warp does 64 iterations of max-finding over 192 elements. - 192/32 = 6 comparisons per thread per iteration. - 64 * (6 + 5) = 704 ops per thread. - 188 blocks on 188 SMs: ~0.5 microseconds (parallel). Pass 2: - 188 blocks * 4 warps * 64 = 48128 candidates. - With 1 block for pass 2, each warp processes tiles of 512 candidates. - 48128 / 512 = 94 tiles. - Each tile: warp-select top-64 from 512 = 64 * (16 + 5) = 1344 ops. - 94 tiles * 1344 = 126336 ops per thread. - But we have 8 warps, so total warp-work = 8 * 126336 ops. - With 1 block on 1 SM: ~0.7 milliseconds. Too slow! Wait, pass 2 also needs to scale. We can use multiple blocks for pass 2 too, doing a tree reduction. Actually, for pass 2, we can use the same multi-block approach: - Launch 188 blocks for pass 2. - Each block reads a subset of candidates and produces top-k. - But then we need another pass to merge the 188 top-k results. This becomes a multi-pass reduction tree. For a tree reduction with 188 inputs: - Level 0: 188 blocks produce 188 top-k arrays. - Level 1: 94 blocks merge pairs -> 94 top-k. - Level 2: 47 blocks -> 47. - Level 3: 24 -> 24. - Level 4: 12 -> 12. - Level 5: 6 -> 6. - Level 6: 3 -> 3. - Level 7: 2 -> 2. - Level 8: 1 -> 1. Total passes: 9 kernel launches. Each launch has a small number of blocks. The kernel launch overhead might add up. Alternatively, we can do a single large pass 2: - Launch many blocks, each reading candidates and writing to a smaller set of candidate buffers. - This is essentially a parallel reduction. Actually, there's a simpler approach. What if pass 1 writes `num_blocks * k` candidates (instead of `num_blocks * threads * k`)? We can do a block-level merge in pass 1 to reduce the output size. In pass 1, after each warp selects top-k from its elements, we do a block-level merge of all warp results. This produces exactly k candidates per block. Then pass 2 only needs to merge `num_blocks * k` candidates. For num_blocks=188 and k=64: 188 * 64 = 12032 candidates. Pass 2 with 12032 candidates: - Launch 1 block with 256 threads. - Load candidates in tiles of 1024. - 12 tiles. - Each tile: bitonic sort 1024 in shared memory, then merge top-64 with running top-64. - Sort 1024: 25600 comparators. With 256 threads, ~100 comparators per thread. - 12 tiles: 1200 comparators per thread. - Plus 12 merges of 64+64: ~768 comparisons (with one thread). - Total: ~2000 ops per thread = ~1.3 microseconds. This is fast! So the key is: in pass 1, each block must produce only k candidates (not threads * k). This requires a block-level merge within each block. And block-level merge of `num_warps * k` candidates into k candidates can be done with bitonic sort in shared memory. For num_warps=4, k=64: block_cand = 256. Bitonic sort 256 in shared memory: 256 * 8 * 8 / 4 = 4096 comparators. With 128 threads: ~32 comparators per thread. Very fast. Shared memory for block-level merge: 256 * 8 = 2KB. And for warp-level selection, shared memory: 4 warps * 192 elements * 8 bytes = 6KB. Total per block: 8KB. Fits in 48KB easily. And with 188 blocks, we use 188 SMs! Let me now formalize the complete design. **Pass 1: Multi-block warp-select + block-merge** Kernel: `topk_pass1` - Launch `batch * num_blocks` blocks. - Each block has `THREADS` threads (e.g., 128). - Elements per block: `block_n = ceil(n / num_blocks)`. - `EPT = ceil(block_n / THREADS)` = elements per thread. - Warp array size `M = 32 * EPT`. - Shared memory: `num_warps * M * 8` (warp areas) + `num_warps * k * 8` (block candidate area). Wait, for block-level merge, we need to merge `num_warps * k` candidates. We can do this by: 1. Each warp writes its top-k to shared memory. 2. Block-level bitonic sort on `num_warps * k` candidates. 3. Thread 0 writes top-k to global memory. Block candidate area: `num_warps * k * 8` bytes. For num_warps=4, k=64: 2048 bytes. Warp areas: `num_warps * M * 8` bytes. For M=192, num_warps=4: 4 * 192 * 8 = 6144 bytes. Total: 8192 bytes = 8KB. Algorithm per block: 1. Initialize warp shared memory arrays to -inf. 2. Each thread loads `EPT` elements into scalar registers. 3. Write elements to warp shared memory. 4. Warp-level selection: k iterations of max-finding, write top-k to a warp output array in shared memory. 5. After all warps done, block-level bitonic sort on `num_warps * k` candidates. 6. Thread 0 writes top-k to global memory at `row * num_blocks + block_idx`. Wait, step 3: "Write elements to warp shared memory." But we have multiple rounds if EPT > M? No, M = 32 * EPT, so each thread writes exactly EPT elements. One shot. Actually, let me reconsider. If block_n = 697 and THREADS = 128, then EPT = ceil(697 / 128) = 6. Each thread loads 6 elements. Warp array size = 32 * 6 = 192. Total elements loaded per block = 128 * 6 = 768. But block_n is only 697. So some threads load dummy elements (-inf) for the last few positions. For the last block, we need to handle the boundary. Some elements might be past n. We can pad with -inf. OK, the algorithm is: 1. For thread i in [0, THREADS): - For j in [0, EPT): - idx = block_start + i + j * THREADS - val = (idx < n) ? row_x[idx] : -FLT_MAX - Write to warp shared memory at appropriate position. Wait, for coalesced access, thread i should read consecutive elements. With `i + j * THREADS`, thread 0 reads 0, 128, 256, ... which are NOT consecutive. For coalesced access, thread i should read i, i+THREADS, i+2*THREADS... but with blockDim.x = 128, these are 128 floats apart = 512 bytes apart. Not coalesced within a warp! Actually, in a warp (32 threads), threads 0-31 read elements 0, 1, 2, ..., 31 if we use `idx = block_start + i + j * 32`. But with blockDim.x = 128, consecutive threads in the block are not necessarily in the same warp. Warp 0: threads 0-31 Warp 1: threads 32-63 Warp 2: threads 64-95 Warp 3: threads 96-127 For coalesced access within a warp, threads in the same warp should read consecutive memory. So thread i should read `block_start + (i & ~31) * EPT + (i & 31) + j * 32`. This is complex. Let me simplify by having each warp load its own chunk of elements. Actually, a simpler approach: assign each warp a contiguous chunk of the block's elements. - Block has `block_n` elements. - Warp w processes elements `[w * block_n / num_warps, (w+1) * block_n / num_warps)`. - Each thread in warp w processes a strided subset of its warp's chunk. For block_n=697 and 4 warps: each warp gets ~174 elements. Each thread in a warp processes 174/32 = 5.4 elements, round up to 6. Total elements per warp = 32 * 6 = 192. For memory access within a warp: - Thread lane reads at `warp_start + lane, warp_start + lane + 32, warp_start + lane + 64, ...` - These are 32 floats apart = 128 bytes apart. - A warp reads 128-byte chunks. With 128-byte memory transactions, this is perfectly coalesced! Wait, consecutive threads in a warp read addresses 128 bytes apart. Memory transactions are 128 bytes. So each thread gets its own 128-byte segment? No, that would be 32 * 128 = 4096 bytes per warp read. That's too much. Actually, for coalesced access, consecutive threads should read consecutive elements (or consecutive 4/8/16-byte words). Here, thread lane reads at offset `lane * 128`. These are 128 bytes apart. This is NOT coalesced! Each thread accesses a different 128-byte segment. The hardware can handle this (one transaction per thread), but it's not efficient. For coalesced float4 access: - Thread lane should read at `warp_start + lane * 4`. - This gives consecutive 128-byte chunks per warp. Perfect coalescing. But each thread needs to read 6 elements. With float4, each thread reads one float4 (4 elements) and then 2 individual floats. Actually, for EPT=6, we can read: - One float4 at `warp_start + lane * 4` - Then two floats at `warp_start + 128 + lane * 2`? No, that might not be coalesced. This is getting complicated. Let me just use simple sequential loads per thread for now. The warp-level selection is compute-bound, so memory access might not be the bottleneck. For warp_start + lane, warp_start + lane + 32, etc.: - Thread 0 reads at warp_start + 0 - Thread 1 reads at warp_start + 32 - These are 128 bytes apart. - The warp issues 32 separate 128-byte transactions. Total 4096 bytes. - At 1.8 TB/s, reading 4096 bytes takes ~2.3 nanoseconds. Actually, reading 4096 bytes with 32 separate transactions is terrible. The L2 cache line is 128 bytes, so each transaction reads one cache line. 32 cache lines = 4096 bytes. But the latency of each transaction is ~100-200 cycles. Without enough warps to hide latency, this is slow. For better memory access, we should have consecutive threads read consecutive elements. With 32 threads in a warp, if thread lane reads `warp_start + lane + j * 32`, then for j=0: - Thread 0: warp_start + 0 - Thread 1: warp_start + 1 - ... - Thread 31: warp_start + 31 These are consecutive floats = 4 bytes apart. The warp reads a contiguous 128-byte segment in ONE transaction! Perfect! For j=1: - Thread 0: warp_start + 32 - Thread 1: warp_start + 33 - ... - Thread 31: warp_start + 63 Another contiguous 128-byte segment. One transaction. This is the standard strided access pattern and it's perfectly coalesced! So the load pattern should be: ```cuda for (int j = 0; j < EPT; j++) { int idx = warp_start + lane + j * 32; float val = (idx < warp_end) ? row_x[idx] : -FLT_MAX; // Store to warp shared memory s_warp_vals[warp_id * M + j * 32 + lane] = val; s_warp_idxs[warp_id * M + j * 32 + lane] = (idx < warp_end) ? idx : -1; } ``` Wait, but `lane + j * 32` is the position within the warp's shared memory array. Thread lane writes at position `j * 32 + lane`. After the loop, the warp array contains elements in the order they were loaded. For coalesced global memory access, thread lane reads `row_x[warp_start + lane + j * 32]`. This is perfect. For shared memory access, thread lane writes `s_warp[warp_id * M + j * 32 + lane]`. Bank = `(j * 32 + lane) % 32 = lane % 32`. So thread lane writes to bank lane. No bank conflicts! Great! This is perfect. Now, for the warp-level selection, the elements are stored in shared memory in the order they were loaded. The max-finding loop accesses them with stride 32: ```cuda for (int i = lane; i < M; i += 32) { if (s_warp_vals[i] > local_max) { ... } } ``` Bank = `i % 32 = lane % 32`. Each thread accesses its own bank. No conflicts! This is excellent. Let me now think about the overall kernel launch configuration. For each shape, I want to choose `num_blocks` to maximize parallelism while keeping shared memory reasonable. Actually, instead of a fixed num_blocks, I can compute it based on the input size. For batch=1, n=131072, k=64: - We want many blocks to saturate SMs. Let's use num_blocks = 128 or 256. - With 256 blocks, each block processes 512 elements. - With 128 threads per block, each warp processes 128 elements. - EPT = 128 / 32 = 4. - M = 32 * 4 = 128. - Warp area = 4 * 128 * 8 = 4096 bytes. - Block merge = 4 * 64 * 8 = 2048 bytes. - Total = 6KB. Fits easily. - 256 blocks on 188 SMs: all SMs are busy, some have 2 blocks. Pass 2: 256 blocks * 64 = 16384 candidates. Pass 2 with 1 block: 16384 candidates in tiles of 1024 = 16 tiles. Sort 1024: ~100 comparators per thread. 16 tiles: 1600 comparators per thread. Plus 16 merges: ~1000 comparisons (1 thread). Total: ~2600 ops per thread = ~1.7 microseconds. Or we can use multiple blocks for pass 2 too. With 16 blocks, each block processes 1024 candidates and produces top-64. Then 1 block merges 16 * 64 = 1024 candidates. But kernel launch overhead might be noticeable. Let me just use a single block for pass 2 since it's fast enough. Wait, 1.7 microseconds for pass 2 is great. But pass 1 with 256 blocks: - Each block: 128 threads * 4 loads = 512 elements. - Warp select from 128 elements: 64 iterations * (4 + 5) = 576 ops per thread. - Block merge: sort 256 elements: 256 * 8 * 8 / 4 = 4096 comparators. 128 threads: 32 comparators per thread. - Total per block: ~600 ops per thread. - 256 blocks on 188 SMs: some SMs run 2 blocks. - Time: ~0.4 microseconds. Total: ~2.1 microseconds. That's 25x faster than torch.topk! This seems too good to be true. Let me verify by considering memory bandwidth. - 512KB read in pass 1. - 256 blocks * 64 * 8 = 131KB written in pass 1. - 131KB read in pass 2. - 512 bytes written in pass 2. - Total memory: ~775KB. - At 1.8 TB/s: 0.43 microseconds. Compute time: ~2 microseconds. Total: ~2.4 microseconds. Compare to torch.topk: 55 microseconds. So we'd be ~23x faster. That's incredible! But I need to be careful. The warp-level selection might be slower than I estimated because of warp divergence and shared memory latency. For warp-level selection from 128 elements (M=128): - Each thread processes 4 elements per iteration. - 64 iterations * 4 comparisons = 256 comparisons per thread. - 64 iterations * 5 shuffles = 320 shuffles per thread. - Total: ~576 ops per thread. - At 1.5GHz: ~0.4 microseconds per thread. - With 4 warps per block: 0.4 microseconds (parallel). Block merge (sort 256): 32 comparators per thread. Negligible. So per block: ~0.4 microseconds. 256 blocks: limited by memory bandwidth and SM count. Memory read per block: 512 elements * 4 bytes = 2KB. 256 blocks * 2KB = 512KB total. At 1.8 TB/s: 0.28 microseconds. With 188 SMs, peak throughput is 512KB / 188 SMs = 2.7KB per SM. At 1.8 TB/s per SM (theoretical), 2.7KB takes 1.5 nanoseconds. But with latency and contention, maybe 10-20 nanoseconds per SM. Total: ~0.02 microseconds. Wait, that doesn't make sense. Actually, memory bandwidth is shared across all SMs. The total bandwidth is 1.8 TB/s. Reading 512KB takes 512KB / 1.8TB/s = 0.28 microseconds regardless of how many SMs participate. The only question is whether the GPU can sustain this bandwidth with 256 small blocks. For 256 blocks of 128 threads = 32768 threads. At 1 thread per 4 bytes read, and each thread reads 4 floats (16 bytes) via 4 separate loads... actually, each thread does 4 separate loads, not coalesced into one vector load. Wait, with the strided pattern, threads in a warp read 32 consecutive floats in one transaction. Each thread gets one float. So 4 transactions per warp for 4 iterations. 256 blocks * 4 warps = 1024 warps. Each warp issues 4 memory transactions. Total 4096 transactions. At 128 bytes each = 512KB. With enough warps in flight, the memory controller can batch these efficiently. 1024 warps should easily saturate bandwidth. So pass 1 should take roughly the memory time: ~0.3-0.5 microseconds. Pass 2: 1 block, 256 threads. Read 131KB. At 1.8 TB/s: 0.07 microseconds. Compute: ~2 microseconds. So pass 2 dominates. Total: ~2.5 microseconds. Still 20x faster than torch.topk! This is a very exciting design. Let me implement it! Actually, wait. For pass 2, reading 131KB with 1 block of 256 threads = 8 warps. Only 8 warps on 1 SM. Can 8 warps saturate L2 cache bandwidth for 131KB? 131KB / 8 warps = 16KB per warp. Each warp has 32 threads, each reading 4 floats sequentially (with strided pattern). Each warp does 4 memory transactions of 128 bytes = 512 bytes. Total 16KB = 32 transactions per warp. With memory latency of ~200 cycles and 4 warp schedulers per SM, the SM can issue one transaction per cycle (approx). 32 transactions * 200 cycles = 6400 cycles per warp. But with 8 warps, the latency is hidden. Total time ~6400 cycles / 8 warps = 800 cycles = 0.5 microseconds for memory. Plus compute: ~2600 ops / 256 threads = ~10 ops per thread = ~7 microseconds? Wait, I said 1600 comparators per thread for pass 2. That's a lot. Let me recalculate pass 2: - 16384 candidates, tiles of 1024. 16 tiles. - Each tile: bitonic sort 1024 in shared memory. - 55 stages. Each stage: 1024/2 = 512 comparators. - With 256 threads: 2 comparators per thread per stage. - 55 * 2 = 110 comparators per thread. - Each comparator: 2 reads, 1 compare, possible 2 writes. ~5 ops. - 110 * 5 = 550 ops per thread. - 16 tiles: 16 * 550 = 8800 ops per thread. - Plus 16 merges of 64+64: 16 * 128 = 2048 comparisons. If done by 1 thread, that's 2048 ops. But we can parallelize. - Total: ~10000 ops per thread. - At 1.5GHz: ~6.7 microseconds. Hmm, 6.7 microseconds is still good, but not the 1.7 microseconds I estimated earlier. I underestimated the sort cost. Actually, for the merge in pass 2, instead of sorting each tile and then merging, we can just use warp-level selection on each tile! For a tile of 1024 candidates: - 8 warps, each warp gets 128 candidates. - Each warp selects top-64 from 128 candidates using the warp-select algorithm. - Then block merges 8 * 64 = 512 candidates using bitonic sort. Warp select from 128 candidates (M=128, k=64): - 64 iterations * (4 comparisons + 5 shuffles) = 576 ops per thread. - 8 warps: 576 ops per thread (parallel). - Then block sort 512: 512 * 9 * 9 / 4 = 10368 comparators. - With 256 threads: ~40 comparators per thread. - Total per tile: ~600 ops per thread. 16 tiles: 16 * 600 = 9600 ops per thread. Similar to the sort approach. But we can optimize: instead of doing 16 separate tiles, stream through candidates and maintain a running top-64. Running top-64 in shared memory. For each new tile of 1024: - Load 1024 candidates into shared memory. - Find top-64 from the tile using warp-select. - Merge the tile's top-64 with the running top-64. For finding top-64 from 1024 candidates: - 8 warps, each gets 128 candidates. - Each warp selects top-64 from 128: 576 ops per thread. - Total: 576 ops per thread. For merging two sorted 64-element lists into one 64-element list: - Sequential merge: 128 comparisons. 1 thread. Negligible. 16 tiles: 16 * 576 = 9216 ops per thread. Plus some overhead. At 1.5GHz: ~6 microseconds. This is similar. The pass 2 time is dominated by the warp-select operations. Can we make pass 2 faster? What if we use more blocks for pass 2? With 16 blocks for pass 2: - Each block processes 1024 candidates and produces top-64. - 16 blocks on 16 SMs: parallel. - Time per block: 576 ops per thread = 0.4 microseconds. - Then a final block merges 16 * 64 = 1024 candidates. - Load 1024 into shared memory. - 8 warps select top-64 from 128 each: 576 ops per thread. - Block sort 512: 40 ops per thread. - Total: ~600 ops per thread = 0.4 microseconds. Total pass 2: ~0.4 + 0.4 = 0.8 microseconds. Plus kernel launch overhead: maybe 0.5-1 microseconds per kernel launch. With 2 launches for pass 2, ~1-2 microseconds overhead. Total time: pass 1 (~0.5 us) + pass 2 (~0.8 us) + overhead (~2 us) = ~3.3 microseconds. Still 15x faster than torch.topk. And much of the overhead is from Python/kernel launches. But with the two-pass approach, we need to launch 2 kernels from Python. The `load_inline` function call overhead is small, but the CUDA kernel launch overhead from Python might be 5-10 microseconds. This could dominate! Actually, for PyTorch CUDA kernels, the launch overhead is typically 3-5 microseconds. With 2 kernels, that's 6-10 microseconds. Then our total would be ~10-15 microseconds. Still faster than torch.topk's 55 microseconds, but not by 15x. Hmm, but torch.topk also launches kernels internally. Its 55 microseconds includes kernel execution, not just launch overhead. So our ~10-15 microseconds would still be 3-5x faster. Wait, but my original single-kernel approach took 1.7ms, which is much slower. The two-pass approach should be much faster. Actually, let me think about this differently. What if I stay with a single kernel but increase block size and reduce local memory usage? The issue with the single kernel was: 1. Local memory for per-thread top-k 2. Lack of parallelism for batch=1 For batch=1, even with 1024 threads, we only have 1 block. But if we can make that 1 block extremely fast, it might compete. With 1024 threads, each thread processes 128 elements for n=131072. - Each thread loads 128 elements and sorts them using a bitonic sort in local memory. - But local memory sorting is slow. What if each thread just loads elements and keeps the max? Then we have 1024 max values. We need top-64 from these 1024. A bitonic sort of 1024 in shared memory is fast. But as I proved, keeping only top-1 per thread is NOT sufficient. But what if each thread keeps top-64 from its 128 elements using a register-based approach? For 128 elements, we can load them in chunks of 16, sort each chunk in registers (scalar variables), and merge into a running top-64 in local memory. The chunk sort is in registers (fast). The merge with top-64 is in local memory (slow). But we only do 128/16 = 8 merges. Each merge: 64 + 16 = 80 comparisons. 8 merges = 640 comparisons. Plus bitonic sort of 16 * 8 = 288 comparisons. Total ~900 comparisons per thread. With 1024 threads: 900 ops per thread. At 1.5GHz: 0.6 microseconds. But local memory accesses make it slower. Say 5x slower = 3 microseconds. Then block-level sort of 1024 * 64 = 65536 candidates in shared memory. 65536 * 16 * 16 / 4 = 4M comparators. With 1024 threads: 4K comparators per thread. At 1.5GHz: 2.7 microseconds. Total: ~6 microseconds. Plus memory reads: 512KB / 1.8TB/s = 0.3 microseconds. Total single-kernel time: ~6.5 microseconds. With 1 block on 1 SM. But 1024 threads might not all fit if shared memory is large. Shared memory for block-level sort of 65536 candidates: 65536 * 8 = 512KB. Way too much. So single-kernel with 1024 threads doesn't work because shared memory is too small. What if we do the block-level merge in stages? - Warp 0 merges 32 top-64 arrays into one top-64. - Then warp 1 merges its 32, etc. - Then block merges 32 warp results. For 1024 threads = 32 warps. Each warp merges 32 * 64 = 2048 candidates into top-64. Shared memory per warp: 2048 * 8 = 16KB. For 32 warps: 512KB. Too much. Unless warps do it sequentially. But then we lose parallelism. OK, the two-pass approach really does seem like the way to go. Let me implement it. But wait - I just realized that for the two-pass approach, I need to handle multiple batch sizes. The first pass produces `batch * num_blocks * k` candidates. The second pass reads these and produces `batch * k` output. Actually, let me reconsider whether I really need two passes. What if I use a single kernel with persistent threads that process multiple rows? For batch=128, n=4096, k=1: we launch 128 blocks. My current kernel is already fast (0.009ms vs eager 0.022ms). For batch=64, n=8192, k=8: 64 blocks of 256 threads. My current kernel takes 0.059ms vs eager 0.049ms. Only 20% slower. This is acceptable. For batch=32, n=16384, k=32: 32 blocks of 256 threads. Current kernel: 0.278ms vs eager 0.051ms. 5.4x slower. Bad. For batch=16, n=12000, k=16: 16 blocks of 256 threads. Current: 0.138ms vs eager 0.047ms. 2.9x slower. Bad. For batch=1, n=131072, k=64: 1 block of 128 threads. Current: 1.709ms vs eager 0.055ms. 31x slower. Terrible. So the performance degrades as batch size decreases (fewer blocks) and k increases (more local memory ops). What if I just change the algorithm for small batch sizes? For batch <= 32, I can use multiple blocks per row. For batch > 32, one block per row is fine (already decent performance for batch=64). So I can have a hybrid approach: - If batch is large (>= 32), use single-kernel with 1 block per row. - If batch is small (< 32), use two-pass with multiple blocks per row. But this adds complexity. Let me think if there's a simpler unified approach. Actually, what if I use a single kernel but with `gridDim.x = batch * num_blocks_per_row`? Each block processes a chunk of one row and writes its top-k to a global memory buffer. Then I use a second kernel to merge. This is the two-pass approach, and it works for all batch sizes. For large batch sizes, I can set `num_blocks_per_row = 1`, which degenerates to the single-block approach. But the single-block approach with local memory heap is slow for k=32,64. Hmm. What if in pass 1, each block always uses the warp-select approach (no local memory heap)? For warp-select with EPT elements per thread: - Each thread loads EPT elements into registers. - Writes to warp shared memory. - Warp selects top-k from M = 32 * EPT candidates. - Block merges warp results. The key is that the element loading and warp-select are fast, with minimal local memory usage. The only shared memory is for warp arrays and block merge. And this works for ANY number of blocks per row! If num_blocks=1, each block processes the entire row. If num_blocks=32, each block processes 1/32 of the row. So I can use a single kernel design that works for all configurations: **Unified kernel: `topk_kernel`** Parameters: - `x`: input array - `out_vals`, `out_idxs`: output arrays - `n`: elements per row - `k`: top-k - `block_start`, `block_end`: this block's element range Actually, for simplicity, let me just use the two-pass approach with a fixed `num_blocks_per_row` computed based on n and k. For each row, `num_blocks = max(1, min(MAX_BLOCKS, n / MIN_ELEMENTS_PER_BLOCK))`. Where `MIN_ELEMENTS_PER_BLOCK` is chosen so that each thread processes a reasonable number of elements. Let me set: - Threads per block = 128 - EPT = 8 (so M = 256) - Shared memory per block = 4 * 256 * 8 + 4 * 64 * 8 = 8KB + 2KB = 10KB. Wait, for 4 warps and k=64. - Warp area: 4 * 256 * 8 = 8KB. - Block merge: 4 * 64 * 8 = 2KB. - Total: 10KB. With 10KB per block, we can run many blocks per SM. An SM has 100KB shared memory, so ~10 blocks per SM. For batch=1, n=131072: - We want many blocks. If each block processes 1024 elements, num_blocks = 128. - 128 blocks. Each SM runs ~1 block (since 128 blocks / 188 SMs < 1). - Wait, 128 < 188, so some SMs are idle. Let's use num_blocks = 256. - 256 blocks / 188 SMs ≈ 1.36 blocks per SM. Good utilization. For batch=64, n=8192: - 64 rows * num_blocks_per_row. - If num_blocks_per_row = 1: 64 blocks total. OK. - If num_blocks_per_row = 4: 256 blocks. Even better. Actually, for simplicity, let me just set `num_blocks_per_row = max(1, 256 / batch)`. This ensures we always have at least 256 blocks total. Wait, for batch=128: 256/128 = 2 blocks per row. 256 blocks total. For batch=1: 256 blocks per row. 256 blocks total. This gives consistent total parallelism! And for each block, elements per block = n / num_blocks_per_row. For n=4096 and num_blocks=2: 2048 elements per block. For n=8192 and num_blocks=2: 4096 elements per block. For n=131072 and num_blocks=256: 512 elements per block. With 128 threads per block: - For 512 elements: EPT = 512 / 128 = 4. M = 32 * 4 = 128. - For 2048 elements: EPT = 2048 / 128 = 16. M = 32 * 16 = 512. - For 4096 elements: EPT = 4096 / 128 = 32. M = 32 * 32 = 1024. Shared memory for M=1024, k=64: - Warp area: 4 * 1024 * 8 = 32KB. - Block merge: 4 * 64 * 8 = 2KB. - Total: 34KB. Still under 48KB! Great! This works for all shapes. For n=131072 with 256 blocks: - Each block: 512 elements, EPT=4, M=128. - Warp select from 128: 64 * (4 + 5) = 576 ops per thread. - Block sort 256: ~32 ops per thread. - Total per block: ~600 ops. - 256 blocks on 188 SMs: ~1 microsecond (limited by memory: 512KB / 1.8TB/s = 0.3 us, plus compute). For pass 2: - 256 blocks * 64 = 16384 candidates per row. - For batch=1: 16384 candidates total. - Pass 2 with 1 block: 16 tiles of 1024. - Each tile: warp select from 1024 = 8 warps * 128 candidates each. - Warp select from 128: 576 ops per thread. - 8 warps parallel. - Block sort 512: 40 ops per thread. - Total per tile: ~600 ops per thread. - 16 tiles: ~9600 ops per thread. - At 1.5GHz: ~6.4 microseconds. Total: ~7.5 microseconds. Still 7x faster than torch.topk. But can we make pass 2 faster by using more blocks? Yes! For pass 2, we can launch `batch * num_blocks2` blocks. - Each block reads `total_candidates / num_blocks2` candidates and produces top-k. - Then another pass to merge. For batch=1, total_candidates=16384. - Pass 2a: 16 blocks, each reads 1024 candidates, produces top-64. - Each block: warp select from 1024 = 576 ops per thread. - 16 blocks on 16 SMs: ~0.4 microseconds. - Pass 2b: 1 block merges 16 * 64 = 1024 candidates. - Warp select from 1024 (8 warps, 128 each): 576 ops per thread. - Block sort 512: 40 ops. - Total: ~0.4 microseconds. Total pass 2: ~0.8 microseconds. Plus kernel launch overhead: 2 launches * 5us = 10us. Total with overhead: ~11 microseconds. Hmm, the kernel launch overhead dominates! A single kernel that does everything in ~50 microseconds might actually be faster than 3 kernels with 15 microseconds compute but 15 microseconds launch overhead. Wait, but torch.topk launches at least 1 kernel too. And it takes 55 microseconds. So our 11 microseconds total is still 5x faster. Actually, let me reconsider. The kernel launch overhead in PyTorch is about 3-5 microseconds on modern GPUs. For 3 kernels, that's 9-15 microseconds. Plus compute ~8 microseconds. Total ~20 microseconds. Still 2.7x faster than torch.topk. But there's also Python function call overhead for each kernel launch. With `load_inline`, calling the C++ function from Python has some overhead. But the C++ function can launch all 3 kernels in one call! Yes! In the C++ wrapper, I can launch pass 1, pass 2a, and pass 2b all in one C++ function call. The Python only calls the C++ function once. The kernel launch overhead is just the CUDA driver overhead, which is ~3-5 microseconds per kernel. With 3 kernels: ~10-15 microseconds total launch overhead. Plus compute: ~8 microseconds. Total: ~20 microseconds. For comparison, torch.topk takes ~55 microseconds. So we'd be ~2.7x faster. Can we reduce to 2 kernels? Yes, if pass 1 does block-level merge (which it does), then pass 2 only needs one kernel for the final merge. Wait, pass 1 produces `batch * num_blocks * k` candidates. For batch=1, num_blocks=256, that's 16384 candidates. Can a single block handle 16384 candidates? With 256 threads and tiles of 1024: 16 tiles. Each tile does warp-select (576 ops) + block merge (40 ops) = 600 ops. 16 tiles = 9600 ops. At 1.5GHz: 6.4 microseconds. A single block on 1 SM takes 6.4 microseconds. That's fine. With 2 kernels total: - Kernel 1: 256 blocks. ~0.5 microseconds (memory bound). - Kernel 2: 1 block. ~6.4 microseconds (compute bound). - Launch overhead: 2 * 5us = 10us. - Total: ~17 microseconds. Still 3x faster than torch.topk. But wait, for batch=64, n=8192, k=8: - Pass 1: 64 rows * 4 blocks = 256 blocks. - Each block: 2048 elements, EPT=16, M=512. - Warp select from 512: 64 * (16 + 5) = 1344 ops per thread. - Block sort 256 (4 warps * 8 = 32 candidates): 32 * 5 * 5 / 4 = 200 comparators. Very fast. - Total per block: ~1400 ops. - 256 blocks: ~1 microsecond. - Pass 2: 64 rows * 256 blocks * 8 = 131072 candidates total. - Launch 64 blocks (1 per row), each processes 2048 candidates. - Each block: 256 threads. Tiles of 1024. 2 tiles. - Each tile: warp select from 1024 (128 per warp): 64 * (16 + 5) = 1344 ops. - Block sort 512 (8 warps * 8 = 64 candidates): 64 * 6 * 6 / 4 = 576 comparators. - Total per block: 2 * 1344 + 576 = 3264 ops. - 64 blocks: ~2.2 microseconds. - Total compute: ~3.2 microseconds. - Launch overhead: 10 microseconds. - Total: ~13 microseconds. Compare to eager torch.topk: 49 microseconds. 3.8x faster. This is looking very good! Now, I need to implement this. The key components are: 1. **Pass 1 kernel**: processes a chunk of elements, produces top-k. - Inputs: `x`, `row`, `start`, `end`, `k` - Outputs: `tmp_vals[(row * num_blocks + block) * k + i]`, `tmp_idxs[...]` 2. **Pass 2 kernel**: merges candidates into final top-k. - Inputs: `tmp_vals`, `tmp_idxs`, `num_candidates`, `k` - Outputs: `out_vals`, `out_idxs` For pass 1, the warp-select algorithm: ```cuda __global__ void topk_pass1(const float* __restrict__ x, float* __restrict__ tmp_vals, int* __restrict__ tmp_idxs, int n, int k, int num_blocks_per_row) { int row = blockIdx.x / num_blocks_per_row; int block_in_row = blockIdx.x % num_blocks_per_row; int block_start = block_in_row * n / num_blocks_per_row; int block_end = (block_in_row + 1) * n / num_blocks_per_row; int block_n = block_end - block_start; const float* row_x = x + row * n + block_start; int tid = threadIdx.x; int lane = tid & 31; int warp_id = tid >> 5; int num_warps = blockDim.x >> 5; // Shared memory extern __shared__ char smem[]; // Warp areas int EPT = (block_n + blockDim.x - 1) / blockDim.x; // elements per thread, rounded up // But EPT might vary. Let's use a fixed EPT computed on host. } ``` Wait, computing EPT dynamically is tricky because shared memory size depends on it. We should pass EPT and M as template parameters or compute shared memory dynamically. Actually, for simplicity, let me fix `THREADS = 128` and compute `EPT` on the host side for each launch. But CUDA kernels can't have dynamic shared memory sizes that depend on runtime variables easily... well, they can via the `<<<..., smem>>>` parameter. But the shared memory layout also depends on EPT. Let me use a simpler approach: fix `EPT` based on the maximum block size. For the largest block size: when batch=1, n=131072, num_blocks=256. - block_n = 131072 / 256 = 512. - EPT = ceil(512 / 128) = 4. - M = 32 * 4 = 128. For the smallest block size: batch=128, n=4096, num_blocks=2. - block_n = 4096 / 2 = 2048. - EPT = ceil(2048 / 128) = 16. - M = 32 * 16 = 512. For n=12000, batch=16, num_blocks=16. - block_n = 12000 / 16 = 750. - EPT = ceil(750 / 128) = 6. - M = 32 * 6 = 192. So M ranges from 128 to 512. And k ranges from 1 to 64. Shared memory = num_warps * M * 8 + num_warps * k * 8. For num_warps=4, M=512, k=64: 4 * 512 * 8 + 4 * 64 * 8 = 16384 + 2048 = 18432 bytes = 18KB. For num_warps=4, M=128, k=64: 4096 + 2048 = 6144 bytes = 6KB. All well under 48KB. For pass 2, M (tile size) can be fixed. Let's use TILE = 1024. Shared memory = TILE * 8 + num_warps * k * 8 + k * 8. For TILE=1024, k=64: 8192 + 2048 + 512 = 10752 bytes = 10.5KB. This is getting complex to implement. Let me start coding it. Actually, wait. I should also consider a simpler optimization to my original kernel that might fix the performance issue without a complete rewrite. The original kernel's problem was local memory for the heap arrays. What if I replace the heap with a **simple linear scan** that finds the top-k by doing a partial sort? For K=64, instead of a heap, each thread can: 1. Load all its elements into local memory. 2. Use `std::nth_element` or quickselect to find the k-th largest. 3. Then filter and sort the top-k. But `std::nth_element` isn't available in CUDA device code, and implementing quickselect with dynamic indexing would also use local memory. What if each thread uses **radix sort** on its elements? For fp32, a 1-pass radix sort by the sign+exponent might be fast. But implementing radix sort per thread is complex. Actually, I just realized: the ptxas output showed 0 spill stores/loads and a stack frame. The stack frame IS local memory. But it says 0 spill stores/loads because the compiler isn't spilling registers - it's deliberately allocating local memory for the arrays. The performance impact of local memory depends on how well the L1 cache handles it. With 128 threads and 512 bytes each = 64KB local memory, this fits in L1 (128KB on Blackwell). So local memory should be L1 cached and relatively fast. But the actual time is 1.7ms. That's 1700 microseconds. For 1024 elements per thread and 6 heap levels: 6144 local memory accesses per thread. At 128 threads: 786K accesses total. At 20 cycles per access: 15.7M cycles. At 1.5GHz: 10.5 milliseconds. Wait, that's even worse! Actually, L1 hit latency is ~10-20 cycles, but throughput is much higher. An SM can do many L1 accesses per cycle. The bottleneck might be something else. Let me run `ncu` to profile the kernel and see what's actually happening. Actually, before I completely rewrite the kernel, let me run a quick profile to confirm the bottleneck.

I don't have permission to run ncu. Let me try with nsys or use torch.profiler instead. Or I can just proceed with the algorithmic redesign since I know the issue. Actually, let me try running with sudo or checking if there's an alternative profiling method. But even without profiling, the ptxas output clearly shows local memory usage, and the performance numbers clearly show terrible scaling for small batch sizes. The fix is clear: I need to increase parallelism for small batch sizes and reduce reliance on local memory. Let me implement the two-pass approach. This is a significant redesign but should yield much better performance. Key design decisions: 1. Pass 1: Multi-block warp-select kernel - Each block processes a chunk of elements - Threads load elements in a coalesced pattern - Warp-level selection extracts top-k from warp's elements - Block-level bitonic sort merges warp results - Writes top-k to global memory 2. Pass 2: Multi-block or single-block merge kernel - Reads all candidates for a row - Processes in tiles using warp-select - Block-level merge to final top-k Actually, I want to minimize kernel launches. Let me try to do it in 2 kernels total. For pass 2, if I have many candidates, I can use multiple blocks that each produce top-k, then another pass. But to keep it to 2 kernels, I can have pass 2 process all candidates in a single block using the tile-based approach. For up to ~50000 candidates and k=64, this is ~50 tiles of 1024, which is ~30000 ops per thread. At 1.5GHz that's ~20 microseconds. Acceptable. Wait, 20 microseconds is fine. And with 2 kernels, total time would be ~5 + 20 + launch_overhead = ~30 microseconds. Still faster than torch.topk. But can I do pass 2 with multiple blocks to speed it up? Yes, but then I need a 3rd kernel. Actually, for pass 2, the number of candidates is `batch * num_blocks * k`. For batch=1, num_blocks=256, k=64: 16384 candidates. For batch=64, num_blocks=4, k=8: 64 * 4 * 8 = 2048 candidates. With 2048 candidates, a single block can process them very quickly: - 2 tiles of 1024. - Each tile: warp select from 1024 (8 warps * 128). - Warp select from 128: 64 * (4 + 5) = 576 ops per thread. - Block sort 512 (8 warps * 8 = 64 candidates for k=8? No, block merge is 8 warps * 8 = 64 candidates. Bitonic sort 64 = 64 * 6 * 6 / 4 = 576 comparators. With 256 threads: 2 comparators per thread). - Total per tile: ~580 ops. - 2 tiles: ~1160 ops. - At 1.5GHz: ~0.8 microseconds. So pass 2 is very fast for most cases. The only slow case is batch=1 with 16384 candidates: ~10 microseconds. Can I optimize pass 2 for large candidate counts? What if I use larger tiles? With TILE=2048: - Shared memory: 2048 * 8 = 16KB. Still under 48KB. - Warp select from 2048: 8 warps * 256 candidates. - Warp select from 256: 64 * (8 + 5) = 832 ops per thread. - 16384 / 2048 = 8 tiles. - 8 * 832 = 6656 ops per thread. - At 1.5GHz: ~4.4 microseconds. Better than 10 microseconds with TILE=1024. What about TILE=4096? - Shared memory: 4096 * 8 = 32KB. Still under 48KB. - Warp select from 4096: 8 warps * 512 candidates. - Warp select from 512: 64 * (16 + 5) = 1344 ops. - 16384 / 4096 = 4 tiles. - 4 * 1344 = 5376 ops. - ~3.6 microseconds. Even better! TILE=8192? - Shared memory: 8192 * 8 = 64KB. Needs opt-in. - Warp select from 8192: 8 warps * 1024 candidates. - Warp select from 1024: 64 * (32 + 5) = 2368 ops. - 16384 / 8192 = 2 tiles. - 2 * 2368 = 4736 ops. - ~3.2 microseconds. With opt-in shared memory, we can use TILE=8192 and get ~3.2 microseconds for pass 2. But we need opt-in for pass 2. That's fine, we already use it for other things. Actually, wait. For pass 2 with TILE=8192 and 256 threads, each thread processes 8192/256 = 32 candidates per iteration. For 64 iterations: 64 * 32 = 2048 comparisons per thread. Plus shuffles: 64 * 5 = 320. Total: ~2368 ops. Yes. With 2 tiles: ~4700 ops. At 1.5GHz: 3.1 microseconds. This is excellent! For pass 1, with num_blocks=256 and 128 threads: - Each block: 512 elements, EPT = 4, M = 128. - Warp select from 128: 64 * (4 + 5) = 576 ops. - Block sort 256 (4 warps * 64 = 256 candidates): 256 * 8 * 8 / 4 = 4096 comparators. With 128 threads: ~32 comparators per thread. - Total per block: ~600 ops. - Memory: 512 elements * 4 bytes = 2KB. - 256 blocks: 512KB total. At 1.8TB/s: 0.28 microseconds. - But with compute per block: ~600 ops = ~0.4 microseconds. - Total pass 1: ~0.5 microseconds (limited by compute, not memory). So total with 2 kernels: ~0.5 + 3.1 + launch overhead = ~8-10 microseconds. Wait, but pass 1 has 256 blocks. Can 256 blocks all execute their 600 ops in 0.5 microseconds? Each block has 128 threads = 4 warps. The warp select takes 576 ops per thread. With 4 warps, and each warp executing 576 ops, the block takes ~576 ops * 1 cycle/op = 576 cycles = 0.38 microseconds. But there are 256 blocks. The GPU has 188 SMs. So ~1.36 blocks per SM. Each SM has 4 warp schedulers. 1.36 blocks * 4 warps = 5.4 warps per SM. Each warp scheduler can issue 1 instruction per cycle. So 4 warp schedulers can issue 4 instructions per cycle. 5.4 warps * 576 instructions = 3100 instructions per SM. At 4 instructions per cycle: 775 cycles = 0.5 microseconds. So yes, pass 1 takes about 0.5 microseconds. Pass 2 with 1 block and 4700 ops: 4700 / 256 threads = ~18 ops per thread. Wait, no. The ops are per thread, not total. 2368 ops per thread. With 256 threads = 8 warps. 8 warp schedulers on 1 SM can issue 8 instructions per cycle (if there are enough warps). But we only have 8 warps on 1 SM. Each scheduler gets 2 warps, so it can issue 1 instruction per cycle. Total 4 instructions per cycle. 2368 cycles per warp = 1.6 microseconds per tile. 2 tiles = 3.2 microseconds. Total: 0.5 + 3.2 + 10 (launch overhead) = ~14 microseconds. Still 4x faster than torch.topk! OK, I'm convinced. Let me implement this two-pass approach. For simplicity, let me fix some parameters: - THREADS_PASS1 = 128 - THREADS_PASS2 = 256 - TILE_PASS2 = 8192 (needs opt-in shared memory of 64KB) Wait, for TILE_PASS2 = 8192, shared memory = 8192 * 8 = 65536 bytes = 64KB. This needs opt-in (limit is 101KB). OK. For pass 1, I need to compute num_blocks per row. Let me use: `num_blocks = max(1, 256 / batch)` for batch >= 1. But if batch=128, num_blocks = max(1, 256/128) = 2. For n=4096, block_n = 2048. EPT = ceil(2048 / 128) = 16. M = 32 * 16 = 512. Shared memory = 4 * 512 * 8 + 4 * 1 * 8 = 16384 + 32 = 16416 bytes. (k=1) For k=1, we can use the specialized argmax kernel instead. So pass 1 for k=1 is different. Actually, for k=1, the argmax reduction is much simpler. I can keep the specialized k=1 kernel for pass 1, or just use a warp-level max reduction. For k=1: - Pass 1: each block finds max of its chunk. - Pass 2: find max of all block maxes. This is very fast. For k > 1, use the warp-select approach. Now let me think about the pass 1 kernel more carefully. **Pass 1 kernel for k > 1:** ```cuda template <int K> __global__ void topk_pass1(const float* __restrict__ x, float* __restrict__ tmp_vals, int* __restrict__ tmp_idxs, int n, int num_blocks_per_row) { int row = blockIdx.x / num_blocks_per_row; int block_in_row = blockIdx.x % num_blocks_per_row; int block_start = block_in_row * n / num_blocks_per_row; int block_end = (block_in_row + 1) * n / num_blocks_per_row; int block_n = block_end - block_start; const float* row_x = x + row * n + block_start; int tid = threadIdx.x; int lane = tid & 31; int warp_id = tid >> 5; int num_warps = blockDim.x >> 5; // Compute elements per thread int EPT = (block_n + blockDim.x - 1) / blockDim.x; int M = 32 * EPT; // Shared memory: warp areas + block candidate area extern __shared__ char smem[]; float* s_warp_vals = reinterpret_cast<float*>(smem); int* s_warp_idxs = reinterpret_cast<int*>(smem + num_warps * M * sizeof(float)); float* s_block_vals = reinterpret_cast<float*>(smem + num_warps * M * (sizeof(float) + sizeof(int))); int* s_block_idxs = reinterpret_cast<int*>(s_block_vals + num_warps * K); // Initialize shared memory to -inf for (int i = tid; i < num_warps * M; i += blockDim.x) { s_warp_vals[i] = -FLT_MAX; } __syncthreads(); // Load elements into warp shared memory int warp_start = warp_id * EPT; for (int j = 0; j < EPT; j++) { int idx = warp_start + j * 32 + lane; float val = -FLT_MAX; int idx_out = -1; if (idx < block_n) { val = row_x[idx]; idx_out = block_start + idx; } s_warp_vals[warp_id * M + j * 32 + lane] = val; s_warp_idxs[warp_id * M + j * 32 + lane] = idx_out; } __syncwarp(); // Warp-level selection of top-K from M candidates float* my_warp_vals = s_warp_vals + warp_id * M; int* my_warp_idxs = s_warp_idxs + warp_id * M; float local_max; int local_sm_idx; for (int iter = 0; iter < K; iter++) { local_max = -FLT_MAX; local_sm_idx = -1; for (int i = lane; i < M; i += 32) { if (my_warp_vals[i] > local_max) { local_max = my_warp_vals[i]; local_sm_idx = i; } } // Warp reduction #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { float other_max = __shfl_down_sync(0xFFFFFFFF, local_max, offset); int other_idx = __shfl_down_sync(0xFFFFFFFF, local_sm_idx, offset); if (other_max > local_max) { local_max = other_max; local_sm_idx = other_idx; } } int winner_sm_idx = __shfl_sync(0xFFFFFFFF, local_sm_idx, 0); float winner_val = __shfl_sync(0xFFFFFFFF, local_max, 0); int winner_idx = my_warp_idxs[winner_sm_idx]; // Write to block candidate area if (lane == 0) { s_block_vals[warp_id * K + iter] = winner_val; s_block_idxs[warp_id * K + iter] = winner_idx; } // Mark as -inf if (lane == (winner_sm_idx & 31)) { my_warp_vals[winner_sm_idx] = -FLT_MAX; } __syncwarp(); } __syncthreads(); // Block-level bitonic sort on num_warps * K candidates int block_cand = num_warps * K; for (int kk = 2; kk <= block_cand; kk <<= 1) { for (int j = kk >> 1; j > 0; j >>= 1) { for (int i = tid; i < block_cand; i += blockDim.x) { int ixj = i ^ j; if (ixj > i) { if ((i & kk) == 0) { if (s_block_vals[i] < s_block_vals[ixj]) { float tv = s_block_vals[i]; s_block_vals[i] = s_block_vals[ixj]; s_block_vals[ixj] = tv; int ti = s_block_idxs[i]; s_block_idxs[i] = s_block_idxs[ixj]; s_block_idxs[ixj] = ti; } } else { if (s_block_vals[i] > s_block_vals[ixj]) { float tv = s_block_vals[i]; s_block_vals[i] = s_block_vals[ixj]; s_block_vals[ixj] = tv; int ti = s_block_idxs[i]; s_block_idxs[i] = s_block_idxs[ixj]; s_block_idxs[ixj] = ti; } } } } __syncthreads(); } } // Write top-K to global memory if (tid < K) { int out_pos = (row * num_blocks_per_row + block_in_row) * K + tid; tmp_vals[out_pos] = s_block_vals[tid]; tmp_idxs[out_pos] = s_block_idxs[tid]; } } ``` Wait, there's an issue. In the warp-level selection, after finding the winner, I write to `s_block_vals[warp_id * K + iter]`. But `iter` goes from 0 to K-1, and `warp_id` is 0-3. So `s_block_vals` needs to be `num_warps * K` elements. Yes. But the bitonic sort operates on `s_block_vals[0..block_cand-1]` where `block_cand = num_warps * K`. After the sort, `s_block_vals[0..K-1]` contains the top-K of the block. We write these to global memory. For k=64 and num_warps=4, block_cand=256. The sort produces top-256, but we only need top-64. Writing top-64 is correct. Actually, we only need to write top-K to global memory, not top-block_cand. So writing `s_block_vals[tid]` for tid < K is correct. But wait, we need `s_block_vals` to be sorted descending for the top-K to be at the beginning. The bitonic sort does sort descending. Good. One issue: in the warp-level selection, we write candidates to `s_block_vals` in the order we find them (largest first). So they're already sorted descending within each warp's segment. But the block-level sort re-sorts everything anyway. Another issue: `s_block_vals` is reused for both warp output and block sorting. That's fine because we write all warp outputs before the block sort, and `__syncthreads()` ensures visibility. Wait, in the code above, I write to `s_block_vals` inside the warp loop (`if (lane == 0) { s_block_vals[...] = ... }`). But `s_block_vals` is in shared memory, and other warps might read it during the block sort. The `__syncthreads()` after all warps finish ensures it's safe. But do I need a `__syncthreads()` between the warp loops? No, because each warp writes to its own segment of `s_block_vals`, and there's no `__syncthreads()` between warps. But warps in a block can diverge and execute at different times. So warp 0 might finish its selection and start the block sort while warp 1 is still doing selection! This is a bug! I need a `__syncthreads()` after ALL warps finish their selection. But I can't put `__syncthreads()` inside the warp loop because not all threads in the block reach it (only lane 0 writes, but all threads in the warp participate in the shuffle and reach the end of the iteration). Actually, ALL threads in ALL warps execute the `for (int iter = 0; iter < K; iter++)` loop. So all threads reach the end of the loop. After the loop, I can put `__syncthreads()`. That's fine. Wait, but inside the loop, there's `if (lane == 0) { s_block_vals[...] = ... }`. Only lane 0 writes. But all threads reach the `__syncwarp()` at the end of each iteration. And all threads reach the end of the `for` loop. So `__syncthreads()` after the loop is safe. Good. Now for pass 2: ```cuda template <int K, int TILE> __global__ void topk_pass2(const float* __restrict__ tmp_vals, const int* __restrict__ tmp_idxs, float* __restrict__ out_vals, int64_t* __restrict__ out_idxs, int num_candidates_per_row) { int row = blockIdx.x; const float* row_vals = tmp_vals + row * num_candidates_per_row; const int* row_idxs = tmp_idxs + row * num_candidates_per_row; int tid = threadIdx.x; int lane = tid & 31; int warp_id = tid >> 5; int num_warps = blockDim.x >> 5; extern __shared__ char smem[]; float* s_tile_vals = reinterpret_cast<float*>(smem); int* s_tile_idxs = reinterpret_cast<int*>(smem + TILE * sizeof(float)); float* s_top_vals = reinterpret_cast<float*>(smem + TILE * (sizeof(float) + sizeof(int))); int* s_top_idxs = reinterpret_cast<int*>(s_top_vals + K); // Initialize top-K to -inf for (int i = tid; i < K; i += blockDim.x) { s_top_vals[i] = -FLT_MAX; s_top_idxs[i] = -1; } __syncthreads(); int num_tiles = (num_candidates_per_row + TILE - 1) / TILE; for (int t = 0; t < num_tiles; t++) { int tile_start = t * TILE; int tile_end = min(tile_start + TILE, num_candidates_per_row); int tile_n = tile_end - tile_start; // Load tile into shared memory for (int i = tid; i < TILE; i += blockDim.x) { if (i < tile_n) { s_tile_vals[i] = row_vals[tile_start + i]; s_tile_idxs[i] = row_idxs[tile_start + i]; } else { s_tile_vals[i] = -FLT_MAX; s_tile_idxs[i] = -1; } } __syncthreads(); // Warp-level selection of top-K from TILE candidates int EPT = TILE / blockDim.x; // elements per thread in tile // Actually, each warp processes TILE/num_warps candidates int warp_tile_start = warp_id * (TILE / num_warps); int warp_tile_size = TILE / num_warps; float* my_tile_vals = s_tile_vals + warp_tile_start; int* my_tile_idxs = s_tile_idxs + warp_tile_start; float local_max; int local_sm_idx; for (int iter = 0; iter < K; iter++) { local_max = -FLT_MAX; local_sm_idx = -1; for (int i = lane; i < warp_tile_size; i += 32) { if (my_tile_vals[i] > local_max) { local_max = my_tile_vals[i]; local_sm_idx = i; } } #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { float other_max = __shfl_down_sync(0xFFFFFFFF, local_max, offset); int other_idx = __shfl_down_sync(0xFFFFFFFF, local_sm_idx, offset); if (other_max > local_max) { local_max = other_max; local_sm_idx = other_idx; } } int winner_sm_idx = __shfl_sync(0xFFFFFFFF, local_sm_idx, 0); float winner_val = __shfl_sync(0xFFFFFFFF, local_max, 0); int winner_idx = my_tile_idxs[winner_sm_idx]; // Write to warp output area if (lane == 0) { s_top_vals[num_warps * K + warp_id * K + iter] = winner_val; s_top_idxs[num_warps * K + warp_id * K + iter] = winner_idx; } if (lane == (winner_sm_idx & 31)) { my_tile_vals[winner_sm_idx] = -FLT_MAX; } __syncwarp(); } __syncthreads(); // Block-level bitonic sort on num_warps * K candidates (from this tile) // Then merge with running top-K int block_cand = num_warps * K; float* s_tile_top_vals = s_top_vals + num_warps * K; int* s_tile_top_idxs = s_top_idxs + num_warps * K; for (int kk = 2; kk <= block_cand; kk <<= 1) { for (int j = kk >> 1; j > 0; j >>= 1) { for (int i = tid; i < block_cand; i += blockDim.x) { int ixj = i ^ j; if (ixj > i) { if ((i & kk) == 0) { if (s_tile_top_vals[i] < s_tile_top_vals[ixj]) { swap } } else { if (s_tile_top_vals[i] > s_tile_top_vals[ixj]) { swap } } } } __syncthreads(); } } // Merge s_tile_top_vals[0..K-1] with s_top_vals[0..K-1] // Both are sorted descending if (tid == 0) { float merged_vals[K]; int merged_idxs[K]; int a = 0, b = 0; for (int i = 0; i < K; i++) { float va = (a < K) ? s_top_vals[a] : -FLT_MAX; float vb = (b < K) ? s_tile_top_vals[b] : -FLT_MAX; if (va >= vb) { merged_vals[i] = va; merged_idxs[i] = s_top_idxs[a]; a++; } else { merged_vals[i] = vb; merged_idxs[i] = s_tile_top_idxs[b]; b++; } } for (int i = 0; i < K; i++) { s_top_vals[i] = merged_vals[i]; s_top_idxs[i] = merged_idxs[i]; } } __syncthreads(); } // Write output if (tid < K) { out_vals[row * K + tid] = s_top_vals[tid]; out_idxs[row * K + tid] = (int64_t)s_top_idxs[tid]; } } ``` Wait, this is getting very complex. And there are shared memory conflicts. The running top-K and tile top-K both need to be in shared memory. Let me simplify the shared memory layout for pass 2: - `s_tile_vals[TILE]` - `s_tile_idxs[TILE]` - `s_running_vals[K]` - `s_running_idxs[K]` - `s_tile_top_vals[num_warps * K]` - `s_tile_top_idxs[num_warps * K]` Total: TILE * 8 + K * 8 + num_warps * K * 8. For TILE=8192, K=64, num_warps=8: 65536 + 512 + 4096 = 70144 bytes = 68KB. Needs opt-in. For TILE=4096: 32768 + 512 + 4096 = 37376 bytes = 36KB. Fits in 48KB! With TILE=4096: - num_tiles = 16384 / 4096 = 4. - Warp select from 4096/8 = 512 candidates: 64 * (16 + 5) = 1344 ops per thread. - 4 tiles: 5376 ops per thread. - Block sort 512: 512 * 9 * 9 / 4 = 10368 comparators. With 256 threads: ~40 comparators per thread. - 4 tiles: 160 ops per thread. - 4 merges: 4 * 128 = 512 comparisons (1 thread). - Total: ~6000 ops per thread. - At 1.5GHz: ~4 microseconds. And shared memory = 36KB. Fits in 48KB! No opt-in needed! This is great. Let me use TILE=4096 for pass 2. For pass 1, shared memory = num_warps * M * 8 + num_warps * K * 8. For the largest M=512 (n=4096, 2 blocks, 128 threads): - num_warps=4, M=512, K=64: 4*512*8 + 4*64*8 = 16384 + 2048 = 18432 bytes = 18KB. Fits easily. For k=1, we can use a simpler pass 1 and pass 2: - Pass 1: each block finds max of its chunk using warp shuffle reduction. - Pass 2: each row finds max of all block maxes. But actually, my original k=1 kernel already works well. Let me keep it. Wait, for k=1 with the multi-block approach: - Pass 1: 256 blocks find local max. - Pass 2: 1 block per row finds global max from 256 candidates. - Time: pass 1 ~0.3us + pass 2 ~0.1us + overhead = ~5-10us. - My original k=1 kernel took 0.009ms = 9us. Similar. So I can keep the original k=1 kernel and use the two-pass approach only for k > 1. Actually, let me just rewrite everything cleanly with the two-pass approach. Let me also think about how to pass `num_blocks_per_row` and `EPT` to the kernel. I can compute them in the C++ wrapper and pass as parameters. For pass 1, EPT depends on block_n: `EPT = (block_n + THREADS - 1) / THREADS` `M = 32 * EPT` But M varies per launch. The shared memory size varies too. I can pass `smem` dynamically. In the kernel, I need to know M to loop over the warp array. I can pass M as a parameter. Actually, instead of passing EPT, I can pass `block_n` and compute EPT inside the kernel: `int EPT = (block_n + blockDim.x - 1) / blockDim.x;` `int M = 32 * EPT;` But `block_n` varies per block! Because `n / num_blocks_per_row` might not be evenly divisible. The last block in a row might have fewer elements. For example, n=12000, num_blocks=16. - block_n = 12000 / 16 = 750. - But 750 * 16 = 12000. Evenly divisible. - For n=131072, num_blocks=256: 131072 / 256 = 512. Evenly divisible. - For n=8192, num_blocks=4: 2048. Even. - For n=16384, num_blocks=8: 2048. Even. - For n=4096, num_blocks=2: 2048. Even. All our n values are divisible by reasonable num_blocks! But to be safe, I'll handle the general case. For the last block in a row: `block_start = block_in_row * n / num_blocks_per_row` `block_end = (block_in_row + 1) * n / num_blocks_per_row` Using integer division: `block_start = block_in_row * n / num_blocks_per_row` This might not be perfectly even. For example, n=100, num_blocks=3: - block 0: 0..33 (33 elements) - block 1: 33..66 (33 elements) - block 2: 66..100 (34 elements) So block_n can vary by 1. EPT would be `(34 + 127) / 128 = 2` for block 2, and `(33 + 127) / 128 = 2` for blocks 0,1. Same EPT! For n=12000, num_blocks=16: 750 each. Same. For n=131072, num_blocks=256: 512 each. So EPT is the same for all blocks in a given launch. Good. But for safety, I'll compute `block_start` and `block_end` per block. Actually, let me compute EPT on the host and pass it as a kernel parameter. This avoids division in the kernel. Wait, if all blocks have the same EPT except possibly the last few, I can just compute EPT from the maximum block size. The extra elements loaded will be -inf (out of bounds), which doesn't affect the result. In the kernel: ```cuda int block_start = block_in_row * block_n; int block_end = min(block_start + block_n, n); int actual_block_n = block_end - block_start; ``` And in the load loop: ```cuda for (int j = 0; j < EPT; j++) { int idx = j * 32 + lane; float val = -FLT_MAX; int idx_out = -1; if (idx < actual_block_n) { val = row_x[idx]; idx_out = block_start + idx; } s_warp_vals[warp_id * M + j * 32 + lane] = val; s_warp_idxs[warp_id * M + j * 32 + lane] = idx_out; } ``` This handles variable block sizes safely. Now, for the warp-level selection, we loop over M elements: ```cuda for (int i = lane; i < M; i += 32) { if (my_warp_vals[i] > local_max) { ... } } ``` If `actual_block_n < M`, the extra elements are -inf, so they don't affect the max-finding. Good. For the block-level sort, we sort `num_warps * K` candidates. Some of these might be -inf if the block has very few elements. But the sort handles -inf correctly. Now let me think about the number of blocks. For each shape, I'll compute `num_blocks_per_row`: ```cpp int num_blocks_per_row = std::max(1, 256 / batch); ``` But for batch=128: num_blocks_per_row = 2. Total blocks = 128 * 2 = 256. For n=4096: block_n = 2048. EPT = (2048 + 127) / 128 = 16. M = 32 * 16 = 512. Shared memory = 4 * 512 * 8 + 4 * 1 * 8 = 16416 bytes. (k=1, but k=1 uses different kernel). For k=8: shared memory = 4 * 512 * 8 + 4 * 8 * 8 = 16384 + 256 = 16640 bytes. For k=64: 16384 + 2048 = 18432 bytes. All fit in 48KB. For batch=64, n=8192, k=8: num_blocks_per_row = max(1, 256/64) = 4. block_n = 2048. EPT = 16. M = 512. Shared memory = 4 * 512 * 8 + 4 * 8 * 8 = 16640 bytes. For batch=32, n=16384, k=32: num_blocks_per_row = max(1, 256/32) = 8. block_n = 2048. EPT = 16. M = 512. Shared memory = 4 * 512 * 8 + 4 * 32 * 8 = 16384 + 1024 = 17408 bytes. For batch=16, n=12000, k=16: num_blocks_per_row = max(1, 256/16) = 16. block_n = 750. EPT = (750 + 127) / 128 = 6. M = 32 * 6 = 192. Shared memory = 4 * 192 * 8 + 4 * 16 * 8 = 6144 + 512 = 6656 bytes. For batch=1, n=131072, k=64: num_blocks_per_row = 256. block_n = 512. EPT = 4. M = 128. Shared memory = 4 * 128 * 8 + 4 * 64 * 8 = 4096 + 2048 = 6144 bytes. All fit in 48KB easily. Great! Now for pass 2: - num_candidates_per_row = num_blocks_per_row * K. - batch=1, k=64: 256 * 64 = 16384. - batch=64, k=8: 4 * 8 = 32. - batch=128, k=1: 2 * 1 = 2. (But k=1 uses different path). For pass 2 with small num_candidates, we can use a smaller TILE. TILE=4096 is the maximum. For 32 candidates, we only need 1 tile of 32. But we can still use TILE=4096; the extra elements are -inf. Actually, to save shared memory for small candidate counts, I can just use `num_candidates_per_row` as the tile size if it's small. Or always use TILE=4096 but only allocate `min(TILE, num_candidates_per_row) * 8` for the tile? No, the kernel needs a fixed shared memory size. I can just use `TILE = min(4096, num_candidates_per_row)` and allocate accordingly. But the template needs a compile-time TILE. Actually, I can make TILE a template parameter and instantiate for different sizes. Or I can just always use TILE=4096 and allocate 4096 * 8 = 32KB for the tile. This fits in 48KB for all k. For k=64: 32KB + 512 + 4096 = 36KB. For k=8: 32KB + 64 + 512 = 32.5KB. All fit. Good. But what about k=1? I'll handle k=1 separately with the argmax kernel. Wait, for pass 2 with num_candidates_per_row = 32 and TILE=4096: - We load 32 candidates and 4064 -infs. - Warp select from 4096/8 = 512 candidates, but only 32 are real. - 64 iterations of max-finding over 512 elements each. - Since most are -inf, after 1 iteration the max is found, and subsequent iterations find -inf. - But the loop still runs 64 times, doing 512 comparisons each time. - 64 * 16 = 1024 comparisons per thread. - With 1 tile: 1024 ops per thread = 0.7 microseconds. Fine. Now let me think about the k=1 case. Should I keep the original argmax kernel or use the two-pass approach? The original argmax kernel takes 0.009ms = 9 microseconds for batch=128, n=4096. With two-pass: - Pass 1: 256 blocks, each finds local max of 16 elements. - 256 blocks * 128 threads = 32768 threads. - Each thread: 16 comparisons. - 256 blocks on 188 SMs: ~0.1 microseconds. - Pass 2: 128 blocks, each merges 2 candidates. - 128 blocks * 256 threads. - But we only need to compare 2 values per row. - This is massive overkill. - Total: ~0.1 + 0.1 + overhead = ~5-10 microseconds. Similar performance. I'll keep the original k=1 kernel since it's simple and already fast. OK, let me now write the complete solution.py with the two-pass approach. One more thing: I need to handle the `tmp_vals` and `tmp_idxs` allocation. I'll allocate them in the C++ wrapper. For tmp size: `batch * num_blocks_per_row * k` elements. For batch=1, num_blocks=256, k=64: 16384 elements. For batch=64, num_blocks=4, k=8: 2048 elements. For batch=128, num_blocks=2, k=1: 256 elements. (but k=1 is handled separately) Actually, for k=1, I don't need tmp arrays. Let me write the code now. This is going to be long, but let me be careful. Wait, I should also think about whether the warp-level selection with `__shfl_down_sync` is correct for the case where multiple elements have the same value. If two elements have the same max value, the shuffle reduction will pick one based on the lane ordering. That's fine - any max is acceptable. And when we mark the winner as -inf, if there are duplicates, we only mark one instance. The next iteration will find the other duplicate. This is correct. One subtle bug: in the warp-level selection, `winner_sm_idx` is the index within `my_warp_vals`, which is `warp_id * M + winner_sm_idx` in the full shared memory. When we check `if (lane == (winner_sm_idx & 31))`, we're checking if this thread owns that element in the warp's array. Since the warp array is laid out as `j * 32 + lane`, the element at index `winner_sm_idx` is owned by thread `winner_sm_idx & 31` within the warp. This is correct! But what if `winner_sm_idx >= M`? That can't happen because `winner_sm_idx` is found by iterating `i = lane; i < M; i += 32`. So `winner_sm_idx < M`. OK, the code should be correct. Let me also think about the block-level bitonic sort in pass 1 and pass 2. For small `block_cand`, some threads might have `tid >= block_cand` and do nothing in the inner loop. But they still participate in `__syncthreads()`. This is fine. For pass 2, the merge of `s_top_vals` and `s_tile_top_vals` is done by thread 0 only. After the merge, all threads need the updated `s_top_vals`. So we do `__syncthreads()` after the merge. Wait, but `s_top_vals` is written by thread 0 and read by all threads in the next tile. So `__syncthreads()` after the merge is necessary. I have that. One more thing: in pass 2, `s_tile_top_vals` and `s_tile_top_idxs` are written by the warp-level selection. But `s_top_vals` and `s_top_idxs` are also in shared memory. I need to make sure there's no overlap. Let me design the shared memory layout for pass 2 explicitly: ```cuda // Layout: // [s_tile_vals[TILE]] // [s_tile_idxs[TILE]] // [s_running_vals[K]] // [s_running_idxs[K]] // [s_warp_out_vals[num_warps * K]] // temporary for tile top-k // [s_warp_out_idxs[num_warps * K]] ``` Total: TILE * 8 + K * 8 + num_warps * K * 8. For TILE=4096, K=64, num_warps=8: 32768 + 512 + 4096 = 37376 bytes. Offset calculations: - `s_tile_vals = smem` - `s_tile_idxs = smem + TILE * sizeof(float)` - `s_running_vals = smem + TILE * (sizeof(float) + sizeof(int))` - `s_running_idxs = s_running_vals + K * sizeof(float)` - `s_warp_out_vals = s_running_idxs + K * sizeof(int)` - `s_warp_out_idxs = s_warp_out_vals + num_warps * K * sizeof(float)` Actually, I'll just use a struct-based approach or compute offsets carefully. Let me write the code. For pass 1, I'll create a template kernel `topk_pass1<K>` where K is the compile-time k. For pass 2, I'll create a template kernel `topk_pass2<K, TILE>`. In the C++ wrapper: ```cpp std::vector<torch::Tensor> topk_forward_cuda(torch::Tensor x, int k) { int batch = x.size(0); int n = x.size(1); auto out_vals = torch::empty({batch, k}, x.options()); auto out_idxs = torch::empty({batch, k}, torch::dtype(torch::kInt64).device(x.device())); if (k == 1) { topk1_kernel<<<batch, 256>>>(x.data_ptr<float>(), out_vals.data_ptr<float>(), out_idxs.data_ptr<int64_t>(), n); return {out_vals, out_idxs}; } int num_blocks_per_row = std::max(1, 256 / batch); int total_blocks = batch * num_blocks_per_row; int block_n = (n + num_blocks_per_row - 1) / num_blocks_per_row; int EPT = (block_n + 127) / 128; // THREADS_PASS1 = 128 int M = 32 * EPT; auto tmp_vals = torch::empty({batch, num_blocks_per_row, k}, x.options()); auto tmp_idxs = torch::empty({batch, num_blocks_per_row, k}, torch::dtype(torch::kInt32).device(x.device())); // Pass 1 int threads1 = 128; int num_warps1 = threads1 / 32; size_t smem1 = num_warps1 * M * sizeof(float) + num_warps1 * M * sizeof(int) + num_warps1 * k * sizeof(float) + num_warps1 * k * sizeof(int); switch (k) { case 8: topk_pass1<8> <<<total_blocks, threads1, smem1>>>(...); break; case 16: topk_pass1<16><<<total_blocks, threads1, smem1>>>(...); break; case 32: topk_pass1<32><<<total_blocks, threads1, smem1>>>(...); break; case 64: topk_pass1<64><<<total_blocks, threads1, smem1>>>(...); break; } // Pass 2 int threads2 = 256; int num_warps2 = threads2 / 32; const int TILE = 4096; size_t smem2 = TILE * sizeof(float) + TILE * sizeof(int) + k * sizeof(float) + k * sizeof(int) + num_warps2 * k * sizeof(float) + num_warps2 * k * sizeof(int); int num_candidates = num_blocks_per_row * k; switch (k) { case 8: topk_pass2<8, TILE><<<batch, threads2, smem2>>>(...); break; case 16: topk_pass2<16, TILE><<<batch, threads2, smem2>>>(...); break; case 32: topk_pass2<32, TILE><<<batch, threads2, smem2>>>(...); break; case 64: topk_pass2<64, TILE><<<batch, threads2, smem2>>>(...); break; } return {out_vals, out_idxs}; } ``` Wait, for pass 1, the `block_n` computation: ```cpp int block_n = (n + num_blocks_per_row - 1) / num_blocks_per_row; ``` This gives the maximum block size. Some blocks might be smaller. But as I computed earlier, for our shapes, all blocks are the same size. For n=12000, num_blocks=16: block_n = 750. Exactly even. For n=131072, num_blocks=256: block_n = 512. Even. So it's fine. Let me now write the full CUDA code. One concern: in pass 1, `M = 32 * EPT` might not be a power of 2. But the warp-level selection doesn't require M to be a power of 2. The loop `for (int i = lane; i < M; i += 32)` works for any M. However, for the block-level bitonic sort, `block_cand = num_warps * k` must be a power of 2 for the standard bitonic sort. Let me check: - num_warps = 4, k=64: block_cand = 256. Power of 2. Good. - num_warps = 4, k=32: block_cand = 128. Power of 2. Good. - num_warps = 4, k=16: block_cand = 64. Power of 2. Good. - num_warps = 4, k=8: block_cand = 32. Power of 2. Good. For num_warps=4, block_cand = 4 * k. Since k is 8, 16, 32, 64, block_cand is always a power of 2! Great. For pass 2, block_cand = num_warps * k = 8 * k. - k=8: 64. Power of 2. - k=16: 128. Power of 2. - k=32: 256. Power of 2. - k=64: 512. Power of 2. Great, all powers of 2. For TILE=4096, it's also a power of 2. Good. Now, in the pass 2 kernel, the warp tile size is TILE / num_warps = 4096 / 8 = 512. This is a power of 2. But the warp selection doesn't require it to be a power of 2. Let me write the code now. I'll be very careful with shared memory offsets. Actually, I realize the code is going to be quite long. Let me write it section by section. For the C++ wrapper, I need to compute `EPT` and `M` correctly. `EPT` is the maximum elements any thread might need to process. `block_n = (n + num_blocks_per_row - 1) / num_blocks_per_row;` `EPT = (block_n + 127) / 128;` `M = 32 * EPT;` For n=12000, num_blocks=16: block_n = 750. EPT = (750 + 127) / 128 = 877 / 128 = 6 (integer division). M = 192. For n=131072, num_blocks=256: block_n = 512. EPT = (512 + 127) / 128 = 639 / 128 = 4. M = 128. Good. One more thing: I need to be careful with `cudaFuncSetAttribute` for pass 2 if smem2 exceeds 48KB. smem2 = 4096 * 4 + 4096 * 4 + 64 * 4 + 64 * 4 + 8 * 64 * 4 + 8 * 64 * 4 = 16384 + 16384 + 256 + 256 + 2048 + 2048 = 37376 bytes = 36.5KB. Under 48KB. No opt-in needed. For pass 1, smem1 max = 18432 bytes = 18KB. Also under 48KB. Great! No opt-in needed anywhere. Let me write the full code now. Actually, wait. Let me double-check pass 2 smem2 for k=64: TILE * sizeof(float) = 4096 * 4 = 16384 TILE * sizeof(int) = 4096 * 4 = 16384 K * sizeof(float) = 64 * 4 = 256 K * sizeof(int) = 64 * 4 = 256 num_warps * K * sizeof(float) = 8 * 64 * 4 = 2048 num_warps * K * sizeof(int) = 8 * 64 * 4 = 2048 Total: 16384 + 16384 + 256 + 256 + 2048 + 2048 = 37376 bytes. Yes. For k=8: 16384 + 16384 + 32 + 32 + 256 + 256 = 33344 bytes. Yes. Now let me also think about whether I can fuse pass 1 and pass 2 into a single kernel for the case where num_blocks_per_row = 1. When num_blocks_per_row = 1 (e.g., batch >= 256), pass 1 produces k candidates per row, and pass 2 just reads k candidates and writes them back. This is wasteful. For num_blocks_per_row = 1, we can skip pass 2 entirely and just use the single-block warp-select kernel directly. The block processes the entire row and writes top-k directly to output. This would save one kernel launch. In the C++ wrapper: ```cpp if (num_blocks_per_row == 1) { // Single-pass kernel } else { // Two-pass } ``` The single-pass kernel is just pass 1 but writing to output instead of tmp. For batch=64, num_blocks_per_row = 4. Two-pass. For batch=128, num_blocks_per_row = 2. Two-pass. Actually, for batch=128, num_blocks_per_row = max(1, 256/128) = 2. For batch=256: num_blocks_per_row = 1. Single-pass. But our max batch is 128. For batch=32: num_blocks_per_row = 8. For batch=16: num_blocks_per_row = 16. For batch=1: num_blocks_per_row = 256. So for all our shapes, num_blocks_per_row >= 2. The two-pass approach is always used (except k=1). Wait, is there any shape with batch >= 256? No. So we always use at least 2 blocks per row in pass 1. OK, let me write the code. For the single-pass case, I can just add it as an optimization for future-proofing. Actually, let me keep it simple and just always do two-pass for k > 1. The overhead of pass 2 for small candidate counts is negligible. Let me write the code now. I'll create a new solution.py from scratch.

I wrote the new two-pass solution. Let me test it for correctness first, then benchmark.

The correctness check failed for shape 0 (batch=1, n=131072, k=64). The max absolute difference is 0.508565, which is huge. This means my kernel is returning incorrect values. Let me debug this. I need to figure out what's wrong. Let me run a quick test to compare my output with torch.topk for this specific shape.

The solution is returning duplicate indices (6669 appears 4 times) and missing many large values. This is clearly a bug in the kernel. Let me think about what could cause duplicate indices. In pass 1, the warp-level selection marks the winner as -FLT_MAX after finding it. But maybe the marking isn't working correctly, causing the same element to be selected multiple times. Wait, looking at the pass 1 kernel: ```cuda if (lane == (winner_pos & 31)) { my_warp_vals[winner_pos] = -FLT_MAX; } __syncwarp(); ``` `winner_pos` is the position within `my_warp_vals`, which ranges from 0 to M-1. `winner_pos & 31` gives the lane that "owns" that position. But wait - in our layout, position `i` is written by lane `i & 31`. So this is correct. But what if `winner_pos` is not in the current warp's array? No, `winner_pos` comes from iterating over `my_warp_vals`, so it's in the current warp's array. Actually, wait! `__shfl_down_sync` does a reduction but the `local_pos` is NOT the actual shared memory index - it's the index within the thread's OWN local search range. When we do `__shfl_down_sync` on `local_pos`, we pass the position that the thread found. But after comparing values, the winning thread's `local_pos` is kept. But `local_pos` is relative to `my_warp_vals`, so all threads in the warp use the same `my_warp_vals`. So `winner_pos` is correct. But wait - in the shuffle reduction: ```cuda for (int offset = 16; offset > 0; offset >>= 1) { float other_max = __shfl_down_sync(0xFFFFFFFFu, local_max, offset); int other_pos = __shfl_down_sync(0xFFFFFFFFu, local_pos, offset); if (other_max > local_max) { local_max = other_max; local_pos = other_pos; } } ``` After this loop, `local_max` and `local_pos` in lane 0 are the global max and its position. But what about other lanes? They have the max from their subtree, not the global max. That's fine because we broadcast from lane 0. ```cuda int winner_pos = __shfl_sync(0xFFFFFFFFu, local_pos, 0); ``` This broadcasts lane 0's `local_pos` to all lanes. Good. Then: ```cuda if (lane == (winner_pos & 31)) { my_warp_vals[winner_pos] = -FLT_MAX; } ``` This sets the winning element to -FLT_MAX. Good. But wait, what if there are multiple elements with the same max value? The shuffle reduction might pick one, but then on the next iteration, the other equal element would be found. But in the output, I see the SAME index 6669 appearing 4 times. This means it's being selected 4 times without being marked. Why would marking fail? Let me check: `winner_pos & 31` gives the lane. But `my_warp_vals[winner_pos]` is accessed by that lane. Is it possible that `winner_pos` is out of bounds? No, because `winner_pos < M`. Wait, I see a potential bug. In the shuffle reduction, `local_pos` from a higher lane is shuffled down to a lower lane. But `local_pos` is the index within `my_warp_vals` that the higher lane found. This is correct because all threads in the warp are searching the same `my_warp_vals`. But what about `__shfl_down_sync` semantics? `__shfl_down_sync(mask, var, delta)` returns the value of `var` from the thread `lane + delta`. If `lane + delta >= 32`, it returns the value from thread `lane` (or is undefined?). Actually, `__shfl_down_sync` with `delta` returns `var` from `lane + delta`. For `lane + delta >= warpSize`, the behavior is that it returns the value from the source lane modulo warpSize? No, I think for `lane + delta >= 32`, it wraps around or returns the same lane's value. Wait, actually `__shfl_down_sync` returns the value from `lane + delta`. If `lane + delta` is out of range, it returns the value from the caller's own lane. So for lane 16 and offset 16, `lane + delta = 32`, which is out of range, so it returns lane 16's own value. This is fine for a reduction. Let me trace: - Start: each lane has its local max and position. - offset=16: lane 0 gets lane 16's value. lane 1 gets lane 17's value. ... lane 15 gets lane 31's value. lane 16 gets lane 16's own value (out of range). - Compare and keep max. - offset=8: lane 0 gets lane 8's value. lane 8 gets lane 16's value. etc. - Continue until offset=1. - After all offsets, lane 0 has the global max and position. This is correct. So why are there duplicate indices? Let me check if maybe the block-level bitonic sort is corrupting the indices. Wait, looking at the output more carefully: - sol_idx=6669 appears 4 times with value 3.931298. - This means in pass 2, the same index from pass 1 is appearing multiple times. In pass 1, each block writes its top-64 to tmp. If a block has duplicate indices in its top-64, then pass 2 will see duplicates. Why would pass 1 produce duplicate indices? Oh! I think I see the bug. In pass 1, the warp-level selection uses `__syncwarp()` after marking the winner. But `__syncwarp()` only synchronizes threads within a warp. The `winner_pos` is broadcast to all threads in the warp via `__shfl_sync`. Then the correct lane sets `my_warp_vals[winner_pos] = -FLT_MAX`. Then `__syncwarp()` ensures all threads in the warp see the update before the next iteration. This should work correctly within a warp. But what about the block-level sort? Wait, after the warp-level selection, we write the top-K to `s_block_vals[warp_id * K + iter]`. Each warp writes K elements. Then `__syncthreads()` before the block sort. The block sort sorts `s_block_vals` and `s_block_idxs`. The bitonic sort code looks correct. But what if `s_block_vals` has -FLT_MAX values? For example, if a block has fewer than `num_warps * K` valid candidates. But since each warp selects K elements from M candidates, and M >= K (since EPT >= 1), each warp should find K valid elements. Unless the block has fewer than K elements total. For n=131072, num_blocks=256, block_n=512. M=128. K=64. Each warp has 128 candidates and selects 64. This is fine. Wait, but what if some elements are -FLT_MAX? In the load loop: ```cuda if (idx < block_n) { val = row_x[idx]; idx_out = block_start + idx; } ``` For idx >= block_n, val = -FLT_MAX and idx_out = -1. If block_n = 512 and M = 128 (EPT=4), then all 512 elements are valid. No -FLT_MAX padding. But the warp selects top-64 from 128 elements. If there are no -FLT_MAX values, the top-64 should all be valid and unique (assuming no ties). But wait - in the output, the same index appears 4 times in the SOLUTION, not in a single block. The solution output is after pass 2. So maybe pass 2 is introducing duplicates? In pass 2, the warp-level selection also marks winners as -FLT_MAX. And then the block sort sorts the warp results. Then a merge with the running top-K. Could the merge be introducing duplicates? The merge code: ```cuda bool take_a = (a < K) && (b >= K || s_run_vals[a] >= s_wout_vals[b]); ``` If `s_run_vals[a] == s_wout_vals[b]`, it takes from `s_run_vals`. This shouldn't introduce duplicates, but it could drop valid elements from `s_wout_vals`. Wait, but duplicates in the output mean that some elements from the input appear multiple times. This could happen if: 1. Pass 1 produces duplicate indices across different blocks. 2. Pass 2 doesn't deduplicate. But different blocks process disjoint element ranges. So they can't produce the same index... unless there's a bug in index calculation. Wait! In pass 1: ```cuda int idx = j * 32 + lane; float val = -FLT_MAX; int idx_out = -1; if (idx < block_n) { val = row_x[idx]; idx_out = block_start + idx; } s_warp_vals[warp_off + j * 32 + lane] = val; s_warp_idxs[warp_off + j * 32 + lane] = idx_out; ``` `idx` is the offset within the block. `idx_out = block_start + idx` is the global index within the row. But `block_start` is computed as: ```cuda int block_start = block_in_row * n / num_blocks_per_row; ``` For n=131072, num_blocks=256, block_in_row = 0..255. block_start = block_in_row * 131072 / 256 = block_in_row * 512. For block_in_row = 0: block_start = 0. For block_in_row = 1: block_start = 512. ... So idx_out ranges from block_start to block_start + block_n - 1. These are unique across blocks. Good. But wait, `block_start` uses integer division. Is `block_in_row * n / num_blocks_per_row` exact? For n=131072 and num_blocks=256: block_start = block_in_row * 512. Exact. For n=12000 and num_blocks=16: block_start = block_in_row * 750. Exact. OK, so indices are unique across blocks. So pass 1 shouldn't produce duplicates. Unless... the warp selection is selecting the same element multiple times within a warp. Let me check the marking logic again. ```cuda int winner_pos = __shfl_sync(0xFFFFFFFFu, local_pos, 0); ... if (lane == (winner_pos & 31)) { my_warp_vals[winner_pos] = -FLT_MAX; } __syncwarp(); ``` `winner_pos` is the index within `my_warp_vals`. The owner thread is `winner_pos & 31`. That thread sets the value to -FLT_MAX. But what if `winner_pos` is, say, 33. Then `winner_pos & 31 = 1`. Thread 1 sets `my_warp_vals[33] = -FLT_MAX`. On the next iteration, when threads scan `my_warp_vals`, thread 1 checks `my_warp_vals[33]`. Since thread 1 just set it to -FLT_MAX, it should see -FLT_MAX. But wait - is there a memory consistency issue? `__syncwarp()` should ensure that all threads in the warp see the write. Actually, I just realized: `__syncwarp()` is a memory fence for threads within the warp. But `my_warp_vals` is in shared memory. Shared memory is visible to all threads in the block. Within a warp, `__syncwarp()` should ensure consistency. Hmm, let me think about another possibility. What if the block-level bitonic sort is corrupting the data? Looking at the block sort in pass 1: ```cuda for (int kk = 2; kk <= block_cand; kk <<= 1) { for (int j = kk >> 1; j > 0; j >>= 1) { for (int i = tid; i < block_cand; i += blockDim.x) { int ixj = i ^ j; if (ixj > i) { if ((i & kk) == 0) { if (s_block_vals[i] < s_block_vals[ixj]) { // swap vals and idxs } } else { if (s_block_vals[i] > s_block_vals[ixj]) { // swap vals and idxs } } } } __syncthreads(); } } ``` This is the standard bitonic sort for descending order. I verified it earlier. It should be correct. Wait, there's a potential issue with `block_cand` being small and `tid` being large. For example, if block_cand = 32 and tid = 32..127, the loop `for (int i = tid; i < block_cand; i += blockDim.x)` doesn't execute. But `__syncthreads()` is reached by all threads. That's fine. But what if `block_cand` is larger than expected? No, block_cand = num_warps * K. For num_warps=4, K=64, block_cand=256. This is correct. Let me check if the issue is in pass 2. In pass 2, the warp-level selection also uses marking. Could the same bug appear there? Actually, in pass 2: ```cuda float* my_tile_vals = s_tile_vals + warp_id * warp_tile; int* my_tile_idxs = s_tile_idxs + warp_id * warp_tile; ``` And then the selection loop. This is identical to pass 1's warp selection, but on `s_tile_vals` instead of `s_warp_vals`. Wait, but `warp_tile = TILE / num_warps = 4096 / 8 = 512`. And `my_tile_vals` points to `s_tile_vals + warp_id * 512`. In the load loop: ```cuda for (int i = tid; i < TILE; i += blockDim.x) { if (i < tile_n) { s_tile_vals[i] = row_vals[tile_start + i]; s_tile_idxs[i] = row_idxs[tile_start + i]; } else { s_tile_vals[i] = -FLT_MAX; s_tile_idxs[i] = -1; } } ``` All threads load into `s_tile_vals[i]`. So `s_tile_vals` is a flat array of TILE elements. In the warp selection, `my_tile_vals = s_tile_vals + warp_id * 512`. So warp 0 scans 0..511, warp 1 scans 512..1023, etc. And `winner_pos` is the index within `my_tile_vals`, i.e., 0..511. The owner is `winner_pos & 31`. The actual shared memory index is `warp_id * 512 + winner_pos`. But the marking code is: ```cuda if (lane == (winner_pos & 31)) { my_tile_vals[winner_pos] = -FLT_MAX; } ``` Wait! `my_tile_vals[winner_pos]` is `s_tile_vals[warp_id * 512 + winner_pos]`. This is correct. But what about `__syncwarp()`? After marking, `__syncwarp()` ensures all threads in the warp see the update. But other warps are also running concurrently and might be reading/writing to their own regions of `s_tile_vals`. Since each warp uses a different 512-element region, there's no conflict. So the marking should work. Hmm, let me think about another potential issue. In pass 2, after the warp selection, the block sort operates on `s_wout_vals` and `s_wout_idxs`. These are written by each warp during the selection. But `s_wout_vals` and `s_wout_idxs` are in shared memory. During the warp selection, each warp writes K elements to `s_wout_vals[warp_id * K + iter]`. Since `iter` goes from 0 to K-1, and each warp writes to its own segment, there's no conflict. Then `__syncthreads()` before the block sort. Good. Then the block sort. Good. Then the merge: ```cuda if (tid == 0) { float merged_vals[64]; int merged_idxs[64]; int a = 0, b = 0; for (int i = 0; i < K; ++i) { bool take_a = (a < K) && (b >= K || s_run_vals[a] >= s_wout_vals[b]); if (take_a) { merged_vals[i] = s_run_vals[a]; merged_idxs[i] = s_run_idxs[a]; ++a; } else { merged_vals[i] = s_wout_vals[b]; merged_idxs[i] = s_wout_idxs[b]; ++b; } } for (int i = 0; i < K; ++i) { s_run_vals[i] = merged_vals[i]; s_run_idxs[i] = merged_idxs[i]; } } ``` Wait, `s_wout_vals` is the tile's top-K. But after the block sort, is `s_wout_vals` sorted descending? Yes, the block sort sorts `s_wout_vals` descending. So `s_wout_vals[0..K-1]` are the top-K of the tile, sorted descending. And `s_run_vals` is the running top-K, also sorted descending. The merge should produce the correct top-K from the union. But what if `s_run_vals` contains -FLT_MAX on the first tile? Yes, it's initialized to -FLT_MAX. On the first tile, the merge should just copy `s_wout_vals`. On subsequent tiles, the merge combines the two sorted arrays. This looks correct. So where are the duplicates coming from? Wait, I just realized something. In the pass 1 block sort, after sorting, the top-K elements are at `s_block_vals[0..K-1]`. But what if `s_block_vals` contains duplicate values from different warps? For example, if two different warps both found the same element? No, warps process disjoint element ranges, so they can't find the same element. But what if `s_block_vals` contains duplicate VALUES (different indices with the same float value)? That's possible but shouldn't cause the SAME index to appear multiple times. The indices should be different. Unless... the indices themselves are wrong! Let me check the index computation in pass 1 again. ```cuda int idx = j * 32 + lane; float val = -FLT_MAX; int idx_out = -1; if (idx < block_n) { val = row_x[idx]; idx_out = block_start + idx; } s_warp_vals[warp_off + j * 32 + lane] = val; s_warp_idxs[warp_off + j * 32 + lane] = idx_out; ``` `idx` is the position within the block's element range. `idx_out = block_start + idx` is the global row index. For block 0, block_start=0, idx=0..511. idx_out=0..511. For block 1, block_start=512, idx=0..511. idx_out=512..1023. These are all unique. But wait! The warp selection scans `my_warp_vals` and records `my_warp_idxs[winner_pos]`. Since `my_warp_idxs` contains `idx_out` values, the recorded index should be the global row index. Then the block sort swaps both values and indices. So the top-K should have unique indices. Unless... the block sort has a bug in index swapping! Let me check: ```cuda if (s_block_vals[i] < s_block_vals[ixj]) { float tv = s_block_vals[i]; s_block_vals[i] = s_block_vals[ixj]; s_block_vals[ixj] = tv; int ti = s_block_idxs[i]; s_block_idxs[i] = s_block_idxs[ixj]; s_block_idxs[ixj] = ti; } ``` This swaps both value and index. Correct. And for the other direction: ```cuda if (s_block_vals[i] > s_block_vals[ixj]) { float tv = s_block_vals[i]; s_block_vals[i] = s_block_vals[ixj]; s_block_vals[ixj] = tv; int ti = s_block_idxs[i]; s_block_idxs[i] = s_block_idxs[ixj]; s_block_idxs[ixj] = ti; } ``` Also correct. Hmm, let me check the pass 1 output more carefully. Maybe the issue is that `tmp_vals` and `tmp_idxs` are not being written correctly. ```cuda if (tid < K) { int out_pos = (row * num_blocks_per_row + block_in_row) * K + tid; tmp_vals[out_pos] = s_block_vals[tid]; tmp_idxs[out_pos] = s_block_idxs[tid]; } ``` `tmp_vals` has shape `{batch, num_blocks_per_row, k}`. `out_pos = (row * num_blocks_per_row + block_in_row) * K + tid`. For row=0, num_blocks=256, block=0..255: out_pos = block * 64 + tid. This maps to a flat index in the 2D tensor `{batch * num_blocks, K}`. Since `tmp_vals` is stored as `{batch, num_blocks, K}` in row-major, the offset is: `row * num_blocks * K + block * K + tid = (row * num_blocks + block) * K + tid`. Yes, this is correct. In pass 2: ```cuda const float* row_vals = tmp_vals + row * num_candidates_per_row; const int* row_idxs = tmp_idxs + row * num_candidates_per_row; ``` `num_candidates_per_row = num_blocks_per_row * k = 256 * 64 = 16384`. `row_vals` starts at `row * 16384`. Correct. Then: ```cuda s_tile_vals[i] = row_vals[tile_start + i]; s_tile_idxs[i] = row_idxs[tile_start + i]; ``` For i=0..16383, this reads all candidates for the row. Correct. Wait, could the issue be in the `tmp_vals` allocation? I allocated: ```cpp auto tmp_vals = torch::empty({batch, num_blocks_per_row, k}, x.options()); auto tmp_idxs = torch::empty({batch, num_blocks_per_row, k}, torch::dtype(torch::kInt32).device(x.device())); ``` For batch=1, num_blocks=256, k=64: tmp_vals shape = {1, 256, 64}. tmp_idxs shape = {1, 256, 64}. In pass 1, `tmp_vals.data_ptr<float>()` gives a pointer to the flat array. `out_pos = (row * num_blocks + block) * K + tid`. For row=0, block=0, tid=0: out_pos = 0. For row=0, block=1, tid=0: out_pos = 64. ... For row=0, block=255, tid=63: out_pos = 255 * 64 + 63 = 16383. Total elements = 1 * 256 * 64 = 16384. out_pos ranges from 0 to 16383. Correct. In pass 2, `tmp_vals + row * num_candidates_per_row = tmp_vals + row * 16384`. For row=0: tmp_vals + 0. Correct. For row=0: reads 0..16383. Correct. So the indexing seems correct. Let me think about what else could cause duplicates. Actually, wait. In pass 1, when loading elements: ```cuda int warp_off = warp_id * M; for (int j = 0; j < EPT; ++j) { int idx = j * 32 + lane; ... s_warp_vals[warp_off + j * 32 + lane] = val; s_warp_idxs[warp_off + j * 32 + lane] = idx_out; } ``` `M = 32 * EPT`. So `warp_off + j * 32 + lane` ranges from `warp_off` to `warp_off + M - 1`. But what if `block_n` is not a multiple of 32? For n=12000, num_blocks=16, block_n=750. EPT=6. M=192. Warp 0 processes indices 0..191. Warp 1 processes indices 192..383. Warp 2 processes indices 384..575. Warp 3 processes indices 576..767. But block_n = 750. So warp 3 should only process indices 576..749. Indices 750..767 are out of bounds and should be -FLT_MAX. In the load loop: ```cuda if (idx < block_n) { val = row_x[idx]; idx_out = block_start + idx; } ``` For warp 3, lane=0, j=0: idx=0. But this is relative to the warp's start? NO! `idx = j * 32 + lane;` This is NOT relative to warp_start. It's just `j * 32 + lane`. For warp 3: - lane=0, j=0: idx=0. `idx < block_n` (0 < 750) is true. So it loads `row_x[0]`! - But warp 3 should load `row_x[576]`! THIS IS THE BUG! `idx` is the offset within the WARP's range, not the BLOCK's range. But `row_x` points to `x + row * n + block_start`, which is the start of the block. So `row_x[0]` is the first element of the block, not the first element of warp 3's range! Wait, actually all warps in the block load from the SAME `row_x` pointer. So warp 0 loads elements 0, 1, 2, ..., warp 1 also loads elements 0, 1, 2, ...! They're all loading the SAME elements! No no no, let me re-read: ```cuda const float* row_x = x + row * n + block_start; ``` So `row_x` is the block's start. All warps load from `row_x`. For warp 0: `idx = j * 32 + lane` loads `row_x[0], row_x[1], ...` For warp 1: same `idx` values load `row_x[0], row_x[1], ...` ALL WARPS LOAD THE SAME ELEMENTS! That's the bug! Each warp should load a DIFFERENT portion of the block. Warp 0 should load `row_x[0..191]`, warp 1 should load `row_x[192..383]`, etc. I need to add a warp offset to the load: ```cuda int warp_start = warp_id * (block_n / num_warps); // Or more precisely: int warp_chunk = (block_n + num_warps - 1) / num_warps; int warp_start = warp_id * warp_chunk; int warp_end = min(warp_start + warp_chunk, block_n); ``` But wait, my current code has: ```cuda int warp_off = warp_id * M; for (int j = 0; j < EPT; ++j) { int idx = j * 32 + lane; ... s_warp_vals[warp_off + j * 32 + lane] = val; s_warp_idxs[warp_off + j * 32 + lane] = idx_out; } ``` `warp_off` is used for the SHARED MEMORY offset, not for the GLOBAL MEMORY offset! The global memory index is just `idx`, not `warp_start + idx`. So all warps read the same global memory elements but write to different shared memory regions. This means each warp is processing the SAME elements, not different ones. For example, with block_n=512 and 4 warps: - Warp 0 loads elements 0..127 into shared memory 0..127. - Warp 1 loads elements 0..127 into shared memory 128..255. - Warp 2 loads elements 0..127 into shared memory 256..383. - Warp 3 loads elements 0..127 into shared memory 384..511. Then each warp selects top-64 from its 128 elements. But all warps are selecting from the SAME 128 elements! So the block has 4 copies of the top-64 of the first 128 elements, not the top-64 of all 512 elements. This explains the duplicates! And it explains why the output values are wrong (only using a subset of elements). The fix is to have each warp load a different portion of the block: ```cuda int warp_chunk = (block_n + num_warps - 1) / num_warps; int warp_start = warp_id * warp_chunk; int warp_end = min(warp_start + warp_chunk, block_n); int warp_n = warp_end - warp_start; for (int j = 0; j < EPT; ++j) { int idx = warp_start + j * 32 + lane; float val = -FLT_MAX; int idx_out = -1; if (idx < warp_end) { val = row_x[idx]; idx_out = block_start + idx; } s_warp_vals[warp_off + j * 32 + lane] = val; s_warp_idxs[warp_off + j * 32 + lane] = idx_out; } ``` But wait, `EPT` is computed based on `block_n`, not `warp_n`. So some threads might load elements past `warp_end`. We handle this with the `if (idx < warp_end)` check. And `M = 32 * EPT` is the shared memory size per warp. But `warp_n` might be less than `M`. The extra elements are -FLT_MAX, which is fine. Wait, but `warp_chunk` might not be a multiple of 32. For example, block_n=750, num_warps=4: warp_chunk = (750 + 3) / 4 = 188. warp_start for warp 0: 0, warp 1: 188, warp 2: 376, warp 3: 564. warp_end: 188, 376, 564, 750. For warp 3: warp_n = 750 - 564 = 186. EPT = 6. M = 192. The last 6 elements (186..191) are -FLT_MAX. Fine. For coalesced access: warp 0 reads: 0, 1, 2, ..., 31 (first transaction), 32, 33, ..., 63 (second), etc. This is perfectly coalesced! Great, let me fix this bug. But wait, there's another issue. `EPT = (block_n + 127) / 128` is the same for all warps. But `warp_chunk` might be slightly different for the last warp. However, since `EPT` is the maximum, it's fine. The extra elements are just -FLT_MAX. Actually, to keep things simple, I can just compute `EPT_warp = (warp_n + 31) / 32` and `M_warp = 32 * EPT_warp`. But this makes shared memory layout per-warp non-uniform. It's easier to use a uniform `M` for all warps. So I'll keep `M = 32 * EPT` where `EPT = (block_n + 127) / 128`, but add `warp_start` to the global memory load. Actually, there's a subtlety. `block_n` is the number of elements in the block. `block_n / num_warps` might not be divisible by 32. So `warp_start` might not be aligned to 32. For example, block_n=750, num_warps=4, warp_start=188. For warp_start=188: - Thread 0 reads row_x[188] - Thread 1 reads row_x[189] - ... - Thread 31 reads row_x[219] These are consecutive elements. Coalesced! Perfect. But thread 0 in the second iteration reads row_x[188 + 32] = row_x[220]. Also consecutive. Coalesced. So the load is coalesced regardless of warp_start alignment. Great. Let me fix the code. For pass 2, do I have the same bug? Let me check: ```cuda float* my_tile_vals = s_tile_vals + warp_id * warp_tile; int* my_tile_idxs = s_tile_idxs + warp_id * warp_tile; ``` In pass 2, the load is: ```cuda for (int i = tid; i < TILE; i += blockDim.x) { if (i < tile_n) { s_tile_vals[i] = row_vals[tile_start + i]; s_tile_idxs[i] = row_idxs[tile_start + i]; } else { s_tile_vals[i] = -FLT_MAX; s_tile_idxs[i] = -1; } } ``` This loads ALL candidates into `s_tile_vals` flat. Then each warp scans a different region: `warp_id * warp_tile` to `(warp_id + 1) * warp_tile`. This is correct! No bug in pass 2. So only pass 1 has the bug. Let me fix it. Actually, wait. In pass 1, I also need to make sure the warp selection loop uses `warp_n` instead of `M`. Currently: ```cuda for (int i = lane; i < M; i += 32) { if (my_warp_vals[i] > local_max) { ... } } ``` This scans all M elements, including the -FLT_MAX padding. That's fine - the padding won't be selected. But to be slightly more efficient, I could use `warp_n`. However, `warp_n` varies per warp, and using it in the loop condition is fine. But `M` is the same for all warps, and the extra elements are -FLT_MAX, so using `M` is also fine. Actually, using `M` is simpler and avoids divergence. The -FLT_MAX elements just won't be selected. Let me update the code to add `warp_start` to the global memory load. Wait, I also need to handle the case where `block_n < num_warps`. For example, if block_n=16 and num_warps=4, then warp_chunk=4. Some warps might have warp_n=0 or very small. But with our parameters, block_n is always >= 512 for the large shapes and >= 750 for n=12000. So warp_chunk is at least 128. Fine. But for safety, I'll add `if (warp_start < block_n)` around the load. Actually, the current `if (idx < warp_end)` already handles this. Let me make the fix.

I need to recompile and test again. Let me change the extension name to force recompilation.

Let me run check.py again.

Great, correctness passes now! Let me run the benchmark to see the performance.

The performance is much better than before but still far from ideal. Let me analyze the numbers: 1. shape=0 (1, 131072, 64): solution 0.152ms vs eager 0.054ms. 2.8x slower. 2. shape=1 (64, 8192, 8): solution 0.022ms vs eager 0.048ms. 2.2x FASTER! 3. shape=2 (32, 16384, 32): solution 0.038ms vs eager 0.050ms. 1.3x faster. 4. shape=3 (16, 12000, 16): solution 0.026ms vs eager 0.046ms. 1.8x faster. 5. shape=4 (128, 4096, 1): solution 0.009ms vs eager 0.023ms. 2.6x faster. So for shapes 1-4, we're already faster than eager. But shape 0 is still 2.8x slower. And the overall peak_fraction is only 2.3%, which is LOW. The target is probably something like 50%+ of torch.topk performance, or better. Let me think about how to optimize shape 0. For shape 0 (batch=1, n=131072, k=64): - Pass 1: 256 blocks, each with 128 threads. - block_n = 512, EPT = 4, M = 128. - Each warp loads 128 elements (warp_chunk = 128). - Warp select top-64 from 128: 64 * (4 + 5) = 576 ops per thread. - Block sort 256: ~32 ops per thread. - Total per block: ~600 ops. - 256 blocks: memory 512KB, compute ~600 ops per block. - Time: probably ~50-100 microseconds? But actual is 0.152ms = 152 microseconds. Hmm, 152 microseconds is quite slow. Let me think about what's taking time. Pass 1 with 256 blocks: - Each block does warp select (576 ops/thread) + block sort (32 ops/thread). - But the warp select involves many shared memory reads and warp shuffles. - 64 iterations * 4 shared memory reads per thread = 256 shared memory reads. - Plus 64 * 5 = 320 warp shuffles. - Total per thread: ~576 ops. At 1.5GHz = 0.38 microseconds per thread. - With 128 threads per block: the block takes ~0.38 microseconds (warp select is warp-parallel). - 256 blocks total. With 188 SMs, ~1.36 blocks per SM. - Time: ~0.5 microseconds. Pass 2 with 1 block: - num_candidates = 256 * 64 = 16384. - TILE = 4096. num_tiles = 4. - Each tile: warp select from 512 (warp_tile = 4096/8 = 512). - 64 iterations * (16 + 5) = 1344 ops per thread. - Block sort 512: 512 * 9 * 9 / 4 = 10368 comparators. 256 threads: ~40 comparators per thread. - Merge 64+64: ~128 ops (1 thread). - Per tile: ~1500 ops per thread. - 4 tiles: ~6000 ops per thread. - At 1.5GHz: ~4 microseconds. Total compute: ~4.5 microseconds. But actual time is 152 microseconds. There's a huge gap! Where is the time going? Let me check if maybe the Python overhead or PyTorch compilation is dominating. Or maybe the kernel launches are very slow. Actually, I think the benchmark does 50 iterations and takes the median. Let me check what the benchmark timing function does. Looking at benchmark.py, `time_fn` is imported from `src.eval.timing`. I don't know exactly what it does, but it probably does a warmup and then times multiple iterations. Let me run a quick manual timing to see the breakdown. Actually, wait. The benchmark shows `ms=0.152` for shape 0. That's 152 microseconds. For a single iteration of the model forward pass. With 50 iterations, the total time would be ~7.6ms. But the kernel compute is only ~5 microseconds. Where is the remaining ~147 microseconds going? Possibilities: 1. Memory allocation in the forward pass. 2. Kernel launch overhead. 3. PyTorch tensor operations outside the kernel. 4. The kernels are actually much slower than I estimated. Wait! I just realized: in the C++ wrapper, I allocate `tmp_vals` and `tmp_idxs` on EVERY call to `topk_forward_cuda`! ```cpp auto tmp_vals = torch::empty({batch, num_blocks_per_row, k}, x.options()); auto tmp_idxs = torch::empty({batch, num_blocks_per_row, k}, torch::dtype(torch::kInt32).device(x.device())); ``` These are allocated inside the C++ function, which is called from Python on every forward pass. CUDA memory allocation can take 50-100 microseconds! This is a huge overhead. To fix this, I should pre-allocate the temporary buffers and reuse them. But the Model's forward is called many times during benchmarking. And the C++ function is called from Python each time. I can't easily cache the temp buffers across calls unless I store them in the Model or pass them in. Options: 1. Pre-allocate the temp buffers in the Python Model and pass them to the C++ function. 2. Use a persistent buffer inside the C++ module (e.g., a static tensor cache). 3. Avoid temp buffers entirely by doing everything in one kernel. Option 1 is clean. In the Model's `__init__`, I can compute the max temp buffer size needed and pre-allocate it. But the temp size depends on the input shape. Since the Model is initialized with (batch, n, k), I know the shape. Actually, the temp size for pass 1 is `batch * num_blocks_per_row * k`. I can pre-allocate it in Python. But wait, `num_blocks_per_row` depends on the batch size. For a fixed Model, batch is fixed. So I can compute it in `__init__`. In Python: ```python self.num_blocks_per_row = max(1, 256 // batch) self.tmp_vals = torch.empty(batch, self.num_blocks_per_row, k, dtype=torch.float32) self.tmp_idxs = torch.empty(batch, self.num_blocks_per_row, k, dtype=torch.int32) ``` Then pass these to the C++ function. But the C++ function currently allocates its own temps. I need to modify it to accept temp buffers as arguments. Actually, a simpler approach: use `torch::Tensor` buffers that are cached in the C++ module. Since `load_inline` creates a module, I can maintain a static cache of temp buffers. But `load_inline` doesn't easily support stateful C++ modules. The module just exposes functions. Another option: modify the Python Model to pre-allocate and pass the buffers. Let me modify the interface: ```python class Model(nn.Module): def __init__(self, batch, n, k): ... if k > 1: self.num_blocks_per_row = max(1, 256 // batch) self.register_buffer("_tmp_vals", torch.empty(batch, self.num_blocks_per_row, k, dtype=torch.float32)) self.register_buffer("_tmp_idxs", torch.empty(batch, self.num_blocks_per_row, k, dtype=torch.int32)) def forward(self, x): if self.k == 1: return _topk_cuda.topk_forward_cuda(x, self.k) else: return _topk_cuda.topk_forward_cuda_2pass(x, self.k, self.num_blocks_per_row, self._tmp_vals, self._tmp_idxs) ``` Then in C++: ```cpp std::vector<torch::Tensor> topk_forward_cuda_2pass(torch::Tensor x, int k, int num_blocks_per_row, torch::Tensor tmp_vals, torch::Tensor tmp_idxs) { // ... use tmp_vals and tmp_idxs instead of allocating new ones } ``` This would eliminate the allocation overhead. But wait, `check.py` and `benchmark.py` call `model.forward()` many times. Pre-allocating in `__init__` means the buffers are allocated once. This should significantly reduce overhead. Actually, another source of overhead might be the two kernel launches. Each CUDA kernel launch from Python/C++ takes ~3-5 microseconds. With 2 kernels, that's ~6-10 microseconds. For the fastest shapes, this might be a significant fraction of the total time. For shape 0, the actual kernel execution might be ~20 microseconds, but with launch overhead and allocation overhead, it becomes 152 microseconds. Let me measure the overhead by running a micro-benchmark. Actually, let me first try to eliminate the temp allocation by modifying the C++ code. I'll pass pre-allocated tensors. Wait, there's a simpler approach. I can use CUDA graphs or just persistent allocations in C++ using a static cache keyed by shape. But that adds complexity. The simplest fix is to modify the Python Model to pre-allocate temps and pass them to C++. But actually, I realize that `torch.empty` on CUDA is very fast - it just reserves memory, doesn't actually allocate from the driver. The allocation happens lazily on first use. So the first call might be slow, but subsequent calls should reuse the cached allocator. Wait, `torch.empty` returns a tensor whose storage is allocated from PyTorch's CUDA memory pool. For small allocations, this is very fast (a few microseconds). For larger allocations, it might take longer. For tmp size = batch * num_blocks * k: - shape 0: 1 * 256 * 64 = 16384 floats = 64KB. Very small. - shape 1: 64 * 4 * 8 = 2048 floats = 8KB. - shape 2: 32 * 8 * 32 = 8192 floats = 32KB. - shape 3: 16 * 16 * 16 = 4096 floats = 16KB. - shape 4: k=1, no temp. These are tiny allocations. PyTorch's allocator should handle them in microseconds. So allocation is probably NOT the bottleneck. Let me think again. Maybe the kernels are just slow because of shared memory bank conflicts or warp divergence. For pass 1, warp select from 128 elements with 32 threads: - Each thread does 4 comparisons per iteration. - 64 iterations. - Total: 256 comparisons per thread. - Each comparison reads shared memory. Bank = `i % 32`. Since `i = lane, lane+32, lane+64, lane+96`, all these map to bank `lane % 32`. Each thread accesses its own bank. NO bank conflicts! - But the `__shfl_down_sync` operations have a cost. Each shuffle is ~1-4 cycles. - 64 iterations * 5 shuffles = 320 shuffles. At 4 cycles: 1280 cycles = 0.85 microseconds. - Comparisons: 256 comparisons. At 1 cycle: 256 cycles = 0.17 microseconds. - Marking writes: 64 writes to shared memory. At ~10 cycles: 640 cycles = 0.43 microseconds. - Total per warp: ~1.5 microseconds. With 4 warps per block, and the warp select happening in parallel: ~1.5 microseconds per block. Block sort 256: 256 * 8 * 8 / 4 = 4096 comparators. With 128 threads = 4 warps. Each thread does ~32 comparators. Each comparator: 2 shared memory reads, 1 compare, possible 2 writes. 32 * 5 = 160 operations. ~0.1 microseconds. So pass 1 per block: ~1.6 microseconds. 256 blocks on 188 SMs: ~2.2 microseconds (since some SMs run 2 blocks). Pass 2 with 1 block and 4 tiles: - Each tile: warp select from 512. 64 iterations * 16 comparisons = 1024 comparisons per thread. Plus 320 shuffles. Plus 64 writes. Total: ~1024 + 320*4 + 64*10 = 1024 + 1280 + 640 = 2944 cycles = 2 microseconds per tile. - 4 tiles: ~8 microseconds. - Block sort 512: 512 * 9 * 9 / 4 = 10368 comparators. 256 threads: ~40 comparators per thread. 40 * 5 = 200 ops = 0.13 microseconds. - 4 merges: 4 * 128 = 512 comparisons. 1 thread. Negligible. - Total pass 2: ~8.2 microseconds. Total kernel time: ~10.4 microseconds. But benchmark says 152 microseconds. There must be a HUGE overhead somewhere. Let me check if maybe the C++ function is being called from Python with significant overhead, or if there's some synchronization. Actually, wait. In the benchmark, `time_fn` might be doing something like: ```python def time_fn(fn, inputs, iters=50): # warmup for _ in range(5): fn(*inputs) torch.cuda.synchronize() start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() for _ in range(iters): fn(*inputs) end.record() torch.cuda.synchronize() return start.elapsed_time(end) / iters ``` If so, it measures the average time per iteration, including the forward pass. Let me run a manual timing to verify.

The manual timing confirms ~143 microseconds per iteration. This is much higher than the expected ~10 microseconds for kernel execution. Let me investigate where the time is going. One possibility: the kernel execution is actually much slower than I estimated because of warp divergence or other inefficiencies. Let me run a more detailed timing breakdown. Another possibility: there are implicit synchronizations or host-side operations. Let me measure the time for each kernel separately by using `torch.cuda.Event` around each kernel launch in the C++ code. But I can't easily modify the C++ code for timing. Wait, I can use `torch.profiler` to see the kernel timeline.

The profiler shows: - `topk_pass1`: 21.5 microseconds per call - `topk_pass2`: 121.7 microseconds per call - Total: ~143 microseconds per iteration Pass 1 is 21.5us, which is close to my estimate (~2-5us) but a bit higher. Pass 2 is 121.7us, which is WAY higher than my estimate (~8us). Pass 2 is the bottleneck! Let me analyze why pass 2 is so slow. For pass 2 with batch=1, num_candidates=16384, TILE=4096, num_tiles=4: - Each tile: warp select from 512 candidates. - 64 iterations. - Each iteration: 512/32 = 16 comparisons per thread. - 5 shuffle reductions. - 1 broadcast. - 1 write to s_wout. - 1 write to mark -inf. - Total per iteration: ~16 + 5 + 2 + 2 = 25 ops. - 64 iterations: 1600 ops per thread. - 4 tiles: 6400 ops per thread. - Block sort 512: 512 * 9 * 9 / 4 = 10368 comparators. - With 256 threads: ~40 comparators per thread. - 40 * 5 = 200 ops per thread. - 4 tiles: 800 ops per thread. - 4 merges: each merge is done by 1 thread (tid==0). 4 * 128 = 512 comparisons. - This is negligible for overall time but the thread doing it is busy. Total per thread: ~7200 ops. At 1.5GHz: 7200 / 1.5e9 = 4.8 microseconds. But actual is 121.7 microseconds. That's 25x slower! Why? One big issue: **warp divergence in the merge**. The merge is done by `if (tid == 0)`. All other threads wait at `__syncthreads()`. But the 255 idle threads are still consuming scheduler slots. But that alone doesn't explain 25x. Another issue: **shared memory bank conflicts in the block sort**. For j values that are multiples of 32, there are 2-way intra-thread conflicts. With block_cand = 512 and j = 32, 64, 128, 256: - For j=32: each thread accesses banks `i%32` and `(i^32)%32`. For even i, `i^32 = i+32`. Bank difference: 0. Same bank! 2-way conflict. - For j=64: `i^64`. For i=0: banks 0 and 0. Same. 2-way conflict. - For j=128: same. - For j=256: same. How many stages have j as a multiple of 32? For block_cand = 512: - kk=64: j=32, 16, 8, 4, 2, 1. j=32 is multiple of 32. 1 stage. - kk=128: j=64, 32, 16, 8, 4, 2, 1. j=64, 32. 2 stages. - kk=256: j=128, 64, 32, 16, 8, 4, 2, 1. j=128, 64, 32. 3 stages. - kk=512: j=256, 128, 64, 32, 16, 8, 4, 2, 1. j=256, 128, 64, 32. 4 stages. Total conflicting stages: 1+2+3+4 = 10 stages out of 6+7+8+9 = 30 stages. Each conflicting stage has 2-way intra-thread conflict, so ~2x slower. Overall slowdown: ~10/30 * 2 + 20/30 * 1 = 1.33x. Not enough to explain 25x. Another issue: **the warp select in pass 2 scans 512 elements with 64 iterations. But after the first few iterations, most of the real candidates have been found, and the rest are -FLT_MAX.** Wait, in pass 2, the candidates are the top-64 from each pass-1 block. These are all large values (not -FLT_MAX). There are 16384 candidates, and we need the top-64. So almost ALL candidates are valid and relevant! The warp select marks winners as -FLT_MAX. After 64 iterations, all real candidates in that warp's 512-element tile have been replaced with -FLT_MAX? No! Each warp only selects 64 candidates from its 512-element region. After 64 iterations, there are still 448 candidates left. But they're not needed for this tile. Wait, but the loop still runs 64 times even after finding the top-64. It can't stop early because K is fixed at 64. Actually, for each tile, we only need the top-64 from the tile. The warp select finds exactly 64 candidates. After 64 iterations, we're done with that tile. That's correct. But why is it taking 121 microseconds? Let me think about what else could be slow. Could it be that the `__syncthreads()` between tiles is causing significant overhead? 4 tiles * 4 `__syncthreads()` per tile = ~16 `__syncthreads()` calls. Each `__syncthreads()` takes a few microseconds when there are many warps. With 8 warps, it's fast. Could it be that the kernel is memory-bound? Pass 2 reads 16384 candidates from global memory. That's 16384 * 8 = 131KB. At 1.8 TB/s: 0.07 microseconds. Not the bottleneck. Wait, could the issue be that pass 2 uses only 1 block for batch=1? With 1 block on 1 SM, we have 8 warps. But the SM has 4 warp schedulers, so it can issue 4 instructions per cycle. 8 warps means 2 warps per scheduler, which is good for latency hiding. But 7200 ops per thread * 256 threads = ~1.8M ops total. At 4 instructions per cycle and 1.5GHz: 1.8M / (4 * 1.5e9) = 0.3 milliseconds = 300 microseconds. Wait wait wait! I made a calculation error. 7200 ops per thread is PER THREAD. With 256 threads, the total instruction count is 256 * 7200 = ~1.8 million instructions. But these execute in parallel across the warps. Actually, the SM issues 4 instructions per cycle, and there are 8 warps. So 4 out of 8 warps can issue an instruction each cycle. That's 4 instructions per cycle. But each warp has 32 threads, and each thread executes the same instruction (SIMT). So 4 instructions per cycle means 4 warps * 32 threads = 128 thread-instructions per cycle. Total thread-instructions = 256 threads * 7200 ops = 1,843,200. At 128 thread-instructions per cycle: 1,843,200 / 128 = 14,400 cycles. At 1.5GHz: 14,400 / 1.5e9 = 9.6 microseconds. Hmm, that's closer to my original estimate. But actual is 121 microseconds. Still a 12x gap. Oh! I think I know. The warp select loop has a LOT of warp shuffles. Each `__shfl_down_sync` might take 4-10 cycles. And there are 64 * 5 = 320 shuffles per thread. With 8 warps * 32 threads = 256 threads, total shuffles = 256 * 320 = 81,920. At 4 cycles per shuffle: 327,680 cycles. But shuffles can only be issued by certain execution units, and there might be limited throughput. Also, the loop has many branches and the compiler might not be generating optimal code. But 121 microseconds is still way too high. Let me check if maybe the block sort is the bottleneck. For block sort of 512 with 256 threads: - 30 stages. - Each stage: 512/2 = 256 comparators. - With 256 threads: 1 comparator per thread per stage. - 30 * 1 = 30 comparators per thread. - 4 tiles: 120 comparators per thread. - Very small. So the block sort is definitely not the bottleneck. The bottleneck must be the warp select. Let me count operations more carefully. For warp select from 512 candidates: ```cuda for (int iter = 0; iter < K; ++iter) { float local_max = -FLT_MAX; int local_pos = -1; for (int i = lane; i < warp_tile; i += 32) { if (my_tile_vals[i] > local_max) { local_max = my_tile_vals[i]; local_pos = i; } } #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { float other_max = __shfl_down_sync(0xFFFFFFFFu, local_max, offset); int other_pos = __shfl_down_sync(0xFFFFFFFFu, local_pos, offset); if (other_max > local_max) { local_max = other_max; local_pos = other_pos; } } int winner_pos = __shfl_sync(0xFFFFFFFFu, local_pos, 0); float winner_val = __shfl_sync(0xFFFFFFFFu, local_max, 0); int winner_idx = my_tile_idxs[winner_pos]; if (lane == 0) { s_wout_vals[warp_id * K + iter] = winner_val; s_wout_idxs[warp_id * K + iter] = winner_idx; } if (lane == (winner_pos & 31)) { my_tile_vals[winner_pos] = -FLT_MAX; } __syncwarp(); } ``` Inner loop: `for (int i = lane; i < 512; i += 32)` - Iterations: i = lane, lane+32, lane+64, ..., lane+480. Total 16 iterations. - Each iteration: 1 shared memory read, 1 comparison, possible 2 assignments. - 16 * 4 = 64 ops per thread. Shuffle reduction: - 5 iterations of `__shfl_down_sync` on both `float` and `int`. - Each shuffle: ~4-10 cycles. - 10 shuffles * 5 = 50 shuffle ops. - Plus 5 comparisons and 5 assignments. - Total: ~65 ops. Broadcast: - 2 `__shfl_sync` ops. - 1 shared memory read for index. - Total: ~10 ops. Marking: - 1 comparison, 1 shared memory write. - `__syncwarp()`. - Total: ~10 ops. Total per iteration: ~150 ops. 64 iterations: ~9600 ops per thread. Wait, I previously estimated 25 ops per iteration, but it's more like 150 ops per iteration! Let me recount: - Inner scan: 16 iterations * (1 read + 1 compare + 2 assign) = 64 ops. But assignments are conditional, so branch divergence. - Shuffle reduction: 5 shuffles of float + 5 shuffles of int = 10 shuffles. Each shuffle might be ~1 instruction, but the throughput is limited. Let's count as 10 instructions. - 5 comparisons and assignments in reduction: 10 instructions. - 2 broadcasts: 2 instructions. - 1 shared memory read for index: 1 instruction. - 2 shared memory writes (s_wout): 2 instructions (only lane 0, but all threads evaluate the condition). - 1 comparison for marking: 1 instruction. - 1 shared memory write for marking: 1 instruction. - __syncwarp: ~1 instruction overhead. Total: ~64 + 10 + 10 + 2 + 1 + 2 + 1 + 1 + 1 = 92 instructions per iteration. 64 iterations: ~5888 instructions per thread. For 8 warps * 32 threads = 256 threads: total thread-instructions = 256 * 5888 = ~1.5M. At 128 thread-instructions per cycle: 1.5M / 128 = 11,718 cycles = 7.8 microseconds. Still far from 121 microseconds. Unless... the `__shfl_down_sync` is MUCH slower than I think. On some architectures, warp shuffles have ~20-30 cycle latency. But with multiple warps, the latency is hidden. Actually, let me look at this from a different angle. Maybe the compiler is generating very inefficient code. Let me check the PTX for the pass 2 kernel. Actually, I can't easily check the PTX. But I can add some instrumentation or try to simplify the kernel. One major issue with the current warp select is that it does 64 iterations, each with 5 shuffle reductions. That's 320 shuffles per thread. Even at 10 cycles each, that's 3200 cycles = 2.1 microseconds. With 256 threads, the total shuffle time is still small. But wait, there's another possibility: **the kernel is not launching with 256 threads!** Let me check. In the C++ wrapper: ```cpp int threads2 = 256; ``` Yes, 256 threads. Could it be that `cudaFuncSetAttribute` is needed for pass 2? smem2 max = 36KB, which is under 48KB. So no. Could it be that the kernel is stalling on memory? Pass 2 reads 131KB from global memory. But the global memory reads are coalesced and cached. Wait, I just realized something. In pass 2, I use `__syncthreads()` inside the tile loop: ```cuda for (int t = 0; t < num_tiles; t++) { // load tile __syncthreads(); // warp select __syncwarp(); __syncthreads(); // block sort __syncthreads(); // merge __syncthreads(); } ``` With 8 warps, `__syncthreads()` should be fast. But maybe there's a performance issue with the repeated `__syncthreads()`? Actually, `__syncthreads()` on modern GPUs is very fast (~10-20 cycles) when all warps reach it quickly. But if some warps are delayed, it waits. In the warp select, all warps execute independently and should finish at roughly the same time. So `__syncthreads()` shouldn't be a bottleneck. Let me think about whether there's a simpler algorithm for pass 2. **Alternative pass 2: Bitonic sort the entire candidate array in tiles** Instead of warp select + block sort + merge, what if I just bitonic sort the entire candidate array in tiles, keeping only the top-K after each tile merge? Actually, I can use a simpler approach: 1. Load a tile into shared memory. 2. Bitonic sort the tile descending. 3. The tile's top-K is at `s_tile[0..K-1]`. 4. Merge tile top-K with running top-K. This avoids the expensive warp select entirely! For TILE=4096: - Bitonic sort 4096 in shared memory: 4096 * 12 * 12 / 4 = 147456 comparators. - With 256 threads: 147456 / 256 = 576 comparators per thread. - 4 tiles: 2304 comparators per thread. - Plus 4 merges of 64+64: ~512 comparisons. - Total: ~2800 comparators per thread. At 5 ops per comparator: ~14000 ops per thread = 9.3 microseconds. Wait, that's still slower than my theoretical warp select time. But actual warp select takes 121 microseconds, while bitonic sort might actually take ~10 microseconds. Let me try this approach! Replace the warp select in pass 2 with a simple bitonic sort of the tile. But can I bitonic sort 4096 elements in shared memory with 256 threads? - Shared memory for 4096 elements: 4096 * 8 = 32KB. - Under 48KB. Yes! And 4096 is a power of 2. Perfect for bitonic sort. Let me redesign pass 2: 1. Load 4096 candidates into `s_tile_vals[4096]` and `s_tile_idxs[4096]`. 2. Bitonic sort 4096 in shared memory. 3. The top-64 are at `s_tile_vals[0..63]`. 4. Merge with running top-64. For 16384 candidates: 4 tiles. Each tile: - Sort 4096: ~576 comparators per thread. - Merge 64+64: ~128 comparisons (1 thread). Total: 4 * 576 = 2304 comparators per thread. At 5 ops each: ~11500 ops per thread. At 1.5GHz: ~7.7 microseconds. This is better than warp select's theoretical 7.8 microseconds, and hopefully much better than warp select's actual 121 microseconds. But why was warp select so slow? Maybe because of branch divergence or shuffle overhead. Bitonic sort is much more regular. Let me implement this! Actually, there's an even better optimization. Instead of fully sorting 4096 elements, I can do a **partial bitonic sort** that only sorts the first K elements. For bitonic sort, after the stages where kk <= K, the first K elements are not fully sorted. But we can run a full sort of 2*K elements and then merge with the remaining elements... Actually, a simpler partial sort: 1. Load 4096 elements. 2. Sort in groups of 64. 3. Merge groups progressively, keeping only the top 64 at each step. This is complex. Let me just do a full bitonic sort of 4096. It's still fast. Wait, I can also use a tournament tree approach in pass 2: 1. Load 4096 candidates. 2. Each thread processes 16 candidates and finds its local max. 3. Warp shuffle to find global max. 4. Repeat 64 times. But that's similar to warp select and might have the same performance issues. Let me just try the bitonic sort approach for pass 2. For pass 2 kernel: ```cuda template <int K, int TILE> __global__ void topk_pass2(const float* __restrict__ tmp_vals, const int* __restrict__ tmp_idxs, float* __restrict__ out_vals, int64_t* __restrict__ out_idxs, int num_candidates_per_row) { int row = blockIdx.x; const float* row_vals = tmp_vals + row * num_candidates_per_row; const int* row_idxs = tmp_idxs + row * num_candidates_per_row; int tid = threadIdx.x; extern __shared__ char smem[]; float* s_tile_vals = reinterpret_cast<float*>(smem); int* s_tile_idxs = reinterpret_cast<int*>(smem + TILE * sizeof(float)); float* s_run_vals = reinterpret_cast<float*>(smem + TILE * (sizeof(float) + sizeof(int))); int* s_run_idxs = reinterpret_cast<int*>(s_run_vals + K); for (int i = tid; i < K; i += blockDim.x) { s_run_vals[i] = -FLT_MAX; s_run_idxs[i] = -1; } __syncthreads(); int num_tiles = (num_candidates_per_row + TILE - 1) / TILE; for (int t = 0; t < num_tiles; t++) { int tile_start = t * TILE; int tile_n = min(TILE, num_candidates_per_row - tile_start); // Load tile for (int i = tid; i < TILE; i += blockDim.x) { if (i < tile_n) { s_tile_vals[i] = row_vals[tile_start + i]; s_tile_idxs[i] = row_idxs[tile_start + i]; } else { s_tile_vals[i] = -FLT_MAX; s_tile_idxs[i] = -1; } } __syncthreads(); // Bitonic sort tile descending for (int kk = 2; kk <= TILE; kk <<= 1) { for (int j = kk >> 1; j > 0; j >>= 1) { for (int i = tid; i < TILE; i += blockDim.x) { int ixj = i ^ j; if (ixj > i) { if ((i & kk) == 0) { if (s_tile_vals[i] < s_tile_vals[ixj]) { float tv = s_tile_vals[i]; s_tile_vals[i] = s_tile_vals[ixj]; s_tile_vals[ixj] = tv; int ti = s_tile_idxs[i]; s_tile_idxs[i] = s_tile_idxs[ixj]; s_tile_idxs[ixj] = ti; } } else { if (s_tile_vals[i] > s_tile_vals[ixj]) { float tv = s_tile_vals[i]; s_tile_vals[i] = s_tile_vals[ixj]; s_tile_vals[ixj] = tv; int ti = s_tile_idxs[i]; s_tile_idxs[i] = s_tile_idxs[ixj]; s_tile_idxs[ixj] = ti; } } } } __syncthreads(); } } // Merge s_tile[0..K-1] with s_run[0..K-1] if (tid == 0) { float a_vals[64]; int a_idxs[64]; for (int i = 0; i < K; i++) { a_vals[i] = s_run_vals[i]; a_idxs[i] = s_run_idxs[i]; } int a = 0, b = 0; for (int i = 0; i < K; i++) { bool take_a = (a < K) && (b >= K || a_vals[a] >= s_tile_vals[b]); if (take_a) { s_run_vals[i] = a_vals[a]; s_run_idxs[i] = a_idxs[a]; a++; } else { s_run_vals[i] = s_tile_vals[b]; s_run_idxs[i] = s_tile_idxs[b]; b++; } } } __syncthreads(); } if (tid < K) { out_vals[row * K + tid] = s_run_vals[tid]; out_idxs[row * K + tid] = static_cast<int64_t>(s_run_idxs[tid]); } } ``` Wait, but the bitonic sort of 4096 has `kk` going up to 4096. For j=2048, `i ^ j` might be large. With `TILE=4096` and `blockDim.x=256`, each thread processes 16 elements per stage. For j=2048, each thread compares `i` and `i^2048`. If both are in the tile, it swaps if needed. But `i` and `i^2048` are 2048 apart. Bank conflict: both map to the same bank if 2048 % 32 == 0. Yes! 2048 is a multiple of 32. So j=2048 causes 2-way intra-thread bank conflict. But as I analyzed before, this is only a 2x slowdown for a few stages. Actually, for TILE=4096, the number of stages with j >= 32: - kk=64: j=32. 1 stage. - kk=128: j=64, 32. 2 stages. - kk=256: j=128, 64, 32. 3 stages. - kk=512: j=256, 128, 64, 32. 4 stages. - kk=1024: j=512, 256, 128, 64, 32. 5 stages. - kk=2048: j=1024, 512, 256, 128, 64, 32. 6 stages. - kk=4096: j=2048, 1024, 512, 256, 128, 64, 32. 7 stages. Total conflicting stages: 1+2+3+4+5+6+7 = 28 stages. Total stages: 6+7+8+9+10+11+12 = 63 stages. Fraction: 28/63 = 44%. Slowdown: 0.44 * 2 + 0.56 * 1 = 1.44x. So bitonic sort of 4096 is ~1.44x slower than ideal. Still manageable. For 4 tiles: 4 * 63 stages * 4096/2 / 256 threads = 4 * 63 * 8 = 2016 comparators per thread. At 5 ops: ~10000 ops per thread. At 1.5GHz: ~6.7 microseconds. Plus 4 merges: negligible. Total: ~7 microseconds. This is much better than 121 microseconds! But I need to be careful. The full bitonic sort of 4096 uses a LOT of shared memory accesses. Each comparator does 2 reads and up to 2 writes. With bank conflicts in some stages, the effective time might be higher. Still, even if it's 20 microseconds, that's a huge improvement over 121 microseconds. Let me implement this version of pass 2 and test. Actually, there's another optimization: use a smaller TILE. With TILE=1024: - 16 tiles for 16384 candidates. - Sort 1024: 55 stages. 1024/2/256 = 2 comparators per thread per stage. - 55 * 2 = 110 comparators per thread. - 16 tiles: 1760 comparators per thread. - Much less work than TILE=4096 (4 tiles * 504 comparators per thread = 2016). Wait, 110 * 16 = 1760 vs 504 * 4 = 2016. TILE=1024 is slightly better! And shared memory for TILE=1024: 1024 * 8 = 8KB. Plus K * 8 = 0.5KB. Total ~8.5KB. For TILE=4096: 32KB + 0.5KB = 32.5KB. Both fit in 48KB. But TILE=1024 has less shared memory and slightly fewer comparators. Wait, let me recalculate: TILE=1024: 16 tiles * 55 stages * 512 comparators / 256 threads = 16 * 55 * 2 = 1760 comparators per thread. TILE=4096: 4 tiles * 63 stages * 2048 comparators / 256 threads = 4 * 63 * 8 = 2016 comparators per thread. So TILE=1024 is indeed better! And shared memory is smaller. For TILE=2048: 8 tiles * 58 stages * 1024 comparators / 256 = 8 * 58 * 4 = 1856 comparators per thread. TILE=1024 is the best! Let me use TILE=1024. Actually, for smaller candidate counts: - batch=64, num_candidates=2048: 2 tiles of 1024. 2 * 110 = 220 comparators per thread. Very fast. - batch=32, num_candidates=4096: 4 tiles. 440 comparators per thread. - batch=16, num_candidates=2048: 2 tiles. 220 comparators per thread. All very fast. OK, let me rewrite pass 2 with bitonic sort of tiles. One more thing: I should also optimize pass 1. Currently pass 1 takes 21.5us. Is that reasonable? For pass 1 with 256 blocks and 128 threads: - block_n = 512, warp_chunk = 128, EPT = 4, M = 128. - Warp select from 128: 64 * (4 + 5) = 576 ops per thread. - Block sort 256: ~32 ops per thread. - Total per block: ~608 ops per thread. - 256 blocks on 188 SMs: ~0.8 microseconds. But actual is 21.5us. Still a gap, but smaller than pass 2's gap. Could pass 1 also be slow because of the warp select? Let me think... For pass 1, each warp selects from 128 candidates. 64 iterations * 4 comparisons = 256 comparisons. Plus 320 shuffles. Plus 64 writes. Total ~640 ops per thread. With 4 warps per block: total warp-ops = 4 * 32 * 640 = 81,920. At 128 thread-instructions per cycle: 81920 / 128 = 640 cycles = 0.43 microseconds per block. 256 blocks: ~0.6 microseconds. But actual is 21.5us. That's still a 36x gap! Wait, could the kernel launch overhead be included in the profiler's CUDA time? No, the profiler shows separate CPU and CUDA times. The CUDA time is just the kernel execution. Hmm, let me check if maybe there are other operations happening on the GPU between kernel launches. Or if the kernels are waiting for something. Actually, I wonder if the issue is that the blocks are not running concurrently because of register or shared memory limits. For pass 1: shared memory = 6KB, registers = ? I didn't check the register usage for the new pass 1 kernel. Let me compile with `-Xptxas=-v` to see. Actually, let me just check if occupancy is the issue. With 6KB shared memory per block and 188 SMs with 100KB each, we can run ~16 blocks per SM. With 128 threads per block, that's 2048 threads per SM. The SM can handle 1536 threads (from device properties). So occupancy is limited by threads, not shared memory. 128 threads * 16 blocks = 2048 threads, but max is 1536. So ~12 blocks per SM = 1536 threads. 256 blocks / 188 SMs = 1.36 blocks per SM. So occupancy is not the issue. We have plenty of room. What about registers? If each thread uses many registers, fewer blocks can run concurrently. Let me check the register usage.

Interesting! Pass 2 for k=64 uses 142 registers! That's a lot. With 256 threads and 142 registers per thread, the block uses 256 * 142 = 36,352 registers. An SM has 65,536 registers. So only 1 block can fit per SM (2 blocks would need 72K registers > 65K). With 1 block per SM and only 8 warps, the GPU utilization is limited. But for batch=1, we only have 1 block anyway. So the register count doesn't explain why 1 block takes 121us. Wait, 142 registers per thread! That's huge. The compiler is using many registers for the warp select loops and temporary variables. For pass 1, k=64 uses only 48 registers. Much more reasonable. The high register count in pass 2 is because of the large TILE=4096 and the complex warp select logic. The compiler allocates many registers for loop variables, shared memory pointers, and temporaries. With 142 registers, each thread has a lot of state. But with 8 warps, the SM still has enough registers. The issue is not register spilling (0 spills). So why is it slow? Actually, I think the main issue might be that with TILE=4096, the bitonic sort stages with large j cause massive serialization due to bank conflicts. But I already analyzed that and it should only be ~1.5x slowdown. Wait, but pass 2 with k=64 and TILE=4096 doesn't do a bitonic sort of the tile! It does warp select. The warp select uses many `__shfl_down_sync` operations. Could shuffle operations be very slow on Blackwell? Actually, looking at the profiler output again: - pass1: 21.5us - pass2: 121.7us Pass 2 is dominated by the warp select. With 142 registers, there might be register bank conflicts or other microarchitectural issues. But more importantly, I notice that pass 2 uses 142 registers but 0 stack frame. This means all arrays are either in registers or not used. The shared memory arrays don't use registers directly. Hmm, let me think about this differently. Maybe the issue is that with 142 registers per thread and 256 threads, the total register file usage is 36,352 registers. The SM has 65,536 registers. But due to allocation granularity, the effective limit might be lower. If each block uses 36K registers, then for batch=1 we have 1 block on 1 SM. The SM runs 8 warps. At 4 warp schedulers, 2 warps per scheduler. This should be fine. Wait, I just realized: maybe the 121us is not for kernel execution but includes the time waiting for pass 1 to finish? No, CUDA kernels from the same stream execute sequentially. Pass 1 finishes, then pass 2 starts. The profiler measures each kernel's execution time independently. Could there be a lot of cache misses in pass 2? Pass 2 reads from `tmp_vals` and `tmp_idxs`, which were just written by pass 1. The L2 cache might not have all the data. But 131KB should easily fit in L2 (which is many MB). Actually, let me check if maybe the problem is that pass 2's memory access pattern is bad. In pass 2: ```cuda for (int i = tid; i < TILE; i += blockDim.x) { if (i < tile_n) { s_tile_vals[i] = row_vals[tile_start + i]; s_tile_idxs[i] = row_idxs[tile_start + i]; } } ``` `row_vals` is a flat array of `num_candidates_per_row` floats. For batch=1, num_candidates=16384. Thread tid reads at `tile_start + tid`, then `tile_start + tid + 256`, etc. For tile_start=0 and tid=0: reads 0, 256, 512, etc. These are 256 floats apart = 1024 bytes apart. A warp reads: thread 0 at 0, thread 1 at 256, thread 2 at 512, etc. These are 1024 bytes apart per thread. NOT coalesced! Oh! This is a major issue! In the load loop, threads access memory with stride 256. This means each thread in a warp accesses a different 1024-byte segment. The warp needs 32 separate 128-byte transactions. Total 4096 bytes per warp read. For TILE=4096 and blockDim.x=256, each thread loads 16 elements. 16 transactions per thread. With 256 threads: the total memory bandwidth utilization is terrible. Wait, actually, let me re-read the loop: ```cuda for (int i = tid; i < TILE; i += blockDim.x) { if (i < tile_n) { s_tile_vals[i] = row_vals[tile_start + i]; ... } } ``` Thread 0: i=0, 256, 512, ..., 3840. Reads row_vals[0], [256], [512], ... Thread 1: i=1, 257, 513, ..., 3841. Reads row_vals[1], [257], [513], ... For i=0: warp reads row_vals[0], [1], [2], ..., [255]. These are consecutive! Coalesced! For i=256: warp reads row_vals[256], [257], ..., [511]. Consecutive! Coalesced! Oh! I was wrong. The access IS coalesced! Because within each iteration of the inner loop, consecutive threads read consecutive elements. The `for` loop iterates with stride `blockDim.x = 256`, but within each iteration, threads 0..255 read elements `tile_start + 0..255`, then `tile_start + 256..511`, etc. So the load is perfectly coalesced. My mistake. Then what's causing the slowdown? Actually, I just realized that in pass 2, each tile load does: ```cuda for (int i = tid; i < TILE; i += blockDim.x) { s_tile_vals[i] = row_vals[tile_start + i]; s_tile_idxs[i] = row_idxs[tile_start + i]; } ``` For TILE=4096 and blockDim.x=256, this is 16 iterations per thread. Each iteration does 2 global memory reads and 2 shared memory writes. Total 32 global reads and 32 shared writes per thread. For 4 tiles: 128 global reads and 128 shared writes per thread. At 1 read per 10 cycles: 1280 cycles = 0.85 microseconds. So memory is not the bottleneck. Let me look at the warp select more carefully. Maybe the `__shfl_down_sync` and `__syncwarp()` are the culprits. In Blackwell (SM120), warp shuffle instructions might have different latency/throughput than previous architectures. But they should still be fast. Another thought: maybe the compiler is generating predicated code that's inefficient. The `if (lane == (winner_pos & 31))` branch is highly divergent (only 1 out of 32 threads takes it). This means 31 threads are idle for that instruction. But it's just 1 instruction. Wait, but `__syncwarp()` after the divergent branch requires all threads in the warp to reconverge. On Volta+, independent thread scheduling handles this, but there might be some overhead. For 64 iterations * 1 divergent branch: 64 `__syncwarp()` calls. Each `__syncwarp()` might take ~10-20 cycles. 64 * 15 = 960 cycles = 0.64 microseconds. Not huge. Hmm, I'm really puzzled by the 121us. Let me try a completely different pass 2 implementation: just bitonic sort the tile and merge. No warp select. This avoids all the shuffles and syncwarps. Let me also try reducing TILE to 1024 to reduce the sort cost. Actually, here's another idea: instead of processing tiles sequentially in one block, what if I process tiles in parallel across multiple warps, with each warp sorting a smaller chunk? For example, for 16384 candidates: - 8 warps, each gets 2048 candidates. - Each warp bitonic sorts 2048 candidates in shared memory. - But 8 warps can't sort in the same shared memory area simultaneously. What if each warp sorts 2048 candidates in its own shared memory region? - 8 regions * 2048 * 8 bytes = 131072 bytes = 128KB. Too much. Or use 2048 candidates total per block, with 8 warps each sorting 256 candidates? - 8 * 256 * 8 = 16KB. - Then block merges the 8 sorted 256-element lists into top-64. But merging 8 sorted lists is complex. Actually, a simpler approach: each warp loads and sorts 256 candidates. Then the block does a tournament tree to extract top-64 from 8 * 256 = 2048 candidates. For the tournament tree: - Maintain a heap of 8 elements (one per warp's current best). - For 64 extractions: 1. Find the max of 8 elements using warp shuffle (1 warp does this). 2. Record it. 3. Advance the pointer for the warp that provided it. 4. Update the heap. But this requires all warps to participate in some way. Actually, let me just try the bitonic sort of tiles approach. It's simpler and might be faster because it avoids shuffles. Wait, I just realized something. In the profiler, pass 2 for k=64 takes 121us. But for k=8, how long does pass 2 take? Let me check all shapes with the profiler.

Let me fix the profiler script.

Let me fix the profiler script to use the correct attribute names.

The profiler output is showing 0.0us because I'm not accessing the attributes correctly. Let me use a simpler approach with `torch.cuda.Event` timing around each kernel launch by instrumenting the C++ code. Or I can just use nsys or ncu with sudo. Actually, I can just add timing code to the C++ wrapper. Let me create a version that prints kernel execution times. But first, let me think about whether there's an obvious optimization I'm missing. One thing I notice: pass 2 with TILE=4096 uses 142 registers for k=64. That's extremely high. The high register count limits occupancy but more importantly, it might indicate the compiler is generating very inefficient code. What if I reduce TILE to 1024? This might reduce register pressure and improve performance. Also, I should replace the warp-select in pass 2 with a bitonic sort. The bitonic sort is more regular and might compile to much more efficient code. Let me rewrite pass 2 to use bitonic sort of tiles (TILE=1024) and then merge top-K. This eliminates all the warp shuffles. Actually, I just realized another issue. In pass 2, the warp select does: ```cuda for (int iter = 0; iter < K; ++iter) { ... if (lane == (winner_pos & 31)) { my_tile_vals[winner_pos] = -FLT_MAX; } __syncwarp(); } ``` But `my_tile_vals` is a pointer into shared memory: `s_tile_vals + warp_id * warp_tile`. The pointer arithmetic `my_tile_vals[winner_pos]` computes the address dynamically. This might cause the compiler to keep the pointer in a register and do dynamic addressing. With 142 registers, the compiler is definitely doing something complex. But 0 spills means everything fits. Let me just try the bitonic sort approach for pass 2 with TILE=1024. I think it will be much faster. For TILE=1024: - Shared memory: 1024 * 8 + K * 8 = 8192 + 512 + 64*8 = 8192 + 512 + 512 = 9216 bytes for k=64. - Wait, I also need s_run_vals and s_run_idxs: K * 8 = 512 bytes. - Total: 1024 * 8 + 512 = 8704 bytes? No: - s_tile_vals: 1024 * 4 = 4096 - s_tile_idxs: 1024 * 4 = 4096 - s_run_vals: 64 * 4 = 256 - s_run_idxs: 64 * 4 = 256 - Total: 8704 bytes = 8.5KB. For num_candidates=16384: 16 tiles. Each tile: sort 1024 elements. 1024 * 10 * 10 / 4 = 25600 comparators. With 256 threads: 100 comparators per thread. 16 tiles: 1600 comparators per thread. Plus 16 merges: ~2000 comparisons (1 thread). Total: ~1600 * 5 = 8000 ops per thread. At 1.5GHz: ~5.3 microseconds. This should be much faster than 121 microseconds! Let me implement it. Wait, but for the merge, I need both s_run and s_tile_top to be accessible. After sorting the tile, s_tile_vals[0..K-1] contains the top-K. I can merge directly without copying. The merge can be done in-place or with a temporary array in shared memory. But since I need to preserve s_tile_vals for the next tile, I should merge into a temporary and then copy back. Actually, I can just do: ```cuda if (tid == 0) { float merged_vals[64]; int merged_idxs[64]; int a = 0, b = 0; for (int i = 0; i < K; i++) { bool take_a = (a < K) && (b >= K || s_run_vals[a] >= s_tile_vals[b]); if (take_a) { merged_vals[i] = s_run_vals[a]; merged_idxs[i] = s_run_idxs[a]; a++; } else { merged_vals[i] = s_tile_vals[b]; merged_idxs[i] = s_tile_idxs[b]; b++; } } for (int i = 0; i < K; i++) { s_run_vals[i] = merged_vals[i]; s_run_idxs[i] = merged_idxs[i]; } } ``` This uses a local array `merged_vals[64]` and `merged_idxs[64]`. With K=64, that's 128 values = 512 bytes. The compiler will put this in local memory (stack frame). But local memory for 512 bytes should be fast (L1 cached). Wait, actually for K=64, the arrays are `float merged_vals[64]` and `int merged_idxs[64]`. The indexing is dynamic (`merged_vals[a]` where `a` is computed at runtime). So the compiler will put these in local memory. But 512 bytes of local memory is tiny. It should be very fast. Alternatively, I can do the merge directly without a temporary: ```cuda if (tid == 0) { float run_vals_copy[64]; int run_idxs_copy[64]; for (int i = 0; i < K; i++) { run_vals_copy[i] = s_run_vals[i]; run_idxs_copy[i] = s_run_idxs[i]; } int a = 0, b = 0; for (int i = 0; i < K; i++) { bool take_a = (a < K) && (b >= K || run_vals_copy[a] >= s_tile_vals[b]); if (take_a) { s_run_vals[i] = run_vals_copy[a]; s_run_idxs[i] = run_idxs_copy[a]; a++; } else { s_run_vals[i] = s_tile_vals[b]; s_run_idxs[i] = s_tile_idxs[b]; b++; } } } ``` This copies s_run to local memory first, then merges into s_run. Still uses local memory but only for the copy. For small K (8, 16), the local memory is even smaller. OK, let me rewrite pass 2 with bitonic sort. But wait, I should also think about whether I can optimize pass 1. Pass 1 takes 21.5us for k=64. With 256 blocks, that's actually not bad. 21.5us / 256 blocks = ~0.08us per block. But they run in parallel, so the actual time is ~21.5us total. For pass 1, can I also replace warp select with bitonic sort? For pass 1, M = 128 (for shape 0). Bitonic sort of 128: 128 * 7 * 7 / 4 = 1568 comparators. With 32 threads (1 warp): 1568 / 32 = 49 comparators per thread. 4 warps: 49 comparators per thread (parallel). Total: much faster than warp select (576 ops per thread). And bitonic sort of 128 in shared memory has minimal bank conflicts (only j=64 has conflict? 64 % 32 = 0, yes conflict. 1 stage out of 21). So bitonic sort of 128 would take ~49 * 5 = 245 ops per thread = 0.16 microseconds per warp. vs. warp select: 576 ops per thread = 0.38 microseconds per warp. Bitonic sort is ~2.4x faster for M=128! And for larger M (like 512 for other shapes), bitonic sort is 512 * 9 * 9 / 4 = 10368 comparators. 10368 / 32 = 324 comparators per thread. vs. warp select: 64 * (16 + 5) = 1344 ops per thread. Bitonic sort is ~4x faster! This is a huge insight! I should replace warp select with bitonic sort EVERYWHERE. **Revised design:** - Pass 1: Each warp loads its chunk into shared memory, then bitonic sorts it. Then the block merges the sorted warp results using block-level bitonic sort. - But wait, if each warp sorts its entire M-element array, we get sorted M-element arrays per warp. Then the block needs to merge `num_warps` sorted M-element arrays into top-K. - With block-level bitonic sort on `num_warps * K` candidates... but the warp results are M elements each, not K elements. Actually, after sorting M elements, the top-K of each warp are at positions 0..K-1. So we can just take the first K from each warp and then do a block-level bitonic sort on `num_warps * K` candidates. This is exactly what I already do! But instead of warp select to get top-K, I bitonic sort the entire M-element array and take the first K. For M=128, bitonic sort is 1568 comparators per warp. For M=512, bitonic sort is 10368 comparators per warp. And block sort of `num_warps * K` candidates is still needed. Total pass 1 per block: - Bitonic sort M elements per warp: M * log2(M)^2 / 4 comparators per warp. - Block sort num_warps * K candidates: num_warps * K * log2(num_warps * K)^2 / 4 comparators. For shape 0: M=128, num_warps=4, K=64. - Warp sort: 1568 comparators per warp. 4 warps: 1568 comparators per block. - Block sort: 4 * 64 = 256 candidates. 256 * 8 * 8 / 4 = 4096 comparators. - Total: 1568 + 4096 = 5664 comparators per block. - With 128 threads = 4 warps: 5664 / 128 = 44 comparators per thread. - At 5 ops per comparator: 220 ops per thread. - At 1.5GHz: 0.15 microseconds per block. 256 blocks: ~0.2 microseconds. This is MUCH faster than the warp select approach (600 ops per thread = 0.4us per block)! And for pass 2: - Bitonic sort TILE elements per tile: TILE * log2(TILE)^2 / 4 comparators. - Merge top-K with running top-K. For TILE=1024, 4 tiles (for shape 0): - Sort 1024: 25600 comparators per tile. - 4 tiles: 102400 comparators. - With 256 threads: 400 comparators per thread. - At 5 ops: 2000 ops per thread. - At 1.5GHz: 1.3 microseconds. Plus 4 merges: negligible. Total pass 2: ~1.5 microseconds. Total: ~1.7 microseconds. vs. torch.topk's ~55 microseconds. 32x faster! This is incredibly promising. Let me implement this right away. The key changes: 1. In pass 1, replace warp select with bitonic sort of M elements. 2. In pass 2, replace warp select with bitonic sort of TILE elements. For the bitonic sort of M elements within a warp: - All threads in the warp operate on the same shared memory array `my_warp_vals[0..M-1]`. - But bitonic sort requires all threads in the block to participate, not just one warp. Wait, if I want to sort each warp's array, I can either: a. Have all threads in the block participate in sorting each warp's array sequentially. b. Have only the threads in that warp participate. Option (b) is what I was doing with warp select. For bitonic sort, can a single warp sort 128 elements? Yes! With 32 threads, each thread processes 128/32 = 4 elements per stage. The bitonic sort code is the same, just with `blockDim.x` replaced by `32` and `__syncthreads()` replaced by `__syncwarp()`. But wait, in the current pass 1 kernel, the warp arrays are stored in shared memory. A single warp can sort its own array using the standard bitonic sort algorithm, with `for (int i = lane; i < M; i += 32)` and `__syncwarp()`. This should work! And it's much more efficient than warp select. Let me write the bitonic sort for a single warp: ```cuda // Sort my_warp_vals[0..M-1] descending for (int kk = 2; kk <= M; kk <<= 1) { for (int j = kk >> 1; j > 0; j >>= 1) { for (int i = lane; i < M; i += 32) { int ixj = i ^ j; if (ixj > i) { if ((i & kk) == 0) { if (my_warp_vals[i] < my_warp_vals[ixj]) { // swap } } else { if (my_warp_vals[i] > my_warp_vals[ixj]) { // swap } } } } __syncwarp(); } } ``` After this, `my_warp_vals[0..K-1]` contains the top-K of the warp. Then copy top-K to block candidate area: ```cuda if (lane < K) { s_block_vals[warp_id * K + lane] = my_warp_vals[lane]; s_block_idxs[warp_id * K + lane] = my_warp_idxs[lane]; } __syncthreads(); ``` Then block-level bitonic sort as before. This is much cleaner and faster! For pass 2, same thing: ```cuda // Sort s_tile_vals[0..TILE-1] descending for (int kk = 2; kk <= TILE; kk <<= 1) { for (int j = kk >> 1; j > 0; j >>= 1) { for (int i = tid; i < TILE; i += blockDim.x) { // compare and swap } __syncthreads(); } } ``` Then merge `s_tile_vals[0..K-1]` with `s_run_vals[0..K-1]`. This eliminates all warp shuffles! Let me implement this. I'll rewrite both pass 1 and pass 2. One concern: M might not be a power of 2. For example, M=192 (n=12000, block_n=750, EPT=6). Bitonic sort requires a power of 2. Solutions: 1. Round M up to the next power of 2 and pad with -FLT_MAX. 2. Use a different sorting algorithm for non-power-of-2 sizes. For simplicity, I'll round M up to the next power of 2. Let me call it `M_PAD`. In pass 1: ```cpp int EPT = (block_n + 127) / 128; int M = 32 * EPT; int M_PAD = 1; while (M_PAD < M) M_PAD <<= 1; ``` But wait, the shared memory size is `num_warps * M_PAD`. For M=192, M_PAD=256. Shared memory = 4 * 256 * 8 = 8KB + 2KB = 10KB. Still fine. For the load loop, I can load up to M elements and pad the rest to -FLT_MAX: ```cuda for (int j = 0; j < EPT; ++j) { int idx = warp_start + j * 32 + lane; float val = -FLT_MAX; int idx_out = -1; if (idx < warp_end) { val = row_x[idx]; idx_out = block_start + idx; } if (j * 32 + lane < M_PAD) { my_warp_vals[j * 32 + lane] = val; my_warp_idxs[j * 32 + lane] = idx_out; } } ``` But if M_PAD > M, the extra elements are already -FLT_MAX from initialization. So I just need to make sure I don't access beyond M in the load. Actually, since I initialized all shared memory to -FLT_MAX, I only need to load the valid elements. The padding is already -FLT_MAX. Wait, in my current code, I have: ```cuda for (int i = tid; i < num_warps * M_PAD; i += blockDim.x) { s_warp_vals[i] = -FLT_MAX; } __syncthreads(); ``` Then load valid elements. The extra positions in each warp's array (from M to M_PAD-1) remain -FLT_MAX. For the bitonic sort, I sort M_PAD elements. The padding elements are -FLT_MAX, so they end up at the bottom. The top-K valid elements are at the top. This works! But M_PAD might be larger than necessary. For M=128, M_PAD=128 (power of 2). For M=192, M_PAD=256. For M=512, M_PAD=512. For the block-level sort, `block_cand = num_warps * K`. This is always a power of 2 (since num_warps=4 and K is 8,16,32,64). For pass 2, TILE=1024 is a power of 2. Good. Let me now compute M_PAD for each shape: Shape 0: block_n=512, EPT=4, M=128, M_PAD=128. Shape 1: block_n=2048, EPT=16, M=512, M_PAD=512. Shape 2: block_n=2048, EPT=16, M=512, M_PAD=512. Shape 3: block_n=750, EPT=6, M=192, M_PAD=256. Shape 4: k=1, handled separately. Shared memory for pass 1 with M_PAD: - shape 0: 4 * 128 * 8 + 4 * 64 * 8 = 4096 + 2048 = 6144 bytes. - shape 1: 4 * 512 * 8 + 4 * 8 * 8 = 16384 + 256 = 16640 bytes. - shape 2: 4 * 512 * 8 + 4 * 32 * 8 = 16384 + 1024 = 17408 bytes. - shape 3: 4 * 256 * 8 + 4 * 16 * 8 = 8192 + 512 = 8704 bytes. All fit in 48KB. For pass 2 with TILE=1024: - shape 0: 1024 * 8 + 64 * 8 = 8192 + 512 = 8704 bytes. - shape 1: 1024 * 8 + 8 * 8 = 8192 + 64 = 8256 bytes. - shape 2: 1024 * 8 + 32 * 8 = 8192 + 256 = 8448 bytes. - shape 3: 1024 * 8 + 16 * 8 = 8192 + 128 = 8320 bytes. All fit. Now, for the bitonic sort in pass 2, the number of comparators per thread for TILE=1024: - 55 stages. - Each stage: 1024 / 2 = 512 comparators. - With 256 threads: 2 comparators per thread per stage. - 55 * 2 = 110 comparators per thread. - 4 tiles for shape 0: 440 comparators per thread. - At 5 ops: 2200 ops per thread = 1.5 microseconds. For shape 1 (64 blocks * 8 = 512 candidates): 1 tile. 110 comparators per thread = 0.07 microseconds. This is incredibly fast! Let me now rewrite the kernel with bitonic sort instead of warp select. Actually, one more thing. In pass 1, after the warp-level bitonic sort, I need to copy the top-K to the block candidate area. But the block-level bitonic sort needs all `num_warps * K` candidates. So: ```cuda // After warp sort if (lane < K) { s_block_vals[warp_id * K + lane] = my_warp_vals[lane]; s_block_idxs[warp_id * K + lane] = my_warp_idxs[lane]; } __syncthreads(); ``` Then block sort. Good. For pass 2, after tile sort: ```cuda // Merge s_tile[0..K-1] with s_run[0..K-1] if (tid == 0) { ... } __syncthreads(); ``` Good. Let me write the complete new solution.py.

I rewrote the kernel to use bitonic sort instead of warp select. Let me test for correctness and then benchmark.

Correctness passes. Let me run the benchmark.

The performance got WORSE! Shape 0 went from 0.152ms to 0.525ms. Shape 1 went from 0.022ms to 0.104ms. What happened? The bitonic sort approach is much slower than expected. Why? For shape 0, pass 2 now sorts 1024 elements 16 times. Each sort has 55 stages with `__syncthreads()` between each stage. That's 880 `__syncthreads()` calls in pass 2 alone! `__syncthreads()` is expensive! Each call might take 1-2 microseconds. 880 * 1us = 880us. That's close to the 525us we see. Ah! `__syncthreads()` is the killer. In the bitonic sort, we have log2(N) * (log2(N)+1) / 2 stages, each with a `__syncthreads()`. For TILE=1024, that's 55 `__syncthreads()` calls per tile. For 16 tiles: 880 `__syncthreads()` calls. Each `__syncthreads()` requires all threads in the block to reach the barrier. With 256 threads = 8 warps, if even one warp is delayed, all others wait. On modern GPUs, `__syncthreads()` is not free - it can take hundreds of cycles. For pass 1, M_PAD=128. Bitonic sort of 128 has 21 stages with `__syncwarp()`. `__syncwarp()` is cheaper than `__syncthreads()` because it only synchronizes one warp. 21 * 256 blocks = 5376 `__syncwarp()` calls. Still significant. But in the previous version (warp select), we had 64 `__syncwarp()` calls per warp in pass 1, and 64 `__syncwarp()` calls per warp per tile in pass 2. For pass 2 with 4 tiles and 8 warps: 4 * 64 * 8 = 2048 `__syncwarp()` calls. That's fewer than 880 `__syncthreads()` calls? Hmm, but `__syncthreads()` is much more expensive. Actually, wait. The previous version had 121us for pass 2. The new version has 525us. So the bitonic sort with `__syncthreads()` is 4x slower than warp select with `__syncwarp()`. This means `__syncthreads()` is indeed very expensive. To fix this, I need to reduce the number of synchronizations. Options: 1. Go back to warp select but optimize it. 2. Use a bitonic sort with fewer synchronizations. 3. Use a different algorithm entirely. For option 2, can I reduce synchronizations in bitonic sort? Only by having fewer threads participate, but that increases work per thread. Actually, what if I do the bitonic sort within each warp using registers + shuffles, without any shared memory synchronization? For M=128 in pass 1, I can have each thread hold 4 elements (since 128/32 = 4). Then the warp does a bitonic sort using shuffles. No shared memory needed for the sort! But I still need to load elements from global memory and store the top-K. Wait, a shuffle-based bitonic sort for 128 elements with 32 threads (4 elements per thread) is possible but complex. I'd need to carefully orchestrate the compare-and-swap operations between threads. Actually, there's a simpler approach: **use a selection algorithm instead of full sort**. For finding top-K from M candidates, I can use quickselect. But quickselect requires dynamic indexing and recursion, which is hard on GPU. Another idea: **use a parallel reduction tree to find top-K**. For M=128 and K=64: 1. Compare adjacent pairs, keep larger (64 elements). 2. Compare adjacent pairs again, keep larger (32 elements). 3. Continue until 64 elements remain. Wait, that's just finding the max repeatedly. To get top-64 from 128: - Stage 1: Compare (0,1), (2,3), ..., (126,127). Keep larger of each pair. Result: 64 elements. - Stage 2: These 64 elements are the top-64! Because each pair's larger element is in the top-64 of the 128. Wait, is that true? If I compare (a0, a1) and keep max, (a2, a3) and keep max, etc., the resulting 64 elements are NOT necessarily the top-64 of the original 128. For example, if a0=100, a1=99, and all others are 0, then the top-64 includes a0 and a1, but a1 might be discarded if it's paired with a0. Oh right! Pairwise reduction only keeps one element per pair. So I lose elements. To get top-K from N elements, I need a more sophisticated network. For K=64 from N=128, I can use a bitonic merge network that merges two sorted 64-element sequences. But I need the 128 elements to be in bitonic form first. Actually, if I just fully sort 128 elements, that's 21 stages with `__syncwarp()`. Maybe `__syncwarp()` is not the bottleneck - it's `__syncthreads()` that kills performance. Let me try a hybrid approach: - Pass 1: Use warp select (which uses `__syncwarp()`, not `__syncthreads()`). This was 21.5us. Acceptable. - Pass 2: Instead of bitonic sort with `__syncthreads()`, use warp select on smaller tiles with `__syncwarp()`. But the issue was that pass 2 with warp select took 121us. Why did pass 2 with warp select take 121us? Let me analyze more carefully. In pass 2 with warp select and TILE=4096: - 8 warps, each processes 512 candidates. - Each warp does 64 iterations of max-finding over 512 candidates. - Each iteration: 16 comparisons + 5 shuffles + 2 broadcasts + 1 write. - 64 iterations: ~1500 ops per thread. - 4 tiles: ~6000 ops per thread. - At 1.5GHz: 4 microseconds. But actual was 121us. So there's a 30x gap. Wait, I wonder if the compiler is generating bad code for the inner loop. Maybe the loop is not being unrolled, and there's a lot of loop overhead. Actually, looking at the ptxas output for pass 2 with warp select and TILE=4096, k=64: - Used 142 registers. With 142 registers per thread, the total register file usage for 256 threads is 36,352 registers. An SM has 65,536 registers. So only 1 block fits per SM. But for batch=1, we only have 1 block anyway. But with 142 registers, each warp uses 32 * 142 = 4544 registers. An SM can run up to 32 warps, but with 4544 registers per warp, the SM can only run 65,536 / 4544 = 14.4 warps. We have 8 warps, so this is fine. Could the high register count be causing instruction cache pressure or pipeline stalls? Actually, I think the main issue might be that with 142 registers, the compiler is doing something very inefficient with the inner loops. Let me look at whether I can reduce register usage. One way to reduce registers: use smaller TILE. With TILE=1024, each warp processes 128 candidates instead of 512. This might reduce register pressure because fewer loop iterations. For TILE=1024, warp_tile = 1024/8 = 128. Inner loop: `for (int i = lane; i < 128; i += 32)`. That's 4 iterations. 64 iterations of outer loop * 4 inner iterations = 256 comparisons. Plus 320 shuffles. Plus 64 writes. Total: ~640 ops per thread per tile. For 16 tiles: ~10,000 ops per thread. With TILE=1024, the compiler might use fewer registers. Let me try TILE=1024 with the warp select approach. But wait, I already tried this mentally and the time was similar. Let me just test it. Actually, let me think about this from first principles. The previous version with heap-based single-kernel took 1.7ms. The two-pass warp-select version took 0.143ms. The two-pass bitonic version took 0.525ms. So the warp-select two-pass is the best so far at 0.143ms. But it's still 2.8x slower than torch.topk (0.052ms). Can I optimize the warp-select two-pass further? The bottleneck is pass 2: 121us out of 143us. Pass 1 is 21.5us. If I can speed up pass 2 from 121us to ~30us, total would be ~50us, matching torch.topk. How to speed up pass 2? Current pass 2 does: - For each tile (4096 elements): - Load 4096 candidates. - 8 warps, each warp selects top-64 from 512 candidates. - Block sort 512 candidates. - Merge 64+64. What if I increase the number of blocks in pass 2? Instead of 1 block per row, use multiple blocks. For batch=1, num_candidates=16384: - Use 16 blocks for pass 2. Each block processes 1024 candidates. - Each block: warp select from 1024 (8 warps * 128 candidates). - 64 iterations * (4 comparisons + 5 shuffles) = 576 ops per thread. - Block sort 512: ~40 ops per thread. - Total per block: ~600 ops per thread. - 16 blocks on 16 SMs: ~0.4 microseconds. - Then a final block merges 16 * 64 = 1024 candidates. - Similar to pass 2 but with 1024 candidates. - 2 tiles of 512 or just bitonic sort 1024. - ~1 microsecond. - Total pass 2: ~1.4 microseconds + kernel launch overhead. But with 3 kernels total, launch overhead is ~15us. Total: ~15 + 21.5 + 1.4 = ~38us. Faster than torch.topk! But can I do this without adding another kernel? I can use a single kernel that does everything if I use persistent blocks... but that's complex. Actually, for pass 2, I can just launch `batch * num_pass2_blocks` blocks where each block processes a chunk of candidates and writes top-k to a smaller buffer. Then the same kernel can be called recursively until only 1 block remains. In C++: ```cpp // Pass 2a: reduce candidates by factor of N while (num_candidates > K) { int blocks = max(1, num_candidates / (1024)); // launch kernel that processes chunks and writes top-K per chunk num_candidates = blocks * K; } ``` For num_candidates=16384, K=64: - Iteration 1: 16384 candidates. Launch 16 blocks. Each reads 1024, outputs 64. - Result: 1024 candidates. - Iteration 2: 1024 candidates. Launch 1 block. Reads 1024, outputs 64. - Result: 64 candidates. Done. Total: 2 kernel launches for pass 2. Plus pass 1: 3 kernels total. For each iteration, the kernel time is very small (~1us). Launch overhead dominates at ~5us per kernel. Total: 15us + 21.5us = ~36.5us. This is faster than torch.topk! And for smaller candidate counts: - shape 1: 2048 candidates. - Iteration 1: 2048 candidates. 2 blocks -> 128 candidates. - Iteration 2: 128 candidates. 1 block -> 64 candidates. - Total: 2 iterations. - shape 2: 4096 candidates. - Iteration 1: 4 blocks -> 256 candidates. - Iteration 2: 1 block -> 64 candidates. - shape 3: 2048 candidates. 2 iterations. All very fast. Let me implement this multi-iteration reduction for pass 2! The kernel for pass 2 reduction: ```cuda template <int K, int CHUNK> __global__ void topk_reduce(const float* __restrict__ in_vals, const int* __restrict__ in_idxs, float* __restrict__ out_vals, int* __restrict__ out_idxs, int num_candidates) { int row = blockIdx.y; int block_in_row = blockIdx.x; int chunk_start = block_in_row * CHUNK; int chunk_n = min(CHUNK, num_candidates - chunk_start); const float* row_vals = in_vals + row * num_candidates + chunk_start; const int* row_idxs = in_idxs + row * num_candidates + chunk_start; int tid = threadIdx.x; int lane = tid & 31; int warp_id = tid >> 5; int num_warps = blockDim.x >> 5; extern __shared__ char smem[]; float* s_vals = reinterpret_cast<float*>(smem); int* s_idxs = reinterpret_cast<int*>(smem + CHUNK * sizeof(float)); float* s_out_vals = reinterpret_cast<float*>(smem + CHUNK * (sizeof(float) + sizeof(int))); int* s_out_idxs = reinterpret_cast<int*>(s_out_vals + num_warps * K); // load chunk for (int i = tid; i < CHUNK; i += blockDim.x) { if (i < chunk_n) { s_vals[i] = row_vals[i]; s_idxs[i] = row_idxs[i]; } else { s_vals[i] = -FLT_MAX; s_idxs[i] = -1; } } __syncthreads(); // warp select top-K from CHUNK / num_warps candidates int warp_chunk = CHUNK / num_warps; float* my_vals = s_vals + warp_id * warp_chunk; int* my_idxs = s_idxs + warp_id * warp_chunk; for (int iter = 0; iter < K; iter++) { float local_max = -FLT_MAX; int local_pos = -1; for (int i = lane; i < warp_chunk; i += 32) { if (my_vals[i] > local_max) { local_max = my_vals[i]; local_pos = i; } } #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { float other_max = __shfl_down_sync(0xFFFFFFFFu, local_max, offset); int other_pos = __shfl_down_sync(0xFFFFFFFFu, local_pos, offset); if (other_max > local_max) { local_max = other_max; local_pos = other_pos; } } int winner_pos = __shfl_sync(0xFFFFFFFFu, local_pos, 0); float winner_val = __shfl_sync(0xFFFFFFFFu, local_max, 0); int winner_idx = my_idxs[winner_pos]; if (lane == 0) { s_out_vals[warp_id * K + iter] = winner_val; s_out_idxs[warp_id * K + iter] = winner_idx; } if (lane == (winner_pos & 31)) { my_vals[winner_pos] = -FLT_MAX; } __syncwarp(); } __syncthreads(); // block sort int block_cand = num_warps * K; for (int kk = 2; kk <= block_cand; kk <<= 1) { for (int j = kk >> 1; j > 0; j >>= 1) { for (int i = tid; i < block_cand; i += blockDim.x) { int ixj = i ^ j; if (ixj > i) { if ((i & kk) == 0) { if (s_out_vals[i] < s_out_vals[ixj]) { float tv = s_out_vals[i]; s_out_vals[i] = s_out_vals[ixj]; s_out_vals[ixj] = tv; int ti = s_out_idxs[i]; s_out_idxs[i] = s_out_idxs[ixj]; s_out_idxs[ixj] = ti; } } else { if (s_out_vals[i] > s_out_vals[ixj]) { float tv = s_out_vals[i]; s_out_vals[i] = s_out_vals[ixj]; s_out_vals[ixj] = tv; int ti = s_out_idxs[i]; s_out_idxs[i] = s_out_idxs[ixj]; s_out_idxs[ixj] = ti; } } } } __syncthreads(); } } if (tid < K) { int out_pos = (row * gridDim.x + block_in_row) * K + tid; out_vals[out_pos] = s_out_vals[tid]; out_idxs[out_pos] = s_out_idxs[tid]; } } ``` In the C++ wrapper: ```cpp int num_candidates = num_blocks_per_row * k; auto tmp2_vals = torch::empty({batch, num_candidates, k}, x.options()); auto tmp2_idxs = torch::empty({batch, num_candidates, k}, torch::dtype(torch::kInt32).device(x.device())); const int CHUNK = 1024; int threads_reduce = 256; int num_warps_reduce = threads_reduce >> 5; size_t smem_reduce = CHUNK * sizeof(float) + CHUNK * sizeof(int) + num_warps_reduce * k * sizeof(float) + num_warps_reduce * k * sizeof(int); while (num_candidates > k) { int blocks_per_row = (num_candidates + CHUNK - 1) / CHUNK; dim3 grid(blocks_per_row, batch); switch (k) { case 8: topk_reduce<8, CHUNK><<<grid, threads_reduce, smem_reduce>>>(...); break; ... } // Swap input/output buffers std::swap(tmp_vals, tmp2_vals); std::swap(tmp_idxs, tmp2_idxs); num_candidates = blocks_per_row * k; } // Final output is in tmp_vals // Copy to out_vals, out_idxs ``` Wait, this requires ping-ponging between two temp buffers. That's fine. For num_candidates=16384, k=64: - Iteration 1: blocks_per_row = 16. num_candidates becomes 1024. - Iteration 2: blocks_per_row = 1. num_candidates becomes 64. Done. Each iteration launches a kernel. 2 iterations for pass 2. Kernel time per iteration: - Load 1024 elements: 1024 * 8 = 8KB. Coalesced. - Warp select from 128 (1024/8): 64 * (4 + 5) = 576 ops per thread. - Block sort 512 (8*64): ~40 ops per thread. - Total: ~600 ops per thread = 0.4 microseconds. - 16 blocks on 16 SMs: ~0.4 microseconds. - 1 block on 1 SM: ~0.4 microseconds. Total kernel time for pass 2: ~0.8 microseconds. Kernel launch overhead: 2 * 5us = 10us. Pass 1: 21.5us. Total: ~32us. This would beat torch.topk! But I need to be careful about the memory layout. The temp buffers need to be large enough for all iterations. For the largest temp buffer: - Iteration 0 (pass 1 output): batch * num_blocks_per_row * k = 1 * 256 * 64 = 16384 elements. - Iteration 1 output: batch * 16 * 64 = 1024 elements. - Iteration 2 output: batch * 1 * 64 = 64 elements. So the max temp size is 16384 elements. I can allocate two buffers of this size and reuse them. Let me implement this! Actually, wait. There's a subtlety in the reduction kernel: `gridDim.x` changes each iteration. In the kernel, `out_pos = (row * gridDim.x + block_in_row) * K + tid`. This writes to the output buffer with the correct layout. And in the next iteration, `in_vals + row * num_candidates + chunk_start` reads from the output buffer. Let me verify: - Iteration 1: 16 blocks per row. Output shape: batch * 16 * k. - Block 0 writes at row*16*K + 0*K + tid. - Block 1 writes at row*16*K + 1*K + tid. - ... - Iteration 2: num_candidates = 16 * k = 1024. blocks_per_row = 1. - Block 0 reads chunk_start=0, chunk_n=1024. - Writes at row*1*K + 0*K + tid. - Final output: batch * k elements. This works! One issue: for batch=64, num_blocks_per_row=4, num_candidates=2048. - Iteration 1: blocks_per_row = 2. Output: 64 * 2 * 8 = 1024 elements. - Iteration 2: num_candidates = 16. blocks_per_row = 1. Output: 64 * 8 = 512 elements. Done. For batch=32, num_blocks_per_row=8, num_candidates=4096. - Iteration 1: blocks_per_row = 4. Output: 32 * 4 * 32 = 4096 elements. - Iteration 2: num_candidates = 128. blocks_per_row = 1. Output: 32 * 32 = 1024 elements. Done. For batch=16, num_blocks_per_row=16, num_candidates=2048. - Iteration 1: blocks_per_row = 2. Output: 16 * 2 * 16 = 512 elements. - Iteration 2: num_candidates = 32. blocks_per_row = 1. Output: 16 * 16 = 256 elements. Done. Great! At most 2 iterations for pass 2. Now, for shared memory in the reduce kernel: - CHUNK = 1024 elements = 1024 * 8 = 8KB. - block_cand = num_warps * K = 8 * K. - For K=64: 8 * 64 * 8 = 4KB. - Total: 12KB. Under 48KB. For the loop in C++, I need to pass the correct pointers and `num_candidates` each iteration. Let me also think about whether I can avoid the final copy from tmp to output. After the last iteration, the result is in one of the temp buffers. I can just copy it to the output buffer. Or better, I can have the last iteration write directly to the output buffer! In the last iteration, when `num_candidates <= CHUNK` and `blocks_per_row == 1`, the kernel writes `batch * 1 * k` elements. I can have it write directly to `out_vals` and `out_idxs`. Actually, to keep it simple, I'll just have all iterations write to temp buffers, and then copy the final result. The copy is negligible (batch * k elements). Or, I can use the same output buffer as one of the ping-pong buffers. Let me think... Actually, the simplest is: 1. Allocate `tmp_a` and `tmp_b` both of size `batch * max_candidates * k`. 2. Pass 1 writes to `tmp_a`. 3. Iteration 1 reads `tmp_a`, writes to `tmp_b`. 4. Iteration 2 reads `tmp_b`, writes to `tmp_a`. 5. Copy `tmp_a[0:batch*k]` to output. The copy is `batch * k` elements. For batch=1, k=64: 64 elements. Negligible. Wait, but `tmp_a` has shape `{batch, max_candidates, k}`. After the last iteration, the valid data is at the beginning: `tmp_a[row * k + tid]`. So I can just slice it. In C++: ```cpp auto tmp_a = torch::empty({batch, max_candidates, k}, x.options()); auto tmp_b = torch::empty({batch, max_candidates, k}, torch::dtype(torch::kInt32).device(x.device())); // Pass 1 writes to tmp_a // Iterations read from one buffer and write to the other // Final result is in tmp_a (or tmp_b, depending on iteration count) // Copy to output ``` Actually, I can avoid the copy by using `torch::slice` or by having the last kernel write directly to the output. Let me have the reduction kernel take separate input and output pointers. In the last iteration, I pass the output pointers directly. In the wrapper: ```cpp float* in_v = tmp_vals.data_ptr<float>(); int* in_i = tmp_idxs.data_ptr<int>(); float* out_v = tmp2_vals.data_ptr<float>(); int* out_i = tmp2_idxs.data_ptr<int>(); while (num_candidates > k) { int blocks_per_row = (num_candidates + CHUNK - 1) / CHUNK; dim3 grid(blocks_per_row, batch); // launch kernel reading in_v/in_i, writing out_v/out_i ... num_candidates = blocks_per_row * k; // swap for next iteration std::swap(in_v, out_v); std::swap(in_i, out_i); } // After loop, in_v/in_i point to the final result // Copy to output for (int row = 0; row < batch; row++) { for (int i = 0; i < k; i++) { out_vals[row * k + i] = in_v[row * num_candidates + i]; out_idxs[row * k + i] = static_cast<int64_t>(in_i[row * num_candidates + i]); } } ``` Wait, after the loop, `num_candidates` is the number of candidates in the final buffer. For the last iteration, `num_candidates = blocks_per_row * k = 1 * k = k`. So the final buffer has `batch * k` elements, arranged as `row * k + i`. The copy loop can be replaced with a simple memcpy or tensor slice: ```cpp auto final_vals = torch::from_blob(in_v, {batch, k}, x.options()); auto final_idxs = torch::from_blob(in_i, {batch, k}, torch::dtype(torch::kInt32).device(x.device())); out_vals.copy_(final_vals); out_idxs.copy_(final_idxs.to(torch::kInt64)); ``` Actually, `out_vals` and `out_idxs` are already allocated. I can just do: ```cpp cudaMemcpy(out_vals.data_ptr<float>(), in_v, batch * k * sizeof(float), cudaMemcpyDeviceToDevice); cudaMemcpy(out_idxs.data_ptr<int64_t>(), in_i, batch * k * sizeof(int), cudaMemcpyDeviceToDevice); ``` This is very fast. But wait, there's an even better optimization: I can just have the last reduction kernel write directly to `out_vals` and `out_idxs`. Instead of swapping to a temp buffer for the last iteration, I pass the output pointers. In the loop: ```cpp bool is_last = false; while (num_candidates > k) { int blocks_per_row = (num_candidates + CHUNK - 1) / CHUNK; dim3 grid(blocks_per_row, batch); float* next_v = (blocks_per_row * k == k) ? out_vals.data_ptr<float>() : out_v; int* next_i = (blocks_per_row * k == k) ? out_idxs.data_ptr<int64_t>() : out_i; // launch kernel reading in_v/in_i, writing next_v/next_i ... num_candidates = blocks_per_row * k; in_v = next_v; in_i = next_i; } ``` If the next iteration would produce exactly k candidates per row, write directly to output! For all our shapes: - shape 0: iteration 1 produces 1024, iteration 2 produces 64. Iteration 2 writes to output. - shape 1: iteration 1 produces 128, iteration 2 produces 64. Iteration 2 writes to output. - shape 2: iteration 1 produces 256, iteration 2 produces 64. Iteration 2 writes to output. - shape 3: iteration 1 produces 128, iteration 2 produces 64. Iteration 2 writes to output. So only the last iteration writes to output. This saves one copy! Let me implement this. I'll modify the solution to use the warp-select approach for pass 2, but with multiple iterations. Actually, wait. Let me think about whether I need to keep the bitonic sort in pass 1 or go back to warp select. In the bitonic sort pass 1, we had 21.5us. In the warp select pass 1, we also had 21.5us. So they're similar. But the bitonic sort version had correctness issues with the warp chunk loading. Let me make sure the loading is correct. Actually, I already fixed the loading bug in the bitonic sort version. And it passed correctness. But it was much slower because of `__syncthreads()` in pass 2. So I'll keep the bitonic sort for pass 1 (or warp select, whichever is easier) and use warp-select with multiple iterations for pass 2. Let me write a clean version. For pass 1, I'll use the bitonic sort approach because it's already correct and the time is reasonable. For pass 2, I'll use the iterative reduction with warp select. Let me write the code. Actually, there's a simpler alternative for pass 2. Instead of warp select + block sort in each block, I can just use bitonic sort of CHUNK elements. Since CHUNK=1024 and each block processes one chunk, the bitonic sort of 1024 is fast. But the `__syncthreads()` issue remains. With CHUNK=1024, bitonic sort has 55 stages with `__syncthreads()`. For 16 blocks in iteration 1: 55 * 16 = 880 `__syncthreads()` calls. At ~1us each: 880us total. That's terrible. But if I launch 16 blocks on 16 SMs, the `__syncthreads()` calls happen in parallel across SMs. Each block does 55 `__syncthreads()` in ~10us. 16 blocks in parallel: ~10us total. Wait, that's very different from the single-block case! With 16 blocks on 16 SMs, each SM does its own `__syncthreads()` independently. So the total time is the time for ONE block, not all blocks combined! For a single block sorting 1024 elements with 256 threads: - 55 stages. - Each stage: 2 shared memory reads, 1 comparison, possible 2 writes. - Plus `__syncthreads()`. - Let's say each stage takes 1us (including `__syncthreads()`). - Total: 55us per block. With 16 blocks in parallel: 55us total for iteration 1. Iteration 2: 1 block, 55us. Total pass 2: 110us. Hmm, still slow. But with warp select instead of bitonic sort, each block takes ~0.4us. 16 blocks: ~0.4us. 1 block: ~0.4us. Total: ~0.8us. Much faster! So I definitely want warp select for pass 2, not bitonic sort. OK, let me write the iterative warp-select version. For pass 1, I'll use the bitonic sort version that was correct. Time: ~21us. Acceptable. Wait, but can I also make pass 1 faster? For pass 1 with 256 blocks, each block does bitonic sort of M_PAD=128 elements per warp + block sort of 256 elements. For M_PAD=128, warp sort: 21 stages with `__syncwarp()`. Let's say each `__syncwarp()` takes 0.1us. 21 * 0.1 = 2.1us per warp. With 4 warps: ~2.1us per block (parallel). Block sort 256: 28 stages with `__syncthreads()`. 28 * 1us = 28us per block. Wait, but the 256 blocks run in parallel! 256 blocks / 188 SMs = ~1.36 blocks per SM. Each SM runs 1-2 blocks. If each block takes 30us, and 2 blocks run sequentially on some SMs, the total time is ~60us. But the profiler showed 21.5us. So the blocks must be faster than 30us. Actually, `__syncthreads()` with 4 warps (128 threads) is fast because there are only 4 warps to synchronize. Maybe ~20 cycles = 13 nanoseconds. For block sort 256: 28 stages * 20 cycles = 560 cycles = 0.37 microseconds per block. Warp sort 128: 21 stages * 10 cycles = 210 cycles = 0.14 microseconds per warp. Total per block: ~0.5 microseconds. 256 blocks on 188 SMs: ~0.7 microseconds. But profiler says 21.5us. Where does the extra time come from? I think the main overhead might be: 1. Global memory loads: 512 elements * 4 bytes = 2KB per block. At 1.8TB/s: 1.1 nanoseconds. But with latency: maybe 0.1us. 2. Shared memory initialization: writing num_warps * M_PAD = 512 elements. ~0.05us. 3. Writing output: 64 elements to global memory. ~0.01us. Total per block: ~0.7us. 256 blocks: ~1us. But actual is 21.5us. So there's still a ~20x gap. Could the issue be that many blocks are scheduled on the same SM, causing queuing delays? With 256 blocks and 188 SMs, 68 SMs get 2 blocks. Each SM with 2 blocks runs them sequentially. So some SMs take ~1.4us instead of 0.7us. The overall time is still ~1.4us. Wait, maybe the blocks are not all launched at the same time. The GPU scheduler might launch them in waves. With 256 blocks and each SM holding 1-2 blocks, the total time should be close to the block execution time. I really don't understand where the 21.5us is coming from. Let me try adding some timing prints to the C++ wrapper to measure each kernel launch. Actually, I just realized: the 21.5us includes the kernel LAUNCH overhead from the CPU side! The profiler's "CUDA time" for a kernel includes the time from when the kernel is submitted to when it finishes. But if the CPU launches kernels sequentially with `cudaStreamDefault`, there might be some launch latency. Actually, no. CUDA kernel launches are asynchronous. The CPU queues the kernel and returns immediately. The GPU executes the kernel. The profiler measures the GPU execution time. Unless... PyTorch is doing some synchronization between kernel launches? Let me check. In the C++ wrapper, I launch pass 1, then pass 2. Since they're in the same stream, they execute sequentially. The profiler should show their execution times separately. Let me run the profiler on the bitonic sort version to see the exact times. Actually, let me just implement the iterative warp-select version and see if it improves things. If pass 2 with iterative warp-select is much faster, then the total time should drop significantly. For pass 2 with iterative warp-select: - Iteration 1: 16 blocks, each does warp select from 128 + block sort 512. - Time per block: ~0.4us. - 16 blocks on 16 SMs: ~0.4us. - Iteration 2: 1 block, same work. - Time: ~0.4us. - Total pass 2: ~0.8us. If pass 1 is 21.5us and pass 2 is 0.8us, total is ~22.3us. That's 2.3x faster than torch.topk (52us)! But I need to verify if the warp-select blocks are really that fast. The previous single-block warp-select took 121us. Why would 16 blocks take only 0.4us each? Ah! The previous pass 2 used TILE=4096 and 1 block. Each warp processed 512 candidates. The warp select did 64 iterations over 512 candidates. With CHUNK=1024 and 16 blocks, each block has 1 warp processing 1024/8 = 128 candidates? No, with 256 threads = 8 warps, each warp processes 1024/8 = 128 candidates. The warp select does 64 iterations over 128 candidates. That's 64 * (4 + 5) = 576 ops per thread. But the previous version had 64 * (16 + 5) = 1344 ops per thread. So with smaller CHUNK, the warp select is ~2.3x faster per block. And with 16 blocks instead of 1, we get 16x more parallelism. 16 blocks * 0.4us = 6.4us (if they ran sequentially). But they run in parallel on 16 SMs: ~0.4us. So pass 2 should be ~0.4us. Total: 21.5us + 0.4us + launch overhead = ~27us. This is excellent! Let me implement it. One more thing: for the iterative reduction, I need two temp buffers. I can allocate them once and reuse them. Actually, for pass 1, I still need a temp buffer for the initial output. For pass 2 iterations, I need ping-pong buffers. Let me allocate: - `tmp_a`: size `batch * max_candidates * k` for values (float) - `tmp_b`: size `batch * max_candidates * k` for indices (int32) Wait, pass 1 outputs `batch * num_blocks_per_row * k` elements. The max this can be is for batch=1, num_blocks=256, k=64: 16384 elements. For pass 2 iteration 1: outputs `batch * blocks_per_row * k` elements. Max blocks_per_row = 16. So max output = 1 * 16 * 64 = 1024 elements. So the max temp size is 16384 elements. I'll allocate both buffers with 16384 elements. In the C++ wrapper: ```cpp int max_candidates = num_blocks_per_row * k; // 16384 for shape 0 auto tmp_a_vals = torch::empty({batch, max_candidates}, x.options()); auto tmp_a_idxs = torch::empty({batch, max_candidates}, torch::dtype(torch::kInt32).device(x.device())); auto tmp_b_vals = torch::empty({batch, max_candidates}, x.options()); auto tmp_b_idxs = torch::empty({batch, max_candidates}, torch::dtype(torch::kInt32).device(x.device())); ``` Wait, but `max_candidates` is `num_blocks_per_row * k`. For shape 0: 256 * 64 = 16384. For pass 2 iteration 1, blocks_per_row = 16. Output: 16 * 64 = 1024 elements per row. Total output: 1 * 1024 = 1024 elements. But `tmp_b` is allocated as `{batch, max_candidates}`. So it has 16384 elements. The iteration 1 output only uses the first 1024 elements. For iteration 2, input is 1024 elements, output is 64 elements. I read from `tmp_b` and write to `tmp_a` (or output directly). This works! Let me write the reduction kernel. Actually, for the reduction kernel, I should use `blockIdx.y` for the row and `blockIdx.x` for the block within the row. The grid is `dim3(blocks_per_row, batch)`. For `out_pos`: ```cuda int out_pos = (row * gridDim.x + block_in_row) * K + tid; ``` This writes to a flat buffer of size `batch * gridDim.x * K`. In the next iteration, `num_candidates = gridDim.x * K`. And `in_vals + row * num_candidates + chunk_start` reads from the buffer. This is correct. Now, for the merge at the end: if the last iteration writes directly to output, I don't need a copy. Let me detect when the next iteration would produce exactly K elements: ```cpp int blocks_per_row = (num_candidates + CHUNK - 1) / CHUNK; bool write_to_output = (blocks_per_row * k == k); // i.e., blocks_per_row == 1 ``` Wait, `blocks_per_row` is at least 1. So `blocks_per_row * k == k` means `blocks_per_row == 1`. Yes! When `blocks_per_row == 1`, we can write directly to the output buffer. Let me write the complete C++ wrapper. One optimization: instead of allocating two temp buffers of size `batch * max_candidates`, I can allocate one buffer for values and one for indices, each of size `batch * max_candidates`. But I need two of each for ping-ponging. So 4 buffers total. Actually, I can just allocate 2 value buffers and 2 index buffers. Let me write the code. For the reduction kernel, I should also handle the case where CHUNK > num_candidates. The load loop pads with -FLT_MAX. OK, let me write the complete solution.

Let me test correctness and then benchmark.

There's an index out of range error. The indices are huge (4294967298), which suggests that `out_i` is being interpreted incorrectly. In the C++ wrapper, for the last iteration, I do: ```cpp int* next_i = last_iter ? reinterpret_cast<int*>(out_i) : out_i; ``` But `out_i` is `int*`, and I'm casting it to `int*` again. Wait, `out_i` is already `int*`. And `next_i` is also `int*`. That should be fine. Wait, let me look at the variable declarations: ```cpp float* out_v = tmp_b_vals.data_ptr<float>(); int* out_i = tmp_b_idxs.data_ptr<int>(); ``` And for the last iteration: ```cpp float* next_v = last_iter ? v_ptr : out_v; int* next_i = last_iter ? reinterpret_cast<int*>(out_i) : out_i; ``` Wait, `out_i` is already `int*`. Why am I `reinterpret_cast<int*>(out_i)`? That's redundant but shouldn't cause issues. Actually, I think the problem is that in the last iteration, `next_i` points to `out_i`, which is `tmp_b_idxs.data_ptr<int>()`. But `out_i` is `int*`, and the kernel expects `int*`. So the kernel writes `int` indices to `tmp_b_idxs`. But then `out_idxs` is `int64_t*`. The kernel is supposed to write `int64_t` indices to `out_idxs`. In the last iteration, I pass `next_i = reinterpret_cast<int*>(out_i)`, which is `tmp_b_idxs.data_ptr<int>()`. The kernel writes `int` indices there. But `out_idxs` is never written to directly! Wait, I see the bug. In the last iteration, I should write to `out_idxs` (int64_t), not to `tmp_b_idxs` (int). But the `topk_reduce` kernel writes `int` indices. So I can't write directly to `out_idxs` without modifying the kernel. I have two options: 1. Always write to temp buffers and copy at the end. 2. Modify the kernel to handle int64_t output indices. Option 1 is simpler. Let me just remove the `last_iter` optimization and always write to the ping-pong buffers. Then copy the final result. Actually, I can do the copy with a simple CUDA kernel or just do it on the CPU. But for batch*k elements, it's tiny. Wait, but the copy from `int` to `int64_t` needs to be done. I can write a tiny kernel for this. Or, I can modify `topk_reduce` to take a template parameter for the index type. But that's more complex. Simplest fix: remove the `last_iter` optimization. Always write to `out_v`/`out_i` temp buffers. After the loop, copy `in_v[0:batch*k]` to `v_ptr` and `in_i[0:batch*k]` to `i_ptr` with a cast. In C++: ```cpp // After the reduction loop, copy final result to output for (int row = 0; row < batch; row++) { for (int i = 0; i < k; i++) { v_ptr[row * k + i] = in_v[row * k + i]; i_ptr[row * k + i] = static_cast<int64_t>(in_i[row * k + i]); } } ``` This is a simple CPU loop over batch*k elements. For our shapes, batch*k is at most 128*64 = 8192. But this runs on the CPU, so it reads GPU memory, which is very slow! Better to do the copy on the GPU. I can use `cudaMemcpy` or write a small kernel. Actually, `torch::Tensor` has `copy_` method. But I need to handle the int->int64 cast. I can write a tiny kernel: ```cuda __global__ void copy_output(const float* in_vals, const int* in_idxs, float* out_vals, int64_t* out_idxs, int k) { int idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx < k) { out_vals[idx] = in_vals[idx]; out_idxs[idx] = static_cast<int64_t>(in_idxs[idx]); } } ``` And launch it with `batch * k` threads. But actually, for the last iteration, I can just pass `out_vals` and `out_idxs` directly to the kernel if I modify the kernel to accept int64_t output indices. But that means changing the template. Actually, there's a much simpler fix: change `out_i` to be `int64_t*` and modify `topk_reduce` to write `int64_t` indices. But then the intermediate buffers need to be `int64_t` too, which wastes memory. Alternatively, I can keep the intermediate buffers as `int` and do a final copy with `cudaMemcpy` + kernel. Wait, the simplest fix is to just not do the last-iter optimization. Always write to the `int` temp buffer, and then use `cudaMemcpy` for values and a tiny kernel for indices. Actually, I can just add the copy kernel to the C++ code. Let me rewrite the wrapper to always use ping-pong buffers and then copy at the end.

I need to add the `copy_output` kernel to the CUDA source. Let me do that.

Let me recompile and test.

Let me run check.py.

Correctness passes. Let me run the benchmark to see the performance.

The performance is still poor for most shapes. Shape 0 is 0.083ms vs eager 0.056ms. Shape 1 is 0.091ms vs eager 0.051ms. Shape 4 is fast (0.010ms vs 0.022ms) because it's k=1. The iterative reduction with warp select is still slow. Let me profile with the profiler to see where time is going.

The profiler is showing 0.0us or 0.1us for the kernels, which is clearly wrong. The `cuda_time_total` is deprecated and returning 0. Let me use `device_time_total` instead. But actually, the profiler seems broken for this. Let me try a different approach: use `torch.cuda.Event` timing around the Python forward call. I already did this and got ~143us for the old version, and ~83us for the current version. Wait, the benchmark shows 0.083ms for shape 0, which is 83us. That's better than the original 152us, but still slower than torch.topk's 52us. Let me think about this more carefully. Maybe the overhead is not in the kernels themselves, but in the C++ wrapper function. Every time the Python model's forward is called, it calls the C++ function `topk_forward_cuda`, which: 1. Computes `num_blocks_per_row` 2. Allocates 4 temp buffers with `torch::empty` 3. Launches pass 1 kernel 4. Launches pass 2 reduction kernels 5. Launches copy_output kernel The temp buffer allocation with `torch::empty` on CUDA might be taking time! Even though `torch::empty` doesn't zero the memory, it still needs to allocate from PyTorch's CUDA memory pool. For small buffers (like 64KB), PyTorch's allocator should be fast. But if the allocation requires a CUDA API call or synchronization, it could add microseconds. Actually, `torch::empty` in C++ is the same as `torch.empty` in Python. For CUDA tensors, it allocates from the cached allocator. For small sizes, it's very fast. But there might still be some overhead. What if I pre-allocate the temp buffers in the Python Model and pass them to the C++ function? This would eliminate allocation overhead. In Python: ```python class Model(nn.Module): def __init__(self, batch, n, k): ... if k > 1: num_blocks = max(1, 256 // batch) max_candidates = num_blocks * k self.register_buffer("_tmp_a_vals", torch.empty(batch, max_candidates, dtype=torch.float32)) self.register_buffer("_tmp_a_idxs", torch.empty(batch, max_candidates, dtype=torch.int32)) self.register_buffer("_tmp_b_vals", torch.empty(batch, max_candidates, dtype=torch.float32)) self.register_buffer("_tmp_b_idxs", torch.empty(batch, max_candidates, dtype=torch.int32)) def forward(self, x): return _topk_cuda.topk_forward_cuda(x, self.k, self._tmp_a_vals, self._tmp_a_idxs, self._tmp_b_vals, self._tmp_b_idxs) ``` And in C++: ```cpp std::vector<torch::Tensor> topk_forward_cuda(torch::Tensor x, int k, torch::Tensor tmp_a_vals, torch::Tensor tmp_a_idxs, torch::Tensor tmp_b_vals, torch::Tensor tmp_b_idxs) { // ... no allocations, just use the passed buffers } ``` This would save allocation time on every forward call. But wait, `register_buffer` in Python means the buffers are part of the model state. They get moved to CUDA when `.cuda()` is called. And they're persistent across forward calls. Let me try this! Actually, I also need to pass `num_blocks_per_row` or compute it in Python. In Python: ```python class Model(nn.Module): def __init__(self, batch, n, k): ... self.num_blocks_per_row = max(1, 256 // batch) if k > 1: max_candidates = self.num_blocks_per_row * k self.register_buffer("_tmp_a_vals", torch.empty(batch, max_candidates, dtype=torch.float32)) ... def forward(self, x): if self.k == 1: return _topk_cuda.topk_forward_cuda(x, self.k) return _topk_cuda.topk_forward_cuda_2pass(x, self.k, self.num_blocks_per_row, self._tmp_a_vals, ...) ``` And in C++: ```cpp std::vector<torch::Tensor> topk_forward_cuda_2pass(torch::Tensor x, int k, int num_blocks_per_row, torch::Tensor tmp_a_vals, torch::Tensor tmp_a_idxs, torch::Tensor tmp_b_vals, torch::Tensor tmp_b_idxs) { // ... } ``` Let me implement this. It should reduce the overhead significantly. Also, I realize that the benchmark does 50 iterations, and the median time includes all the forward pass overhead. Reducing overhead is crucial. Let me rewrite solution.py with pre-allocated buffers. Actually, one more optimization: instead of iterative reduction in pass 2 with multiple kernel launches, I can use a single kernel that does everything if I launch enough blocks. For batch=1, num_candidates=16384: - Launch 16 blocks in pass 2. - Each block processes 1024 candidates and produces 64. - But then I need another launch to merge the 16 * 64 = 1024 candidates. Wait, what if I launch 64 blocks? Each block processes 256 candidates and produces 64. - 64 blocks on 64 SMs: ~0.4us. - But then I have 64 * 64 = 4096 candidates left. Need another iteration. The number of iterations is log(CHUNK) / log(K). For K=64 and CHUNK=1024: 16384 -> 1024 -> 64. 2 iterations. What if CHUNK = 16384? Then 1 iteration! But shared memory for 16384 elements = 128KB. Too much. What if CHUNK = 4096? Shared memory = 32KB + 2KB = 34KB (under 48KB). 16384 / 4096 = 4 blocks. 1 iteration: 4 blocks -> 256 candidates. Need another iteration. Still 2 iterations. What if I use global memory for the tile instead of shared memory? Then I can process larger chunks. Actually, for pass 2, I can process the entire candidate array in one kernel with multiple blocks, and use a final block to merge. But that requires global memory atomics or a second kernel. So 2 iterations is fine. The key is to minimize overhead. Let me implement the pre-allocated buffer version and see if it helps.

Let me test correctness and benchmark.

The check.py compares state_dicts between reference and solution models. The reference model doesn't have `_tmp_a_vals`, etc. So I need to make sure my Model's state_dict matches the reference. The reference model has only `_dummy`. My model has `_dummy` plus the temp buffers. The temp buffers should not be in the state_dict. I can either: 1. Not register them as buffers (use regular attributes) 2. Or register them but exclude them from state_dict somehow The simplest is to not use `register_buffer` and instead just assign them as attributes. They'll still be moved to the correct device when `.to(device)` is called if I handle it properly... actually, regular tensors assigned in `__init__` won't be moved by `.to()` or `.cuda()`. I can override `to()` or just manually move them. But that's complex. Actually, a simpler approach: register them as buffers but give them names that start with an underscore and aren't in the reference state_dict. But check.py does `sol_model.load_state_dict(sd, strict=True)`, which requires exact match. The reference state_dict only has `_dummy`. So my model's state_dict must also only have `_dummy`. Solution: don't register the temp buffers. Instead, create them lazily in `forward()` or store them as regular Python attributes. Wait, if I store them as regular attributes in `__init__`: ```python self._tmp_a_vals = torch.empty(...) ``` Then when the model is moved with `.cuda()`, these tensors won't move. But check.py does: ```python sol_model = solution.Model(*init_args).to(device).eval() ``` So `.to(device)` moves all parameters and buffers, but not regular attributes. The temp buffers would stay on CPU, causing errors. I can create the temp buffers inside `forward()` on the same device as `x`: ```python def forward(self, x): if self.k == 1: return _topk_cuda.topk_forward_cuda(x, self.k, 0, ...) if not hasattr(self, '_tmp_a_vals') or self._tmp_a_vals.device != x.device: max_candidates = self.num_blocks_per_row * self.k self._tmp_a_vals = torch.empty(self.batch, max_candidates, dtype=torch.float32, device=x.device) ... return _topk_cuda.topk_forward_cuda(...) ``` This creates them on the correct device lazily. But it allocates them on every call if I check device mismatch. Actually, I can create them once and cache them. Wait, but this brings back the allocation overhead! Though it's only done once (on the first call), so subsequent calls reuse the same buffers. Let me implement this. Actually, another approach: register the buffers with names that ARE in the state_dict but have the same shape/value in both models. But the reference model doesn't have them. The cleanest solution is to not register them and create them lazily in forward. Let me rewrite the Model class.

Let me recompile and test.