reasoning
reasoning (truncated 35175 chars): This is a causal grouped-query attention implementation. Let me analyze what's happening:
1. Query heads outnumber KV heads (8 Q heads vs 2 KV heads)
2. Keys and values are expanded by repeating each KV head across its assigned query-head group
3. For each position t, we compute attention over all positions up to t (causal attention)
The naive PyTorch implementation uses nested loops which is very slow. I need to write a custom CUDA kernel that:
1. Expands K and V by repeating each KV head `group_size` times
2. Computes scaled dot-product attention in a causal manner
3. Uses softmax over the scores
4. Outputs the weighted sum of values
Let me write an optimized CUDA kernel for this.
The key observations:
- batch_size = 2, seq_len = 80, num_q_heads = 8, num_kv_heads = 2, head_dim = 32
- group_size = 4 (each KV head is shared by 4 query heads)
- For causal attention, position t can only attend to positions 0 to t
For efficiency, I should:
1. Parallelize over batch, query head, and sequence position
2. Compute the attention scores and apply softmax
3. Use shared memory for intermediate results if possible
Let me write the kernel:
```cpp
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>
__global__ void causal_gqa_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
float* __restrict__ out,
int batch_size, int seq_len, int num_q_heads, int num_kv_heads, int head_dim
) {
int group_size = num_q_heads / num_kv_heads;
float scale = 1.0f / sqrtf((float)head_dim);
int b = blockIdx.x;
int h = blockIdx.y;
int t = threadIdx.x;
if (b >= batch_size || h >= num_q_heads || t >= seq_len) return;
// Map query head to KV head
int kv_h = h / group_size;
// Compute attention scores for position t
float scores[80]; // max seq_len
float max_score = -INFINITY;
for (int s = 0; s <= t; s++) {
float score = 0.0f;
for (int d = 0; d < head_dim; d++) {
score += q[b * seq_len * num_q_heads * head_dim + t * num_q_heads * head_dim + h * head_dim + d] *
k[b * seq_len * num_kv_heads * head_dim + s * num_kv_heads * head_dim + kv_h * head_dim + d];
}
scores[s] = score * scale;
if (scores[s] > max_score) max_score = scores[s];
}
// Softmax
float sum_exp = 0.0f;
for (int s = 0; s <= t; s++) {
scores[s] = expf(scores[s] - max_score);
sum_exp += scores[s];
}
// Compute output
for (int d = 0; d < head_dim; d++) {
float out_val = 0.0f;
for (int s = 0; s <= t; s++) {
out_val += scores[s] / sum_exp * v[b * seq_len * num_kv_heads * head_dim + s * num_kv_heads * head_dim + kv_h * head_dim + d];
}
out[b * seq_len * num_q_heads * head_dim + t * num_q_heads * head_dim + h * head_dim + d] = out_val;
}
}
```
Wait, this kernel approach has issues. Using threadIdx.x for t means we can't have t up to 80 with typical block sizes. Let me reconsider.
Actually, I should think about this differently. Let me use a more flexible approach:
- Use block dimensions that work with the problem size
- Or use a 1D kernel and compute indices from the global thread index
Let me rewrite:
```cpp
__global__ void causal_gqa_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
float* __restrict__ out,
int batch_size, int seq_len, int num_q_heads, int num_kv_heads, int head_dim,
float scale
) {
int group_size = num_q_heads / num_kv_heads;
// Total threads needed: batch * num_q_heads * seq_len
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = batch_size * num_q_heads * seq_len;
if (idx >= total) return;
int t = idx % seq_len;
int h = (idx / seq_len) % num_q_heads;
int b = idx / (seq_len * num_q_heads);
// Map query head to KV head
int kv_h = h / group_size;
// ... rest of the logic
}
```
This is better. Now each thread handles one (b, h, t) tuple.
But there's a concern about memory access patterns and shared memory. Let me think about optimizing this further.
For shared memory, I could:
1. Load the query vector for this position into registers
2. For each position s <= t, load k[b, s, kv_h] and compute the score
3. Store scores in shared memory or registers
4. Compute softmax
5. Compute weighted sum
Given seq_len = 80, I could store up to 80 float scores in registers (320 bytes per thread, which is reasonable).
Actually, let me think about tiling and shared memory more carefully.
For a given (b, h), we need to compute attention for all positions t from 0 to seq_len-1.
For each t, we compute attention over positions 0 to t.
One optimization: For a given (b, h) pair, we could have threads cooperate by loading K and V values into shared memory.
But for simplicity, let me start with a straightforward kernel and see if it works:
```cpp
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
template <int HEAD_DIM>
__global__ void causal_gqa_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
float* __restrict__ out,
int batch_size, int seq_len, int num_q_heads, int num_kv_heads
) {
constexpr float scale = 1.0f / 5.656854249f; // sqrt(32)
int group_size = num_q_heads / num_kv_heads;
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = batch_size * num_q_heads * seq_len;
if (idx >= total) return;
int t = idx % seq_len;
int h = (idx / seq_len) % num_q_heads;
int b = idx / (seq_len * num_q_heads);
int kv_h = h / group_size;
// Load query vector
float q_vec[HEAD_DIM];
for (int d = 0; d < HEAD_DIM; d++) {
q_vec[d] = q[b * seq_len * num_q_heads * HEAD_DIM + t * num_q_heads * HEAD_DIM + h * HEAD_DIM + d];
}
// Compute scores for positions 0 to t
extern __shared__ float shared_mem[];
float* scores = shared_mem + threadIdx.x * (80 + 1); // Each thread gets up to 80+1 floats
float max_score = -INFINITY;
for (int s = 0; s <= t; s++) {
float score = 0.0f;
for (int d = 0; d < HEAD_DIM; d++) {
score += q_vec[d] * k[b * seq_len * num_kv_heads * HEAD_DIM + s * num_kv_heads * HEAD_DIM + kv_h * HEAD_DIM + d];
}
scores[s] = score * scale;
if (scores[s] > max_score) max_score = scores[s];
}
// Softmax
float sum_exp = 0.0f;
for (int s = 0; s <= t; s++) {
scores[s] = expf(scores[s] - max_score);
sum_exp += scores[s];
}
// Compute output
for (int d = 0; d < HEAD_DIM; d++) {
float out_val = 0.0f;
for (int s = 0; s <= t; s++) {
out_val += (scores[s] / sum_exp) * v[b * seq_len * num_kv_heads * HEAD_DIM + s * num_kv_heads * HEAD_DIM + kv_h * HEAD_DIM + d];
}
out[b * seq_len * num_q_heads * HEAD_DIM + t * num_q_heads * HEAD_DIM + h * HEAD_DIM + d] = out_val;
}
}
```
Wait, the shared memory approach won't work well because each thread needs its own scores array. That would require too much shared memory.
Let me use a simpler approach with registers for the scores. Since seq_len is 80, I need to store up to 80 scores, which is 320 bytes per thread. This might fit in registers if I use a fixed-size array.
Actually, for a more flexible approach, let me just use local arrays (which will be in registers or local memory depending on compiler decisions):
```cpp
#include <torch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void causal_gqa_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
float* __restrict__ out,
int batch_size, int seq_len, int num_q_heads, int num_kv_heads, int head_dim
) {
float scale = 1.0f / sqrtf((float)
...[ omitted 35175 characters ]...
rch/extension.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
__global__ void causal_gqa_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
float* __restrict__ out,
int batch_size, int seq_len, int num_q_heads, int num_kv_heads, int head_dim,
float scale
) {
const int group_size = num_q_heads / num_kv_heads;
const int idx = blockIdx.x * blockDim.x + threadIdx.x;
const int total = batch_size * num_q_heads * seq_len;
if (idx >= total) return;
const int t = idx % seq_len;
const int h = (idx / seq_len) % num_q_heads;
const int b = idx / (seq_len * num_q_heads);
const int kv_h = h / group_size;
// Compute pointer offsets
const int64_t q_offset = (int64_t)b * seq_len * num_q_heads * head_dim
+ (int64_t)t * num_q_heads * head_dim
+ (int64_t)h * head_dim;
const int64_t kv_base_offset = (int64_t)b * seq_len * num_kv_heads * head_dim;
const int64_t kv_stride_s = (int64_t)num_kv_heads * head_dim;
const float* q_ptr = q + q_offset;
const float* k_base = k + kv_base_offset;
const float* v_base = v + kv_base_offset;
// Load query vector
float q_vec[32];
for (int d = 0; d < head_dim; d++) {
q_vec[d] = q_ptr[d];
}
// Online softmax accumulation
float max_score = -INFINITY;
float sum_exp = 0.0f;
float output[32];
for (int d = 0; d < head_dim; d++) {
output[d] = 0.0f;
}
for (int s = 0; s <= t; s++) {
const int64_t kv_offset = (int64_t)s * kv_stride_s + (int64_t)kv_h * head_dim;
const float* k_ptr = k_base + kv_offset;
const float* v_ptr = v_base + kv_offset;
// Compute attention score
float score = 0.0f;
for (int d = 0; d < head_dim; d++) {
score += q_vec[d] * k_ptr[d];
}
score *= scale;
// Online softmax update
const float new_max = fmaxf(max_score, score);
const float exp_old = expf(max_score - new_max);
const float exp_new = expf(score - new_max);
// Rescale and accumulate
for (int d = 0; d < head_dim; d++) {
output[d] = output[d] * exp_old + exp_new * v_ptr[d];
}
sum_exp = sum_exp * exp_old + exp_new;
max_score = new_max;
}
// Normalize output
float* out_ptr = out + q_offset;
for (int d = 0; d < head_dim; d++) {
out_ptr[d] = output[d] / sum_exp;
}
}
torch::Tensor causal_gqa_cuda(
torch::Tensor q,
torch::Tensor k,
torch::Tensor v
) {
const int batch_size = q.size(0);
const int seq_len = q.size(1);
const int num_q_heads = q.size(2);
const int head_dim = q.size(3);
const int num_kv_heads = k.size(2);
const float scale = 1.0f / std::sqrt((float)head_dim);
auto out = torch::zeros_like(q);
const int total = batch_size * num_q_heads * seq_len;
const int threads = 256;
const int blocks = (total + threads - 1) / threads;
causal_gqa_kernel<<<blocks, threads>>>(
q.data_ptr<float>(),
k.data_ptr<float>(),
v.data_ptr<float>(),
out.data_ptr<float>(),
batch_size, seq_len, num_q_heads, num_kv_heads, head_dim,
scale
);
return out;
}
"""
module = load_inline(
name='causal_gqa_ext',
cpp_sources=['torch::Tensor causal_gqa_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v);'],
cuda_sources=[cuda_source],
functions=['causal_gqa_cuda'],
verbose=False,
extra_cuda_cflags=['-O3']
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, q, k, v):
q = q.contiguous()
k = k.contiguous()
v = v.contiguous()
return module.causal_gqa_cuda(q, k, v)
```
This should be correct and reasonably efficient. Let me submit it.
Actually, let me reconsider. I want to make sure I'm handling all edge cases:
1. What if q, k, v are not on the same device? -> The user should handle this
2. What if q, k, v have different shapes? -> The kernel assumes correct shapes
3. What if head_dim > 32? -> The static arrays would overflow
For robustness, I could add a check in the host function:
```cpp
TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, "Inputs must be 4D tensors");
TORCH_CHECK(q.size(3) <= 32, "head_dim must be <= 32");
TORCH_CHECK(q.size(0) == k.size(0) && q.size(0) == v.size(0), "Batch sizes must match");
TORCH_CHECK(q.size(1) == k.size(1) && q.size(1) == v.size(1), "Sequence lengths must match");
TORCH_CHECK(q.size(3) == k.size(3) && q.size(3) == v.size(3), "Head dimensions must match");
```
But for simplicity, I'll assume the inputs are valid (as specified in the problem).
One more thing: I should ensure that the tensor data is on GPU and contiguous. The `contiguous()` call in Python should handle this, but let me add explicit checks in C++:
```cpp
TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "Inputs must be CUDA tensors");
TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(), "Inputs must be contiguous");
```
Actually, `torch::zeros_like(q)` should create an output tensor on the same device as q, so that's fine.
Let me add these checks to make the kernel more robust:
```cpp
torch::Tensor causal_gqa_cuda(
torch::Tensor q,
torch::Tensor k,
torch::Tensor v
) {
TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), "Inputs must be CUDA tensors");
TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(), "Inputs must be contiguous");
TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, "Inputs must be 4D tensors");
TORCH_CHECK(q.size(3) <= 32, "head_dim must be <= 32");
const int batch_size = q.size(0);
const int seq_len = q.size(1);
const int num_q_heads = q.size(2);
const int head_dim = q.size(3);
const int num_kv_heads = k.size(2);
const float scale = 1.0f / std::sqrt((float)head_dim);
auto out = torch::zeros_like(q);
const int total = batch_size * num_q_heads * seq_len;
const int threads = 256;
const int blocks = (total + threads - 1) / threads;
causal_gqa_kernel<<<blocks, threads>>>(
q.data_ptr<float>(),
k.data_ptr<float>(),
v.data_ptr<float>(),
out.data_ptr<float>(),
batch_size, seq_len, num_q_heads, num_kv_heads, head_dim,
scale
);
return out;
}
```
Alright, this should be robust enough. Let me submit the final code.
Wait, I should also consider if I need to synchronize after the kernel launch. In PyTorch, CUDA operations are asynchronous by default, and PyTorch handles synchronization when needed (e.g., when copying data to CPU). So I don't need to explicitly synchronize in the kernel.
However, if there's an error in the kernel, it might not be caught until the next synchronization point. For debugging, I could add:
```cpp
C10_CUDA_CHECK(cudaGetLastError());
```
after the kernel launch. This checks for launch errors. But for the final submission, I'll omit this for simplicity.
Alright, I think the code is ready. Let me finalize it:
Actually, I want to reconsider the memory access patterns once more. The key issue is that each thread reads K and V values in a sequential but strided manner. This could lead to poor cache utilization.
One potential improvement is to have multiple threads cooperate to load K and V values into shared memory. But this adds complexity and might not be necessary for such small problem sizes.
Another improvement is to use shared memory to cache the query vector and output vector across multiple threads. But again, this adds complexity.
For now, I'll stick with the simple scalar version. The performance should be much better than the naive PyTorch implementation regardless.
Let me submit the code.
call → submit_kernel
callsubmit_kernel
{
"kernel_code": "import math\nimport torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\ncuda_source = \"\"\"\n#include <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n#include <cmath>\n\n__global__ void causal_gqa_kernel(\n const float* __restrict__ q,\n const float* __restrict__ k,\n const float* __restrict__ v,\n float* __restrict__ out,\n int batch_size, int seq_len, int num_q_heads, int num_kv_heads, int head_dim,\n float scale\n) {\n const int group_size = num_q_heads / num_kv_heads;\n \n const int idx = blockIdx.x * blockDim.x + threadIdx.x;\n const int total = batch_size * num_q_heads * seq_len;\n if (idx >= total) return;\n \n const int t = idx % seq_len;\n const int h = (idx / seq_len) % num_q_heads;\n const int b = idx / (seq_len * num_q_heads);\n const int kv_h = h / group_size;\n \n // Compute pointer offsets\n const int64_t q_offset = (int64_t)b * seq_len * num_q_heads * head_dim \n + (int64_t)t * num_q_heads * head_dim \n + (int64_t)h * head_dim;\n const int64_t kv_base_offset = (int64_t)b * seq_len * num_kv_heads * head_dim;\n const int64_t kv_stride_s = (int64_t)num_kv_heads * head_dim;\n \n const float* q_ptr = q + q_offset;\n const float* k_base = k + kv_base_offset;\n const float* v_base = v + kv_base_offset;\n \n // Load query vector into registers\n float q_vec[32];\n for (int d = 0; d < head_dim; d++) {\n q_vec[d] = q_ptr[d];\n }\n \n // Online softmax accumulation for numerical stability\n float max_score = -INFINITY;\n float sum_exp = 0.0f;\n float output[32];\n for (int d = 0; d < head_dim; d++) {\n output[d] = 0.0f;\n }\n \n for (int s = 0; s <= t; s++) {\n const int64_t kv_offset = (int64_t)s * kv_stride_s + (int64_t)kv_h * head_dim;\n const float* k_ptr = k_base + kv_offset;\n const float* v_ptr = v_base + kv_offset;\n \n // Compute attention score: q_t . k_s * scale\n float score = 0.0f;\n for (int d = 0; d < head_dim; d++) {\n score += q_vec[d] * k_ptr[d];\n }\n score *= scale;\n \n // Online softmax update (Flash Attention style)\n const float new_max = fmaxf(max_score, score);\n const float exp_old = expf(max_score - new_max);\n const float exp_new = expf(score - new_max);\n \n // Rescale previous accumulation and add new contribution\n for (int d = 0; d < head_dim; d++) {\n output[d] = output[d] * exp_old + exp_new * v_ptr[d];\n }\n \n sum_exp = sum_exp * exp_old + exp_new;\n max_score = new_max;\n }\n \n // Normalize output by sum of exp\n float* out_ptr = out + q_offset;\n for (int d = 0; d < head_dim; d++) {\n out_ptr[d] = output[d] / sum_exp;\n }\n}\n\ntorch::Tensor causal_gqa_cuda(\n torch::Tensor q,\n torch::Tensor k,\n torch::Tensor v\n) {\n TORCH_CHECK(q.is_cuda() && k.is_cuda() && v.is_cuda(), \"Inputs must be CUDA tensors\");\n TORCH_CHECK(q.is_contiguous() && k.is_contiguous() && v.is_contiguous(), \"Inputs must be contiguous\");\n TORCH_CHECK(q.dim() == 4 && k.dim() == 4 && v.dim() == 4, \"Inputs must be 4D tensors\");\n TORCH_CHECK(q.size(3) <= 32, \"head_dim must be <= 32\");\n \n const int batch_size = q.size(0);\n const int seq_len = q.size(1);\n const int num_q_heads = q.size(2);\n const int head_dim = q.size(3);\n const int num_kv_heads = k.size(2);\n const float scale = 1.0f / std::sqrt((float)head_dim);\n \n auto out = torch::zeros_like(q);\n \n const int total = batch_size * num_q_heads * seq_len;\n const int threads = 256;\n const int blocks = (total + threads - 1) / threads;\n \n causal_gqa_kernel<<<blocks, threads>>>(\n q.data_ptr<float>(),\n k.data_ptr<float>(),\n v.data_ptr<float>(),\n out.data_ptr<float>(),\n batch_size, seq_len, num_q_heads, num_kv_heads, head_dim,\n scale\n );\n \n return out;\n}\n\"\"\"\n\nmodule = load_inline(\n name='causal_gqa_ext',\n cpp_sources=['torch::Tensor causal_gqa_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v);'],\n cuda_sources=[cuda_source],\n functions=['causal_gqa_cuda'],\n verbose=False,\n extra_cuda_cflags=['-O3']\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n \n def forward(self, q, k, v):\n q = q.contiguous()\n k = k.contiguous()\n v = v.contiguous()\n return module.causal_gqa_cuda(q, k, v)"
}