reasoning
Now let me analyze the computation more carefully:
For each timestep t:
1. gate_t = gate[:, t].unsqueeze(-2) # shape: (batch, 1, num_heads, head_dim)
2. kv_outer = k[:, t].unsqueeze(-1) * v[:, t].unsqueeze(-2) # shape: (batch, 1, num_heads, head_dim, head_dim) - outer product
3. state = gate_t * state + kv_outer # state shape: (batch, num_heads, head_dim, head_dim)
4. y_t = (q[:, t].unsqueeze(-1) * state).sum(dim=-2) # shape: (batch, 1, num_heads, head_dim)
So the state is a matrix of shape (head_dim, head_dim) for each (batch, num_heads).
The key operations:
- kv_outer: outer product of k and v vectors -> produces head_dim x head_dim matrix
- state update: element-wise multiply with gate broadcast, plus kv_outer
- output: q * state (with sum over one dimension)
Let me think about the parallelization:
- We have batch_size * num_heads independent state recurrences
- Each recurrence is over seq_len steps
- Each step involves:
- outer product: head_dim * head_dim operations
- state update: head_dim * head_dim operations
- output computation: head_dim * head_dim operations
The state matrix is small (head_dim=24), so state for one (batch, head) is 24*24 = 576 floats = 2304 bytes.
Total state for all batches and heads: 2 * 4 * 576 = 4608 floats = ~18KB.
Let me design a kernel where each thread block handles one (batch, num_head) pair, iterating through all timesteps.
Actually, for better parallelism, I could have each warp handle one (batch, head) pair, with multiple warps per block handling different (batch, head) pairs.
Wait, let me reconsider. With head_dim=24, the state matrix is 24x24. Each element needs:
- Load gate value
- Compute kv_outer
- Update state
- Compute output
Let me think about tiling. If we have a thread block with 256 threads:
- We could have each thread handle one or more elements of the state matrix
- 24x24 = 576 elements
- With 256 threads, we need 576/256 ≈ 2.25, so about 2-3 elements per thread on average
Actually, let me think differently. The recurrence is inherently sequential over timesteps, but within each timestep, the operations on the state matrix can be parallelized.
Here's my approach:
1. Each thread block handles one (batch, head) pair
2. All threads in a block cooperatively process the state matrix
3. Iterate through timesteps sequentially within the kernel
Let me design for head_dim=24:
- State matrix: 24x24 = 576 elements
- Use 256 threads per block (one warp can cooperatively do a lot)
- Each thread handles a subset of state elements
Actually, let me simplify and use a 2D thread block layout:
- blockDim = (head_dim, head_dim) = (24, 24) = 576 threads
- But that might be less efficient due to small block size
Let me try another approach:
- Use 128 threads per block
- Handle 576 elements with some threads handling multiple elements
- Or reshape the problem differently
Wait, I should think about what dimensions might vary. The reference mentions batch_size=2, seq_len=128, num_heads=4, head_dim=24. But the kernel should work for arbitrary values.
Let me try a simpler approach first:
- Launch a 2D grid: (batch_size * num_heads) blocks in x, 1 in y
- Each block block handles one (batch, head) pair
- Within each block, use enough threads to parallelize the state operations
For the state operations, since state is head_dim x head_dim, I'll use a 2D thread block:
- blockDim.x = head_dim, blockDim.y = head_dim
- Total threads = head_dim^2
But this might be too small for large head_dim. Let me make it more flexible.
Actually, the problem says head_dim=24, which is fixed. With head_dim=24:
- Thread block: (24, 24) = 576 threads
- This is a reasonable size for modern GPUs
Let me design the kernel:
```cpp
__global__ void gated_delta_attention(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ output,
int batch_size, int seq_len, int num_heads, int head_dim) {
// Each block handles one (batch, head) pair
int batch = blockIdx.x / num_heads;
int head = blockIdx.x % num_heads;
// Thread indices for the state matrix
int row = threadIdx.x;
int col = threadIdx.y;
// State is stored in shared memory for efficiency
extern __shared__ float state[];
// Initialize state to zero
state[row * head_dim + col] = 0.0f;
__syncthreads();
// Iterate through timesteps
for (int t = 0; t < seq_len; t++) {
// Load inputs for this timestep
float q_val = q[batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + col];
float k_val = k[...similar...];
float v_val = v[...similar...];
float gate_val = gate[...similar...];
// Compute outer product contribution
float kv = k_val * v_val; // But this needs to be outer product...
// Actually, each thread handles one element of the state matrix
// The outer product k * v^T at position (i, j) is k[i] * v[j]
// My thread handles position (row, col), so:
// k_row = k[row], v_col = v[col]
// Let me reconsider the indexing...
}
}
```
Hmm, let me think about this more carefully.
The state matrix is head_dim x head_dim. Thread (tx, ty) handles element state[tx, ty].
At each timestep:
- k is a vector of length head_dim
- v is a vector of length head_dim
- gate is a vector of length head_dim (different from what I initially thought!)
- q is a vector of length head_dim
Wait, I need to re-read the code:
```python
gate_t = gate[:, t].unsqueeze(-2) # gate[:, t] has shape (batch, num_heads, head_dim), unsqueeze gives (batch, 1, num_heads, head_dim)
kv_outer = k[:, t].unsqueeze(-1) * v[:, t].unsqueeze(-2) # k[:, t] is (batch, num_heads, head_dim), unsqueeze(-1) gives (batch, num_heads, head_dim, 1), v[:, t].unsqueeze(-2) gives (batch, num_heads, 1, head_dim), product gives (batch, num_heads, head_dim, head_dim)
```
So:
- gate_t has shape (batch, 1, num_heads, head_dim) - this is used for broadcasting
- kv_outer has shape (batch, num_heads, head_dim, head_dim)
- state has shape (batch, num_heads, head_dim, head_dim)
When we do `gate_t * state`:
- gate_t is (batch, 1, num_heads, head_dim)
- state is (batch, num_heads, head_dim, head_dim)
- Broadcasting will expand gate_t to (batch, num_heads, head_dim, head_dim)
- The last dimension of gate_t (head_dim) broadcasts with the last dimension of state (head_dim)
- But wait, there's also a broadcast in the second-to-last dimension?
Let me trace through more carefully:
- gate_t: (batch, 1, num_heads, head_dim)
- state: (batch, num_heads, head_dim, head_dim)
For broadcasting:
- dim 0: batch matches
- dim 1: 1 vs num_heads -> broadcasts from 1 to num_heads
- dim 2: num_heads vs head_dim... wait these are different dimensions!
Hmm, let me look at the shapes again:
- gate_t: (batch, 1, num_heads, head_dim)
- state: (batch, num_heads, head_dim, head_dim)
The unsqueeze(-2) adds a dimension at position -2, so:
- gate[:, t] has shape (batch, num_heads, head_dim)
- gate[:, t].unsqueeze(-2) has shape (batch, num_heads, 1, head_dim)
Let me re-check the reference code:
```python
gate_t = gate[:, t].unsqueeze(-2)
```
gate has shape (batch, seq_len, num_heads, head_dim)
gate[:, t] has shape (batch, num_heads, head_dim)
gate[:, t].unsqueeze(-2) has shape (batch, num_heads, 1, head_dim)
OK so gate_t has shape (batch, num_heads, 1, head_dim).
Now for the state update:
- state has shape (batch, num_heads, head_dim, head_dim)
- gate_t * state: gate_t is (batch, num_heads, 1, head_dim), state is (batch, num_heads, head_dim, head_dim)
- Broadcasting: dim 2 is 1 vs head_dim, so gate_t broadcasts to (batch, num_heads, head_dim, head_dim)
- So each "row" of state (across the last dimension) is multiplied by the corresponding gate value
So gate[:, t, head, i] scales all of state[batch, head, i, :] (one row of the state matrix).
For the output:
```python
y_t = (q[:, t].unsqueeze(-1) * state).sum(dim=-2)
```
- q[:, t] has shape (batch, num_heads, head_dim)
- q[:, t].unsqueeze(-1) has shape (batch, num_heads, head_dim, 1)
- state has shape (batch, num_heads, head_dim, head_dim)
- q[:, t].unsqueeze(-1) * state has shape (batch, num_heads, head_dim, head_dim)
- Broadcasting: dim 3 is 1 vs head_dim, so q broadcasts
- So q[:, t, head, i] * state[batch, head, i, :] scales each row
- sum(dim=-2) sums over the second-to-last dimension (the "row" dimension)
- Result has shape (batch, num_heads, head_dim)
So the output y_t[batch, head, j] = sum over i of q[:, t, head, i] * state[batch, head, i, j]
This is essentially q^T @ state, where q is (head_dim,), state is (head_dim, head_dim), result is (head_dim,)
Now I understand the computation better. Let me design the kernel:
Each (batch, head) pair has:
- A state matrix of shape (head_dim, head_dim)
- At each timestep:
- Load q, k, v, gate vectors of length head_dim
- Update state: state[i, j] = gate[i] * state[i, j] + k[i] * v[j]
- Compute output: y[j] = sum_i (q[i] * state[i, j])
This is essentially:
- state = diag(gate) @ state + k @ v^T
- y = q^T @ state
For the kernel design, I'll have each thread handle one or more elements of the output vector y, and cooperatively update the state matrix.
Actually, let me think about this differently. The state update is:
state[i, j] = gate[i] * state[i, j] + k[i] * v[j]
And the output is:
y[j] = sum_i q[i] * state[i, j]
So if I have threads that each handle one column of the state (fixed j), then:
- The state update for column j: state[i, j] = gate[i] * state[i, j] + k[i] * v[j] for all i
- The output computation for column j: y[j] = sum_i q[i] * state[i, j]
This is nice because v[j] is constant for column j, so:
- For each column j_j, we need to load v[j]
- Then for each row i, update state[i, j]
Let me try a kernel where:
- Thread block handles one (batch, head) pair
- Threads are organized as (head_dim,) - each thread handles one column of state
- Or threads are organized as (head_dim, head_dim) - each thread handles one element of state
Let me try the 2D approach first:
```cpp
template<int HEAD_DIM>
__global__ void gated_delta_attention_kernel(...) {
int batch = blockIdx.x / num_heads;
int head = blockIdx.x % num_heads;
int i = threadIdx.y; // row index
int j = threadIdx.x; // column index
__shared__ float state[HEAD_DIM][HEAD_DIM];
__shared__ float q_vec[HEAD_DIM];
__shared__ float k_vec[HEAD_DIM];
__shared__ float v_vec[HEAD_DIM];
__shared__ float gate_vec[HEAD_DIM];
__shared__ float y_vec[HEAD_DIM];
// Initialize state
state[i][j] = 0.0f;
for (int t = 0; t < seq_len; t++) {
// Cooperatively load vectors
if (i == 0) {
q_vec[j] = q[batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j];
k_vec[j] = ...;
v_vec[j] = ...;
gate_vec[j] = ...;
}
__syncthreads();
// Update state
state[i][j] = gate_vec[i] * state[i][j] + k_vec[i] * v_vec[j];
__syncthreads();
// Compute output (this is a reduction)
// y[j] = sum_i q[i] * state[i][j]
// Each thread computes q[i] * state[i][j] for its (i,j)
// Then threads in a column (same j) reduce
float contrib = q_vec[i] * state[i][j];
// Reduction within column
for (int offset = HEAD_DIM / 2; offset > 0; offset >>= 1) {
if (i < offset) {
contrib += contrib_shuffled...; // Need shared memory for this
}
}
if (i == 0) {
y_vec[j] = contrib; // After reduction
}
__syncthreads();
// Write output
if (i == 0) {
output[batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j] = y_vec[j];
}
}
}
```
This is getting complex. Let me simplify and use a cleaner approach.
Since HEAD_DIM=24 is small, I can use a 1D thread block with 32 threads (one warp), and have each warp handle one (batch, head) pair. Each thread handles a subset of the state matrix.
Actually, with 24*24 = 576 elements, and 32 threads per warp, we need ~18 elements per thread. That's manageable.
Let me redesign:
```cpp
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ output,
int batch_size, int seq_len, int num_heads, int head_dim) {
int batch = blockIdx.x / num_heads;
int head = blockIdx.x % num_heads;
int tid = threadIdx.x;
int warp_size = 32;
extern __shared__ float smem[];
float* state = smem; // head_dim * head_dim
float* q_vec = state + head_dim * head_dim;
float* k_vec = q_vec + head_dim;
float* v_vec = k_vec + head_dim;
float* gate_vec = v_vec + head_dim;
float* y_vec = gate_vec + head_dim;
// Initialize state
for (int idx = tid; idx < head_dim * head_dim; idx += blockDim.x) {
state[idx] = 0.0f;
}
__syncthreads();
// Base indices for this (batch, head)
int base_idx = batch * seq_len * num_heads * head_dim + head * head_dim;
int batch_stride = seq_len * num_heads * head_dim;
int head_stride = head_dim;
for (int t = 0; t < seq_len; t++) {
// Load vectors
for (int i = tid; i < head_dim; i += blockDim.x) {
q_vec[i] = q[batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + i];
k_vec[i] = k[...];
v_vec[i] = v[...];
gate_vec[i] = gate[...];
}
__syncthreads();
// Update state and compute output
for (int i = tid; i < head_dim * head_dim; i += blockDim.x) {
int row = i / head_dim;
int col = i % head_dim;
state[i] = gate_vec[row] * state[i] + k_vec[row] * v_vec[col];
}
__syncthreads();
// Compute output: y[j] = sum_i q[i] * state[i, j]
for (int j = tid; j < head_dim; j += blockDim.x) {
float sum = 0.0f;
for (int i = 0; i < head_dim; i++) {
sum += q_vec[i] * state[i * head_dim + j];
}
y_vec[j] = sum;
}
__syncthreads();
// Write output
for (int i = tid; i < head_dim; i += blockDim.x) {
output[batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + i] = y_vec[i];
}
}
}
```
This is cleaner. Let me now write the full implementation.
For optimization:
1. Use enough threads per block to hide latency
2. Use shared memory efficiently
3. Consider vectorization for loads/stores
Let me start with a working version and optimize from there.
Actually, I realize I should think about the memory access patterns more carefully. The input tensors are:
- q: (batch, seq, head, dim) - strided access as we iterate through seq
- k: (batch, seq, head, dim) - same
- v: (batch, seq, head, dim) - same
- gate: (batch, seq, head, dim) - same
- output: (batch, seq, head, dim) - same
For a given (batch, head), we access q[batch, :, head, :], k[batch, :, head, :], etc. This is not contiguous in memory because of the interleave between seq and head dimensions.
Memory layout: batch * seq * head * dim
For fixed (batch, head), elements are at indices: batch * seq * head * dim + t * head * dim + head * dim + i
Which simplifies to: batch * seq * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + i
So the stride between consecutive timesteps for the same (batch, head, dim_idx) is num_heads * head_dim.
The stride between consecutive dim_idx for the same (batch, seq, head) is 1.
If we use a larger thread block, we can have multiple warps handle different (batch, head) pairs. This could improve memory throughput.
Let me design with a thread block size of 128 threads:
- Each warp (32 threads) handles one (batch, head) pair
- With 128 threads, we have 4 warps, so each block handles 4 (batch, head) pairs
Wait, but then shared memory becomes more complex because each warp needs its own state matrix.
Actually, let me simplify and have each block handle one (batch, head) pair. The shared memory usage is:
- state: head_dim * head_dim = 576 floats = 2304 bytes
- vectors: 4 * head_dim = 96 floats = 384 bytes
- output: head_dim = 24 floats = 96 bytes
Total: ~2.8 KB per block, well within the 227 KB limit.
With 132 SMs (H100), and each block using 2.8 KB, we could have up to 80 blocks per SM theoretically, but limited by other resources.
For batch=2, num_heads=4, we have 8 blocks total, which is very small. We need to parallelize differently.
Let me reconsider. Maybe I should process multiple timesteps in parallel? But that would require parallelizing the recurrence, which is not straightforward.
Alternatively, I could have each thread handle one (batch, head) pair and process all timesteps sequentially. This would give us batch * num_heads = 8 threads, which is not enough to saturate the GPU.
Hmm, this is a challenge. The problem is that the recurrence is inherently sequential over timesteps, but we have only 8 independent recurrences.
Let me look at this from a different angle. What if we parallelize over the head_dim dimension?
For each (batch, head, dim_idx), we can compute the contribution to the output:
- y[t, batch, head, dim_idx] depends on all q[t, batch, head, :] and all state[batch, head, :, dim_idx]
- The state column [:, dim_idx] evolves as:
state[i, dim_idx] = gate[i] * state[i, dim_idx] + k[i] * v[dim_idx]
So if we fix dim_idx, the state column evolution involves:
- Loading v[dim_idx] at each timestep
- Updating state[:, dim_idx] using gate[:] and k[:]
Let me try a different kernel design:
- Grid: (batch * num_heads * head_dim) blocks, each handling one (batch, head, dim_idx)
- Each thread updates one column of the state matrix
Actually, this might work better:
- Grid: (batch * num_heads) blocks in x
- Thread block: head_dim threads
- Each thread handles one column of state (one output dimension)
```cpp
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ output,
int batch_size, int seq_len, int num_heads, int head_dim) {
int batch = blockIdx.x / num_heads;
int head = blockIdx.x % num_heads;
int j = threadIdx.x; // output dimension (column of state)
// Each thread has its own state column
float state_col[HEAD_DIM]; // This is local memory, might be slow
// Initialize state
for (int i = 0; i < head_dim; i++) {
state_col[i] = 0.0f;
}
// Base indices
int base = batch * seq_len * num_heads * head_dim + head * head_dim;
for (int t = 0; t < seq_len; t++) {
// Load v[j] - this is constant for my column
int v_idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j;
float v_j = v[v_idx];
// Load q, k, gate vectors (need all of them for my column update)
// ... this requires loading head_dim elements each, which is a lot of work per thread
// Update state column: state[i, j] = gate[i] * state[i, j] + k[i] * v[j]
for (int i = 0; i < head_dim; i++) {
int idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + i;
float k_i = k[idx];
float gate_i = gate[idx];
state_col[i] = gate_i * state_col[i] + k_i * v_j;
}
// Compute output: y[j] = sum_i q[i] * state[i, j]
float y_j = 0.0f;
for (int i = 0; i < head_dim; i++) {
int idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + i;
float q_i = q[idx];
y_j += q_i * state_col[i];
}
// Write output
int out_idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j;
output[out_idx] = y_j;
}
}
```
But this approach has each thread doing O(head_dim) work for loading vectors, which is inefficient.
Let me go back to the cooperative approach within a thread block. Let me use shared memory to load vectors cooperatively.
Actually, I think the key insight is that we need to balance parallelism:
- We have only 8 independent recurrences (batch * num_heads)
- Each recurrence involves O(seq_len * head_dim^2) work
If we launch 8 thread blocks, each with many threads, we can:
1. Cooperatively load the input vectors
2. Parallelize the state update and output computation
Let me design a kernel with:
- 8 thread blocks (one per (batch, head) pair)
- Each block has enough threads to parallelize the work within each timestep
With head_dim = 24:
- State has 576 elements
- With 256 threads per block, each thread handles ~2 elements
- This gives good parallelism within each timestep
Actually, I realize we can also parallelize across head_dim at the grid level:
- Launch grid with (batch * num_heads * head_dim) blocks
- Each block handles one output dimension for one (batch, head) pair
- Thread block size can be small (e.g., 32 threads = one warp)
- Each thread handles part of the state column
Let me try this:
Grid: (batch * num_heads, head_dim) = (8, 24) = 192 blocks
Block size: 32 threads
Each block handles:
- One (batch, head, output_dim) triple
- Maintains and updates one column of the state matrix
- Computes one element of the output vector
```cpp
__global__ void gated_delta_attention_kernel(...) {
int batch = blockIdx.x / num_heads;
int head = blockIdx.x % num_heads;
int j = blockIdx.y; // output dimension
int tid = threadIdx.x;
// State column - distributed across threads
// Each thread handles head_dim / blockDim.x elements
int num_elements_per_thread = (head_dim + blockDim.x - 1) / blockDim.x;
// ... but this requires cooperation between threads for the reduction
}
```
This is getting complicated. Let me just go with a simpler approach:
- Each thread handles one (batch, head, output_dim)
- Process all timesteps sequentially
- For each timestep, update the state column and compute output
This gives us batch * num_heads * head_dim = 2 * 4 * 24 = 192 threads, which is still not great but better.
Actually, let me try a different approach. What if we parallelize over timesteps?
We can use a parallel scan/prefix sum approach:
- For each (batch, head), the state at time t is: state[t] = gate[t] * state[t-1] + k[t] * v[t]^T
- This is a linear recurrence: state[t] = A[t] * state[t-1] + B[t]
- We can parallelize using prefix scan with matrix multiplication
But this requires implementing matrix prefix scan, which is complex and might not be worth it for small head_dim.
Let me just implement a straightforward kernel first:
```cpp
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ output,
int batch_size, int seq_len, int num_heads, int head_dim) {
// Each thread handles one output dimension for one (batch, head)
int batch = blockIdx.x / num_heads;
int head = blockIdx.x % num_heads;
int j = threadIdx.x; // output dimension
if (j >= head_dim) return;
// State column - we maintain all of it in registers/local memory
// Actually, we need all state_col[i] for the update and output
// Let's use a different approach: each thread handles one (batch, head) entirely
// and iterates through seq_len
}
```
Actually, I think I'm overcomplicating this. Let me just implement a kernel where:
- Each thread handles one (batch, head) pair
- We have batch * num_heads threads (might need more blocks if > max threads per block)
- Each thread processes all timesteps and all state elements
For 8 threads, this is not great, but let's start here and optimize.
Actually, since we need more threads, let me organize as:
- Grid: (batch, num_heads)
- Block: (head_dim, head_dim) = (24, 24) = 576 threads
Each block handles one (batch, head) pair, and threads cooperatively process the state matrix.
Let me write this:
```cpp
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ output,
int batch_size, int seq_len, int num_heads, int head_dim) {
int batch = blockIdx.y;
int head = blockIdx.x;
int i = threadIdx.y; // row of state
int j = threadIdx.x; // column of state
if (i >= head_dim || j >= head_dim) return;
// State in shared memory
extern __shared__ float smem[];
float* state = smem;
float* scratch = state + head_dim * head_dim;
// Initialize state
state[i * head_dim + j] = 0.0f;
__syncthreads();
for (int t = 0; t < seq_len; t++) {
// Load inputs
float q_val, k_val, v_val, gate_val;
// Only load if within bounds
if (i < head_dim) {
int idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + i;
q_val = q[idx];
k_val = k[idx];
gate_val = gate[idx];
}
if (j < head_dim) {
int idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j;
v_val = v[idx];
}
// Share the loaded values
scratch[i] = q_val; // q[i]
scratch[head_dim + i] = k_val; // k[i]
scratch[2 * head_dim + i] = gate_val; // gate[i]
scratch[3 * head_dim + j] = v_val; // v[j]
__syncthreads();
// Now all threads can access the vectors
float q_i = scratch[i];
float k_i = scratch[head_dim + i];
float gate_i = scratch[2 * head_dim + i];
float v_j = scratch[3 * head_dim + j];
// Update state
state[i * head_dim + j] = gate_i * state[i * head_dim + j] + k_i * v_j;
__syncthreads();
// Compute output: y[j] = sum_i q[i] * state[i, j]
// Each thread computes q[i] * state[i, j]
float contrib = q_i * state[i * head_dim + j];
// Use parallel reduction across rows (same j)
// This requires threads with the same j to reduce
// ... this is tricky with 2D thread layout
// Alternative: each column (fixed j) has threads with different i
// We need to reduce across i values for each j
// Store contribution in scratch
scratch[i * head_dim + j] = contrib;
__syncthreads();
// One thread per column does the sum
if (i == 0) {
float sum = 0.0f;
for (int row = 0; row < head_dim; row++) {
sum += scratch[row * head_dim + j];
}
// Write output
int out_idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j;
output[out_idx] = sum;
}
}
}
```
Hmm, the reduction is a bottleneck. Let me use warp-level primitives for better efficiency.
Actually, since head_dim = 24, the reduction is over only 24 values, which is small. The simple loop approach should be fine.
But wait, I have an issue. When I do:
```cpp
scratch[i * head_dim + j] = contrib;
```
This is using head_dim * head_dim = 576 elements of scratch, but I haven't allocated that much.
Let me redesign the shared memory layout:
- state: head_dim * head_dim = 576 floats
- q_vec: head_dim = 24 floats
- k_vec: head_dim = 24 floats
- v_vec: head_dim = 24 floats
- gate_vec: head_dim = 24 floats
Total: 576 + 96 = 672 floats = 2688 bytes
Let me rewrite:
```cpp
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ output,
int batch_size, int seq_len, int num_heads, int head_dim) {
int batch = blockIdx.y;
int head = blockIdx.x;
int i = threadIdx.y; // row of state
int j = threadIdx.x; // column of state
// Shared memory
extern __shared__ float smem[];
float* state = smem; // [head_dim, head_dim]
float* q_vec = state + head_dim * head_dim; // [head_dim]
float* k_vec = q_vec + head_dim; // [head_dim]
float* v_vec = k_vec + head_dim; // [head_dim]
float* gate_vec = v_vec + head_dim; // [head_dim]
// Initialize state
if (i < head_dim && j < head_dim) {
state[i * head_dim + j] = 0.0f;
}
__syncthreads();
for (int t = 0; t < seq_len; t++) {
// Load vectors cooperatively
// Each thread loads one element
if (i < head_dim && j < head_dim) {
int idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j;
v_vec[j] = v[idx];
}
if (i < head_dim) {
int idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + i;
q_vec[i] = q[idx];
k_vec[i] = k[idx];
gate_vec[i] = gate[idx];
}
__syncthreads();
// Update state
if (i < head_dim && j < head_dim) {
state[i * head_dim + j] = gate_vec[i] * state[i * head_dim + j] + k_vec[i] * v_vec[j];
}
__syncthreads();
// Compute output: y[j] = sum_i q[i] * state[i, j]
// Thread (i, j) computes contribution for its element
if (i < head_dim && j < head_dim) {
float contrib = q_vec[i] * state[i * head_dim + j];
// Reduce across rows (same j)
// Use atomic add or shared memory reduction
// Let's use shared memory reduction
// Store in a temporary array
// Actually, we can reuse state for this... no, that's not safe
// Let me allocate another array in shared memory
}
}
}
```
Actually, I think I'm overcomplicating the reduction. Since head_dim is small (24), let me just have one thread per column do the entire sum:
```cpp
// Compute output: y[j] = sum_i q[i] * state[i, j]
if (i == 0 && j < head_dim) {
float sum = 0.0f;
for (int row = 0; row < head_dim; row++) {
sum += q_vec[row] * state[row * head_dim + j];
}
int out_idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j;
output[out_idx] = sum;
}
```
This is simple and correct. The thread (0, j) will compute the output for dimension j. The inefficiency is that most threads are idle during this phase.
But for head_dim = 24, this is only 24 values to sum, which is fast.
Let me write the complete kernel:
```cpp
__global__ void __launch_bounds__(256) gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ output,
int batch_size, int seq_len, int num_heads, int head_dim) {
int batch = blockIdx.y;
int head = blockIdx.x;
int i = threadIdx.y; // row of state
int j = threadIdx.x; // column of state
// Shared memory
extern __shared__ float smem[];
float* state = smem; // [head_dim, head_dim]
float* q_vec = state + head_dim * head_dim; // [head_dim]
float* k_vec = q_vec + head_dim; // [head_dim]
float* v_vec = k_vec + head_dim; // [head_dim]
float* gate_vec = v_vec + head_dim; // [head_dim]
// Initialize state
if (i < head_dim && j < head_dim) {
state[i * head_dim + j] = 0.0f;
}
__syncthreads();
for (int t = 0; t < seq_len; t++) {
// Load vectors cooperatively
if (j < head_dim) {
int v_idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j;
v_vec[j] = v[v_idx];
}
if (i < head_dim) {
int idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + i;
q_vec[i] = q[idx];
k_vec[i] = k[idx];
gate_vec[i] = gate[idx];
}
__syncthreads();
// Update state
if (i < head_dim && j < head_dim) {
state[i * head_dim + j] = gate_vec[i] * state[i * head_dim + j] + k_vec[i] * v_vec[j];
}
__syncthreads();
// Compute output: y[j] = sum_i q[i] * state[i, j]
if (i == 0 && j < head_dim) {
float sum = 0.0f;
for (int row = 0; row < head_dim; row++) {
sum += q_vec[row] * state[row * head_dim + j];
}
int out_idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j;
output[out_idx] = sum;
}
__syncthreads(); // For next iteration
}
}
```
Wait, I have an issue with the thread indexing. I'm using threadIdx.y for row and threadIdx.x for column, but I need to make sure the thread block dimensions match.
If head_dim = 24, I need blockDim = (24, 24) = 576 threads. But 576 threads might not be a power of 2, which could cause issues.
Let me adjust:
- Use blockDim = (32, 32) = 1024 threads (or smaller)
- Check bounds: if (i < head_dim && j < head_dim)
But with blockDim = (32, 32), we have 1024 threads per block, and only 576 are doing useful work (if head_dim = 24). That's inefficient.
Let me try a different approach:
- Use 1D thread blocks
- Reshape the 2D state into 1D and have each thread handle multiple elements
```cpp
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ output,
int batch_size, int seq_len, int num_heads, int head_dim) {
int batch = blockIdx.x / num_heads;
int head = blockIdx.x % num_heads;
int tid = threadIdx.x;
int state_size = head_dim * head_dim;
int threads_per_block = blockDim.x;
// Shared memory
extern __shared__ float smem[];
float* state = smem;
float* q_vec = state + state_size;
float* k_vec = q_vec + head_dim;
float* v_vec = k_vec + head_dim;
float* gate_vec = v_vec + head_dim;
// Initialize state
for (int idx = tid; idx < state_size; idx += threads_per_block) {
state[idx] = 0.0f;
}
__syncthreads();
for (int t = 0; t < seq_len; t++) {
// Load vectors cooperatively
for (int i = tid; i < head_dim; i += threads_per_block) {
int idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + i;
q_vec[i] = q[idx];
k_vec[i] = k[idx];
v_vec[i] = v[idx];
gate_vec[i] = gate[idx];
}
__syncthreads();
// Update state
for (int idx = tid; idx < state_size; idx += threads_per_block) {
int i = idx / head_dim; // row
int j = idx % head_dim; // column
state[idx] = gate_vec[i] * state[idx] + k_vec[i] * v_vec[j];
}
__syncthreads();
// Compute output: y[j] = sum_i q[i] * state[i, j]
for (int j = tid; j < head_dim; j += threads_per_block) {
float sum = 0.0f;
for (int i = 0; i < head_dim; i++) {
sum += q_vec[i] * state[i * head_dim + j];
}
int out_idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j;
output[out_idx] = sum;
}
__syncthreads();
}
}
```
This looks cleaner. Let me now think about optimization:
1. Block size: With head_dim = 24, state_size = 576. A block size of 128 or 256 would give good coverage.
2. Shared memory: state_size + 4 * head_dim = 576 + 96 = 672 floats = 2688 bytes
3. Grid: batch_size * num_heads = 2 * 4 = 8 blocks
With 8 blocks, we're not using many SMs. But since each block processes 128 timesteps, the total work is significant.
Let me refine and test this. I'll use:
- blockDim = 256 threads
- Shared memory = 2688 bytes (rounded up to a safe size)
- Grid = (batch_size * num_heads) blocks
Now let me write the full Python code:
```python
import torch
import torch.nn as nn
from torch.utils.cpp_extension import load_inline
cuda_source = """
__global__ void gated_delta_attention_kernel(
const float* __restrict__ q,
const float* __restrict__ k,
const float* __restrict__ v,
const float* __restrict__ gate,
float* __restrict__ output,
int batch_size, int seq_len, int num_heads, int head_dim) {
int batch = blockIdx.x / num_heads;
int head = blockIdx.x % num_heads;
int tid = threadIdx.x;
int state_size = head_dim * head_dim;
int threads_per_block = blockDim.x;
// Shared memory
extern __shared__ float smem[];
float* state = smem;
float* q_vec = state + state_size;
float* k_vec = q_vec + head_dim;
float* v_vec = k_vec + head_dim;
float* gate_vec = v_vec + head_dim;
// Initialize state
for (int idx = tid; idx < state_size; idx += threads_per_block) {
state[idx] = 0.0f;
}
__syncthreads();
for (int t = 0; t < seq_len; t++) {
// Load vectors cooperatively
for (int i = tid; i < head_dim; i += threads_per_block) {
int idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + i;
q_vec[i] = q[idx];
k_vec[i] = k[idx];
v_vec[i] = v[idx];
gate_vec[i] = gate[idx];
}
__syncthreads();
// Update state
for (int idx = tid; idx < state_size; idx += threads_per_block) {
int i = idx / head_dim;
int j = idx % head_dim;
state[idx] = gate_vec[i] * state[idx] + k_vec[i] * v_vec[j];
}
__syncthreads();
// Compute output
for (int j = tid; j < head_dim; j += threads_per_block) {
float sum = 0.0f;
for (int i = 0; i < head_dim; i++) {
sum += q_vec[i] * state[i * head_dim + j];
}
int out_idx = batch * seq_len * num_heads * head_dim + t * num_heads * head_dim + head * head_dim + j;
output[out_idx] = sum;
}
__syncthreads();
}
}
torch::Tensor gated_delta_attention_cuda(
torch::Tensor q,
torch::Tensor k,
torch::Tensor v,
torch::Tensor gate) {
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 threads = 256;
int blocks = batch_size * num_heads;
int shared_mem = (head_dim * head_dim + 4 * head_dim) * sizeof(float);
gated_delta_attention_kernel<<<blocks, threads, shared_mem>>>(
q.data_ptr<float>(),
k.data_ptr<float>(),
v.data_ptr<float>(),
gate.data_ptr<float>(),
output.data_ptr<float>(),
batch_size, seq_len, num_heads, head_dim);
return output;
}
"""
cpp_source = """
torch::Tensor gated_delta_attention_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor gate);
"""
gated_attention_module = load_inline(
name='gated_attention',
cpp_sources=cpp_source,
cuda_sources=cuda_source,
functions=['gated_delta_attention_cuda'],
extra_cuda_cflags=['-O3'],
verbose=True,
)
class ModelNew(nn.Module):
def __init__(self):
super().__init__()
self.gated_attention = gated_attention_module.gated_delta_attention_cuda
def forward(self, q, k, v, gate):
return self.gated_attention(q, k, v, gate)
```
Let me compile and test this.