reasoning
Let me analyze this problem step by step.
The model is a scaled low-precision attention mechanism. The inputs are quantized tensors (int8) along with per-head scales (float32). The forward pass:
1. Dequantizes q, k, v by converting to float and multiplying by scales
2. Computes attention scores for each position in the sequence (causal attention)
3. Applies softmax and computes weighted sum of values
The key insight is that this is a causal attention mechanism - at position t, we only attend to positions 0 through t.
Let me understand the shapes:
- q_q, k_q, v_q: (batch_size, seq_len, num_heads, head_dim) - int8
- q_scale, k_scale, v_scale: (num_heads,) - float32
- Output: (batch_size, seq_len, num_heads, head_dim) - float32
Given:
- batch_size = 2
- seq_len = 72
- num_heads = 4
- head_dim = 32
So the total work is quite small, which means overhead will be significant.
The reference implementation has nested loops over batch, heads, and sequence positions. This is very inefficient in PyTorch. We can parallelize this significantly.
A custom CUDA kernel approach:
1. Dequantize q, k, v into float
2. For each (batch, head, seq_pos) combination, compute the causal attention
Actually, let me think about a more efficient approach. Since seq_len is small (72) and head_dim is small (32), we can use shared memory efficiently.
Let me design a kernel where each thread block handles one (batch, head) pair, and threads within the block handle different sequence positions.
Actually, let me think more carefully. The attention computation is:
- For each position t, compute scores = k[0:t+1] @ q[t]^T * scale
- Apply softmax to scores
- Output[t] = softmax(scores) @ v[0:t+1]
This is essentially a causal self-attention but with Q, K, V being different tensors (cross-attention style, but causal).
One approach: Use matrix operations to compute all scores at once, then apply a causal mask for softmax.
But the reference computes it sequentially - for position t, we only look at positions 0 to t. This means we can't compute all scores at once in a simple way.
Actually, we CAN compute all scores at once:
- scores = k @ q.transpose(-2, -1) * scale
- This gives (batch, seq_len, num_heads, seq_len)
- But we need to apply causal mask before softmax
Let me think about a CUDA implementation:
Option 1: Flash Attention style
- Compute attention in tiles
- Use online softmax for numerical stability
- But this might be overkill for small sizes
Option 2: Simple parallelization
- Each thread handles one (b, h, t) tuple
- Compute scores for that position
- Apply softmax
- Compute output
Option 3: Vectorized approach
- Use tensor cores or optimized matmul
- Compute full attention matrix
- Apply causal mask and softmax
- Compute output
Given the small sizes (batch=2, seq_len=72, heads=4, head_dim=32), I think option 2 or 3 might work well.
Let me try option 2 first - a simple CUDA kernel that parallelizes over (b, h, t).
Actually, for better memory access patterns, let me think about tiling.
Let me look at the memory access:
- q: (batch, seq, heads, head_dim) - accessed as q[b, t, h, :]
- k: (batch, seq, heads, head_dim) - accessed as k[b, 0:t+1, h, :]
- v: (batch, seq, heads, head_dim) - accessed as v[b, 0:t+1, h, :]
If we parallelize over (b, h, t), each thread needs to access k[0:t+1, :] and v[0:t+1, :], which are scattered in memory.
A better approach might be to have a thread block handle one (b, h) pair, with threads cooperating to compute attention for all t.
Or even better: use shared memory to load k and v for a (b, h) pair, then have threads compute attention for different t.
Let me design this:
```cuda
// Each thread block handles one (b, h) pair
// Threads within block handle different t
__global__ void attention_kernel(
const int8_t* q_q, const int8_t* k_q, const int8_t* v_q,
const float* q_scale, const float* k_scale, const float* v_scale,
float* out,
int batch_size, int seq_len, int num_heads, int head_dim
) {
int b = blockIdx.x / num_heads;
int h = blockIdx.x % num_heads;
int t = threadIdx.x;
if (t >= seq_len) return;
// Load q for this position
float q_local[head_dim];
float qs = q_scale[h];
for (int d = 0; d < head_dim; d++) {
q_local[d] = (float)q_q[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d] * qs;
}
// Compute scores for positions 0 to t
float ks = k_scale[h];
float scores[seq_len];
float max_score = -INFINITY;
for (int i = 0; i <= t; i++) {
float dot = 0.0f;
for (int d = 0; d < head_dim; d++) {
float k_val = (float)k_q[b * seq_len * num_heads * head_dim + i * num_heads * head_dim + h * head_dim + d] * ks;
dot += k_val * q_local[d];
}
scores[i] = dot / sqrtf((float)head_dim);
if (scores[i] > max_score) max_score = scores[i];
}
// Softmax
float sum_exp = 0.0f;
for (int i = 0; i <= t; i++) {
scores[i] = expf(scores[i] - max_score);
sum_exp += scores[i];
}
for (int i = 0; i <= t; i++) {
scores[i] /= sum_exp;
}
// Compute output
float vs = v_scale[h];
float out_local[head_dim];
for (int d = 0; d < head_dim; d++) {
out_local[d] = 0.0f;
}
for (int i = 0; i <= t; i++) {
float attn = scores[i];
for (int d = 0; d < head_dim; d++) {
float v_val = (float)v_q[b * seq_len * num_heads * head_dim + i * num_heads * head_dim + h * head_dim + d] * vs;
out_local[d] += attn * v_val;
}
}
// Store output
for (int d = 0; d < head_dim; d++) {
out[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d] = out_local[d];
}
}
```
This is a simple approach but has some issues:
1. Each thread loads k and v multiple times (once for each t >= position)
2. Memory access patterns are not coalesced
Let me think about a better approach using shared memory:
```cuda
// Each thread block handles one (b, h) pair
// Use shared memory to cache k and v
__global__ void attention_kernel(...) {
__shared__ float k_shared[seq_len * head_dim];
__shared__ float v_shared[seq_len * head_dim];
int b = blockIdx.x / num_heads;
int h = blockIdx.x % num_heads;
// Cooperatively load k and v into shared memory
for (int i = threadIdx.x; i < seq_len * head_dim; i += blockDim.x) {
int s = i / head_dim;
int d = i % head_dim;
k_shared[i] = (float)k_q[...] * k_scale[h];
v_shared[i] = (float)v_q[...] * v_scale[h];
}
__syncthreads();
// Now each thread handles one or more t
...
}
```
But wait, seq_len * head_dim = 72 * 32 = 2304, and we need both k and v in shared memory, so 4608 floats = 18KB. That's less than the typical 48KB shared memory limit, so this should work.
Let me refine this approach. But actually, for such small sizes, the launch overhead and synchronization might dominate. Let me check GPU specs first.
Actually, let me start with a simpler approach that might be faster due to lower overhead:
Use a single kernel that processes everything in a more vectorized way.
Actually, thinking about this more, since the sizes are small:
- batch_size = 2
- seq_len = 72
- num_heads = 4
- head_dim = 32
The total number of output elements is 2 * 72 * 4 * 32 = 18432.
If we launch with blockDim.x = 256 (typical), we'd have 18432 / 256 = 72 thread blocks. That's not many.
Let me think about a different approach: process multiple (b, h, t) per thread to increase work per thread.
Or, let's use a tiled approach where each thread handles multiple dimensions of the output.
Actually, let me just write a straightforward kernel first and see if it beats PyTorch:
```python
import torch
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
__global__ void attention_kernel(
const int8_t* __restrict__ q_q,
const int8_t* __restrict__ k_q,
const int8_t* __restrict__ v_q,
const float* __restrict__ q_scale,
const float* __restrict__ k_scale,
const float* __restrict__ v_scale,
float* __restrict__ out,
int batch_size, int seq_len, int num_heads, int head_dim
) {
// Each thread handles one (b, h, t) tuple
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = batch_size * num_heads * seq_len;
if (idx >= total) return;
int t = idx % seq_len;
int h = (idx / seq_len) % num_heads;
int b = idx / (seq_len * num_heads);
float scale = 1.0f / sqrtf((float)head_dim);
float qs = q_scale[h];
float ks = k_scale[h];
float vs = v_scale[h];
int base_idx = b * seq_len * num_heads * head_dim + h * head_dim;
// Load q for this position
// q is indexed as [b, t, h, d]
// In memory: b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d
// Compute scores for positions 0 to t
float max_score = -INFINITY;
float scores[72]; // Max seq_len
for (int i = 0; i <= t; i++) {
float dot = 0.0f;
for (int d = 0; d < head_dim; d++) {
int q_idx = b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d;
int k_idx = b * seq_len * num_heads * head_dim + i * num_heads * head_dim + h * head_dim + d;
float q_val = (float)q_q[q_idx] * qs;
float k_val = (float)k_q[k_idx] * ks;
dot += q_val * k_val;
}
scores[i] = dot * scale;
if (scores[i] > max_score) max_score = scores[i];
}
// Softmax
float sum_exp = 0.0f;
for (int i = 0; i <= t; i++) {
scores[i] = expf(scores[i] - max_score);
sum_exp += scores[i];
}
float inv_sum = 1.0f / sum_exp;
for (int i = 0; i <= t; i++) {
scores[i] *= inv_sum;
}
// Compute output
for (int d = 0; d < head_dim; d++) {
float sum = 0.0f;
for (int i = 0; i <= t; i++) {
int v_idx = b * seq_len * num_heads * head_dim + i * num_heads * head_dim + h * head_dim + d;
float v_val = (float)v_q[v_idx] * vs;
sum += scores[i] * v_val;
}
int out_idx = b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d;
out[out_idx] = sum;
}
}
void attention_cuda(
const torch::Tensor& q_q,
const torch::Tensor& k_q,
const torch::Tensor& v_q,
const torch::Tensor& q_scale,
const torch::Tensor& k_scale,
const torch::Tensor& v_scale,
torch::Tensor& out
) {
int batch_size = q_q.size(0);
int seq_len = q_q.size(1);
int num_heads = q_q.size(2);
int head_dim = q_q.size(3);
int total = batch_size * num_heads * seq_len;
int threads = 256;
int blocks = (total + threads - 1) / threads;
attention_kernel<<<blocks, threads>>>(
q_q.data_ptr<int8_t>(),
k_q.data_ptr<int8_t>(),
v_q.data_ptr<int8_t>(),
q_scale.data_ptr<float>(),
k_scale.data_ptr<float>(),
v_scale.data_ptr<float>(),
out.data_ptr<float>(),
batch_size, seq_len, num_heads, head_dim
);
}
"""
# Load the CUDA extension
attention_extension = load_inline(
name='attention_extension',
cpp_sources='',
cuda_sources=cuda_source,
functions=['attention_cuda'],
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, q_q, k_q, v_q, q_scale, k_scale, v_scale):
# Create output tensor
batch_size = q_q.size(0)
seq_len = q_q.size(1)
num_heads = q_q.size(2)
head_dim = q_q.size(3)
out = torch.empty(batch_size, seq_len, num_heads, head_dim,
dtype=torch.float32, device=q_q.device)
attention_extension.attention_cuda(q_q, k_q, v_q, q_scale, k_scale, v_scale, out)
return out
```
Wait, I need to be careful about the indexing. Let me check the tensor layout again.
Looking at the reference:
- q_q.shape = (batch_size, seq_len, num_heads, head_dim)
- q_scale.shape = (num_heads,)
In PyTorch, tensors are stored in row-major order, so:
- q_q[b, t, h, d] is at offset b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d
Let me verify this is correct.
Actually, wait - in the reference, the scales are indexed as:
```python
q = q_q.float() * q_scale.view(1, 1, -1, 1)
```
So q_scale[h] is used for all positions of head h. That matches my understanding.
Let me also double-check the indexing for q_q, k_q, v_q. These are int8 tensors, and PyTorch should store them in the same memory layout as float tensors.
Now let me think about optimization. The kernel I wrote has several inefficiencies:
1. Repeated memory loads - each thread loads q[b,t,h,:] and then k[b,i,h,:] and v[b,i,h,:] for each i
2. The inner loops over head_dim could be optimized
One optimization is to use shared memory to cache k and v for each (b, h) pair, then have all threads for that pair share the cached data.
But given the small problem size, let me first check if the simple kernel works, then optimize.
Let me also think about using tensor cores or other hardware features. With int8 inputs, we could potentially use tensor cores for the matmul operations. But that's more complex.
For now, let me write a clean implementation:
```python
import math
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
__global__ void attention_kernel(
const int8_t* __restrict__ q_q,
const int8_t* __restrict__ k_q,
const int8_t* __restrict__ v_q,
const float* __restrict__ q_scale,
const float* __restrict__ k_scale,
const float* __restrict__ v_scale,
float* __restrict__ out,
int batch_size, int seq_len, int num_heads, int head_dim
) {
// Each thread handles one (b, h, t) tuple
int idx = blockIdx.x * blockDim.x + threadIdx.x;
int total = batch_size * num_heads * seq_len;
if (idx >= total) return;
int t = idx % seq_len;
int h = (idx / seq_len) % num_heads;
int b = idx / (seq_len * num_heads);
float rsqrt_head_dim = 1.0f / sqrtf((float)head_dim);
float qs = q_scale[h];
float ks = k_scale[h];
float vs = v_scale[h];
// Strides
int stride_b = seq_len * num_heads * head_dim;
int stride_t = num_heads * head_dim;
int stride_h = head_dim;
int base_b = b * stride_b;
extern __shared__ float shared_mem[];
float* local_scores = shared_mem + threadIdx.x * (seq_len + 1); // Max seq_len scores per thread
// Load q for this position
int q_base = base_b + t * stride_t + h * stride_h;
// Compute scores for positions 0 to t
float max_score = -INFINITY;
for (int i = 0; i <= t; i++) {
float dot = 0.0f;
int k_base = base_b + i * stride_t + h * stride_h;
for (int d = 0; d < head_dim; d++) {
float q_val = static_cast<float>(q_q[q_base + d]) * qs;
float k_val = static_cast<float>(k_q[k_base + d]) * ks;
dot += q_val * k_val;
}
local_scores[i] = dot * rsqrt_head_dim;
if (local_scores[i] > max_score) max_score = local_scores[i];
}
// Softmax
float sum_exp = 0.0f;
for (int i = 0; i <= t; i++) {
local_scores[i] = expf(local_scores[i] - max_score);
sum_exp += local_scores[i];
}
float inv_sum = 1.0f / sum_exp;
for (int i = 0; i <= t; i++) {
local_scores[i] *= inv_sum;
}
// Compute output
for (int d = 0; d < head_dim; d++) {
float sum = 0.0f;
for (int i = 0; i <= t; i++) {
int v_base = base_b + i * stride_t + h * stride_h;
float v_val = static_cast<float>(v_q[v_base + d]) * vs;
sum += local_scores[i] * v_val;
}
int out_base = base_b + t * stride_t + h * stride_h;
out[out_base + d] = sum;
}
}
void attention_cuda(
const torch::Tensor& q_q,
const torch::Tensor& k_q,
const torch::Tensor& v_q,
const torch::Tensor& q_scale,
const torch::Tensor& k_scale,
const torch::Tensor& v_scale,
torch::Tensor& out
) {
int batch_size = q_q.size(0);
int seq_len = q_q.size(1);
int num_heads = q_q.size(2);
int head_dim = q_q.size(3);
int total = batch_size * num_heads * seq_len;
int threads = 256;
int blocks = (total + threads - 1) / threads;
// Shared memory for scores (max seq_len per thread)
size_t shared_mem_size = threads * 72 * sizeof(float); // Assuming max seq_len = 72
attention_kernel<<<blocks, threads, shared_mem_size>>>(
q_q.data_ptr<int8_t>(),
k_q.data_ptr<int8_t>(),
v_q.data_ptr<int8_t>(),
q_scale.data_ptr<float>(),
k_scale.data_ptr<float>(),
v_scale.data_ptr<float>(),
out.data_ptr<float>(),
batch_size, seq_len, num_heads, head_dim
);
}
"""
attention_extension = load_inline(
name='attention_extension',
cpp_sources='',
cuda_sources=cuda_source,
functions=['attention_cuda'],
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
def forward(self, q_q, k_q, v_q, q_scale, k_scale, v_scale):
batch_size = q_q.size(0)
seq_len = q_q.size(1)
num_heads = q_q.size(2)
head_dim = q_q.size(3)
out = torch.empty(batch_size, seq_len, num_heads, head_dim,
dtype=torch.float32, device=q_q.device)
attention_extension.attention_cuda(q_q, k_q, v_q, q_scale, k_scale, v_scale, out)
return out
```
Hmm, but the shared memory usage might be too much - 256 threads * 72 floats = 72KB, which exceeds the typical 48KB limit.
Let me not use shared memory for now and just use local arrays. With seq_len=72, each thread needs 72 floats = 288 bytes for scores. With 256 threads, that's 72KB of local memory, but local memory is per-thread and not shared, so it might still work (local memory is backed by global memory but accessed through a cache).
Actually, let me simplify and not use any shared memory, just use local arrays for scores.
But wait, I realize there's a better approach. Instead of having each thread compute all scores, we can have a warp or block cooperate.
Actually, let me think about this differently. The problem is:
- We need to compute scores[i] = (k[i] @ q[t]) * scale for i = 0 to t
- Then softmax
- Then weighted sum of v
This is essentially a causal attention computation. The standard approach is:
1. Compute all scores at once: scores = q @ k.T * scale
2. Apply causal mask
3. Apply softmax
4. Multiply by v
But in CUDA, we can be more efficient by using a tiled approach.
For the given sizes (batch=2, seq=72, heads=4, head_dim=32), let me think about the optimal thread configuration.
Total work: 2 * 72 * 4 = 576 positions to compute, each producing 32 output values.
If we use 256 threads per block, we need 576 / 256 = 2.25 -> 3 blocks. That's very few blocks.
Actually, wait. Let me reconsider the indexing. I had:
```
int t = idx % seq_len;
int h = (idx / seq_len) % num_heads;
int b = idx / (seq_len * num_heads);
```
This maps idx to (b, h, t) in a specific order. Let me verify:
- idx = 0 -> t=0, h=0, b=0
- idx = 71 -> t=71, h=0, b=0
- idx = 72 -> t=0, h=1, b=0
- idx = 287 -> t=71, h=3, b=0
- idx = 288 -> t=0, h=0, b=1
- idx = 575 -> t=71, h=3, b=1
Total = 576 threads needed. With 256 threads per block, we need 576/256 = 2.25 -> 3 blocks.
This seems reasonable. Let me write the kernel more carefully.
Actually, I just realized I should check the GPU specs first to understand the hardware better.