reasoning
reasoning (truncated 37470 chars): Let me analyze this task. I need to implement a custom CUDA kernel for Rotary Position Embedding (RoPE).
The reference implementation:
1. Takes query `q` and key `k` tensors of shape (B, H, N, D)
2. Builds cosine and sine embeddings based on position
3. Applies rotation to both q and k using:
- `q_rot = q * cos + rotate_half(q) * sin`
- where `rotate_half` splits x into two halves and concatenates [-x2, x1]
Let me write an optimized CUDA kernel for this:
1. The `inv_freq` is precomputed and registered as a buffer
2. For each forward pass:
- Build cos/sin based on sequence length N
- Apply the rotation formula
The rotation formula for position p and dimension pair (2i, 2i+1):
- For a pair (x_{2i}, x_{2i+1}):
- x'_{2i} = x_{2i} * cos(θ_{i,p}) - x_{2i+1} * sin(θ_{i,p})
- x'_{2i+1} = x_{2i} * sin(θ_{i,p}) + x_{2i+1} * cos(θ_{i,p})
The `rotate_half` function takes x and returns [-x2, x1] where x1 is the first half and x2 is the second half.
So `q * cos + rotate_half(q) * sin` effectively applies:
- For first half: q1 * cos - q2 * sin (since rotate_half gives -q2 for first half)
- For second half: q1 * sin + q2 * cos (since rotate_half gives q1 for second half)
This is the rotation formula.
Let me write an optimized CUDA kernel:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
# CUDA kernel code
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
__global__ void rope_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ inv_freq,
float* __restrict__ q_out,
float* __restrict__ k_out,
int batch_size, int num_heads, int seq_len, int dim, int half_dim
) {
// Each thread handles one element
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = batch_size * num_heads * seq_len * dim;
if (idx >= total) return;
// Decode indices
int d = idx % dim;
int n = (idx / dim) % seq_len;
// int h = (idx / dim / seq_len) % num_heads; // Not needed since cos/sin don't depend on h
// int b = idx / dim / seq_len / num_heads; // Not needed for same reason
// Build frequency for this position and dimension
int freq_idx = d % half_dim;
float freq = n * inv_freq[freq_idx];
float cos_val = cosf(freq);
float sin_val = sinf(freq);
// Apply rotation
float val_q = q[idx];
float val_k = k[idx];
// Determine the other half's index
int other_d;
float sign;
if (d < half_dim) {
// First half: pair with d + half_dim
other_d = d + half_dim;
sign = -1.0f; // rotate_half gives -x2 for first half
// q_rot[idx] = q[idx] * cos + (-q[other]) * sin
} else {
// Second half: pair with d - half_dim
other_d = d - half_dim;
sign = 1.0f; // rotate_half gives x1 for second half
// q_rot[idx] = q[idx] * cos + q[other] * sin
}
// Need to compute the other index
int other_idx = idx - d + other_d;
float val_q_other = q[other_idx];
float val_k_other = k[other_idx];
q_out[idx] = val_q * cos_val + sign * val_q_other * sin_val;
k_out[idx] = val_k * cos_val + sign * val_k_other * sin_val;
}
void rope_forward(
const torch::Tensor& q,
const torch::Tensor& k,
const torch::Tensor& inv_freq,
torch::Tensor& q_out,
torch::Tensor& k_out
) {
int batch_size = q.size(0);
int num_heads = q.size(1);
int seq_len = q.size(2);
int dim = q.size(3);
int half_dim = dim / 2;
int total = batch_size * num_heads * seq_len * dim;
int threads = 256;
int blocks = (total + threads - 1) / threads;
rope_kernel<<<blocks, threads>>>(
q.data_ptr<float>(),
k.data_ptr<float>(),
inv_freq.data_ptr<float>(),
q_out.data_ptr<float>(),
k_out.data_ptr<float>(),
batch_size, num_heads, seq_len, dim, half_dim
);
}
"""
cpp_source = """
void rope_forward(
const torch::Tensor& q,
const torch::Tensor& k,
const torch::Tensor& inv_freq,
torch::Tensor& q_out,
torch::Tensor& k_out
);
"""
rope_extension = load_inline(
name='rope_extension',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['rope_forward'],
verbose=True
)
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: torch.Tensor, k: torch.Tensor) -> tuple:
q_out = torch.empty_like(q)
k_out = torch.empty_like(k)
rope_extension.rope_forward(q, k, self.inv_freq, q_out, k_out)
return q_out, k_out
```
Wait, I need to think more carefully about memory access patterns. The current approach has each thread access two elements that might be far apart in memory, which is not ideal.
Let me reconsider. A better approach might be to:
1. Precompute cos and sin values on GPU
2. Then apply the rotation
Or, process pairs of elements together in one thread.
Actually, let me look at the memory layout more carefully:
- q shape: (B, H, N, D)
- For a given (b, h, n), all D elements are contiguous in memory
- For rotation, we need pairs (d, d+half_dim) for d < half_dim
A better kernel would process pairs together:
```cuda
__global__ void rope_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ inv_freq,
float* __restrict__ q_out,
float* __restrict__ k_out,
int batch_size, int num_heads, int seq_len, int dim
) {
int half_dim = dim / 2;
// Each thread handles one pair (d, d+half_dim) for one (b, h, n)
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total_pairs = batch_size * num_heads * seq_len * half_dim;
if (idx >= total_pairs) return;
// Decode indices
int d = idx % half_dim;
int n = (idx / half_dim) % seq_len;
int h = (idx / half_dim / seq_len) % num_heads;
int b = idx / half_dim / seq_len / num_heads;
// Base index for this (b, h, n)
int base = ((b * num_heads + h) * seq_len + n) * dim;
// Indices for the pair
int idx1 = base + d; // first half
int idx2 = base + d + half_dim; // second half
// Build frequency for this position and dimension pair
float freq = n * inv_freq[d];
float cos_val = cosf(freq);
float sin_val = sinf(freq);
// Apply rotation
float q1 = q[idx1], q2 = q[idx2];
float k1 = k[idx1], k2 = k[idx2];
// q1' = q1 * cos - q2 * sin
// q2' = q1 * sin + q2 * cos
q_out[idx1] = q1 * cos_val - q2 * sin_val;
q_out[idx2] = q1 * sin_val + q2 * cos_val;
k_out[idx1] = k1 * cos_val - k2 * sin_val;
k_out[idx2] = k1 * sin_val + k2 * cos_val;
}
```
This is cleaner - each thread handles a pair, reads both q and k values, and writes both outputs.
But wait, I can also precompute cos/sin once and reuse. Let me think about optimizations:
1. Precompute cos/sin embeddings once per forward call - this avoids redundant trig computation
2. Use vectorized loads/stores (float4) for better memory throughput
3. Process multiple pairs per thread for better ILP
Actually, let me look at the problem again. The cos and sin values depend only on n (position) and d (dimension), not on b or h. So we can precompute them and reuse.
Let me write a version that:
1. Precomputes cos/sin in a separate kernel (or on CPU if small enough)
2. Uses a main kernel that applies the rotation with the precomputed values
But for simplicity and to ensure correctness first, let me start with the basic version and optimize later.
Actually, the simplest approach that should work and be faster than PyTorch:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_ext
...[ omitted 37470 characters ]...
(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: torch.Tensor, k: torch.Tensor) -> tuple:
q_out = torch.empty_like(q)
k_out = torch.empty_like(k)
rope_extension.rope_forward(q, k, self.inv_freq, q_out, k_out)
return q_out, k_out
```
This looks good. Let me think if there are any other optimizations I can add...
One potential optimization: use more threads per block for better occupancy. 256 threads per block is standard, but we could try 512 or 1024 if the registers allow.
Another optimization: process multiple pairs per thread. For example, each thread could handle 4 or 8 pairs, which reduces the number of threads and increases work per thread.
Let me try a version where each thread handles 4 pairs:
```cuda
__global__ void rope_kernel_unrolled(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ inv_freq,
float* __restrict__ q_out,
float* __restrict__ k_out,
const int batch_size,
const int num_heads,
const int seq_len,
const int dim
) {
const int half_dim = dim / 2;
const int pairs_per_thread = 4;
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
const int total_threads = batch_size * num_heads * seq_len * (half_dim / pairs_per_thread);
if (idx >= total_threads) return;
const int d_base = (idx % (half_dim / pairs_per_thread)) * pairs_per_thread;
const int n = (idx / (half_dim / pairs_per_thread)) % seq_len;
const int h = (idx / (half_dim / pairs_per_thread) / seq_len) % num_heads;
const int b = idx / (half_dim / pairs_per_thread) / seq_len / num_heads;
const int base = ((b * num_heads + h) * seq_len + n) * dim;
#pragma unroll
for (int i = 0; i < pairs_per_thread; i++) {
const int d = d_base + i;
const int idx1 = base + d;
const int idx2 = base + d + half_dim;
const float freq = static_cast<float>(n) * inv_freq[d];
const float cos_val = cosf(freq);
const float sin_val = sinf(freq);
const float q1 = q[idx1];
const float q2 = q[idx2];
const float k1 = k[idx1];
const float k2 = k[idx2];
q_out[idx1] = q1 * cos_val - q2 * sin_val;
q_out[idx2] = q1 * sin_val + q2 * cos_val;
k_out[idx1] = k1 * cos_val - k2 * sin_val;
k_out[idx2] = k1 * sin_val + k2 * cos_val;
}
}
```
This version processes 4 pairs per thread, which should reduce kernel launch overhead and improve instruction-level parallelism.
But wait, for half_dim = 32, pairs_per_thread = 4 means we process all 32 pairs with 8 threads per (b, h, n). This might reduce parallelism too much for small problem sizes.
Let me stick with the simpler 1 pair per thread version for now. It should be fast enough.
Actually, let me reconsider. The problem has:
- batch_size = 4
- num_heads = 8
- seq_len = 256
- dim = 64, half_dim = 32
Total pairs = 4 * 8 * 256 * 32 = 262,144
With 256 threads per block and 1 pair per thread, we need 1024 blocks. This is a decent amount of parallelism.
With 4 pairs per thread, we'd need 256 blocks. This might be too few for some GPUs.
Let me stick with 1 pair per thread for simplicity and good parallelism.
Alright, I think my implementation is ready. Let me finalize it.
Actually, I just realized I need to handle the case where `inv_freq` might not be on the GPU. Let me add a check:
```python
def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple:
inv_freq = self.inv_freq.to(q.device)
q_out = torch.empty_like(q)
k_out = torch.empty_like(k)
rope_extension.rope_forward(q, k, inv_freq, q_out, k_out)
return q_out, k_out
```
This ensures inv_freq is on the same device as q and k.
Actually, since inv_freq is registered as a buffer, it should be moved automatically when the model is moved. But adding `.to(q.device)` is a safe practice.
Let me finalize:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = r"""
#include <cuda.h>
#include <cuda_runtime.h>
#include <ATen/cuda/CUDAContext.h>
#include <cmath>
__global__ void rope_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ inv_freq,
float* __restrict__ q_out,
float* __restrict__ k_out,
const int batch_size,
const int num_heads,
const int seq_len,
const int dim
) {
const int half_dim = dim / 2;
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
const int total_pairs = batch_size * num_heads * seq_len * half_dim;
if (idx >= total_pairs) return;
const int d = idx % half_dim;
const int n = (idx / half_dim) % seq_len;
const int h = (idx / half_dim / seq_len) % num_heads;
const int b = idx / half_dim / seq_len / num_heads;
const int base = ((b * num_heads + h) * seq_len + n) * dim;
const int idx1 = base + d;
const int idx2 = base + d + half_dim;
const float freq = static_cast<float>(n) * inv_freq[d];
const float cos_val = cosf(freq);
const float sin_val = sinf(freq);
const float q1 = q[idx1];
const float q2 = q[idx2];
const float k1 = k[idx1];
const float k2 = k[idx2];
q_out[idx1] = q1 * cos_val - q2 * sin_val;
q_out[idx2] = q1 * sin_val + q2 * cos_val;
k_out[idx1] = k1 * cos_val - k2 * sin_val;
k_out[idx2] = k1 * sin_val + k2 * cos_val;
}
void rope_forward(
const torch::Tensor& q,
const torch::Tensor& k,
const torch::Tensor& inv_freq,
torch::Tensor& q_out,
torch::Tensor& k_out
) {
const int batch_size = static_cast<int>(q.size(0));
const int num_heads = static_cast<int>(q.size(1));
const int seq_len = static_cast<int>(q.size(2));
const int dim = static_cast<int>(q.size(3));
const int half_dim = dim / 2;
const int total_pairs = batch_size * num_heads * seq_len * half_dim;
const int threads = 256;
const int blocks = (total_pairs + threads - 1) / threads;
rope_kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(
q.data_ptr<float>(),
k.data_ptr<float>(),
inv_freq.data_ptr<float>(),
q_out.data_ptr<float>(),
k_out.data_ptr<float>(),
batch_size,
num_heads,
seq_len,
dim
);
}
"""
cpp_source = r"""
void rope_forward(
const torch::Tensor& q,
const torch::Tensor& k,
const torch::Tensor& inv_freq,
torch::Tensor& q_out,
torch::Tensor& k_out
);
"""
rope_extension = load_inline(
name='rope_extension',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['rope_forward'],
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: torch.Tensor, k: torch.Tensor) -> tuple:
inv_freq = self.inv_freq.to(q.device)
q_out = torch.empty_like(q)
k_out = torch.empty_like(k)
rope_extension.rope_forward(q, k, inv_freq, q_out, k_out)
return q_out, k_out
```
This looks complete. Let me double-check the logic:
1. In the kernel, each thread handles one pair (d, d+half_dim) for one (b, h, n).
2. The cos and sin values are computed directly from n * inv_freq[d].
3. The rotation formula is applied: q_rot[d] = q[d] * cos - q[d+half_dim] * sin, q_rot[d+half_dim] = q[d] * sin + q[d+half_dim] * cos.
4. The same is done for k.
This matches the reference implementation.
I'm confident this implementation is correct and should be faster than the PyTorch reference.
Let me submit it!