reasoning
Let me analyze this task. I need to implement a custom CUDA kernel for a "scalar-gated delta-rule linear attention recurrence" operation.
Looking at the reference implementation:
1. It maintains a state matrix of shape (batch_size, num_heads, head_dim, head_dim)
2. For each timestep t:
- gate_t: shape (batch_size, num_heads, 1, 1)
- beta_t: shape (batch_size, num_heads, 1, 1)
- kv_outer: shape (batch_size, num_heads, head_dim, head_dim) - outer product of k and v
- state = gate_t * state + beta_t * kv_outer
- y_t = sum over dim=-2 of (q[t] * state), resulting in shape (batch_size, num_heads, head_dim)
The problem is that this is a sequential operation - each step depends on the previous state. This makes it challenging to parallelize across the sequence dimension directly.
However, I can parallelize across batch and head dimensions. Let me think about the CUDA implementation:
For each (batch, head) pair, we need to:
1. Initialize state to zeros
2. For t = 0 to seq_len-1:
- state = gate[b,h,t] * state + beta[b,h,t] * (k[b,t,h,:] outer v[b,t,h,:])
- output[b,t,h,:] = q[b,t,h,:] @ state
The key insight is that for each (batch, head) pair, the computation is independent. So I can assign each (batch, head) pair to a thread block or warp.
The inner loop across time steps involves:
- Reading k, v, gate, beta for that timestep
- Computing kv outer product
- Updating state
- Computing query output
Let me think about memory access patterns and optimization:
- State is head_dim x head_dim = 24 x 24 = 576 floats, about 2.3KB
- This fits well in shared memory, which should reduce DRAM bandwidth pressure
- The sequential timestep dependency means I need to process timesteps in order within each block, but can parallelize the matrix-vector operations inside each step
For the kernel design, I'll assign each thread block to a single (batch, head) pair. Each block keeps its state matrix in registers or shared memory. The outer product computation of k[t] and v[t] produces a head_dim x head_dim matrix, which can be distributed across threads. The output computation y[t] = q[t] @ state involves matrix-vector multiplication that all threads participate in. The state update then applies gate and beta scaling with the outer product addition. I need to reconsider the approach here. With head_dim = 24, a single warp of 32 threads could handle each (batch, head) pair effectively. Each thread could compute a row of the state matrix, where it processes k[t, row] and v[t, :] to form the outer product row, then applies the gate and beta scaling. For the query output, each thread would need to access all columns of state to compute the dot product with the query vector, but that creates a synchronization challenge since different threads own different rows.
Actually, the output computation is simpler: each thread computes one element of the output by taking the dot product between q[t, row] and the corresponding row of state. The thread that owns state[row, :] can directly compute this without cross-thread dependencies, making it embarrassingly parallel.
For the state update, each thread also handles its assigned row: scaling the previous state row by gate, then adding beta multiplied by the outer product row. This keeps the computation per thread self-contained.
The approach scales naturally—each thread manages one row of the state matrix, computes one output element per timestep, and updates its row independently. This eliminates the need for synchronization or shared memory entirely, though register pressure becomes the main constraint with the full state matrix stored across threads.
For head_dim = 24, each thread maintains 24 floats in registers along with several scalar values like gate_t, beta_t, k_t, and v_t, totaling roughly 30 floats or 120 bytes per thread. The full block with 96 threads would consume around 11.5KB of register space, which fits reasonably in GPU register files. The computation itself stays lightweight: the kv_outer row requires just 24 multiplications, the state update needs 24 multiply-accumulate operations, and the query output row takes 24 additional operations per timestep.
The memory access pattern involves reading inputs from global memory for each timestep and writing outputs back, keeping all intermediate state entirely in registers. This creates a very clean access pattern where computation happens on fast register storage while only the necessary input/output data moves through global memory.
For larger head dimensions, the approach might need adjustment, but the provided head_dim of 24 works well within register capacity. Alternatively, shared memory could house the state matrix with each thread processing multiple rows, though this introduces synchronization complexity.
The implementation will follow the thread-per-row model: a kernel processes all timesteps sequentially with each thread responsible for one output element per timestep, loading query, key, and value components as needed. float* output,
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 row = threadIdx.x;
if (row >= head_dim) return;
// State: each thread holds one row
// But wait, we need to share state across threads in the same warp...
// Actually, let me reconsider. Each thread computes output[b, t, h, row]
// which requires state[row, :] dotted with q[t, h, row]... no wait
// y[t, h, :] = q[t, h, :] @ state
// y[t, h, j] = sum_i q[t, h, i] * state[i, j]
// So each thread should compute one column of output, not one row!
// Thread j computes y[t, h, j] = sum_i q[t, h, i] * state[i, j]
}
```
Actually, I need to think about this more carefully.
y_t = (q[:, t].unsqueeze(-1) * state).sum(dim=-2)
q[:, t] has shape (batch_size, num_heads, head_dim)
state has shape (batch_size, num_heads, head_dim, head_dim)
q[:, t].unsqueeze(-1) has shape (batch_size, num_heads, head_dim, 1)
(q[:, t].unsqueeze(-1) * state) has shape (batch_size, num_heads, head_dim, head_dim)
.sum(dim=-2) gives shape (batch_size, num_heads, 1, head_dim) -> effectively (batch_size, num_heads, head_dim)
So y_t[h, d] = sum over i of q[t, h, i] * state[h, i, d]
This means each output element y[h, d] depends on one column of state (state[:, d]) and all elements of q.
Now the question is: how should I organize threads? If each thread computes one output element y[d], then it needs access to the d-th column of state. But state is being updated over time, and different threads need different columns. The outer product computation means each thread contributes to one column of the state update.
Actually, if I think about it differently: the kv outer product produces a matrix where each row comes from k values and each column from v values. When updating state, I'm adding this outer product scaled by beta. For computing the output, I need to sum q weighted by each column of state. So thread j would need access to state[:, j] to compute y[j], but it would need k[t, :] to update its column of state, and the outer product structure means each column j gets contributions from v[t, j] scaled by all k values. So each thread computing output element j can independently update state[:, j] by computing k[t, :] * v[t, j] and adding it to its column. This means no synchronization is needed between threads. Each thread maintains its own portion of the state vector and processes all timesteps independently.
For the kernel, I'll assign one thread per output element per head. Thread j will handle state[j] and output[j], computing the weighted sum at each timestep while updating its state component with the appropriate scaling factors. I need to reconsider the computation flow. The state matrix gets updated first with the gate decay and beta-weighted outer product, then the query reads from that updated state. So the loop should be: compute the kv outer product, decay the state by gate, add the scaled outer product, then use the query to produce the output. Each timestep follows this sequence—update state, then read from it. I should verify whether the PyTorch implementation matches this order or if there's a discrepancy I'm missing.
For the CUDA kernel, I'm setting up the indexing so each block handles one batch-head pair, with threads assigned to different dimensions of the head. The kernel needs to iterate through timesteps, maintaining state across iterations while applying the recurrent update rule. }
}
// Store output
output[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d] = y;
}
The kernel processes each timestep by reading query, key, and value elements for the current position, then updating the recurrent state through a gated mechanism. The state accumulates the outer product of key-value pairs weighted by beta, while the gate provides decay. Finally, the output is computed as the dot product between the query and the updated state.
I need to reconsider the state storage approach. Storing `state[head_dim]` in registers should work fine— compilers handle small fixed-size arrays well, and with 96 total threads per block, register usage stays reasonable at around 11.5KB per block.
The memory access pattern reveals inefficiencies though. Each thread reads k[t, h, :] repeatedly and scatters output writes. A better approach would use shared memory to load k[t, h, :] once per block, letting all threads reuse it. Coalesced writes to the output array would also improve throughput. I should restructure the kernel to have threads compute different head dimensions while sharing the key vector across the block.
Actually, switching to column-major state storage makes the math cleaner. Each thread maintains one column of the state matrix (the portion needed for one output element). For each timestep, the thread loads the query and value element, then broadcasts the key vector through shared memory so all threads can access it efficiently. This way, each thread updates its state column by scaling the key vector with its value element, then applies the gate for the next iteration.
```cpp
#include <cuda.h>
#include <cuda_runtime.h>
__global__ void attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
const float* __restrict__ beta,
float* __restrict__ output,
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 d = threadIdx.x;
if (d >= head_dim) return;
// State: one column per thread
float state[24]; // head_dim = 24
for (int i = 0; i < head_dim; i++) {
state[i] = 0.0f;
}
for (int t = 0; t < seq_len; t++) {
float g = gate[b * seq_len * num_heads + t * num_heads + h];
float b_val = beta[b * seq_len * num_heads + t * num_heads + h];
float v_t = v[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d];
// Update state: state *= g, then add beta * k * v_t
float y = 0.0f;
for (int i = 0; i < head_dim; i++) {
float k_t = k[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + i];
state[i] = g * state[i] + b_val * k_t * v_t;
float q_t = q[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + i];
y += q_t * state[i];
}
output[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d] = y;
}
}
```
The inner loop handles both the state update (decaying by gate, then adding the key-value contribution) and the output computation as a dot product between query and state vectors. Each thread writes its output element directly to global memory. I see the pattern now - after updating state[i] with the new value from the recurrence, I multiply it by q_t and accumulate into y. So for each timestep, I'm computing the sum over all dimensions of q_t * state[i], where state[i] has already been updated for that position.
The indexing is consistent across the kernel - I'm reading q, k, and v at the same linear offset within each timestep, and storing the accumulated result y at the corresponding output position. The state update happens in-place before being used in the dot product calculation.
For the wrapper, I need to set up the kernel launch parameters. The grid should have one block per (batch, head) pair, so grid.x = batch_size * num_heads. Each block gets head_dim threads to process all dimension elements in parallel. Then I call the kernel with the input and output tensors along with their dimensions.
One thing to check - the tensor strides should be contiguous since they're coming from `torch.randn`, but I should verify that assumption holds. The expected layout is [batch_size, seq_len, num_heads, head_dim] with the standard PyTorch contiguity pattern, so as long as the inputs match that, the kernel indexing should work correctly.
For the kernel itself, I'm using `state[24]` which the compiler should handle efficiently - likely keeping it in registers rather than local memory since the size is small and fixed. The real concern is the inner loop where each iteration reads k and q values that aren't coalesced across threads. Loading those values into shared memory first would let all threads in a warp access them more efficiently without each thread doing separate global memory reads. const int head_dim) {
int b = blockIdx.x / num_heads;
int h = blockIdx.x % num_heads;
int d = threadIdx.x;
if (d >= head_dim) return;
// State in registers
float state[24];
for (int i = 0; i < head_dim; i++) {
state[i] = 0.0f;
}
// Shared memory for k and q at each timestep
extern __shared__ float shared[];
float* k_shared = shared;
float* q_shared = shared + head_dim;
for (int t = 0; t < seq_len; t++) {
float g = gate[b * seq_len * num_heads + t * num_heads + h];
float b_val = beta[b * seq_len * num_heads + t * num_heads + h];
float v_t = v[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d];
// Load k and q into shared memory
if (d < head_dim) {
k_shared[d] = k[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d];
q_shared[d] = q[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d];
}
__syncthreads();
float y = 0.0f;
for (int i = 0; i < head_dim; i++) {
state[i] = g * state[i] + b_val * k_shared[i] * v_t;
y += q_shared[i] * state[i];
}
output[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d] = y;
}
}
The kernel maintains state in registers while using shared memory to cache the key and query values for each timestep. Each thread processes one dimension, updating its state through the recurrence relation and computing the output as a dot product between the query and updated state vectors. , y = %.6f\n", t, d, g, b_val, v_t, k_shared[0], q_shared[0], state[0], y);
}
output[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d] = y;
}
}
torch::Tensor attention_cuda(
torch::Tensor q,
torch::Tensor k,
torch::Tensor v,
torch::Tensor gate,
torch::Tensor beta) {
int batch_size = q.size(0);
int seq_len = q.size(1);
int num_heads = q.size(2);
int head_dim = q.size(3);
auto output = torch::empty({batch_size, seq_len, num_heads, head_dim}, q.options());
int blocks = batch_size * num_heads;
int threads = head_dim;
int shared_mem = 2 * head_dim * sizeof(float);
attention_kernel<<<blocks, threads, shared_mem>>>(
q.data_ptr<float>(),
k.data_ptr<float>(),
v.data_ptr<float>(),
gate.data_ptr<float>(),
beta.data_ptr<float>(),
output.data_ptr<float>(),
batch_size, seq_len, num_heads, head_dim);
return output;
}
```
I'm also loading q into shared memory alongside k for the same efficiency gains. The function signature is corrected to match the expected interface. I should verify the kernel works correctly now and check if there are any remaining issues with the indexing or shared memory setup.
call → compile_kernel
callcompile_kernel
{
"kernel_code": "import torch\nimport torch.nn as nn\nfrom torch.utils.cpp_extension import load_inline\n\n# Define the CUDA kernel\ncuda_source = \"\"\"\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n__global__ void attention_kernel(\n const float* __restrict__ q,\n const float* __restrict__ k,\n const float* __restrict__ v,\n const float* __restrict__ gate,\n const float* __restrict__ beta,\n float* __restrict__ output,\n int batch_size, int seq_len, int num_heads, int head_dim) {\n \n int b = blockIdx.x / num_heads;\n int h = blockIdx.x % num_heads;\n int d = threadIdx.x;\n \n if (d >= head_dim) return;\n \n // State: one column per thread (head_dim elements)\n // Thread d computes output[..., d] using state[:, d]\n float state[24]; // head_dim = 24, fixed for register allocation\n \n for (int i = 0; i < head_dim; i++) {\n state[i] = 0.0f;\n }\n \n // Shared memory for k and q at each timestep\n extern __shared__ float shared[];\n float* k_shared = shared;\n float* q_shared = shared + head_dim;\n \n for (int t = 0; t < seq_len; t++) {\n float g = gate[b * seq_len * num_heads + t * num_heads + h];\n float b_val = beta[b * seq_len * num_heads + t * num_heads + h];\n float v_t = v[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d];\n \n // Coalesced load k and q into shared memory\n k_shared[d] = k[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d];\n q_shared[d] = q[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d];\n __syncthreads();\n \n // Compute output and update state\n float y = 0.0f;\n for (int i = 0; i < head_dim; i++) {\n // Update state: state[:, d] = g * state[:, d] + b_val * k[:] * v[d]\n state[i] = g * state[i] + b_val * k_shared[i] * v_t;\n // Accumulate output: y = q[:] @ state[:, d]\n y += q_shared[i] * state[i];\n }\n \n output[b * seq_len * num_heads * head_dim + t * num_heads * head_dim + h * head_dim + d] = y;\n }\n}\n\ntorch::Tensor attention_cuda(\n torch::Tensor q,\n torch::Tensor k,\n torch::Tensor v,\n torch::Tensor gate,\n torch::Tensor beta) {\n \n int batch_size = q.size(0);\n int seq_len = q.size(1);\n int num_heads = q.size(2);\n int head_dim = q.size(3);\n \n auto output = torch::empty({batch_size, seq_len, num_heads, head_dim}, q.options());\n \n int blocks = batch_size * num_heads;\n int threads = head_dim;\n int shared_mem = 2 * head_dim * sizeof(float);\n \n attention_kernel<<<blocks, threads, shared_mem>>>(\n q.data_ptr<float>(),\n k.data_ptr<float>(),\n v.data_ptr<float>(),\n gate.data_ptr<float>(),\n beta.data_ptr<float>(),\n output.data_ptr<float>(),\n batch_size, seq_len, num_heads, head_dim);\n \n return output;\n}\n\"\"\"\n\ncpp_source = \"torch::Tensor attention_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor gate, torch::Tensor beta);\"\n\n# Load the CUDA extension\nattention_extension = load_inline(\n name='attention_extension',\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=['attention_cuda'],\n extra_cuda_cflags=['-O3'],\n verbose=True\n)\n\nclass ModelNew(nn.Module):\n def __init__(self):\n super().__init__()\n \n def forward(self, q, k, v, gate, beta):\n return attention_extension.attention_cuda(q, k, v, gate, beta)"
}