We need to write a custom CUDA kernel for a specific PyTorch operation: Given: - token_hidden: [num_tokens, hidden_dim] - router_logits: [num_tokens, num_experts] - expert_ground: [num_experts, hidden_dim] - alpha: scalar float Operation: 1. Compute grounded = router_logits + alpha * (token_hidden @ expert_ground.t()) - token_hidden @ expert_ground.t() yields [num_tokens, num_experts] 2. Compute top-2 values and indices per row of grounded. 3. Compute softmax over top-2 values per row: top_weights = softmax(top_vals, dim=-1) 4. Return stack of (top_idx as float32, top_weights) along last dimension, shape [num_tokens, 2, 2] Reference output shape: [num_tokens, 2, 2] where for each token, first row is two top expert indices (as float), second row is two weights. We need to write a CUDA kernel that does this fused, minimizing memory bandwidth and computation. Key operations: - A GEMV or GEMM per token: token_hidden dot each expert_ground. That's hidden_dim=128, num_experts=16. Small matrix multiplication. Could compute on the fly in the kernel to avoid extra memory reads/writes. - Add router_logits and alpha*dot product. - Find top-2 values and indices per row. Since num_experts=16, it's tiny. We can unroll a simple linear scan per thread/warp. Since num_experts is small (16), we could even have each thread handle one token and store the 16 floats in registers or shared memory. For num_tokens=4096, one block per token with 1 thread is too many threads? 4096 blocks is okay, but we can also process multiple tokens per block. Let's think about occupancy. GPU specs unknown yet. Need to get them. The kernel should be robust for arbitrary dimensions within reason, but we know dims. Potential kernel design: - Each thread processes one token (row). There are 4096 tokens, which is enough to fill GPU. - For each token, compute dot product of token_hidden[hidden_dim] with each expert_ground[expert][hidden_dim]. hidden_dim=128, num_experts=16. - For each expert e from 0 to num_experts-1: - Compute sum_{d=0}^{hidden_dim-1} token_hidden[token, d] * expert_ground[e, d] - Could load token_hidden into register array of size 128 (float) per thread. That's 128 registers * 4 bytes = 512 bytes. That's a lot. Alternatively, each thread loads one element and we use shared memory? But if each thread handles one token, shared memory is per block. We could have each block handle a chunk of tokens, say 32 or 64, and share expert_ground in shared memory. Actually, expert_ground is [num_experts, hidden_dim], which is 16*128 = 2048 floats = 8KB. That's small enough to load into shared memory once per block. token_hidden is [num_tokens, hidden_dim] = 4096*128 = 524288 floats = 2MB. We need to stream through it. Better approach: - Use 1D grid of blocks, each block has e.g., 128 threads. Each block processes a chunk of tokens, e.g., blockDim.x tokens (one per thread). Wait, if blockDim.x = 128, then block handles 128 tokens. - Load expert_ground into shared memory: size 16*128 floats = 8KB. That's fine. - Each thread loads its token_hidden row from global memory. 128 floats per thread = 512 bytes. We can load it into registers by unrolling a loop. Or we can keep it in global memory and stream through it while computing dot products. Since we need to multiply with each expert, and expert_ground is in shared memory, each thread can compute the 16 dot products sequentially. For each expert e, sum over d: token_hidden[token,d] * shared_expert[e,d]. Access pattern: token_hidden is coalesced if threads in a warp read same d for different tokens? Actually, if thread t in block reads token_hidden[token_start + t, d], memory layout is row-major, so adjacent threads read adjacent addresses: coalesced. Good. So algorithm per block: - Load expert_ground[e,d] into shared memory `s_expert[num_experts][hidden_dim]`. Use all threads to copy. We have 128 threads in block, total elements 2048, so each thread copies 16 elements. Use coalesced access: threadIdx.x copies elements with stride blockDim.x. Or just a simple loop. - __syncthreads() - For each thread (corresponding to token index `token = blockIdx.x * blockDim.x + threadIdx.x`): - If token < num_tokens: - Compute 16 scores: score[e] = router_logits[token, e] + alpha * sum_{d} token_hidden[token,d] * s_expert[e,d] - Top-2 search over e in [0,15]. Since 16 is tiny, just do a simple linear scan. - After computing both top values and indices, compute softmax: w0 = exp(v0) / (exp(v0) + exp(v1)), w1 = exp(v1) / (exp(v0) + exp(v1)). For numerical stability, subtract max(v0,v1) before exp. Actually only two values, so easy. - Write to output: out[token, 0, 0] = (float)idx0, out[token, 0, 1] = (float)idx1, out[token, 1, 0] = w0, out[token, 1, 1] = w1. Wait, the output shape is [num_tokens, 2, 2]. The order from torch.stack((top_idx.to(torch.float32), top_weights), dim=-1) means: - top_idx shape: [num_tokens, 2] - top_weights shape: [num_tokens, 2] - stack on dim=-1 => shape [num_tokens, 2, 2]. The last dimension has top_idx and top_weights. - So out[token, e, 0] = top_idx[token, e] (as float), out[token, e, 1] = top_weights[token, e]. - Wait, stack on last dim means: for each token and top position j: - out[token, j, 0] = top_idx[token, j] - out[token, j, 1] = top_weights[token, j] - So the returned shape is [num_tokens, 2, 2]. Let's verify: torch.stack(tensors, dim=-1). If tensors are [N,2], stacking on last dim gives [N,2,2]. Yes, last dim is the stack dim, so index 0 is first tensor (top_idx), index 1 is second tensor (top_weights). So out[token, j, 0] = float(top_idx[token, j]), out[token, j, 1] = top_weights[token, j]. Thus we need to write output accordingly. We need to consider fp32 precision. All inputs are fp32. Use float throughout. Use expf. Potential issues: - register pressure: each thread stores token_hidden[128] in registers? Actually we don't need to store the full row. We can compute each score on the fly: for e in 0..num_experts-1: float dot = 0; for d in 0..hidden_dim-1: dot += token_hidden[token*hidden_dim + d] * s_expert[e*hidden_dim + d]; score = router_logits[token*num_experts + e] + alpha * dot; // update top-2 This means we read token_hidden row 16 times (once per expert). That's 16*128 = 2048 reads per thread. But token_hidden is in global memory. However, token_hidden row fits in L1 cache if tile is reused. For 128 threads in a warp reading same d for different tokens? Actually, for a given expert e, thread t reads token_hidden[t, d]. Adjacent threads read adjacent memory: coalesced. The line will be cached. Then for next expert e+1, same threads read same token_hidden[t, d] again. Since it's the same address, it will hit L1. So it's fine. Alternatively, we could load token_hidden into a register array of size 128 to avoid re-reading, but 128 registers might be okay? 128 floats = 128 registers. That might be high but manageable if occupancy isn't paramount. However, we also have other variables. Each SM has 64K registers. With 128 threads per block, 128 regs per thread = 16384 regs per block. An SM can run 2 blocks if 32K per block? Actually max registers per SM is 64K (for compute capability >= 7.0). So 16384 per block means up to 4 blocks per SM. With 128 threads per block, that's 512 threads per SM. Occupancy = 512 / 2048 = 25% on Ampere? Actually warp count: 512 threads = 16 warps. Max warps per SM is 64 (for Ampere). So 25% occupancy. That's acceptable but maybe not optimal. However, given the tiny computation, memory bandwidth might dominate. And if we keep token_hidden in registers, we reduce L1 bandwidth and maybe help. But maybe a simpler kernel: each warp (32 threads) processes one token collaboratively. That might be better for register usage but more complex. Let's consider. Option A: 1 thread per token. - Block size: 128 or 256 threads. - expert_ground in shared memory. - Compute 16 dot products per thread. - top-2 reduction is trivial per thread. - Write output. This is very simple and likely fast enough due to high parallelism (4096 threads). But memory bound. Let's compute arithmetic intensity: - Reads: token_hidden row: 128 f32 = 512 bytes. router_logits row: 16 f32 = 64 bytes. expert_ground: 16*128 f32 = 8KB (shared, loaded once per block but amortized). Output: 4 f32 = 16 bytes. Total reads per token ~ 592 bytes. Writes ~16 bytes. - Compute: 16*128 = 2048 FMAs = 4096 FLOPs. Plus some adds, top-2, exp, etc. ~4200 FLOPs. - AI = 4200 / 608 ≈ 6.9 FLOPs/byte. HBM bandwidth is maybe 1-2 TB/s. So peak FLOPs needed: 6.9 * 2TB/s = 13.8 TFLOPs. GPU compute is >10 TFLOPs for FP32 (maybe 19 TFLOPs for A100). So compute and memory are comparable. Fusing helps by reducing intermediate writes (e.g., writing the full [4096,16] grounded matrix). Reference does: matmul (produces [4096,16] in global mem), add (another kernel), topk (another kernel, reads/writes), softmax (another kernel), stack (another kernel). That's a lot of memory traffic. Fusing everything into one kernel will be a huge win. Even a naive kernel will be much faster. But we should still aim for good performance. Option B: warp per token. - 32 threads cooperate to compute 16 dot products. Each thread computes dot product for some experts? Or each thread computes partial sums across hidden_dim and then warp reduce. Since hidden_dim=128, each thread can sum 4 elements per expert. But 16 experts * 128 hidden = 2048 multiplications. 32 threads => 64 muls per thread. That's low. But register usage per thread is tiny. - However, warp-level top-2 reduction requires shuffle. Since 16 values per token, each thread could hold 2 values, then do warp reduce. But it's more complex. Given simplicity and likely sufficient speedup, I'll go with 1 thread per token. To improve, we can process 2 or 4 tokens per thread to amortize block overhead and increase block work. Wait, blockDim.x = 128. If we have 4096 tokens, we need 32 blocks. That's only 32 blocks! If GPU has many SMs (e.g., 108 SMs on A100), that's terrible occupancy. We need more blocks. So better to have more threads per token? Or have each thread process multiple tokens? Or have smaller block size. Let's think occupancy. - 4096 tokens. - If block size = 64, blocks = 64. Still maybe not enough for 108 SMs. - If block size = 32, blocks = 128. Better. - Or each thread processes multiple tokens: e.g., block size = 128, each thread processes 4 tokens => 8 blocks. That's worse. So to get enough blocks, we want block size small, like 64 or 128? Actually 4096/128 = 32 blocks. On Ampere with 108 SMs, 32 blocks will utilization be limited. But each block has 128 threads. 32*128 = 4096 threads. The GPU can run thousands of threads concurrently. But the number of blocks per SM is limited by registers/shared memory. If we have 32 blocks total, we can only distribute them across SMs, at most 1 block per SM for many SMs, leaving others idle. Even if each SM can run multiple blocks, total blocks = 32. So only 32 SMs can be active at once (if 1 per SM). That's about 30% utilization on A100. Not great. So we need to either: 1. Reduce block size to increase total blocks. E.g., block size = 32 => 128 blocks. That's okay, maybe 1 block per SM. 2. Use warp per token: each block has e.g., 4 warps = 128 threads, handling 4 tokens (32 threads per token). Then blocks = 1024. That would be great. But code is more complex. 3. Each thread handles 1 token but block size = 16? Then 256 blocks. Or block size = 8 => 512 blocks. But block size too small reduces memory coalescing efficiency? Actually warps of 32 threads still coalesce if threads are contiguous. However, we need to consider that the workload per thread is small. If block size is 32, each warp handles one token? No, each thread handles one token. So warp handles 32 tokens. That's fine. Block size 32 means 1 warp per block. This is okay for occupancy because blocks are the scheduling unit. 128 blocks of 32 threads each = 4096 threads. On a GPU with 108 SMs and max 2048 threads/SM, that's fine. But warps per SM: 128 warps total. Max 64 warps/SM => at least 2 SMs needed, but more importantly, 128 blocks can be distributed across SMs. If each SM runs e.g., 4 blocks (32 warps), then 32 SMs active. Still not all SMs. Actually 128 blocks total / 108 SMs ≈ 1.2 blocks per SM. That's okay. But can we do better? Block size = 16 => 256 blocks. Then ~2.4 blocks per SM. Better. But block size less than warp size (32) means some threads in warp are inactive, but that's fine because each block is a separate warp? Actually a block with 16 threads still gets a full warp of 32 threads, with 16 masked. So it's fine but slightly wasteful. But we can have block size = 64 => 64 blocks. That's maybe okay. Wait, is 32 blocks really that bad? On some GPUs, a block executes on a single SM, but an SM can execute multiple blocks concurrently. So 32 blocks can run on 32 SMs, and if each SM can run e.g., 4 blocks, we need 8 SMs. But we have many more SMs. The remaining SMs are idle because there are no more blocks. So throughput is limited by the number of blocks. Thus we want number of blocks >= number of SMs * desired concurrency per SM. For max occupancy, we might want ~2-4 blocks per SM. 108 SMs * 2 = 216 blocks. So block size ≈ 4096/216 ≈ 19. So block size of 16 or 32 is reasonable. But block size too small can have overhead. Alternatively, we can have each thread process multiple tokens. E.g., block size = 256, each thread processes 4 tokens, blocks = 4. That's terrible. Better: block size = 128, each thread processes 1 token, but we launch more blocks by processing a smaller chunk? No, tokens = 4096 fixed. But note: we don't need each thread to handle exactly one token. We can have each thread handle e.g., 2 tokens. Then block size = 128 => 16 blocks. Even worse. So to increase block count, we must decrease threads per block or have multiple threads per token. Let's analyze multiple threads per token. Warp per token (32 threads per token): - Number of tokens per block = blockDim.x / 32. If blockDim.x = 128, then 4 tokens per block. - Number of blocks = 4096 / 4 = 1024 blocks. That's great. - Each warp collaborates on computing dot products for one token. - How to partition work: hidden_dim = 128. 32 threads. Each thread computes partial sum of 4 elements per expert? That's 16 experts * 4 elements = 64 multiplications per thread. After computing partial sums, we need to reduce across warp. - However, we need 16 scores per token. We could have each thread compute 1 score (or 2 scores) using all hidden dims. Since there are 32 threads and 16 experts, we can assign 2 threads per expert, each computing half the dot product and then adding via shfl. Or assign 1 thread per expert (16 threads active, 16 idle). But that's okay. - Better: each thread computes the full dot product for one expert. 16 threads in warp do work, 16 idle. But then we need to compute top-2 across the 16 scores. We can use warp-level reduce. All 32 threads can hold a score? No, only 16 have valid scores. We can have lane i compute score for expert i if i < 16, else 0. Then do a warp-level parallel reduction to find top-2. There are known warp-level top-k algorithms. For k=2 and 16 elements, it's manageable. Alternatively, we can have each warp process 2 tokens: 16 threads per token. But 16 threads can't do a warp-level reduce because they are not a full warp. Actually they can, but shfl within a warp works for all 32 lanes; lanes 0-15 handle token 0, lanes 16-31 handle token 1. Then we can do independent shfl within each half-warp? __shfl_sync with lane masks? We can do shfl within a warp for all lanes. But if we have two tokens in one warp, the reduction must not cross tokens. We can do shfl and condition on (lane/16 == other_lane/16). But simpler: one token per warp. Let's explore one token per warp, block size = 128 (4 warps, 4 tokens per block). Blocks = 1024. - Shared memory for expert_ground: 8KB. All warps in block use it. - Each warp processes one token. Algorithm per warp: - Determine token index: token = blockIdx.x * (blockDim.x / 32) + warp_id. - Each lane i (0..31) loads part of token_hidden for that token. Specifically, each lane loads token_hidden[token, d] for d in [i, i+32, i+64, i+96]? Actually hidden_dim=128, so each lane can load 4 elements. We can load into registers. Or each lane can compute partial dot products for a subset of experts and dims. But maybe simpler: each lane computes the dot product for expert e = lane_id if lane_id < num_experts (16). For that expert, it needs to sum over d=0..127 of token_hidden[token,d] * expert_ground[e,d]. Since hidden_dim=128, and 128 threads aren't available per token, each lane must loop over d=0..127 alone. That's 128 loads from global memory (token_hidden row) and 128 loads from shared memory (expert_ground). But all 32 lanes in warp load the same token_hidden elements (since each lane loops over d). However, lanes 0..15 load token_hidden[t,d] for each d, while lanes 16..31 are idle (or also load but not needed). This means the warp issues 16 loads per d. But they are all the same address? No, each lane loads token_hidden[t,d] (same address). On modern GPUs, broadcasting the same address to multiple threads in a warp is efficient: it reads one cache line and broadcasts. So it's okay. But we have 16 active lanes. Actually, even simpler: Have all 32 lanes cooperate to compute one score at a time. For expert e: - lane i loads token_hidden[t, i] and token_hidden[t, i+32] etc. Let's partition hidden_dim=128 across 32 lanes: each lane loads 4 elements. It also loads corresponding expert_ground[e, i], expert_ground[e,i+32], etc. from shared memory. - Each lane computes partial sum of its 4 products. - Then warp reduce (shfl) to sum the 32 partial sums. - The result is the dot product for expert e. Lane 0 (or all lanes) stores it maybe in register. - Repeat for e=0..15. This takes 16 iterations per token. Each iteration: 32 loads from token_hidden (coalesced? token_hidden[t, i] for i = lane_id, t fixed. Adjacent lanes read adjacent elements: coalesced). 32 loads from shared memory: s_expert[e, i] for i=lane_id. Coalesced in shared memory. 4 FMAs per lane. Then warp reduce (5 shuffles). Total work per token: 16*4=64 FMAs per lane, plus loads. This is very lightweight but serializes across 16 experts. However, 1024 blocks * 4 warps = 4096 warps. With 108 SMs, each SM gets ~38 warps. Max warps/SM = 64. So maybe okay. But we can do better: compute all experts in parallel. Since num_experts=16 and warp_size=32, we can have each lane compute a full dot product for one expert. 16 lanes active, 16 idle. The 16 active lanes each load token_hidden[t,d] for all d. The token_hidden row is broadcast. Then each lane loads expert_ground[e,d] from shared memory and accumulates. This is more parallel. After computing 16 scores in registers (one per active lane), we need to find top-2 across the warp. We can do a warp-level top-2 reduction. Since there are only 16 values, we can do it with shuffles. Let's detail this: - Warp size 32. token = blockIdx.x * warps_per_block + warp_id. - lane = threadIdx.x % 32. - Active if lane < num_experts (16). Or we can have all lanes compute but only first 16 matter; but we want to use all 32 lanes? We could assign 2 experts per lane: lane i handles experts i and i+16? But num_experts=16. So 16 lanes idle. That's okay. Actually, we can have 2 tokens per warp: lane 0-15 for token 0, lane 16-31 for token 1. Then each half-warp computes 16 scores independently. But warp-level top-2 would need to be done per half-warp. We can use __shfl_sync with a mask? Actually shfl sync works on the whole warp, but we can have each half-warp only consider its own values by starting with pair comparison within the half-warp and ignoring the other half. But __shfl_sync by default reads any lane. We can use __shfl_sync with lane offset but the data from other half-warp would be irrelevant unless we mask it. However, we can simply treat lane 16-31 as computing scores for token 1 and doing their own top-2. To avoid cross-talk, we can ensure that when lane 0 does shfl from lane 1, lane 1 is in same token group. Since lane 0 and lane 1 are both in token 0, fine. But lane 0 reading lane 16 would be wrong if we don't mask. So we need to ensure the reduction only accesses lanes within the same token. That's a bit more complex because warp shuffle instructions take a lane offset or lane index, and we can conditionally call them or mask the values. For example, we can do a butterfly reduction where at each step, the offset is 1,2,4,8,16. When offset=16, lane 0 would read lane 16. If we mask: if (lane_id < 16) then the partner is lane_id+16 which belongs to other token. We can avoid the offset=16 step by doing a 16-element reduction using 4 steps (1,2,4,8). Because 16 elements, max offset 8 is enough to propagate across 16 lanes. Actually to reduce 16 values in a warp, we can do it within a 16-lane subgroup if we restrict shfl to modulo 16. But the GPU warp shuffle doesn't have a width parameter unless we use __shfl_down_sync with width (from CUDA 9). __shfl_down_sync(mask, var, delta, width) allows width=16. However, width must be a power of two. So we can set width=16. Then lanes 0-15 act as one warp of 16, lanes 16-31 as another warp of 16. That's perfect! Then each warp can handle 2 tokens. This doubles parallelism. Let's see: - Each warp handles 2 tokens. - token0 = base_token + warp_id*2 + 0. - token1 = base_token + warp_id*2 + 1. - For lane in 0..15: compute score for expert = lane for token0. - For lane in 16..31: compute score for expert = lane-16 for token1. - All lanes load token_hidden for their respective token. - token0 lane loads token_hidden[token0, d]. token1 lane loads token_hidden[token1, d]. - Coalescing: adjacent lanes read adjacent tokens, good. - All lanes load s_expert[e,d] where e = lane%16. Since e same for lane i and i+16, they both load the same shared memory address. Broadcast, okay. - Each lane accumulates dot product. - After loop, each lane has a score. We need top-2 per token. - Use __shfl_down_sync with width=16 within each half-warp. - For token0 (lanes 0-15), we can do a reduction to find top-2. - We need to compute both max values and their indices. - A simple approach: we can do a tree reduction for max, then a second pass to find the second max (excluding the max lane). But since k=2, we can do pairwise compare-swap. - A known efficient approach for small k: do a reduction where we keep two best values. - Initialize `float my_val = score; int my_idx = lane%16;` - For offset = 1,2,4,8: - Receive `other_val = __shfl_down_sync(0xFFFF, my_val, offset, 16);` - Receive `other_idx = __shfl_down_sync(0xFFFF, my_idx, offset, 16);` - Compare and keep top 2? Actually shfl_down reduction is good for single max. For top-2, we need a more custom reduction. - Alternatively, we can simply do a serial loop within one lane after gathering all 16 scores to that lane. Since 16 scores, we can have lane 0 collect all 16 scores via shfl, then do serial top-2. But that's 16 shfls + 16 comparisons on one lane, while others idle. Might be okay because it's tiny. - Better: Have lane 0 and lane 1 cooperate. - Actually, we can just have each active lane compute its score, then write score and idx to shared memory within the warp? But shared memory per block is limited. Or we can just have each thread (lane) do an independent serial top-2 after computing its own scores? Wait, each lane only computed one score. To know the other scores, it needs to communicate. Alternative: Use warp-level primitives to broadcast all scores to all lanes in the half-warp, then each lane does its own serial top-2? That's wasteful. Better to have one lane in each half-warp gather the 16 scores and do the top-2. The time is negligible compared to dot product. Let's do this per half-warp: - After computing scores, each lane has `float score` and `int expert = lane % 16`. - Use a ballot or shfl to gather to lane 0: - We can have an array `float scores[16];` in registers for lane 0. For i=0..15: `scores[i] = __shfl_sync(0xFFFF, score, i, 16);` and similarly `int idxs[i] = __shfl_sync(0xFFFF, expert, i, 16);` but since expert = i, idxs are just i. But we can shfl the idx too. - Then lane 0 does a simple linear scan for top-2 over 16 values. - Then lane 0 computes softmax and writes output for its token. - Meanwhile, lane 16 does the same for token1. - Lanes 1-15 and 17-31 are idle during this part. Is 16 shfl instructions expensive? Not really. Shfl is fast. But wait, we need to gather 16 floats. That's fine. But maybe we can avoid gathering all by doing a tree of compare-swaps. Let's consider the tree reduction for top-2: - Each lane holds a pair `(v1, i1, v2, i2)` where v1 is best, v2 is second best in its sub-group. - Initially `v1 = score, i1 = expert, v2 = -inf, i2 = -1`. - For offset = 1,2,4,8: - Get partner's pair via shfl_down with width=16. - Merge the two pairs into a new pair of top-2 from the combined 4 values. - Merging: we have up to 4 values: my v1, my v2, partner v1, partner v2. Find top 2 among them. Since it's just 4, can do a few comparisons. - After offset=8, lane 0 has the top-2 for the 16 lanes. - This is elegant and parallel. But the merge logic might add registers. However, it's tiny. However, writing and debugging this might be error-prone. Since the top-2 over 16 floats is a tiny part, doing it on a single lane is fine. But maybe simplest: each thread handles one token alone. Let's revisit occupancy concern. 4096 tokens. If block size = 64, blocks = 64. Is 64 blocks enough? On a GPU with e.g., 80 SMs (V100) or 108 SMs (A100), 64 blocks < SMs. But an SM can run multiple blocks concurrently. Let's see how many blocks per SM we can get. Our kernel uses shared memory: expert_ground [16][128] floats = 8192 bytes. Registers per thread: let's estimate. If we load token_hidden row into registers: 128 floats = 128 registers. Plus loop variables, etc. Maybe ~140 registers. At 140 registers, max threads per SM = floor(65536/140) = 468 threads. For block size 64, max blocks per SM = floor(468/64) = 7 blocks. But warps per SM: floor(468/32) = 14 warps. Since max warps/SM = 64, register limit is binding. So each SM can run 7 blocks of 64 threads = 448 threads. Total blocks = 64. Total threads = 4096. On 108 SMs, we can run at most 7*108 = 756 blocks, but we only have 64. So number of SMs used = ceil(64/7) = 10 SMs. That's terrible. Wait, is my register estimate correct? If we don't preload token_hidden into registers, but rather stream through it in the loop, we only need a few registers. Let's compute: float score0=-INFINITY, score1=-INFINITY; int idx0=-1, idx1=-1; for (int e=0; e<num_experts; ++e) { float dot = 0; for (int d=0; d<hidden_dim; ++d) { dot += token_hidden[token*hidden_dim + d] * s_expert[e*hidden_dim + d]; } float val = router_logits[token*num_experts + e] + alpha * dot; // update top-2 } Here, token_hidden is reloaded each inner iteration. The compiler might keep the address in a register and load into a temporary register. So per thread: dot, val, maybe some temps, loop counters. That's maybe ~20 registers. With 20 registers, max threads per SM = 2048 (warp limit). Max blocks per SM for block size 64 = 32 (warp limit). But shared memory: 8192 bytes per block. Shared memory per SM = maybe 164KB or 100KB. So max blocks per SM = floor(100KB/8KB) ≈ 12. So shared memory is the limiting factor if it's 100KB. If it's 164KB, max blocks = 20. So blocks per SM limited by shared mem to maybe 12 or 20. But we only have 64 blocks total. If each SM can run 12 blocks, then 64 blocks / 12 ≈ 6 SMs active. Still bad. So 1 thread per token with small block size leads to low occupancy. Actually occupancy is the ratio of active warps to max warps per SM. If only 6 SMs are active, the rest are idle, so effective utilization is low. But if each block takes long enough, maybe the GPU can't hide latency anyway. However, for memory-bound kernels, we want many warps per SM to hide latency. So we need many blocks or many warps per block. Thus, using 1 thread per token is bad because the total number of warps is only 4096 / 32 = 128 warps. Max warps per SM is 64 (on Ampere). So we need at least 2 SMs? Actually 128 warps total, 64 max per SM, means we could saturate 2 SMs fully. But we have 108 SMs. So to saturate all SMs, we'd need at least 108 * some warps. But total work is small: 4096 tokens * 16 experts * 128 hidden = 8,388,608 FMAs. That's not a huge amount. On A100 (19.5 TFLOPS FP32), this is 0.4 ms at peak. But memory bandwidth will dominate. The reference will be slower because it writes intermediates. Even with low occupancy, the fused kernel could be much faster than reference. But can we do even better by increasing parallelism within each block? We can have each thread process multiple tokens sequentially. That doesn't increase block count. Alternatively, we can have each warp process one token (32 threads per token). Then total warps = 4096. Blocks can be e.g., 256 threads (8 warps) per block, processing 8 tokens. Then blocks = 512. That's excellent occupancy. 512 blocks > 108 SMs, so all SMs will be busy. Given that, I think warp-per-token or half-warp-per-token is a better design for occupancy. Let's design warp-per-token with 32 threads per token. Block size: 128 threads (4 warps). Each block handles 4 tokens. Number of blocks = 1024. That's good. Shared memory: load expert_ground [num_experts][hidden_dim] = 8KB. Also maybe a workspace for warp reduction? Not needed if using shuffles. Algorithm per warp: - token = blockIdx.x * 4 + warp_id (where warp_id = threadIdx.x / 32). - lane = threadIdx.x % 32. - Load part of token_hidden for this token: - token_hidden[token * hidden_dim + lane] and token_hidden[token * hidden_dim + lane + 32], token_hidden[token * hidden_dim + lane + 64], token_hidden[token * hidden_dim + lane + 96]. Since hidden_dim=128, lane 0..31 loads 4 elements. We can load them into registers: `float x[4];` Actually 4 floats per thread, fine. - Also load corresponding expert_ground elements for each expert? Since we need to compute dot product for each expert, and each expert has 128 elements, we can process experts sequentially or in parallel. Option 1: Each warp computes all 16 experts sequentially. - For e=0..15: - Each lane loads `s_expert[e*hidden_dim + lane]`, `s_expert[e*hidden_dim + lane+32]`, `s_expert[e*hidden_dim + lane+64]`, `s_expert[e*hidden_dim + lane+96]`. That's 4 loads per lane per expert. - Each lane computes `dot_partial = x[0]*w[0] + x[1]*w[1] + x[2]*w[2] + x[3]*w[3];` - Warp reduce sum across 32 lanes: `dot = warp_sum(dot_partial);` - All lanes get the dot product (broadcast). Then each lane computes `val = router_logits[token*num_experts + e] + alpha * dot;` - Now we need top-2 across experts. Since only one expert value per iteration, we can keep top-2 in registers per thread (all threads hold the same top-2 because they all see the same val). But to avoid divergence, each thread can independently maintain top-2. After 16 iterations, all lanes in warp have identical top-2. - After 16 iterations, compute softmax on lane 0 (or all lanes) and write output. This is straightforward. - Number of iterations: 16. - Per iteration: 4 loads from shared memory (w), 4 FMAs, warp reduce sum (5 shfls), 1 load from global memory (router_logits), 1 FMA. - Total per token: 16*4 = 64 FMAs for dot products + 16 FMAs for alpha*dot + 16 adds. Plus loads. - Warp reduce: 16 * 5 = 80 shfls per token. That's not huge. - Each warp handles one token. 4096 warps total. How long does this take? Each warp has 64 useful FMAs. That's tiny. But the whole kernel has 4096 warps * 64 FMAs = 262k FMAs? Wait no: 4096 tokens * 16 experts * 128 hidden_dim = 8,388,608 multiplications. In warp-per-token with 32 threads per token, each thread does 4 multiplications per expert, 16 experts = 64 multiplications per thread. Total threads = 4096*32 = 131072. 131072*64 = 8,388,608. Good. Can we reduce memory traffic? token_hidden is loaded once per token (128 floats). router_logits is loaded 16 times sequentially, but each load is same for all threads in warp (broadcast). Actually `router_logits[token*num_experts + e]` is same for all lanes. The compiler may cache it in registers or L1. We can have lane 0 load it and broadcast via shfl to save loads. But since it's only 16 floats per token, it's negligible. This design seems robust and simple. Let's code it. CUDA kernel details: ```cuda template <int HIDDEN_DIM, int NUM_EXPERTS> __global__ void grounded_top2_kernel( const float* __restrict__ token_hidden, const float* __restrict__ router_logits, const float* __restrict__ expert_ground, float alpha, float* __restrict__ output, int num_tokens) { // Load expert_ground into shared memory // total size = NUM_EXPERTS * HIDDEN_DIM floats __shared__ float s_expert[NUM_EXPERTS * HIDDEN_DIM]; int tid = threadIdx.x; int total = NUM_EXPERTS * HIDDEN_DIM; for (int i = tid; i < total; i += blockDim.x) { s_expert[i] = expert_ground[i]; } __syncthreads(); // Each warp handles one token int warps_per_block = blockDim.x / 32; int warp_id = tid / 32; int lane = tid % 32; int token = blockIdx.x * warps_per_block + warp_id; if (token >= num_tokens) return; // Load token_hidden for this token (HIDDEN_DIM floats) // Each lane loads 4 elements float x[4]; int base = token * HIDDEN_DIM + lane; x[0] = token_hidden[base]; x[1] = token_hidden[base + 32]; x[2] = token_hidden[base + 64]; x[3] = token_hidden[base + 96]; float top1_val = -CUDART_INF_F; // or -1e30 float top2_val = -CUDART_INF_F; int top1_idx = -1; int top2_idx = -1; // Loop over experts for (int e = 0; e < NUM_EXPERTS; ++e) { // Load expert weights for this expert float w[4]; int ebase = e * HIDDEN_DIM + lane; w[0] = s_expert[ebase]; w[1] = s_expert[ebase + 32]; w[2] = s_expert[ebase + 64]; w[3] = s_expert[ebase + 96]; float partial = x[0]*w[0] + x[1]*w[1] + x[2]*w[2] + x[3]*w[3]; // Warp reduce sum #pragma unroll for (int offset = 16; offset > 0; offset /= 2) { partial += __shfl_down_sync(0xffffffff, partial, offset); } float dot = partial; // all lanes have same dot float logit = router_logits[token * NUM_EXPERTS + e]; float val = logit + alpha * dot; // Update top-2 (all lanes do identical work) if (val > top1_val) { top2_val = top1_val; top2_idx = top1_idx; top1_val = val; top1_idx = e; } else if (val > top2_val) { top2_val = val; top2_idx = e; } } // Compute softmax over top2 // For numerical stability subtract max float max_val = top1_val; // top1 >= top2 float exp0 = expf(top1_val - max_val); float exp1 = expf(top2_val - max_val); float sum = exp0 + exp1; float w0 = exp0 / sum; float w1 = exp1 / sum; // Write output // output shape: [num_tokens, 2, 2] // For each token, out[token, 0, 0] = top1_idx, out[token, 0, 1] = w0 // out[token, 1, 0] = top2_idx, out[token, 1, 1] = w1 int out_base = token * 4; // 2*2 // Wait, careful: output shape [num_tokens, 2, 2]. So linear index: // out[token, j, k] = token*4 + j*2 + k // j=0: top1, j=1: top2 // k=0: idx, k=1: weight // So: // out[token*4 + 0*2 + 0] = (float)top1_idx; // out[token*4 + 0*2 + 1] = w0; // out[token*4 + 1*2 + 0] = (float)top2_idx; // out[token*4 + 1*2 + 1] = w1; // Alternatively, we can write with two threads to avoid divergence? But only lane 0 writes to avoid redundant stores. // Actually all lanes have same values, so we should have only one lane write. if (lane == 0) { output[out_base + 0] = (float)top1_idx; output[out_base + 1] = w0; output[out_base + 2] = (float)top2_idx; output[out_base + 3] = w1; } } ``` Wait, top-2 order: `torch.topk(grounded, k=2, dim=-1)` returns top-2 sorted in descending order. So top1_val is the largest, top2_val is second largest. `torch.softmax(top_vals, dim=-1)` applies softmax to these two values, so weights correspond to same order. So output order is correct. Potential issues: - If top2_val is -inf because all values are -inf? Not possible with normal inputs. - Numerical stability: subtract max_val which is top1_val. - For top-2 update: we need to make sure we handle ties correctly? The reference topk sorting is stable but order of equal values may depend on index. Using strict > and >= vs > can affect tie-breaking. The reference `torch.topk` with default largest=True and sorted=True returns indices in descending order of values. If values are equal, the lower index comes first? Actually, PyTorch's topk is not guaranteed stable for ties? Let's check. In PyTorch, `torch.topk` returns the k largest elements. For ties, the order of indices is not guaranteed to be stable? Actually, CUDA topk uses sorting that may not be stable for equal values. But it's deterministic for a given PyTorch version. The evaluation likely runs reference and compares with tolerance for float values, but the indices must match exactly. If there are ties, our algorithm's tie-breaking might differ. How to match reference tie-breaking exactly? We need to know how PyTorch CUDA topk breaks ties. For small sizes (16), it's likely a bitonic sort or similar. But to be safe, we should try to match the exact behavior. The reference does `torch.topk(grounded, k=2, dim=-1)`. If there are ties in the top two, the exact choice of which index is first/second depends on the sorting algorithm. However, since inputs are random, ties are extremely unlikely (probability ~0 because floats are continuous). So tie-breaking is not a concern in practice. But for correctness with any inputs, we should be careful. But to be safe, we can mimic PyTorch's topk tie-breaking: stable sort would keep lower index first for equal values. Our simple linear scan uses `>` and `>=` to decide: ``` if (val > top1_val) { shift top1 to top2 top1 = val, idx = e } else if (val > top2_val) { top2 = val, idx = e } ``` If a later value equals top1_val, it goes to `else if`, and if greater than top2_val, becomes top2. But if we want lower index to win for ties, we should use `>=` for the first comparison? Let's think: if we encounter a value equal to current top1, we might want to keep the earlier (smaller) index as top1, or maybe the smaller index as top1 and the equal older one as top2? Actually, stable topk would preserve original order. But since we process e from 0 to 15, stable topk would keep the first encountered as top1 if values equal. In our scan, if we use `>` for top1 (not `>=`), then when equal value appears later, it does not replace top1, so the earlier index stays as top1. Good. For top2, if we use `>` (not `>=`), then when a value equal to top2 appears later, it does not replace top2, so earlier index stays as top2. Is that consistent with stable sort? Usually stable sort preserves order of equal elements, so earlier index appears first. Our scan does that for top1 and top2. However, consider a case: e=0 val=1.0, e=1 val=2.0, e=2 val=2.0. Top1 should be e=1 (first 2.0), top2 should be e=0 (1.0) or e=2 (2.0)? Stable topk for k=2 would return [2.0 (e=1), 2.0 (e=2)] because those are the two largest, and stable sort would preserve order e=1 then e=2. Our scan: - e=0: top1=1.0(e=0), top2=-inf - e=1: val=2.0 > top1 -> top2=1.0(e=0), top1=2.0(e=1) - e=2: val=2.0 > top1? No (equal). > top2? Yes (2.0 > 1.0) -> top2=2.0(e=2). Result: top1=e=1, top2=e=2. That's correct for stable topk? Actually, we want top2 to be the second largest. Since e=2 has value 2.0 which is larger than current top2 1.0, it should replace top2. That's correct. So our scan with strict `>` works for stable topk in terms of values. But what if e=2 val=1.0 (equal to top2), then `> top2` fails, so top2 stays e=0. That's stable: earlier index first. Good. So using `>` (strictly greater) for both is a reasonable stable top-2. One more subtlety: the reference topk returns `top_vals, top_idx` sorted by value descending. Our scan returns top1 > top2 by value because we always compare and shift. Good. But what if top2 is never updated (e.g., num_experts=1)? Here num_experts=16, so no issue. Now, we need to ensure we handle arbitrary `num_tokens`, but hidden_dim and num_experts are fixed at compile time via template parameters for performance. The problem statement gives fixed dims, but the kernel should be general enough? We can read hidden_dim and num_experts from Python and pass them, but using template specialization for 128 and 16 is fine because the problem uses those constants. But to be safe, we can pass as arguments and unroll nothing. However, loops with known bounds allow compiler unrolling. We can pass as template arguments. Actually, `num_tokens` is variable, but hidden_dim=128 and num_experts=16 are constants from the problem, but the kernel could be used for any. We can just template them or pass as arguments. For simplicity and performance, template them. But we need to match the shapes in the C++ wrapper. In the wrapper function, we can call a templated kernel with 128 and 16. It will work for these sizes. Wait, what if the evaluator tests with different sizes? The problem statement defines `num_tokens = 4096`, `hidden_dim = 128`, `num_experts = 16`. Usually the evaluator uses these exact sizes, but we should be robust. We can make the kernel handle arbitrary hidden_dim and num_experts passed as arguments. But for performance, we want to unroll the inner loop. Since hidden_dim=128 is moderate, we can just loop. But we can also have the C++ wrapper launch a kernel that accepts them as template parameters from the torch tensor sizes. Let's see: we can get sizes from tensors and launch kernel with dynamic bounds. The compiler can unroll loops if bounds are compile-time constants. But if we template them, we need to instantiate for each possible size. That's impossible. So better to pass as runtime arguments but rely on the compiler to optimize. For 128 and 16, the loop overhead is tiny. Alternatively, we can do `if constexpr` or macro but since we compile at runtime using nvcc, we can pass as template parameters from Python. In the Python wrapper, we know the sizes from the input tensors. But the same compiled module is used for all calls. However, in PyTorch's `load_inline`, the C++ code is compiled once at module load. So we can't instantiate templates per call unless we have a switch. But we can just pass `hidden_dim` and `num_experts` as function arguments and loop. Given that the sizes are fixed, we can just use `#define` or template defaults. But to be general, I'll pass as arguments and loop. The compiler might not unroll the outer loop over 16 experts, but 16 is small. For the inner loop over hidden_dim=128, the compiler can unroll partially. It's fine. Actually, we can keep them dynamic but mark the loops with `#pragma unroll`? That requires compile-time constant trip count. If dynamic, it won't unroll. But we can unroll manually with a `while` loop? Not needed. The memory bandwidth will dominate. Loop overhead is negligible. Let's write the kernel with dynamic sizes to be general. But wait, the token_hidden loading: we assumed hidden_dim=128 and each lane loads 4 elements. If hidden_dim is dynamic, we need a loop. We can have each lane load `hidden_dim/32` elements. But we should handle cases where hidden_dim < 32 or not divisible. Let's just write a generic load loop: ```cuda int token_base = token * hidden_dim; float partial = 0.0f; for (int d = lane; d < hidden_dim; d += 32) { partial += token_hidden[token_base + d] * s_expert[e * hidden_dim + d]; } // warp reduce for (int offset = 16; offset > 0; offset /= 2) { partial += __shfl_down_sync(0xffffffff, partial, offset); } float dot = partial; ``` This works for any hidden_dim. For hidden_dim=128, each lane does 4 iterations. This avoids needing to store x[4] and w[4] in registers (the compiler might still store them). It's simpler. But we need s_expert size. For dynamic sizes, we need to allocate shared memory at launch or use a large fixed size. With `load_inline`, we can declare extern shared memory or use a fixed max size. Since we know max is 16*128=2048, we can declare `__shared__ float s_expert[2048];` and pass dynamic count. But if we want to support arbitrary sizes up to that max, we can do: ```cuda const int MAX_EXPERTS = 16; const int MAX_HIDDEN = 128; __shared__ float s_expert[MAX_EXPERTS * MAX_HIDDEN]; ``` Actually, if the evaluator might test larger sizes, we should not hardcode max. But the problem gives fixed sizes. Better to allocate shared memory dynamically: `extern __shared__ float s_expert[];` and pass shared memory size in bytes at launch: `kernel<<<grid, block, num_experts*hidden_dim*sizeof(float)>>>`. That's the standard way. But with `load_inline`, we pass the kernel in the CUDA code string. We can just use `extern __shared__ float s_expert[];`. However, each block handles multiple warps. All warps in block load s_expert. Then they use it. Let's design the kernel generically: ```cuda __global__ void grounded_top2_kernel( const float* __restrict__ token_hidden, const float* __restrict__ router_logits, const float* __restrict__ expert_ground, float alpha, float* __restrict__ output, int num_tokens, int hidden_dim, int num_experts) { extern __shared__ float s_expert[]; int tid = threadIdx.x; int total = hidden_dim * num_experts; for (int i = tid; i < total; i += blockDim.x) { s_expert[i] = expert_ground[i]; } __syncthreads(); int warps_per_block = blockDim.x >> 5; int warp_id = tid >> 5; int lane = tid & 31; int token = blockIdx.x * warps_per_block + warp_id; if (token >= num_tokens) return; float top1_val = -CUDART_INF_F; float top2_val = -CUDART_INF_F; int top1_idx = -1; int top2_idx = -1; int token_base = token * hidden_dim; int logits_base = token * num_experts; for (int e = 0; e < num_experts; ++e) { float partial = 0.0f; int ebase = e * hidden_dim; for (int d = lane; d < hidden_dim; d += 32) { partial += token_hidden[token_base + d] * s_expert[ebase + d]; } // warp sum #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { partial += __shfl_down_sync(0xffffffff, partial, offset); } float dot = partial; float val = router_logits[logits_base + e] + alpha * dot; if (val > top1_val) { top2_val = top1_val; top2_idx = top1_idx; top1_val = val; top1_idx = e; } else if (val > top2_val) { top2_val = val; top2_idx = e; } } // compute softmax float maxv = top1_val; float exp0 = expf(top1_val - maxv); float exp1 = expf(top2_val - maxv); float sum = exp0 + exp1; float w0 = exp0 / sum; float w1 = exp1 / sum; if (lane == 0) { int out_base = token * 4; output[out_base + 0] = (float)top1_idx; output[out_base + 1] = w0; output[out_base + 2] = (float)top2_idx; output[out_base + 3] = w1; } } ``` Wait, `__shfl_down_sync` requires compute capability >= 7.0? Actually it's available from Pascal and later. We'll target modern GPUs. We can include `#include <cuda_runtime.h>` and maybe `#include <math_constants.h>` for `CUDART_INF_F`. Or just use a large negative number. Let's use `-1e30f` or `-(float)1e30`. Actually `CUDART_INF_F` is defined in `math_constants.h`. To be safe, just define `const float NEG_INF = -1e30f;`. Or `-CUDART_INF_F`? Let's use `-1e30f` since topk on finite floats won't have values less than that. Potential issue: `expf(top2_val - maxv)` where maxv = top1_val >= top2_val. If top2_val is -inf (shouldn't happen), expf(-inf) = 0. If top1_val is large, expf(0) = 1. Good. Now, the C++ wrapper using `torch.utils.cpp_extension.load_inline`: We need to define a host function `grounded_top2_forward` that takes torch tensors and alpha. We need to launch the kernel. C++ wrapper: ```cpp #include <torch/extension.h> #include <cuda.h> #include <cuda_runtime.h> #include <c10/cuda/CUDAStream.h> __global__ void kernel(...) { ... } torch::Tensor grounded_top2_forward( torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, double alpha) { int num_tokens = token_hidden.size(0); int hidden_dim = token_hidden.size(1); int num_experts = router_logits.size(1); TORCH_CHECK(token_hidden.dim() == 2); TORCH_CHECK(router_logits.dim() == 2); TORCH_CHECK(expert_ground.dim() == 2); TORCH_CHECK(expert_ground.size(0) == num_experts); TORCH_CHECK(expert_ground.size(1) == hidden_dim); TORCH_CHECK(router_logits.size(0) == num_tokens); auto output = torch::empty({num_tokens, 2, 2}, token_hidden.options()); const int threads = 128; // 4 warps per block const int warps_per_block = threads / 32; int blocks = (num_tokens + warps_per_block - 1) / warps_per_block; const size_t shared_mem = hidden_dim * num_experts * sizeof(float); grounded_top2_kernel<<<blocks, threads, shared_mem, at::cuda::getCurrentCUDAStream()>>>( token_hidden.data_ptr<float>(), router_logits.data_ptr<float>(), expert_ground.data_ptr<float>(), static_cast<float>(alpha), output.data_ptr<float>(), num_tokens, hidden_dim, num_experts ); return output; } ``` Wait, if num_tokens is not a multiple of warps_per_block, the last block's extra warps must be masked: `if (token >= num_tokens) return;`. Good. Output dtype: token_hidden is float32. top_idx converted to float32. top_weights are float32. So output should be float32. `torch::empty(..., token_hidden.options())` ensures float32. We need to ensure the kernel is compiled with appropriate architecture. `load_inline` uses default flags. Should be fine. Now, we must define `ModelNew` class in Python. It accepts same init args as `Model` (none). In `forward`, it calls the compiled extension. ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline cuda_source = """ ... CUDA/C++ code ... """ cpp_sources = """ torch::Tensor grounded_top2_forward(...); """ # Actually, we need to include the full kernel code in cpp_sources or cuda_sources? # load_inline takes cpp_sources and cuda_sources. The kernel is CUDA code, so it goes in cuda_sources. # The wrapper function is C++ code that calls the CUDA kernel; it can be in cuda_sources as well, # because nvcc compiles .cu files which can contain both device and host code. # But load_inline uses nvcc for cuda_sources and host compiler for cpp_sources. # We can put everything in cuda_sources. # Then functions=['grounded_top2_forward'] ``` From PyTorch docs: ```python module = load_inline( name='grounded_top2_cuda', cpp_sources='', # can be empty cuda_sources=cuda_source, functions=['grounded_top2_forward'], extra_cuda_cflags=['-O3', '--use_fast_math'], verbose=False, ) ``` Then `ModelNew`: ```python class ModelNew(nn.Module): def __init__(self): super().__init__() def forward(self, token_hidden, router_logits, expert_ground, alpha): return grounded_top2_forward(token_hidden, router_logits, expert_ground, alpha) ``` Wait, we need to call `module.grounded_top2_forward`. So we need to store the module or access it. Since `load_inline` is at module level, we can do: ```python _grounded_top2_cuda = load_inline(...) ``` Then in forward: ```python return _grounded_top2_cuda.grounded_top2_forward(token_hidden, router_logits, expert_ground, alpha) ``` But the `load_inline` call must be at top-level. It will compile on import. That's fine. One thing: the first time it runs, it compiles. But evaluation might import once and run forward many times. That's okay. We should perhaps include `-O3` and `--use_fast_math` for performance. `--use_fast_math` makes expf faster, slight precision difference but within tolerance (1e-4). Actually, `softmax` uses expf. `--use_fast_math` might reduce precision but usually within 1e-4 for typical values. To be safe, we can omit `--use_fast_math` or include it. The tolerance is 1e-4. `expf` fast math is usually accurate to within 1-2 ULP, which for softmax weights might be okay. But to be absolutely safe with correctness, I might skip `--use_fast_math` or at least test first. We can always add it later. Let's start without it, then if we need more speed, add it. Let's write the full kernel code now. But wait, top-2 per thread with a small block size may be simpler and just as fast if we process multiple tokens per block. However, we decided warp-per-token is better for occupancy. Let's stick with warp-per-token. One potential issue: `router_logits` loading: each warp loads `router_logits[token*num_experts + e]` in a loop. Since all threads in warp load the same address, it's a broadcast. But we do it 16 times per warp. That's 16 loads per token. With 4096 tokens, that's 65536 loads. Very small. The compiler might optimize it by keeping in a register after first load? No, it's loaded each iteration. Since token_base and e are known, it will load from global memory each time. But L1 cache will likely keep the row. It's fine. Alternatively, lane 0 can load it and broadcast: ```cuda float logit; if (lane == 0) logit = router_logits[logits_base + e]; logit = __shfl_sync(0xffffffff, logit, 0); ``` This reduces redundant loads. Might help slightly. Let's do it to be clean. Actually, we can also do the same for `token_hidden`? No, each lane needs different d. Let's refine the kernel: ```cuda __global__ void grounded_top2_kernel( const float* __restrict__ token_hidden, const float* __restrict__ router_logits, const float* __restrict__ expert_ground, float alpha, float* __restrict__ output, int num_tokens, int hidden_dim, int num_experts) { extern __shared__ float s_expert[]; int tid = threadIdx.x; int total_elems = hidden_dim * num_experts; for (int i = tid; i < total_elems; i += blockDim.x) { s_expert[i] = expert_ground[i]; } __syncthreads(); const int warp_id = tid >> 5; const int lane = tid & 31; const int warps_per_block = blockDim.x >> 5; const int token = blockIdx.x * warps_per_block + warp_id; if (token >= num_tokens) return; const int token_base = token * hidden_dim; const int logits_base = token * num_experts; float top1_val = -1e30f; float top2_val = -1e30f; int top1_idx = -1; int top2_idx = -1; for (int e = 0; e < num_experts; ++e) { float partial = 0.0f; int ebase = e * hidden_dim; for (int d = lane; d < hidden_dim; d += 32) { partial += token_hidden[token_base + d] * s_expert[ebase + d]; } #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { partial += __shfl_down_sync(0xffffffff, partial, offset); } float logit; if (lane == 0) { logit = router_logits[logits_base + e]; } logit = __shfl_sync(0xffffffff, logit, 0); float val = logit + alpha * partial; if (val > top1_val) { top2_val = top1_val; top2_idx = top1_idx; top1_val = val; top1_idx = e; } else if (val > top2_val) { top2_val = val; top2_idx = e; } } float maxv = top1_val; float exp0 = expf(top1_val - maxv); float exp1 = expf(top2_val - maxv); float sum = exp0 + exp1; float w0 = exp0 / sum; float w1 = exp1 / sum; if (lane == 0) { int out_idx = token * 4; output[out_idx + 0] = (float)top1_idx; output[out_idx + 1] = w0; output[out_idx + 2] = (float)top2_idx; output[out_idx + 3] = w1; } } ``` Wait, is there any divergence in the `if (val > top1_val)` block? All threads in a warp execute the same instructions because `val` is identical for all lanes. So no divergence. Good. One more optimization: The inner loop over `d` with step 32: for hidden_dim=128, it runs 4 times. The compiler should unroll it if hidden_dim is known at compile time, but it's dynamic. However, we can add `#pragma unroll 4`? But `hidden_dim` is dynamic. The compiler doesn't know the trip count. But the step size is 32, and if hidden_dim is 128, it runs exactly 4 times. But since the bound is a variable, the compiler can't fully unroll without runtime check. But it might still partially unroll. We can help by adding a compile-time specialization if possible. But it's probably fine. Actually, we can pass hidden_dim as a template parameter to allow unrolling. Since the Python wrapper sees the actual dimensions at runtime, we can use a switch statement: ```cpp if (hidden_dim == 128 && num_experts == 16) { kernel<128, 16><<<...>>>(...); } else { kernel<<<...>>>(..., hidden_dim, num_experts); } ``` But that requires two kernels: one templated and one dynamic. Or we can just use the dynamic version. The performance difference is likely small for this memory-bound kernel. But let's consider: the inner loop runs 4 times. If unrolled, it avoids loop overhead and enables better instruction scheduling. The loop overhead is tiny compared to FMAs and loads. But we can try to add a `#pragma unroll 8` before the loop: `#pragma unroll 8` tells the compiler to unroll up to 8 iterations. Since the compiler doesn't know the bound, it might still unroll if the trip count is small? Typically `#pragma unroll` without a number means full unroll if constant. With a number, it unrolls by that factor, but the trip count must still be known? Actually `#pragma unroll 4` before a loop with dynamic bound will unroll the body 4 times and insert a runtime check for the remainder. That could help. Let's add `#pragma unroll 4` for the inner `d` loop. But wait, the step is 32, and the number of iterations is `ceil(hidden_dim/32)`. For hidden_dim=128, it's exactly 4 iterations. So `#pragma unroll 4` is appropriate. Similarly, the outer loop over `e` has 16 iterations. We can add `#pragma unroll 16` before it if num_experts is compile-time constant? It's dynamic. But we can add `#pragma unroll 16` before the outer loop to unroll up to 16 times. Since the compiler knows it's a loop variable, it will unroll and insert a dynamic check. That could be beneficial. However, 16 unrolled iterations with a shared memory load per iteration (expert weights) might increase register pressure. But since each iteration uses the same registers (`partial`), it might be okay. Let's try to include `#pragma unroll` for both loops. But `#pragma unroll 16` for the outer loop might generate a lot of code. Let's think: the outer loop body includes a `#pragma unroll 4` inner loop and a shfl reduction. That's not huge. Unrolling 16 times might be okay. But if num_experts varies, the generated code is large. However, with `num_experts=16`, it's fine. But we don't know if the evaluator uses other sizes. Usually they do not. I'd rather keep it dynamic to be safe, but we can unroll by a factor. Actually, to keep code size manageable and occupancy high, I will not unroll the outer loop. Register pressure is more important. The performance is likely limited by memory and warp scheduling, not instruction throughput. Let's also check the output layout. The reference uses `torch.stack((top_idx.to(torch.float32), top_weights), dim=-1)`. So shape is [N, 2, 2]. We need to make sure the output memory layout is correct. PyTorch stack on last dim creates a new dim at the end: for two tensors of shape [N,2], the result shape is [N,2,2]. Strides: the output is contiguous. `top_idx` is converted to float. So output is float32. Layout: for token i, expert j (0 or 1), channel k (0=idx, 1=weight). Our writes: - out[i,0,0] = top1_idx - out[i,0,1] = w0 - out[i,1,0] = top2_idx - out[i,1,1] = w1 This matches because `topk` returns sorted in descending order. So `top_vals[:,0]` is largest weight, `top_vals[:,1]` is second largest. `top_weights` correspond: `top_weights[:,0]` is weight for largest, `top_weights[:,1]` for second largest. So j=0 is top1, j=1 is top2. Good. One more thing: The reference computes `grounded = router_logits + alpha * (token_hidden @ expert_ground.t())`. The matmul is standard: [N, H] @ [H, E] = [N, E]. We compute the same dot product. Precision: We use float accumulation. The reference uses torch.matmul which accumulates in float32. Our partial sums are done per thread with 4 elements, then warp-reduced with shfl additions. This should be numerically very close. The only difference is order of summation. The reference likely uses cuBLAS which may do tiled summation or tensor cores. The difference should be within 1e-4, especially for random inputs. Tolerance is 1e-4 for fp32. We should be fine. Let's consider adding `-Xptxas -v`? Not needed. Now, let's write the Python file. ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline CUDA_SOURCE = r''' #include <torch/extension.h> #include <cuda.h> #include <cuda_runtime.h> #include <c10/cuda/CUDAStream.h> __global__ void grounded_top2_kernel( const float* __restrict__ token_hidden, const float* __restrict__ router_logits, const float* __restrict__ expert_ground, float alpha, float* __restrict__ output, int num_tokens, int hidden_dim, int num_experts) { extern __shared__ float s_expert[]; int tid = threadIdx.x; int total_elems = hidden_dim * num_experts; for (int i = tid; i < total_elems; i += blockDim.x) { s_expert[i] = expert_ground[i]; } __syncthreads(); const int warp_id = tid >> 5; const int lane = tid & 31; const int warps_per_block = blockDim.x >> 5; const int token = blockIdx.x * warps_per_block + warp_id; if (token >= num_tokens) return; const int token_base = token * hidden_dim; const int logits_base = token * num_experts; float top1_val = -1e30f; float top2_val = -1e30f; int top1_idx = -1; int top2_idx = -1; for (int e = 0; e < num_experts; ++e) { float partial = 0.0f; int ebase = e * hidden_dim; #pragma unroll 4 for (int d = lane; d < hidden_dim; d += 32) { partial += token_hidden[token_base + d] * s_expert[ebase + d]; } #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { partial += __shfl_down_sync(0xffffffff, partial, offset); } float logit; if (lane == 0) { logit = router_logits[logits_base + e]; } logit = __shfl_sync(0xffffffff, logit, 0); float val = logit + alpha * partial; if (val > top1_val) { top2_val = top1_val; top2_idx = top1_idx; top1_val = val; top1_idx = e; } else if (val > top2_val) { top2_val = val; top2_idx = e; } } float maxv = top1_val; float exp0 = expf(top1_val - maxv); float exp1 = expf(top2_val - maxv); float sum = exp0 + exp1; float w0 = exp0 / sum; float w1 = exp1 / sum; if (lane == 0) { int out_idx = token * 4; output[out_idx + 0] = (float)top1_idx; output[out_idx + 1] = w0; output[out_idx + 2] = (float)top2_idx; output[out_idx + 3] = w1; } } torch::Tensor grounded_top2_forward( torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, double alpha) { int num_tokens = token_hidden.size(0); int hidden_dim = token_hidden.size(1); int num_experts = router_logits.size(1); auto output = torch::empty({num_tokens, 2, 2}, token_hidden.options().dtype(torch::kFloat32)); const int threads = 128; const int warps_per_block = threads / 32; int blocks = (num_tokens + warps_per_block - 1) / warps_per_block; size_t shared_mem = static_cast<size_t>(hidden_dim) * num_experts * sizeof(float); grounded_top2_kernel<<<blocks, threads, shared_mem, at::cuda::getCurrentCUDAStream()>>>( token_hidden.data_ptr<float>(), router_logits.data_ptr<float>(), expert_ground.data_ptr<float>(), static_cast<float>(alpha), output.data_ptr<float>(), num_tokens, hidden_dim, num_experts ); return output; } ''' _grounded_top2_cuda = load_inline( name='grounded_top2_cuda', cpp_sources='', cuda_sources=CUDA_SOURCE, functions=['grounded_top2_forward'], extra_cuda_cflags=['-O3'], verbose=False, ) class ModelNew(nn.Module): def __init__(self): super().__init__() def forward(self, token_hidden, router_logits, expert_ground, alpha): return _grounded_top2_cuda.grounded_top2_forward( token_hidden, router_logits, expert_ground, alpha ) ``` Wait, `token_hidden.options().dtype(torch::kFloat32)` vs `token_hidden.options()`. Since token_hidden is float32, `.options()` already is float32. But top_idx.to(torch.float32) implies output is float32. So if input is fp16, output would still be fp16 if we use options. But problem says fp32. We can just use `token_hidden.options()`. Or better, explicitly `torch::kFloat32`. But using `options()` is fine. Also, in C++ we need to include `<torch/extension.h>`. In `load_inline`, it's automatically included? Usually it's okay to include it in cuda_sources. But `torch/extension.h` includes pybind11 and might conflict? Actually, in `cuda_sources`, including `<torch/extension.h>` is standard and works. However, in `load_inline`, the generated pybind11 code uses `cpp_sources`. Since our host function `grounded_top2_forward` is in `cuda_sources`, it's compiled by nvcc and then linked. The `functions=['grounded_top2_forward']` tells `load_inline` to generate a binding for this function. So it should work. One issue: `load_inline` with `cpp_sources=''` and `cuda_sources` containing the host function. The auto-generated pybind11 code might be in the cpp_sources compilation unit, and it needs the declaration of the function. Since `cpp_sources` is empty, the generated code might just reference the symbol. As long as the symbol is exported from the CUDA compilation, it should link. Usually nvcc compiles host code with external linkage. So it should work. But sometimes `load_inline` requires the function declaration in `cpp_sources`. Let's check: In examples, they put the host function in `cpp_sources` and the kernel in `cuda_sources`. Or everything in `cpp_sources` if it's `.cpp`? Actually, for CUDA, you put kernel and host wrapper in `cuda_sources`. The binding is generated in a separate `.cpp` file that includes a header with declarations? I'm not entirely sure. From PyTorch docs for `load_inline`: ``` load_inline(name, cpp_sources, cuda_sources=None, functions=None, extra_cflags=None, extra_cuda_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False) ... cpp_sources: A string, or list of strings, containing C++ source code. cuda_sources: A string, or list of strings, containing CUDA source code. functions: A list of function names for which to generate Python bindings. ``` If we put everything in `cuda_sources`, does it work? I think yes, because the generated pybind11 code just calls `m.def("grounded_top2_forward", &grounded_top2_forward, ...)` in the auto-generated `cpp_sources`. It relies on `grounded_top2_forward` being declared in `cuda_sources` and having C++ linkage. Since nvcc compiles it as a `.cu` file and then links, it should be fine. Alternatively, we can put a forward declaration in `cpp_sources`: ```cpp torch::Tensor grounded_top2_forward(torch::Tensor, torch::Tensor, torch::Tensor, double); ``` But `cpp_sources` will be compiled by the host compiler (g++), which doesn't know about `torch::Tensor` unless we include headers. `load_inline` includes `<torch/extension.h>` automatically? I think it might. But to be safe, we can just put everything in `cuda_sources`. I recall examples where everything is in `cuda_sources` and `functions` references it. So let's do that. Wait, `torch/utils/cpp_extension/load_inline.py` auto-generates a `.cpp` that includes declarations for the functions. It generates something like: ```cpp #include <torch/extension.h> extern "C" torch::Tensor grounded_top2_forward(...); PYBIND11_MODULE(...) ``` Actually, it generates `extern "C"`? I'm not sure. Let's check the source of `load_inline` mentally: It builds a `.cpp` file that calls `m.def()` for each function. To do that, it needs the function to be declared in a header or known. I think it just emits the function call directly, which implies it must see the declaration at compile time. But since the generated `.cpp` and the `.cu` are compiled separately and linked, the `.cpp` must have a declaration. I believe `load_inline` actually places the provided `cpp_sources` into the generated `.cpp`. So if we provide the declaration in `cpp_sources`, it will be in the same compilation unit as the binding. So to be safe, we can do: ```python cpp_sources = ''' torch::Tensor grounded_top2_forward(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, double alpha); ''' cuda_sources = ''' #include <torch/extension.h> #include <cuda.h> #include <cuda_runtime.h> #include <c10/cuda/CUDAStream.h> __global__ void ... { ... } torch::Tensor grounded_top2_forward(...) { ... } ''' ``` This is a common pattern. The `cpp_sources` has the forward declaration. The `cuda_sources` has the definition. The auto-generated binding code in the `.cpp` will see the declaration because `cpp_sources` is included at the top of the generated file? Actually, `load_inline` creates a separate `.cpp` for the bindings and compiles `cpp_sources` as another `.cpp` or inserts it? Let me recall: `load_inline` uses `torch.utils.cpp_extension.CppExtension` under the hood? It writes the `cpp_sources` and `cuda_sources` to files and builds them. The `cpp_sources` are compiled as `.cpp` files. The generated binding code might be in a separate `.cpp` file that also includes some auto-generated declarations. But the binding `.cpp` needs to see the function declarations. I think `load_inline` automatically extracts or includes them. Actually, let's look at the PyTorch source: In `torch/utils/cpp_extension.py`, `load_inline` calls `_write_inline` which writes sources to files. Then it creates a `CppExtension` or `CUDAExtension`. For `CUDAExtension`, the generated `main.cpp` includes the functions definitions? Wait, I think for `load_inline`, it uses `torch.utils.cpp_extension.load` with `sources` built from the inline strings. It may create a single `.cpp` or `.cu` file containing a wrapper. I found a reliable pattern online: ```python module = load_inline( name='...', cpp_sources='torch::Tensor my_fn(torch::Tensor x);', cuda_sources=''' #include <torch/extension.h> __global__ void kernel(...) {} torch::Tensor my_fn(torch::Tensor x) { ... } ''', functions=['my_fn'], ) ``` This works because the generated binding code is in a `.cpp` that includes the `cpp_sources` as declarations. Actually, the generated binding code calls `m.def("my_fn", &my_fn, ...);`. Since `cpp_sources` is compiled into a `.o`, the linker resolves it. But the `.cpp` file with `m.def` might not see the declaration, causing a compilation error. However, the generated binding code might include a forward declaration automatically: `load_inline` uses ninja to compile, and it might generate a header with `extern` declarations? Let's look up: In PyTorch 1.x, `load_inline` generates a `.cpp` file that looks like: ```cpp #include <torch/extension.h> // cpp_sources inserted here? extern "C" { torch::Tensor my_fn(torch::Tensor x); } PYBIND11_MODULE(...) { m.def("my_fn", &my_fn, ...); } ``` Actually, I think the `cpp_sources` string is placed verbatim into a `.cpp` file. Then another `.cpp` or the same `.cpp` contains the bindings. I need to check. From the PyTorch source code (approx): ```python def load_inline(name, cpp_sources, cuda_sources=None, functions=None, ...): build_directory = ... cpp_source_path = os.path.join(build_directory, 'main.cpp') with open(cpp_source_path, 'w') as f: f.write(cpp_sources) ... return load(name, sources=..., ...) ``` If `cpp_sources` is just declarations, then `main.cpp` only has declarations. The actual implementation is in a `.cu` file. The binding code is auto-generated in yet another `.cpp`? I think `load` with `functions` generates a `bindings.cpp`. But does it include the header? It might not. However, the standard practice is to put the binding manually in `cpp_sources` if using `CppExtension`. But `load_inline` is supposed to generate the binding for you. So how does it find the declaration? Actually, reading the source: `load_inline` in PyTorch creates a temporary file for `cpp_sources` and `cuda_sources`. Then it calls `load(name, sources=[cpp_source_path, cuda_source_path], extra_cflags=..., ...)`. The `functions` argument is passed to `load`. But `load` doesn't generate bindings; it just compiles the sources and returns a shared library. Wait, does `load` generate bindings? Yes, `torch.utils.cpp_extension.load` can generate Python bindings if `functions` is provided? Actually, `load` just compiles and loads a C++/CUDA extension. The extension must define `PYBIND11_MODULE`. If `functions` is provided, `load_inline` generates a `main.cpp` that defines `PYBIND11_MODULE` and calls `m.def` for each function. But for that to compile, the generated `main.cpp` must include declarations of those functions. I think `load_inline` extracts the declarations from `cpp_sources` or simply generates `extern` declarations. Let me check: I remember seeing `load_inline` usage where `cpp_sources` is empty and `cuda_sources` contains both the kernel and the host function, and `functions` lists the host function. That worked. For example: ```python custom_cpu = load_inline( name="custom_cpu", cpp_sources="""... host code ...""", functions=["custom_function"], ) ``` But for CUDA, people often do: ```python module = load_inline( name='mymodule', cpp_sources='torch::Tensor myfunc(torch::Tensor x);', cuda_sources='... implementation ...', functions=['myfunc'], ) ``` To be absolutely safe, I'll provide the forward declaration in `cpp_sources` and the full definition in `cuda_sources`. This is the recommended pattern. Let's check the exact mechanics: The auto-generated pybind11 module will be in a `.cpp` file. That file will include the `cpp_sources`? No, I think `cpp_sources` is compiled as a separate translation unit. The generated module `.cpp` will need declarations. But `load_inline` might generate `extern` declarations automatically based on the signatures? It likely uses the `cpp_sources` to figure out the signatures? I'm not sure. Wait, I recall that `load_inline` does not parse signatures. It simply generates: ```cpp #include <torch/extension.h> // If cpp_sources is a string, it might be included here verbatim? Actually no. // The generated file is just: extern "C" { // maybe nothing } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("func_name", &func_name, "doc"); } ``` If that's the case, the generated `.cpp` needs the function declaration. To provide it, we can put the declaration in `cpp_sources` and the generated file might `#include` the `cpp_sources` file? Or it might compile `cpp_sources` separately. But separate compilation means the generated `.cpp` still needs a declaration. I think `load_inline` may actually prepend the `cpp_sources` to the generated `.cpp`? Let me think. Looking at the actual implementation: There is a function `torch.utils.cpp_extension._write_inline` that writes `cpp_sources` to `main.cpp`. Then `load` is called with `sources=[main.cpp, ...]`. The `load` function creates a `CppExtension` or `CUDAExtension`. For `CUDAExtension`, the `main.cpp` is just one of the sources. The bindings are not auto-generated into `main.cpp`; rather, the user is expected to provide `PYBIND11_MODULE` in `cpp_sources`? Wait, but `load_inline` says "List the host-side wrapper function names in `functions=[...]` instead." So it explicitly says not to write `PYBIND11_MODULE`. That means `load_inline` generates the module. So it must generate a file containing `PYBIND11_MODULE`. And that generated file must see the function declarations. I found a reference: In PyTorch's `cpp_extension.py`, the `load_inline` function generates a `bindings.cpp` (or similar) that includes the declarations by parsing the `cpp_sources`? Actually Python can't parse C++. But I think it just writes `extern "C" { ... }` with the function names from `functions` and uses `decltype`? No. Maybe the generated code is: ```cpp #include <torch/extension.h> // forward declarations are generated by including a header that is built from the function names and signatures? ``` Wait, actually I think `load_inline` doesn't generate a separate binding file for `functions`. Instead, it uses `pybind11` to create bindings at runtime? No, that's `load_inline` from some other library. Let me recall: `torch.utils.cpp_extension.load` (without `functions`) requires the source to define `PYBIND11_MODULE`. `torch.utils.cpp_extension.load_inline` (with `functions`) generates the module automatically. The way it does it is by creating a `main.cpp` that includes the user's `cpp_sources` and then adds `PYBIND11_MODULE`. Actually, looking at the code: ```python def load_inline(name, cpp_sources, cuda_sources=None, functions=None, extra_cflags=None, extra_cuda_cflags=None, extra_ldflags=None, extra_include_paths=None, build_directory=None, verbose=False): ... cpp_source_path = os.path.join(build_directory, 'main.cpp') with open(cpp_source_path, 'w') as f: f.write(cuda_sources or cpp_sources) ... ``` No, that doesn't sound right. I think for `load_inline`, the `cpp_sources` is the main source. If `functions` is provided, it appends the pybind11 module code to `cpp_sources` or creates a separate file. Actually, to avoid this uncertainty, it's safer to put everything in `cpp_sources` and use the `CUDA` launcher? No, CUDA kernel must be in `.cu` compilation unit (nvcc). But we can put the CUDA kernel in a string that is compiled as `.cu` via `cuda_sources`, and put a forward declaration plus maybe a small wrapper in `cpp_sources`? But the wrapper needs to call the kernel, which requires the kernel declaration. If we put the kernel declaration in `cpp_sources` and definition in `cuda_sources`, that works. But we want the host function that does the tensor checks and launching to be bound. We can put that host function in `cpp_sources` but then it can't call the kernel because kernel is in another translation unit. However, CUDA allows host code in `.cpp` to call kernels if the kernel is declared with `cudaError_t` wrapper? Actually, calling a kernel from a `.cpp` file compiled by the host compiler is not allowed unless you use the CUDA runtime API to launch via function pointer or `cuLaunchKernel`. So it's much easier to put the host launcher in a `.cu` file compiled by nvcc. Thus, the host launcher must be in `cuda_sources`. And `functions` references it. Given the uncertainty about binding generation, I will look for a minimal example from memory: ```python custom_add = load_inline( name='custom_add', cpp_sources='', cuda_sources=''' #include <torch/extension.h> #include <cuda.h> #include <cuda_runtime.h> __global__ void add_kernel(const float* a, const float* b, float* out, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; if (i < n) out[i] = a[i] + b[i]; } torch::Tensor add_cuda(torch::Tensor a, torch::Tensor b) { auto out = torch::empty_like(a); int n = a.numel(); add_kernel<<<(n+255)/256, 256>>>(a.data_ptr<float>(), b.data_ptr<float>(), out.data_ptr<float>(), n); return out; } ''', functions=['add_cuda'], ) ``` This is a known working pattern. So we can safely put the host function in `cuda_sources` with `cpp_sources=''` and `functions=['grounded_top2_forward']`. The generated bindings will compile because the generated `.cpp` file might just have `extern "C"` or rely on the linker? Wait, in C++, you need a declaration before use. But the generated `.cpp` might `#include` the `cuda_sources`? No, it can't include a `.cu` file in host compilation. Actually, maybe `load_inline` compiles everything as a single translation unit? No, `.cpp` and `.cu` are separate. But there is a known trick: `load_inline` generates a `main.cpp` that contains: ```cpp #include <torch/extension.h> // forward declarations torch::Tensor grounded_top2_forward(torch::Tensor, torch::Tensor, torch::Tensor, double); PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("grounded_top2_forward", &grounded_top2_forward, "..."); } ``` How does it know the signature? It doesn't. It might not need the full signature if it uses some macro? Actually, it must know the signature to generate the declaration. How does it get it? Let's actually check the PyTorch source code for `load_inline` (I can approximate). In `torch/utils/cpp_extension.py`: ```python def load_inline(name, cpp_sources, cuda_sources=None, functions=None, ...): ... build_directory = ... main_source_file = os.path.join(build_directory, 'main.cpp') with open(main_source_file, 'w') as f: f.write(cpp_sources) ... if functions is not None: # Generate bindings main_source_file = _import_module_from_library(...) ``` Wait, maybe `load_inline` simply takes `cpp_sources` and `cuda_sources`, and then uses `torch.utils.cpp_extension.load` with `is_python_module=True` and `functions=functions`. `load` then generates the bindings. In `load`, there is a `_write_ninja_file` step. For `functions`, it creates a separate `.cpp` file that includes a header generated from the function signatures. But where do the signatures come from? Actually, I think `load_inline` uses `torch.utils.cpp_extension.load` but does not generate declarations. Instead, the generated `bindings.cpp` includes the user's `cpp_sources` as a header. Let me think: If `cpp_sources` is written to `main.cpp`, then `bindings.cpp` can `#include "main.cpp"`? That would be weird. Alternatively, `load_inline` builds a single shared library from `main.cpp` (from `cpp_sources`) and `cuda.cu` (from `cuda_sources`). The `PYBIND11_MODULE` is generated in a separate `.cpp` called `python_bindings.cpp` or similar. This file contains `extern` declarations. But the `extern` declarations might just be generated using template metaprogramming or `decltype`? No. Wait, maybe `load_inline` does not generate a separate file. It appends the `PYBIND11_MODULE` to the `cpp_sources` before writing it. That would make sense! It says: "Do not write your own PYBIND11_MODULE... load_inline auto-generates one". If it auto-generates it, it likely appends the module code to the `cpp_sources` string (or writes it into the same file). So the final `main.cpp` contains both the user's `cpp_sources` and the auto-generated `PYBIND11_MODULE`. Then if we put our forward declarations in `cpp_sources`, the module code will see them because they're in the same file. And the definitions in `cuda_sources` will be linked in. That's the correct mechanism. Let's verify: If `cpp_sources` is empty, and `cuda_sources` has the function definition, then the auto-generated module code appended to `cpp_sources` would not see the declaration (since it's in another TU). So that would fail. Therefore, either: 1. `load_inline` also generates forward declarations for functions listed in `functions` by parsing the code (unlikely). 2. Or it puts the auto-generated module in a file that includes the `cpp_sources` (if `cpp_sources` is a file) and also maybe includes `cuda_sources`? No. 3. Or `load_inline` actually compiles `cpp_sources` and `cuda_sources` together using `nvcc` when generating the bindings? No, binding code is host code. Actually, now I recall: In `torch.utils.cpp_extension`, `load_inline` writes the `cpp_sources` to a file. It also writes `cuda_sources` to a file. Then it invokes `load` with both files as sources. When `functions` is provided, `load` auto-generates a file with a `PYBIND11_MODULE` that calls `m.def` for each function. This generated file includes the `cpp_sources` file via `#include`. Wait, does it? Let me look at the actual code (from memory of reading it): ```python # In load_inline: with open(os.path.join(build_dir, 'main.cpp'), 'w') as source_file: source_file.write(cpp_sources) # ... if functions is not None: with open(os.path.join(build_dir, 'bindings.cpp'), 'w') as bindings_file: bindings_file.write('#include <torch/extension.h>\n') bindings_file.write('#include "main.cpp"\n') # maybe? # define PYBIND11_MODULE for f in functions: bindings_file.write('m.def("{}", &{}, ...);'.format(f,f)) ``` If it includes `main.cpp` in `bindings.cpp`, then `bindings.cpp` sees all declarations in `main.cpp`. And if `main.cpp` has forward declarations for the functions defined in `cuda_sources`, then `bindings.cpp` can call them. This makes sense! So `main.cpp` (our `cpp_sources`) should contain the forward declarations. The definitions are in `cuda_sources` and are linked later. Alternatively, some versions of PyTorch generate a header from `cpp_sources` using regex. But since C++ signatures are hard to parse, including the file is plausibly the approach. Actually, looking at the PyTorch source, there is a function `_write_ninja_file` that builds a library. For `load_inline`, I think it writes a `cpu_ops.cpp` and `cuda_ops.cu`. Then it generates `python_bindings.cpp` with `m.def()`. The `python_bindings.cpp` might include `<torch/extension.h>` and then declare the functions as `extern` by generating them from the user-provided `cpp_sources`? I'm not sure. To be completely safe, I will put the forward declaration in `cpp_sources` and the definition in `cuda_sources`. This is the standard pattern I've seen in multiple tutorials. ```python cpp_sources = ''' torch::Tensor grounded_top2_forward(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, double alpha); ''' ``` But wait, does the auto-generated binding file include `cpp_sources`? If it does not, then the binding file won't see the declaration. But the linker doesn't need declarations for linking, only for compilation. The binding file is compiled separately, so it does need declarations. Therefore, `load_inline` must somehow make the declarations available to the binding file. Usually, it does this by including `cpp_sources` in the same compilation unit as the bindings, or by parsing them. Given the instruction: "List the host-side wrapper function names in functions=[...] instead." and "Do not write your own PYBIND11_MODULE(...)", it implies that `load_inline` knows how to create the bindings. It likely parses the `cpp_sources` for function signatures? Or maybe it just generates `extern "C"` declarations for the function names and uses a generic wrapper? Hmm. Actually, maybe `load_inline` uses `torch::cpp_extension::bind` or something? No. Let's consider a different strategy: put everything in `cpp_sources` and use `aten` to launch CUDA kernel from C++? Actually, you can declare the kernel as `void kernel(...)` in `cpp_sources` and define it in `cuda_sources`. The host compiler (g++) doesn't compile the kernel launch syntax `<<< >>>` in `.cpp`. So you need nvcc for the launcher. So launcher must be in `.cu`. But there is a way: define a `void launch_kernel(...)` function in `cuda_sources` that does the launch, and declare it in `cpp_sources`. Then the host function in `cpp_sources` calls `launch_kernel`. Then we bind the host function in `cpp_sources`. Like: ```cpp // cpp_sources void launch_grounded_top2(...); // defined in cuda_sources torch::Tensor grounded_top2_forward(...) { ... launch_grounded_top2(...); return output; } ``` Then `functions=['grounded_top2_forward']`. This way, all binding is for the C++ function in `cpp_sources`, which is compiled by host compiler and only needs the declaration of `launch_grounded_top2`. The definition is in `cuda_sources`. This pattern definitely works because the host compiler sees the declaration, the binding file sees the declaration via `cpp_sources` (if included), and linking resolves it. But we still rely on the binding file seeing `cpp_sources`. If `load_inline` appends `PYBIND11_MODULE` to `cpp_sources`, then `cpp_sources` is the binding file. In that case, we can put everything there. Wait, but the host compiler cannot compile kernel launch syntax. So we cannot put `grounded_top2_forward` in `cpp_sources` if it uses `<<< >>>`. Unless we make it call an external launcher. So robust pattern: - `cpp_sources`: contains a pure C++ host function that calls an externally declared launcher. Also includes any necessary headers. - `cuda_sources`: contains the kernel and the launcher function. But if `load_inline` appends module to `cpp_sources`, then `cpp_sources` must contain the function definition (the host wrapper) so that `m.def` can reference it. And the host wrapper can call an external function declared in `cpp_sources`. Alternatively, if `load_inline` generates a separate binding file and includes `cpp_sources`, then `cpp_sources` should just have declarations and maybe the host wrapper? No, if binding file includes `cpp_sources`, then `cpp_sources` could contain the host wrapper (with external launcher declaration). That works. Given all this, I'll use the approach where the host wrapper is in `cpp_sources` with an `extern` declaration for the launcher, and the launcher+kernel are in `cuda_sources`. Wait, but the instruction says: "Do not write your own PYBIND11_MODULE(...)." This implies we don't need to worry about the module because `load_inline` handles it. It probably appends to `cpp_sources` or includes it. Let's verify by looking at a known working example from PyTorch docs. I recall: ```python custom_gather = load_inline( name="custom_gather", cpp_sources=""" at::Tensor gather_forward(at::Tensor input, at::Tensor index); """, cuda_sources=""" #include <torch/extension.h> #include <cuda.h> ... at::Tensor gather_forward(at::Tensor input, at::Tensor index) { ... gather_kernel<<<...>>>(...); return output; } """, functions=["gather_forward"], ) ``` This works because `cpp_sources` is treated as a header included by the binding code. The definition is in `cuda_sources` and linked in. The binding code calls `gather_forward` as declared in `cpp_sources`. Yes, this is the pattern. So I will do: ```python cpp_sources = ''' torch::Tensor grounded_top2_forward(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, double alpha); ''' ``` And in `cuda_sources`, define the kernel and the host wrapper function. Wait, but does `cpp_sources` need to include `<torch/extension.h>`? The binding code will likely include it. But to be safe, we can include it in `cpp_sources`. If the binding code also includes it, no harm. Let's just include it. ```python cpp_sources = ''' #include <torch/extension.h> torch::Tensor grounded_top2_forward(torch::Tensor token_hidden, torch::Tensor router_logits, torch::Tensor expert_ground, double alpha); ''' ``` Actually, we might not need `cpp_sources` at all. Some people write: ```python cpp_sources='', cuda_sources='...', functions=['grounded_top2_forward'] ``` and it works. Because the generated binding code might not need a declaration if it uses some macro magic? Or maybe it forwards declares the function with `decltype`? No, `decltype` needs an expression. Wait, `load_inline` might generate: ```cpp #include <torch/extension.h> torch::Tensor grounded_top2_forward(torch::Tensor, torch::Tensor, torch::Tensor, double); ``` by using `torch::cpp_function`? Actually no, it can't know the signature. Let's just search memory: In `torch.utils.cpp_extension.load_inline`, the generated code includes `#include <torch/extension.h>` and then a `PYBIND11_MODULE`. The user's `cpp_sources` is compiled as a separate translation unit and then linked. The binding file is generated automatically. How does it know the signatures? I think it might not. Actually, maybe `functions` is just a list of function names, and the generated code uses `pybind11::cpp_function` with some template deduction? Wait, `m.def("name", &name)` requires the function to be declared. If the binding file is separate, it needs a declaration. PyTorch's `CppExtension` has a mechanism to include `cpp_sources` in the binding file? I think it might simply `#include` your `cpp_sources` in the generated binding file. That would explain why `cpp_sources` can be just declarations: the binding file includes them and links against the definitions in `cuda_sources`. Yes, that makes sense: the generated binding file starts with `#include <torch/extension.h>` and then `#include "main.cpp"` (your `cpp_sources`), then defines `PYBIND11_MODULE`. So any declarations in `cpp_sources` are visible. The definitions are in `cuda_sources` (compiled by nvcc) and linked. So our plan is correct. Wait, but if the binding file `#include`s `main.cpp`, and `main.cpp` contains only a forward declaration, then the binding file will have the declaration. The binding file then calls `m.def(...)`. The linker links the binding `.o` with `main.o` and `cuda.o`. Wait, `main.o` will contain nothing if `main.cpp` only has a declaration. But `cuda.o` contains the definition. So linking works. But what if `cpp_sources` is empty? Then the generated binding file includes nothing. It tries to reference the function name but the function is only declared in `cuda_sources` (a `.cu` file). The linker would still resolve it because C++ linkage is by name, but the compilation of the binding `.cpp` would fail because the function is not declared in that translation unit. So `cpp_sources` must contain the declaration. Unless `load_inline` parses the `cuda_sources` and extracts declarations? I doubt it. Actually, some versions of `load_inline` do parse signatures! There is a function in PyTorch that uses regex to extract function declarations from sources. But it's fragile. To be safe, we should provide the declaration in `cpp_sources`. So final approach: `cpp_sources` contains the forward declaration. `cuda_sources` contains the full implementation. But wait, if `cpp_sources` is included in the binding file, and the binding file is compiled as `.cpp`, then `cpp_sources` cannot contain CUDA-specific syntax (like `__global__`). So our forward declaration is safe. Now, let's also consider if we can improve performance further. Given our kernel: - Each warp processes 1 token. - Block size 128 => 4 warps => 4 tokens per block. - Blocks = ceil(4096 / 4) = 1024 blocks. - Shared memory per block = 8KB. - Registers per thread: partial (1), top1/top2 values (4), loop vars (3). Maybe ~12 registers. So max blocks per SM limited by shared memory. With 164KB shared memory, 164/8 = 20 blocks per SM. 1024 blocks / 20 = ~52 SMs used. That's good! Actually 52 SMs active out of 108 means about 48% utilization? Wait, if each SM can hold 20 blocks, then we can run 20 blocks on each of 52 SMs, but there are 108 SMs. The scheduler distributes blocks across all SMs. As long as there are enough blocks (1024 > 108), all SMs will get some blocks. Since each SM can hold multiple blocks concurrently (up to 20 due to shared mem), the actual concurrency is 108 SMs * up to 20 blocks = 2160 blocks capacity. We have 1024 blocks. So we can have 1024/108 ≈ 9.5 blocks per SM. That's well within the 20 limit. So occupancy is high. Great. Wait, but warps per block = 4. Max warps per SM = 64. So blocks per SM limited by warps = 64/4 = 16. Shared memory limit = 20. Register limit = maybe 64 (assuming low reg count). So blocks per SM = 16. Total capacity = 108*16 = 1728 blocks. We have 1024. So ~9.5 blocks per SM. Good. Latency hiding: 9.5 blocks per SM means 9.5*4 = 38 warps per SM. 38 warps out of 64 is about 59% occupancy. That's decent. Can we improve by increasing warps per block? If block size = 256 (8 warps), blocks = 512. Blocks per SM = 512/108 ≈ 4.7. Warps per SM = 4.7*8 ≈ 37.6. Similar occupancy. Block size 128 might be slightly better due to more blocks. If block size = 64 (2 warps), blocks = 2048. Blocks per SM ≈ 19. Warps per SM ≈ 38. Same. But smaller blocks have some overhead. Block size = 32 (1 warp), blocks = 4096. Blocks per SM ≈ 38. But shared memory of 8KB and 1 warp per block: max blocks per SM by warps = 64. Shared memory limit = 20. So blocks per SM = 20. Total blocks capacity = 2160. We have 4096, so all SMs can run 20 blocks. Warps per SM = 20. That's 31% occupancy. Worse than 128 block size because warps per block is tiny. There is overhead to launching blocks. So 128 is good. Thus block size 128 seems good. Can we increase block size to 256 and have more shared memory reusability? 8KB is small, so no issue. Actually, we can also consider caching router_logits in shared memory per token? No, it's 16 floats per token, tiny. Another optimization: Instead of having each warp do a warp reduction for each expert, we could have each warp compute multiple experts in parallel by having each thread hold one expert's partial and then reduce. But that only reduces the outer loop iterations from 16 to 1, but requires storing 16 partials and doing 16 reductions. The current design with 16 iterations and a single partial is register-light and instruction-light. Since memory bandwidth is the bottleneck, probably not much gain. One thing: `token_hidden` is accessed in the inner loop: `token_hidden[token_base + d]`. For a given warp, threads access `d = lane, lane+32, lane+64, lane+96`. But adjacent warps in the same block access different tokens. So the access pattern for `token_hidden` across the block: thread 0 accesses token0_base+0, thread 1 accesses token0_base+1, ..., thread 31 accesses token0_base+31. Thread 32 accesses token1_base+0, etc. This is perfectly coalesced across the block. So global memory loads are efficient. `expert_ground` is loaded into shared memory in the first loop. This is coalesced as well. `s_expert[ebase + d]` access: within a warp, lane accesses `ebase + lane`, `ebase + lane + 32`, etc. Shared memory banks: on modern GPUs, 32 banks, 4 bytes each. Adjacent threads access adjacent addresses (stride 1), so no bank conflicts. For the second iteration (d += 32), threads access `ebase + lane + 32`. Again stride 1 within the warp. Perfect. `router_logits` access: broadcast from lane 0. `output` write: lane 0 writes 4 floats. Coalescing doesn't matter because it's just 4 floats per token. Memory traffic analysis: - Load expert_ground once per block: 8KB / block. Amortized per token: 8KB / 4 = 2KB per token. Wait, 2048 floats * 4 bytes = 8192 bytes per block. 4 tokens per block => 2048 bytes per token. That's a lot! But it's only loaded once from global memory and reused across all threads in the block. Global memory bandwidth is consumed for 2KB per token. However, the reference implementation loads expert_ground in the matmul too. The difference is that in our kernel, it's loaded into shared memory and then read many times from shared memory (fast). Global memory read of expert_ground is unavoidable. The reference does `token_hidden @ expert_ground.t()`, which also reads expert_ground from global memory. So global memory traffic for expert_ground is comparable. - Load token_hidden: 128 floats = 512 bytes per token. - Load router_logits: 16 floats = 64 bytes per token. - Write output: 4 floats = 16 bytes. - Intermediate traffic avoided: reference writes [4096,16] grounded matrix (256 KB), reads it for topk and softmax, writes top_vals [4096,2] (32 KB), writes top_weights [4096,2] (32 KB), writes output stack [4096,2,2] (64 KB). Our kernel has no intermediate global writes. So total global reads per token: 512 + 64 + 2048 = 2624 bytes. Writes: 16 bytes. Reference reads per token: token_hidden 512, router_logits 64, expert_ground 2048 (in matmul), plus reads/writes of grounded 256K (63 bytes per token amortized? Wait, grounded is [4096,16] = 262144 bytes. Total for all tokens. So per token: 512 + 64 + 2048 + 64 (grounded read by topk) = 2688 bytes. Plus writes: grounded 64 bytes, top_vals 8 bytes, top_weights 8 bytes, output 16 bytes. But topk and softmax are separate kernels with their own reads. The fused approach saves the writes and reads of grounded, top_vals, top_weights. So we save ~64+8+8+64+8+8 = 160 bytes of traffic per token? Actually, reference: 1. Matmul: reads token_hidden (512/token), expert_ground (2048/token once? Actually it's a regular GEMM, so expert_ground is read multiple times by the tiled algorithm, maybe more than once from HBM if not cached). Writes grounded (64/token). 2. Add: reads router_logits (64/token), reads grounded (64/token), writes grounded (64/token). Since add is fused into matmul in cuBLAS? Actually, `router_logits + alpha * (token_hidden @ expert_ground.t())` might be computed as `torch.matmul` then `torch.add`. PyTorch might fuse them? Actually, `torch.matmul` returns a tensor, then `+` is an add operation. If matmul and add are not fused (unless using torch.compile, but reference is plain PyTorch), then yes, separate kernel for add. But wait, `alpha * (matmul)` is a scale, and then addition. Scale might be fused with add? Actually, `a + b * scalar` might be fused into one kernel `addcmul` or similar? In PyTorch, `router_logits + alpha * mm_result` does: mm_result = matmul(...); scaled = alpha * mm_result (maybe fused as a single elementwise kernel? Actually `torch.mul` may be fused with `torch.add` if using JIT, but not in eager mode). So there might be multiple intermediate tensors. 3. topk: reads grounded (64/token) and produces top_vals and top_idx. 4. softmax: reads top_vals (8/token) and writes top_weights (8/token). 5. stack: reads top_idx (8/token) and top_weights (8/token) and writes output (16/token). So reference total HBM traffic per token is high: maybe 512 + 2048 + 64*3 + 8*2 + 8*2 + 16 = ... Actually, let's count precisely (assuming no fusion): - Matmul: read token_hidden (512), read expert_ground (2048), write mm_out (64). - Scale+Add (alpha*mm_out + router_logits): read mm_out (64), read router_logits (64), write grounded (64). - Topk: read grounded (64), write top_vals (8), write top_idx (8) (as int64 maybe 16 bytes? Actually top_idx is int64 by default for torch.topk? No, torch.topk returns LongTensor. But is it converted to float32 in stack. Let's assume int64: 16 bytes). - Softmax: read top_vals (8), write top_weights (8). - Stack: read top_idx (16), read top_weights (8), write output (16). Total read: 512 + 2048 + 64 + 64 + 64 + 8 + 8 + 16 + 8 = 2792 bytes. Total write: 64 + 64 + 8 + 16 + 8 + 16 = 176 bytes. Our kernel read: 512 + 2048 + 64 = 2624 bytes. Write: 16 bytes. We saved ~168 bytes of write and ~168 bytes of read per token. That's about 6% reduction? Wait, 2792 vs 2624 is 168 bytes difference. That's small because expert_ground dominates traffic. But we also avoid kernel launch overhead and avoid writing/reading the 64-byte grounded tensor. 168 bytes is about 6% of 2792. Not huge in terms of bandwidth. But we also avoid launching 5 kernels and their associated overhead. For small problem size (4096 tokens), kernel launch overhead is significant. Also, the reference matmul for [4096,128] @ [128,16] is a small GEMM. cuBLAS might have overhead for such a small matrix. In fact, 4096 independent dot products of length 128 is more like a batched dot product. Our kernel directly computes them. So our kernel should be significantly faster. But we can do even better. Is there a way to avoid reading expert_ground from global memory entirely? No, it's an input. But we can keep it in L2 cache across blocks. Since it's only 8KB, it will fit entirely in L2 and likely L1. So after the first block loads it, subsequent blocks will hit L2. The shared memory load per block is still needed. Actually, we can just load expert_ground into shared memory once per block. If multiple blocks run on the same SM sequentially, they each reload from L2 to shared. That's fine. Another optimization: We can process 2 tokens per warp to reduce the number of warps and increase work per warp? No, that would reduce block count. Keep as is. Another thought: We can compute dot products using vector loads. For example, load `float4` from `token_hidden` and `expert_ground`. This increases memory throughput. Since hidden_dim=128 and each thread loads 4 elements, we can load a `float4` (4 floats aligned) in one instruction. The addresses `token_base + lane`, `token_base + lane + 32`, etc. Are they aligned for `float4`? `token_base` is a multiple of 128 floats = 512 bytes. `lane` is 0..31. `token_base + lane` is aligned to 4 bytes (since `lane` is integer). `float4` requires 16-byte alignment. `lane` must be multiple of 4 for alignment. Since lane varies across threads, each thread loads its own `float4` at address `token_base + lane*4`. If we reorganize data so that each thread loads a contiguous chunk of 4 floats, that is `float4` aligned if `lane*16` is aligned. Since `token_base` is 512-byte aligned, `token_base + lane*16` is 16-byte aligned for any lane. Yes! Because `token_base % 16 == 0` (since 512 % 16 = 0), and `lane*16 % 16 == 0`. So we can have each thread load a `float4` from `token_hidden[token_base + lane*4]`. Wait, that means each thread loads 4 consecutive floats. Then for expert_ground, each thread loads `expert_ground[e*hidden_dim + lane*4]` also as `float4`. This is great because: - Total loads per warp: 128/4 = 32 float4s for token_hidden, and 32 float4s for expert_ground per expert. - This uses vectorized loads (LDG.128), which are more efficient. - The data layout is naturally row-major: `token_hidden[token, d]` is contiguous in d. So `token_hidden[token_base + lane*4 + 0..3]` is contiguous. - For a warp, thread 0 loads d=0..3, thread 1 loads d=4..7, ..., thread 31 loads d=124..127. Perfect coalescing and alignment. - The inner loop over d disappears; it's just one load per thread. Actually, hidden_dim=128, each thread loads 4 floats = one `float4`. So exactly one load per token per thread for token_hidden. And one load per expert per thread for expert_ground. This reduces instructions and improves memory throughput. We should do this. Implementation with `float4`: ```cuda float4 x = reinterpret_cast<const float4*>(token_hidden + token_base)[lane]; float4 w = reinterpret_cast<const float4*>(s_expert + e * hidden_dim)[lane]; float partial = x.x * w.x + x.y * w.y + x.z * w.z + x.w * w.w; ``` But we need to ensure the pointers are aligned. `token_hidden` is allocated by PyTorch, which ensures 256-byte alignment for CUDA tensors. `s_expert` is shared memory; its start address is aligned to at least 16 bytes if we request it? In CUDA, `extern __shared__ float s_expert[];` start is aligned to 16 bytes? Actually, shared memory is aligned to 16 bytes for `float4` memory operations. So it's fine. However, using `float4` in shared memory might require `s_expert` to be declared as `float` and cast to `float4*`. The address `s_expert + e*hidden_dim` must be aligned to 16 bytes. `e*hidden_dim` is a multiple of 128 floats = 512 bytes, which is divisible by 16. So `s_expert + e*hidden_dim` is 16-byte aligned if `s_expert` is. Good. But wait, `s_expert` is `extern __shared__ float s_expert[];`. The base address is aligned to 16 bytes only if the shared memory size is requested appropriately? Actually, shared memory arrays start at a 16-byte boundary for compute capability >= 2.0. So it's fine. Let's use `float4`. But note: `__shfl_down_sync` with `float4` is not needed because we reduce the `partial` float. So code becomes: ```cuda int vec_lane = lane; // 0..31 int token_vec_base = token * hidden_dim / 4; // not needed. Better: const float4* token_hidden_vec = reinterpret_cast<const float4*>(token_hidden + token_base); float4 x = token_hidden_vec[vec_lane]; for (int e=0; e<num_experts; ++e) { const float4* w_vec = reinterpret_cast<const float4*>(s_expert + e * hidden_dim); float4 w = w_vec[vec_lane]; float partial = x.x*w.x + x.y*w.y + x.z*w.z + x.w*w.w; // warp reduce sum #pragma unroll for (int offset = 16; offset > 0; offset >>= 1) { partial += __shfl_down_sync(0xffffffff, partial, offset); } ... } ``` Wait, we need to load `token_hidden` per token. But token_hidden is in global memory. Using `float4` load means each thread reads 16 bytes. For a warp: 32*16 = 512 bytes. Exactly one cache line per warp. Good. For `s_expert` per expert: same pattern. One `float4` load per thread per expert. 32*16 = 512 bytes per warp per expert. Perfect. This is very efficient. One potential issue: The compiler may not like aliasing with `float4*` if the original pointer is not `float4*`. We can use `__ldg` for global memory? Actually `reinterpret_cast` is fine. But we might want to use `const float4*` but ensure `__restrict__`. Note: If `hidden_dim` is not a multiple of 4, vector loads could read out of bounds. But `hidden_dim=128`, so it's fine. For generality, we can handle remainder, but not needed. We can assert `hidden_dim % 4 == 0` in the host wrapper. For safety. Another optimization: In the warp reduction, we can use `__shfl_xor_sync`? `__shfl_down_sync` is fine. What about using `warpReduceSum` with `__shfl_down_sync`? That's what we have. Now, there is one more thing: topk tie-breaking. As discussed, ties are unlikely. Let's also consider using `__expf` (intrinsic) instead of `expf`. `expf` is standard. With `--use_fast_math`, `expf` maps to `__expf`. Without it, it's accurate but slower. For fp32 tolerance 1e-4, we might be okay with fast math, but let's first try without. Wait, is there any chance of numerical mismatch due to different summation order? The reference does matmul via cuBLAS, which may use TF32 on Ampere if allowed. Actually, TF32 is default for matmul on Ampere if the input is float32. Wait, PyTorch's matmul uses TF32 by default on Ampere GPUs for float32 tensors? Yes, since PyTorch 1.7, `torch.backends.cuda.matmul.allow_tf32` defaults to True on Ampere. This means cuBLAS uses TF32 tensor cores, which have lower precision (10-bit mantissa) compared to FP32 (23-bit). This can lead to differences in the matmul result. Our kernel does FP32 FMAs in registers/shuffle, no TF32. So our result will be more accurate than the reference if TF32 is enabled. But the tolerance is 1e-4. Will the TF32 difference exceed that? For random inputs of size 128, the error from TF32 is usually small, maybe ~1e-5 relative. But absolute tolerance 1e-4 might be okay. However, the evaluator runs the reference and our kernel on the same GPU, same PyTorch settings. The reference will use TF32 if enabled. Our kernel won't. So there may be a small discrepancy. To ensure correctness, we might want to disable TF32 for the reference? We can't modify the reference. But we can check what the tolerance allows. The tolerance is relative and absolute: `atol=rtol=1e-4` for fp32. That means `abs(a-b) <= 1e-4 + 1e-4 * abs(b)`. For values around 1.0, tolerance is 2e-4. TF32 error is usually much smaller than that for dot products of length 128. The maximum rounding error per FMA is about 2^-10 relative for TF32 vs 2^-23 for FP32. Summing 128 terms, worst-case error is maybe 128 * 2^-10 * product magnitude ≈ 0.125 * magnitude. For random values ~1, that could be ~0.1? Actually, typical error is much smaller due to random signs canceling. But worst-case could be large. However, for random inputs, it's very unlikely to exceed 1e-4. Still, to be safe, maybe we can make our kernel emulate TF32? That's complex and not necessary. I think FP32 will be within tolerance. Wait, what if the evaluator uses `torch.backends.cuda.matmul.allow_tf32 = False`? We don't know. But if they set it to False, then reference uses FP32 too, and our results should match very closely. The order of summation in cuBLAS may differ from ours, but the difference should be tiny (1e-6 typically). So we're safe. Let's also think about `top_idx` type. The reference returns `top_idx.to(torch.float32)`. In PyTorch, topk on CUDA returns LongTensor (int64). We cast to float32. Our kernel uses `int` for indices. Since `num_experts=16`, `int` is fine. We cast to `float` for output. One subtlety: The reference topk returns indices as int64. When converted to float32, values like 15 are exactly representable. So no precision issue. Now, we need to compile and test. Since the environment is provided by the evaluator, we must ensure the code compiles with PyTorch's nvcc. We should avoid using very new CUDA features. `__shfl_down_sync` is available for compute capability >= 3.0? Actually `__shfl_down_sync` was introduced in CUDA 9.0 for all architectures? It's a PTX instruction that requires sm_30 or higher. Modern PyTorch supports it. Use it. `reinterpret_cast` of shared memory to float4: ensure alignment. We can add `__align__(16)` to shared memory? For `extern __shared__`, we can't add alignment attribute directly. But the base is aligned to 16 bytes. We ensure `e*hidden_dim` is a multiple of 4 floats (16 bytes). Since `hidden_dim % 4 == 0`, it is. Wait, if `hidden_dim` is dynamic and odd, alignment fails. But we can add a check in host code: `TORCH_CHECK(hidden_dim % 4 == 0, "hidden_dim must be multiple of 4");`. This is fine for the given problem. Alternatively, we can handle non-multiple-of-4 generically by not using float4. But since hidden_dim is 128, we can just assert. Let's write the final code with the float4 optimization. But before finalizing, let's think about the `extern __shared__` usage. In the kernel, we need the size of shared memory per block. We pass it at launch: `hidden_dim * num_experts * sizeof(float)`. Since `hidden_dim=128`, `num_experts=16`, it's 8192 bytes. That's less than the default shared memory per block (48KB or more). So no need to request more dynamic shared memory. Good. What if we don't use `extern __shared__` but instead allocate a fixed size `__shared__ float s_expert[2048];`? 2048 floats = 8192 bytes. That avoids passing shared memory size and dynamic indexing issues. But if hidden_dim*num_experts > 2048, we would overflow. We can allocate a large enough fixed size, e.g., 4096, to be safe. But since sizes are known, 2048 is enough. Using fixed size allows the compiler to know the size and maybe optimize. But for simplicity, dynamic is fine. Actually, with fixed size we can avoid the launch parameter. But dynamic is fine. However, there is a subtle issue: If we allocate fixed size inside the kernel, it goes into static shared memory, which counts against the 48KB limit (or 96KB depending on config). We are fine. But if we use `extern __shared__`, the size is dynamic and passed at launch. Also fine. I'll use `extern __shared__ float s_expert[];` and pass size at launch. This is generic. Now, let's think about compile flags. We want `-O3`. We might want `-lineinfo` for debugging but not needed. `extra_cuda_cflags=['-O3']`. Also, we should consider adding `--ftz=true --prec-div=false --prec-sqrt=false`? That's part of fast math. Let's avoid for correctness. Wait, what if the evaluator runs on a GPU with compute capability that doesn't support `__shfl_down_sync`? PyTorch's minimum is sm_50? Actually, `__shfl_down_sync` requires sm_30+. PyTorch 2.0 supports down to sm_50? Shfl is supported on sm_50. So fine. Now, potential correctness issue: The reference `torch.topk` with `k=2` on a tensor of shape [N, 16] returns the top 2 sorted. Our algorithm finds top 1 and top 2 using a linear scan. Since we process e from 0 to num_experts-1, and use strict `>` for replacement, if there are ties where a later element equals the current top1, it will not replace it. This preserves the earlier index as top1, which is stable. For top2, if a later element equals current top2, it will not replace it. So we keep the earliest index among ties. Is this what PyTorch does? PyTorch's CUDA topk implementation uses a bitonic sort or radix sort. For equal values, the sort is not guaranteed to be stable. However, for small sizes like 16, the implementation might use a sorting network that is not stable. So ties could be broken arbitrarily depending on the algorithm, not necessarily preserving original order. Given that inputs are random, ties are extremely unlikely. But to be absolutely safe, we could match PyTorch's exact tie-breaking. But we don't know which CUDA topk algorithm is used. However, `torch.topk` uses a sorting-based approach for small k and small dimensions? Actually, it uses a radix sort for large dimensions and a bitonic sort for small. For 16 elements, it might use a warp-level bitonic sort. In a bitonic sort, the relative order of equal elements depends on the comparator. A typical comparator is `>` (not `>=`) for bitonic sort, which makes it unstable. Wait, bitonic sort with `>` is stable? No, bitonic sort is not stable in general. But for a sorting network, stability depends on the comparator. If comparator uses `>` instead of `>=`, equal elements are not swapped, preserving their relative order? Actually, in a sorting network, if comparator is "if a > b then swap", then when a == b, no swap occurs. Since the network routes elements through fixed paths, the relative order of equal elements is preserved because they are never swapped. So sorting networks with strict `>` are stable! Wait, is that true? Consider a sorting network with compare-swap(a,b): if a > b swap. If a == b, no swap. Then the order of equal elements is whatever order they entered the network. However, the network shuffles elements through wires, so two equal elements might take different paths? In a sorting network, each comparator connects two specific wires. If a == b, no swap, so the element on the top wire stays on top. Thus relative order is preserved if they are on different wires. But if they start on different wires, they will end up sorted according to the network's wire positions. Actually, stability means that if two elements are equal, their original order is preserved. In a sorting network with `>`, equal elements never swap, so the element that started on the upper wire will remain on the upper wire at each comparator they meet. Since the network is deterministic, the final order of equal elements is determined by their initial positions. This is indeed stable with respect to the initial array order, because no swap occurs for equal elements. Wait, but does PyTorch's topk use `>` or `>=`? If it uses `>=`, then equal elements can be swapped, making it unstable. I'm not sure. But I suspect it uses `>` for the value and handles index tie-breaking separately? Actually, the comparator for topk might be `val > other.val || (val == other.val && idx < other.idx)` to ensure deterministic order. To know for sure is hard. But again, ties are random floats. The probability of exact equality is essentially zero. So we can ignore this. Let's move on to code. One more thing: In the C++ host wrapper, we need to get `num_tokens`, `hidden_dim`, `num_experts`. We must verify shapes. `expert_ground` shape: [num_experts, hidden_dim]. `router_logits` shape: [num_tokens, num_experts]. Everything should be contiguous? We should handle non-contiguous tensors? The reference uses standard matmul which handles strides. If inputs are not contiguous, our kernel will fail because we use pointer arithmetic assuming contiguous. The evaluator likely provides contiguous tensors from `torch.randn`. But to be safe, we can call `.contiguous()` in Python before passing. Actually, we can do: ```python token_hidden = token_hidden.contiguous() router_logits = router_logits.contiguous() expert_ground = expert_ground.contiguous() ``` before calling the C++ function. That adds negligible overhead. Alternatively, in C++ we can check strides and call `.contiguous()` there. Let's just do it in Python to keep C++ simple. Also, we need to ensure `token_hidden` is float32. It's generated as float32. What about `alpha`? It's a Python float. We pass it as double to C++ and cast to float. Now, let's write the code. I will also include `TORCH_CHECK` for shape and contiguity? Or just call `.contiguous()` in Python. Let's do that. Potential issue: The compiled extension module name must be unique. `_grounded_top2_cuda` is fine. Let's assemble the final `kernel_code` string. Wait, we need to call `compile_kernel` first. Let's do that. But before writing, let's double-check the output indices. The reference does: `torch.stack((top_idx.to(torch.float32), top_weights), dim=-1)` `top_idx` shape: [num_tokens, 2] `top_weights` shape: [num_tokens, 2] Stack on dim=-1 gives shape [num_tokens, 2, 2]. The last dimension index 0 corresponds to top_idx, index 1 to top_weights. So `output[token, j, 0] = top_idx[token, j]` `output[token, j, 1] = top_weights[token, j]` For j=0, it's the top-1 expert (largest value). For j=1, it's the top-2 expert (second largest value). Our kernel writes: `out[token*4 + 0] = top1_idx` => `output[token, 0, 0]` `out[token*4 + 1] = w0` => `output[token, 0, 1]` `out[token*4 + 2] = top2_idx` => `output[token, 1, 0]` `out[token*4 + 3] = w1` => `output[token, 1, 1]` Yes, correct. Now, let's consider if topk returns indices in descending order. Yes. Now, let's write the code and compile. We'll need to use the tools. But first, we need to write a complete Python file. We can test with compile_kernel.