Now let me write an initial CUDA kernel for RoPE. Let me analyze the reference implementation: 1. `_build_cos_sin`: Creates cos/sin values based on position - `t = torch.arange(seq_len)` - position indices - `freqs = torch.outer(t, inv_freq)` - outer product - `emb = torch.cat([freqs, freqs], dim=-1)` - duplicate along last dim - Return `emb.cos()` and `emb.sin()` 2. `_rotate_half`: Split last dim in half, swap and negate - `x1, x2 = x.chunk(2, dim=-1)` - split into two halves - Return `torch.cat([-x2, x1], dim=-1)` - concatenate with negation 3. `forward`: - `q_rot = q * cos + rotate_half(q) * sin` - Same for k For the CUDA kernel, I'll: 1. Pre-compute cos/sin on host or in a separate kernel 2. Apply RoPE transformation in parallel for each element Let me write the kernel: ```python import torch import torch.nn as nn from torch.utils.cpp_extension import load_inline # CUDA kernel cuda_source = """ #include <cuda.h> #include <cuda_runtime.h> #include <math.h> __global__ void rope_kernel( const float* __restrict__ input, const float* __restrict__ cos, const float* __restrict__ sin, float* __restrict__ output, int batch_size, int num_heads, int seq_len, int dim ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = batch_size * num_heads * seq_len * dim; if (idx < total) { // Compute indices int d = idx % dim; // dimension index int n = (idx / dim) % seq_len; // sequence position int h = (idx / (dim * seq_len)) % num_heads; int b = idx / (dim * seq_len * num_heads); float val = input[idx]; float c = cos[n * dim + d]; float s = sin[n * dim + d]; // rotate_half logic int half_dim = dim / 2; int pair_idx = (d < half_dim) ? (d + half_dim) : (d - half_dim); float pair_val = input[b * num_heads * seq_len * dim + h * seq_len * dim + n * dim + pair_idx]; if (d < half_dim) { // x1 part: -x2 * sin + x1 * cos output[idx] = val * c - pair_val * s; } else { // x2 part: x1 * sin + x2 * cos output[idx] = pair_val * s + val * c; } } } // Combined kernel for building cos/sin and applying RoPE __global__ void rope_build_and_apply_kernel( const float* __restrict__ input, const float* __restrict__ inv_freq, float* __restrict__ output, int batch_size, int num_heads, int seq_len, int dim ) { int idx = blockIdx.x * blockDim.x + threadIdx.x; int total = batch_size * num_heads * seq_len * dim; if (idx < total) { int d = idx % dim; int n = (idx / dim) % seq_len; int h = (idx / (dim * seq_len)) % num_heads; int b = idx / (dim * seq_len * num_heads); int half_dim = dim / 2; int freq_idx = d % half_dim; // Which inv_freq to use // Compute freq = n * inv_freq[freq_idx] float freq = n * inv_freq[freq_idx]; // cos and sin values float c = cosf(freq); float s = sinf(freq); float val = input[idx]; int pair_idx = (d < half_dim) ? (d + half_dim) : (d - half_dim); float pair_val = input[b * num_heads * seq_len * dim + h * seq_len * dim + n * dim + pair_idx]; if (d < half_dim) { output[idx] = val * c - pair_val * s; } else { output[idx] = pair_val * s + val * c; } } } std::vector<torch::Tensor> rope_cuda( torch::Tensor q, torch::Tensor k, torch::Tensor inv_freq ) { int batch_size = q.size(0); int num_heads = q.size(1); int seq_len = q.size(2); int dim = q.size(3); auto q_out = torch::empty_like(q); auto k_out = torch::empty_like(k); int total = batch_size * num_heads * seq_len * dim; int threads = 256; int blocks = (total + threads - 1) / threads; rope_build_and_apply_kernel<<<blocks, threads>>>( q.data_ptr<float>(), inv_freq.data_ptr<float>(), q_out.data_ptr<float>(), batch_size, num_heads, seq_len, dim ); rope_build_and_apply_kernel<<<blocks, threads>>>( k.data_ptr<float>(), inv_freq.data_ptr<float>(), k_out.data_ptr<float>(), batch_size, num_heads, seq_len, dim ); return {q_out, k_out}; } """ rope_extension = load_inline( name='rope_extension', cpp_sources='', cuda_sources=cuda_source, functions=['rope_cuda'], extra_cuda_cflags=['-O3'], verbose=False ) class ModelNew(nn.Module): def __init__(self, dim, max_seq_len=2048, base=10000.0): super().__init__() inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq) self.max_seq_len = max_seq_len self.dim = dim def forward(self, q, k): return rope_extension.rope_cuda(q, k, self.inv_freq) ``` Wait, I need to think more carefully about the cos/sin computation. In the original code: - `freqs = torch.outer(t, inv_freq)` gives shape (seq_len, dim/2) - `emb = torch.cat([freqs, freqs], dim=-1)` gives shape (seq_len, dim) - So for each position n and dimension d, the angle is `n * inv_freq[d % (dim/2)]` Actually, looking more carefully: - `inv_freq` has shape (dim/2,) with values `1.0 / (base ** (i / dim))` for i in [0, 2, 4, ..., dim-2] - `freqs = torch.outer(t, inv_freq)` gives shape (seq_len, dim/2), where `freqs[n, i] = n * inv_freq[i]` - `emb = torch.cat([freqs, freqs], dim=-1)` gives shape (seq_len, dim), where `emb[n, d] = n * inv_freq[d]` for d < dim/2, and `emb[n, d] = n * inv_freq[d - dim/2]` for d >= dim/2 So `cos[n, d] = cos(n * inv_freq[d % (dim/2)])` and same for sin. This matches my implementation above. Let me compile and test.

The profiling reveals significant performance issues: 1. **Memory coalescing is terrible**: 8 sectors/request means 8x wasted bandwidth 2. **Warp stalls**: 809% long_scoreboard stalls - waiting for memory 3. **Uncoalesced global loads**: The current access pattern is not coalesced The issue is that when `idx` goes across threads in a warp, `d = idx % dim` varies across all values 0..dim-1 within a single warp (for the first dim threads). This means threads in a warp are accessing scattered locations in memory, causing poor coalescing. Let me rethink the access pattern. A better approach would be to: 1. Have threads in a warp access consecutive elements in the dim dimension 2. Process multiple elements per thread (vectorized) Actually, the main issue is that threads within a warp are accessing indices that are dim apart in the input. When `idx` increases by 1, we're accessing positions that are dim apart in the original tensor layout. Let me reconsider the indexing. The tensor q is (B, H, N, D) and stored contiguously as [B*H*N*D] elements. So element (b, h, n, d) is at position `b*H*N*D + h*N*D + n*D + d`. If I index with `idx = b*H*N*D + h*N*D + n*D + d`, then: - `d = idx % D` varies quickly - `n = (idx / D) % N` varies more slowly So actually, threads with consecutive `idx` values access consecutive `d` values, which should be coalesced! Let me look more carefully... Wait, the issue is: 1. I'm reading `input[idx]` which is coalesced (consecutive threads read consecutive memory) 2. But I'm also reading `input[idx - d + pair_idx]` which is NOT coalesced - different threads read from different positions The second read `input[idx - d + pair_idx]` is problematic because `pair_idx = d + half_dim` (for d < half_dim) or `d - half_dim` (for d >= half_dim). So for d < half_dim: - `input[idx - d + (d + half_dim)]` = `input[idx + half_dim]` For d >= half_dim: - `input[idx - d + (d - half_dim)]` = `input[idx - half_dim]` So actually, for the first half of threads (d < half_dim), they read from idx + half_dim, and for the second half, they read from idx - half_dim. This should still be somewhat coalesced... threads 0-31 in a warp would be: - Threads 0-31 (d = 0-31 for dim=64) read from idx+half_dim - But wait, this means thread 0 reads from idx+32, thread 1 reads from idx+33, etc. which IS coalesced! Hmm, but the profiling shows 8 sectors/request which indicates poor coalescing. Let me think more... Actually, I think the issue is that each thread reads two separate locations, and those locations don't form a contiguous pattern. Let me try a different approach - using vectorized loads and shared memory to improve memory access patterns. Actually, a simpler optimization: compute cos/sin once per (n, freq_idx) pair instead of per element, and use shared memory for the cos/sin table. Let me rewrite with a more optimized approach: 1. Use vectorized loads (float4) to load input data 2. Pre-compute cos/sin in shared memory 3. Ensure coalesced memory access Actually, let me try a simpler approach first - just vectorizing the loads and stores using float4, and see if that helps.

The error is higher but still within fp32 tolerance. Let me profile this version.

Still having 8x wasted bandwidth and 1441% stalls. The problem is the memory access pattern - each thread accesses `input[base_idx + d]` and `input[base_idx + d + half_dim]` which are scattered. Let me think about a better approach. Since dim=64, each row has 64 elements. I should have threads process consecutive elements in the d dimension to get coalesced access. The key insight: for the reference implementation, cos/sin are computed as: - freqs = outer(positions, inv_freq) -> shape (N, D/2) - emb = cat([freqs, freqs], dim=-1) -> shape (N, D) - cos, sin = emb.cos(), emb.sin() So for each position n, the first half and second half of cos/sin are IDENTICAL: - cos[n, d] = cos[n, d + half_dim] for d in [0, half_dim) - sin[n, d] = sin[n, d + half_dim] for d in [0, half_dim) The rotation formula is: - x_rot[n, d] = x[n, d] * cos[n, d] + rotate_half(x)[n, d] * sin[n, d] - Where rotate_half swaps first/second halves: - rotate_half(x)[n, d] = -x[n, d + half_dim] for d < half_dim - rotate_half(x)[n, d] = x[n, d - half_dim] for d >= half_dim So: - For d < half_dim: x_rot = x[d] * cos[d] - x[d + half_dim] * sin[d] - For d >= half_dim: x_rot = x[d] * cos[d] + x[d - half_dim] * sin[d] Since cos[d] = cos[d % half_dim] and sin[d] = sin[d % half_dim]: - For d < half_dim: x_rot = x[d] * c - x[d + half_dim] * s - For d >= half_dim: x_rot = x[d] * c + x[d - half_dim] * s But wait, the issue is that I need x[d] and x[d + half_dim] in each thread, which are half_dim apart. This creates the non-coalesced access pattern. A better approach: 1. Use vectorized loads - but x[d] and x[d+half_dim] are not adjacent in memory 2. Instead, process consecutive rows (same n, different b,h), so that consecutive threads access consecutive d values 3. Use shared memory to cache the input row, then process pairs Actually, let me try a different approach - have each thread process consecutive d values for the SAME row, but use float4 loads. Actually wait - the real issue is that the accesses to x[d] and x[d+half_dim] for different threads are not contiguous. Let me try processing all d for a single (b,h,n) tuple. Actually, a better approach: use vectorized loads to load chunks of 4 consecutive d values at a time, for the first half, then the second half. This way threads can process 4 pairs each. Let me restructure: - Each thread processes 4 pairs (4 elements from first half, 4 from second half) - This reduces divergence and the number of distinct memory transactions Actually, let me try yet another approach - have consecutive threads process consecutive elements in memory, and compute the rotation pairs using a more clever indexing. The memory layout is [B, H, N, D]. If I want contiguous access, threads should process elements in the last dimension consecutively. Let me try: each thread processes one element, and uses adjacent threads' data to avoid fetching distant pairs. Actually, the simplest improvement is to just process elements linearly and use vectorized loads. Let me try: For each (b,h,n,d), we need to compute: - If d < half_dim: result[d] = x[d] * c - x[d+half_dim] * s - If d >= half_dim: result[d] = x[d] * c + x[d-half_dim] * s To make this coalesce, we want consecutive threads to access consecutive memory. So the linear index should map to (b,h,n,d) where d varies fastest. But the problem is we need to fetch both x[d] and x[d +/- half_dim] - they are 32 positions apart. Let me try using float4 to load 4 elements at a time from both halves: Each thread loads float4 from first half and float4 from second half -> 8 elements total -> computes 8 outputs. Let me try this:

Let me try a different approach - precompute the cos/sin values and use shared memory. Actually, the main issue is the non-contiguous memory access pattern. Let me try restructuring to have contiguous access. Actually, the core issue is that we're accessing x1 and x2 which are separated by half_dim elements. Even if threads access contiguous d values, they need to read from two locations half_dim apart. Let me try a different approach: compute the cos/sin table once on the host or in a separate kernel, then use it in the main kernel. This avoids computing cos/sin repeatedly. Actually, for this small problem size (4 * 8 * 256 * 64 = 524288 elements), the overhead of multiple kernels may be too high. Let me focus on the single kernel approach. The real issue is the memory access pattern. Let me try: 1. Using shared memory to cache inv_freq 2. Process multiple elements per thread with better coalescing Actually, let me look at this more carefully. The reference computes cos/sin based on outer product of positions with inv_freq. If I can precompute the cos/sin for all positions once, then the kernel just needs to apply them. Let me try precomputing cos/sin on the CPU and passing them in.

Good, correctness passes. Let me profile this version and see if memory coalescing improved.

The kernel is correct. Let me profile it to see if there's room for improvement, or if it's already performing well enough to submit.

The profiling shows the main issue is still memory coalescing - 12x wasted bandwidth. The issue is that when each thread processes one pair (x1 at position d, x2 at position half_dim+d), the two loads are separated by half_dim in memory, causing uncoalesced access patterns. Let me think about a better approach: 1. The input tensor is (B, H, N, D) - contiguous in memory with D being the fastest-varying dimension 2. When threads in a warp access input[base_idx + d] with consecutive d values, this is good coalescing for x1 3. But x2 = input[base_idx + half_dim + d] - these are also consecutive, just offset by half_dim Actually, wait - the problem might be that each thread has TWO loads that are far apart. Let me try a different approach where each thread handles just one element, and we use vectorized loads to get both x1 and x2 at once. Actually, a better approach would be to have threads within a warp all process the same (b, h, n) but different d values. This way, all threads in the warp read from contiguous locations. Let me try this approach: reorganize the thread indexing so that threads in a warp all work on the same (b, h, n) but different d values. This should improve coalescing since threads 0-31 would read from positions 0-31 within the same row. The problem is I'm running out of time and have too many constraints. Let me try one more optimization - use float4 loads for x1 values (contiguous first half) and another float4 load for x2 values (contiguous second half). This should give better coalescing. Actually, I think I should submit now since I have a working kernel and I'm running low on turns. The kernel is correct and should be faster than the reference PyTorch implementation which involves multiple kernel launches for cos/sin computation, chunking, and concatenation.