reasoning
Let me analyze this problem carefully.
The reference implementation is a channel-wise gated delta attention recurrence. Let me break down what it does:
1. For each timestep t in the sequence:
- gate_t = gate[:, t].unsqueeze(-2) # shape: [batch, num_heads, 1, head_dim]
- kv_outer = k[:, t].unsqueeze(-1) * v[:, t].unsqueeze(-2) # outer product, shape: [batch, num_heads, head_dim, head_dim]
- state = gate_t * state + kv_outer # state accumulates with decay
- y_t = (q[:, t].unsqueeze(-1) * state).sum(dim=-2) # shape: [batch, num_heads, head_dim]
- outputs.append(y_t)
Let me trace the shapes more carefully:
- q, k, v: [batch_size, seq_len, num_heads, head_dim]
- gate: [batch_size, seq_len, num_heads, head_dim]
- state: [batch_size, num_heads, head_dim, head_dim]
For each timestep t:
- gate[:, t]: [batch, num_heads, head_dim]
- gate[:, t].unsqueeze(-2): [batch, num_heads, 1, head_dim]
- k[:, t]: [batch, num_heads, head_dim]
- k[:, t].unsqueeze(-1): [batch, num_heads, head_dim, 1]
- v[:, t]: [batch, num_heads, head_dim]
- v[:, t].unsqueeze(-2): [batch, num_heads, 1, head_dim]
- kv_outer = k[:, t].unsqueeze(-1) * v[:, t].unsqueeze(-2): [batch, num_heads, head_dim, head_dim]
- state = gate_t * state + kv_outer: [batch, num_heads, head_dim, head_dim]
The decay mechanism multiplies the previous state by the gate, then adds the new key-value outer product. For the query computation:
- q[:, t]: [batch, num_heads, head_dim]
- q[:, t].unsqueeze(-1): [batch, num_heads, head_dim, 1]
- q[:, t].unsqueeze(-1) * state: [batch, num_heads, head_dim, head_dim]
- .sum(dim=-2): [batch, num_heads, head_dim]
The final output stacks all timesteps to produce [batch, seq_len, num_heads, head_dim].
Now I need to implement this as a CUDA kernel. The sequential nature across timesteps means I can't parallelize that dimension, but I can parallelize across batch and heads. For each (batch, head) combination, a single thread or warp should process all timesteps sequentially.
The state has dimensions [head_dim, head_dim], and the operations at each timestep involve the gate decay, key-value outer product, and query-state interaction. For a given (batch, head) pair, I need to iterate through the sequence while maintaining the state matrix. I should structure this with each thread handling one (batch, head) combination, processing timesteps sequentially. The key operations are loading k and v to compute their outer product, applying the gate decay to the accumulated state, then querying against the updated state.
For memory layout, assuming row-major ordering, I need to correctly index into the tensors. The query tensor has shape [batch, seq_len, num_heads, head_dim], so for a specific (batch, t, head, d) position, the index is computed based on this ordering.
The gate tensor follows the same indexing pattern as query. For the state matrix, which has shape [batch, num_heads, head_dim, head_dim], accessing element (batch, head, i, j) requires accounting for all four dimensions in the stride calculation.
The output tensor matches the query shape at [batch, seq_len, num_heads, head_dim], so the indexing scheme is consistent with how query is accessed.
When computing the outer product for kv, I'm storing it in state[head_idx, row, col] where row corresponds to k and col corresponds to v, giving me kv[row, col] = k[row] * v[col]. The gate broadcast multiplies across the state rows. For the y_t computation, I need y_t[i] = sum_j(q[i] * state[i, j]), which the reference code handles by multiplying q.unsqueeze(-1) with state and summing along the appropriate dimension.
When q has shape [batch, num_heads, head_dim, 1] and state has shape [batch, num_heads, head_dim, head_dim], the element-wise product yields [batch, num_heads, head_dim, head_dim], and summing over dimension -2 produces [batch, num_heads, head_dim]. For each output index i, the sum accumulates q[i] * state[i, j] across all j values.
I should parallelize across batch and head pairs, assigning each to a separate thread block, then process the sequence steps within each block. With small head_dim (24), a single warp per block could handle the state matrix operations efficiently, using shared memory for temporary storage and applying parallel reduction when computing y_t. Alternatively, I could have each thread manage one (batch, head) pair and loop through the sequence while maintaining local state in registers or shared memory, which would be simpler to implement.
Let me try assigning one thread per (batch, head) combination. For each sequence step, I'd compute the outer product of k and v into a temporary matrix, decay the state by the gate factor, accumulate the outer product into state, then compute y_t by reducing across the head dimension. The tricky part is that with 192 threads total (2 batch × 4 heads × 24 head_dim), I need to figure out how to map threads to this work efficiently.
A cleaner approach would use one thread per (batch, head) pair—that's 8 threads total—each handling all 24 dimensions of the state matrix. With blockDim = (8, 1, 1), every thread loops through all timesteps and computes the state updates and outputs for its assigned (batch, head) combination.
For each timestep, the thread computes the outer product of k and v, applies the gate decay to the accumulated state, adds the new key-value contribution, then computes the output by summing the weighted state. The state itself is a 24×24 matrix that I'd need to store either in registers if the dimensions are fixed, or in shared memory for flexibility.
Given the small batch size and number of heads, 8 threads would be quite limited. I could expand by having each thread handle one output element instead—so one thread per (batch, num_heads, head_dim) combination, giving 2 × 4 × 24 = 192 threads in a single block. This way, each thread computes its own output element while coordinating through shared memory for the state updates.
Actually, with head_dim = 24, I could use cooperative groups to have warps work together. A single warp of 32 threads could handle each (batch, head) pair, with threads collaborating on the 24×24 state matrix. For 2 batches × 4 heads = 8 pairs, I'd need 8 warps total, which fits in a block.
Let me consider the simplest approach first: using one thread per (batch, head) pair, storing state in shared memory, and sequentially processing timesteps. With batch=2 and num_heads=4, that's 8 threads per block. I'd need 4 floats per thread for the 24×24 state matrix.
Looking at the test parameters, batch_size is 2, seq_len is 128, num_heads is 4, and head_dim is 24. Given these fixed dimensions, I can write a kernel optimized for this exact configuration rather than generalizing.
The most straightforward implementation would assign one thread per (batch, head) pair, iterate through the 128 timesteps, update the state matrix in shared memory, and compute outputs at each step.
call → submit_kernel
callsubmit_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 <torch/extension.h>\n#include <cuda.h>\n#include <cuda_runtime.h>\n\n// Optimized kernel for channel-wise gated delta attention recurrence\n// Each thread handles one (batch, head) pair, processing all timesteps sequentially\n\n__global__ void gated_delta_attention_kernel(\n const float* __restrict__ q,\n const float* __restrict__ k,\n const float* __restrict__ v,\n const float* __restrict__ gate,\n float* __restrict__ output,\n int batch_size, int seq_len, int num_heads, int head_dim\n) {\n // Shared memory for state matrices\n // Each (batch, head) pair needs head_dim * head_dim floats for state\n extern __shared__ float shared_state[];\n \n int batch_idx = blockIdx.x / num_heads;\n int head_idx = blockIdx.x % num_heads;\n int tid = blockIdx.x; // One block per (batch, head) pair\n \n if (batch_idx >= batch_size) return;\n \n // Each thread within the block handles part of the state\n int lane_id = threadIdx.x;\n int warp_size = 32;\n \n // State pointer for this (batch, head) pair\n float* state = shared_state; // All threads in block share the same state\n \n // Initialize state to zero\n for (int i = lane_id; i < head_dim * head_dim; i += blockDim.x) {\n state[i] = 0.0f;\n }\n __syncthreads();\n \n // Pre-compute strides\n // q, k, v, gate: [batch, seq_len, num_heads, head_dim]\n long long qkv_stride_batch = (long long)seq_len * num_heads * head_dim;\n long long qkv_stride_seq = num_heads * head_dim;\n long long qkv_stride_head = head_dim;\n \n // output: [batch, seq_len, num_heads, head_dim]\n // Same strides as q\n \n int num_warps = (blockDim.x + 31) / 32;\n \n // Process each timestep\n for (int t = 0; t < seq_len; t++) {\n // Compute base pointer for this (batch, t, head)\n long long base_offset = (long long)batch_idx * qkv_stride_batch + \n (long long)t * qkv_stride_seq + \n (long long)head_idx * qkv_stride_head;\n \n const float* q_t = q + base_offset;\n const float* k_t = k + base_offset;\n const float* v_t = v + base_offset;\n const float* gate_t = gate + base_offset;\n float* out_t = output + base_offset;\n \n // Load k and v into registers (cooperative load)\n // Each thread loads a portion\n float k_local[24], v_local[24], gate_local[24], q_local[24];\n \n // Cooperative load of k, v, gate, q\n for (int i = lane_id; i < head_dim; i += blockDim.x) {\n k_local[i] = k_t[i];\n v_local[i] = v_t[i];\n gate_local[i] = gate_t[i];\n q_local[i] = q_t[i];\n }\n \n // Need to broadcast k, v, gate, q to all threads in the block\n // Using shared memory for broadcast\n __shared__ float shared_k[24], shared_v[24], shared_gate[24], shared_q[24];\n \n for (int i = lane_id; i < head_dim; i += blockDim.x) {\n shared_k[i] = k_local[i];\n shared_v[i] = v_local[i];\n shared_gate[i] = gate_local[i];\n shared_q[i] = q_local[i];\n }\n __syncthreads();\n \n // Update state: state = gate * state + kv_outer\n // Cooperative update of state matrix\n for (int idx = lane_id; idx < head_dim * head_dim; idx += blockDim.x) {\n int row = idx / head_dim;\n int col = idx % head_dim;\n // gate[row] broadcasts across columns of state[row, :]\n // kv_outer[row, col] = k[row] * v[col]\n float gate_val = shared_gate[row];\n float kv_val = shared_k[row] * shared_v[col];\n state[idx] = gate_val * state[idx] + kv_val;\n }\n __syncthreads();\n \n // Compute output: y[i] = sum_j(q[i] * state[i, j])\n // Each thread computes one or more output elements\n for (int i = lane_id; i < head_dim; i += blockDim.x) {\n float q_val = shared_q[i];\n float sum = 0.0f;\n for (int j = 0; j < head_dim; j++) {\n sum += q_val * state[i * head_dim + j];\n }\n out_t[i] = sum;\n }\n __syncthreads();\n }\n}\n\n// Alternative: simple kernel with one thread per (batch, head) pair\n__global__ void gated_delta_attention_simple_kernel(\n const float* __restrict__ q,\n const float* __restrict__ k,\n const float* __restrict__ v,\n const float* __restrict__ gate,\n float* __restrict__ output,\n int batch_size, int seq_len, int num_heads, int head_dim\n) {\n int batch_idx = blockIdx.y;\n int head_idx = blockIdx.x;\n \n if (batch_idx >= batch_size) return;\n if (head_idx >= num_heads) return;\n \n // State matrix stored in registers (head_dim x head_dim)\n // For head_dim=24, this is 576 floats - too large for registers\n // Use a smaller tile approach or shared memory\n \n extern __shared__ float state[];\n \n // Initialize state to zero\n for (int i = threadIdx.x; i < head_dim * head_dim; i += blockDim.x) {\n state[i] = 0.0f;\n }\n __syncthreads();\n \n // Strides\n long long qkv_stride_batch = (long long)seq_len * num_heads * head_dim;\n long long qkv_stride_seq = num_heads * head_dim;\n long long qkv_stride_head = head_dim;\n \n // Process timesteps\n for (int t = 0; t < seq_len; t++) {\n long long base_offset = (long long)batch_idx * qkv_stride_batch + \n (long long)t * qkv_stride_seq + \n (long long)head_idx * qkv_stride_head;\n \n const float* q_t = q + base_offset;\n const float* k_t = k + base_offset;\n const float* v_t = v + base_offset;\n const float* gate_t = gate + base_offset;\n float* out_t = output + base_offset;\n \n // Load k, v, gate, q cooperatively\n __shared__ float shared_k[24], shared_v[24], shared_gate[24], shared_q[24];\n \n for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {\n shared_k[i] = k_t[i];\n shared_v[i] = v_t[i];\n shared_gate[i] = gate_t[i];\n shared_q[i] = q_t[i];\n }\n __syncthreads();\n \n // Update state\n for (int idx = threadIdx.x; idx < head_dim * head_dim; idx += blockDim.x) {\n int row = idx / head_dim;\n int col = idx % head_dim;\n state[idx] = shared_gate[row] * state[idx] + shared_k[row] * shared_v[col];\n }\n __syncthreads();\n \n // Compute output\n for (int i = threadIdx.x; i < head_dim; i += blockDim.x) {\n float sum = 0.0f;\n for (int j = 0; j < head_dim; j++) {\n sum += shared_q[i] * state[i * head_dim + j];\n }\n out_t[i] = sum;\n }\n __syncthreads();\n }\n}\n\n// Fused kernel optimized for the specific problem dimensions\n// Uses warp-level cooperation with fixed head_dim=24\ntemplate <int HEAD_DIM>\n__global__ void __launch_bounds__(128) gated_delta_attention_fused_kernel(\n const float* __restrict__ q,\n const float* __restrict__ k,\n const float* __restrict__ v,\n const float* __restrict__ gate,\n float* __restrict__ output,\n int batch_size, int seq_len, int num_heads\n) {\n // One block per (batch, head) pair\n int batch_idx = blockIdx.y;\n int head_idx = blockIdx.x;\n \n if (batch_idx >= batch_size) return;\n if (head_idx >= num_heads) return;\n \n const int head_dim = HEAD_DIM;\n \n // Shared memory for state (head_dim x head_dim) and broadcast buffers\n __shared__ float state[24 * 24]; // Fixed size for head_dim=24\n __shared__ float shared_k[24], shared_v[24], shared_gate[24], shared_q[24];\n \n int lane_id = threadIdx.x & 31;\n int warp_id = threadIdx.x >> 5;\n int num_threads = blockDim.x;\n \n // Initialize state to zero\n #pragma unroll\n for (int i = threadIdx.x; i < head_dim * head_dim; i += num_threads) {\n state[i] = 0.0f;\n }\n __syncthreads();\n \n // Strides\n long long qkv_stride_batch = (long long)seq_len * num_heads * head_dim;\n long long qkv_stride_seq = num_heads * head_dim;\n long long qkv_stride_head = head_dim;\n \n // Base offset for this (batch, head)\n long long batch_head_offset = (long long)batch_idx * qkv_stride_batch + \n (long long)head_idx * qkv_stride_head;\n \n // Process timesteps\n for (int t = 0; t < seq_len; t++) {\n const float* q_t = q + batch_head_offset + (long long)t * qkv_stride_seq;\n const float* k_t = k + batch_head_offset + (long long)t * qkv_stride_seq;\n const float* v_t = v + batch_head_offset + (long long)t * qkv_stride_seq;\n const float* gate_t = gate + batch_head_offset + (long long)t * qkv_stride_seq;\n float* out_t = output + batch_head_offset + (long long)t * qkv_stride_seq;\n \n // Cooperative load into shared memory\n for (int i = threadIdx.x; i < head_dim; i += num_threads) {\n shared_k[i] = k_t[i];\n shared_v[i] = v_t[i];\n shared_gate[i] = gate_t[i];\n shared_q[i] = q_t[i];\n }\n __syncthreads();\n \n // Update state matrix: state[r,c] = gate[r]*state[r,c] + k[r]*v[c]\n // Process in tiles for better cache utilization\n for (int idx = threadIdx.x; idx < head_dim * head_dim; idx += num_threads) {\n int row = idx / head_dim;\n int col = idx % head_dim;\n state[idx] = shared_gate[row] * state[idx] + shared_k[row] * shared_v[col];\n }\n __syncthreads();\n \n // Compute output: out[i] = sum_j(q[i] * state[i,j])\n for (int i = threadIdx.x; i < head_dim; i += num_threads) {\n float q_val = shared_q[i];\n float sum = 0.0f;\n \n // Manual unroll for head_dim=24\n #pragma unroll 8\n for (int j = 0; j < head_dim; j++) {\n sum += q_val * state[i * head_dim + j];\n }\n out_t[i] = sum;\n }\n __syncthreads();\n }\n}\n\n// Host function\ntorch::Tensor gated_delta_attention_cuda(\n torch::Tensor q,\n torch::Tensor k,\n torch::Tensor v,\n torch::Tensor gate\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_like(q);\n \n // Use the fused kernel with template parameter\n dim3 grid(num_heads, batch_size);\n int threads = 128; // Multiple warps for cooperative loading\n \n size_t shared_mem_size = 24 * 24 * sizeof(float) + 4 * 24 * sizeof(float);\n \n gated_delta_attention_fused_kernel<24><<<grid, threads, shared_mem_size>>>(\n q.data_ptr<float>(),\n k.data_ptr<float>(),\n v.data_ptr<float>(),\n gate.data_ptr<float>(),\n output.data_ptr<float>(),\n batch_size, seq_len, num_heads\n );\n \n return output;\n}\n\"\"\"\n\ncpp_source = \"\"\"\ntorch::Tensor gated_delta_attention_cuda(torch::Tensor q, torch::Tensor k, torch::Tensor v, torch::Tensor gate);\n\"\"\"\n\n# Load the module\nmodule = load_inline(\n name=\"gated_delta_attention\",\n cpp_sources=cpp_source,\n cuda_sources=cuda_source,\n functions=[\"gated_delta_attention_cuda\"],\n verbose=False,\n extra_cuda_cflags=[\"-O3\", \"--use_fast_math\"],\n)\n\n\nclass ModelNew(nn.Module):\n \"\"\"\n Channel-wise gated delta attention recurrence with custom CUDA kernel.\n \"\"\"\n\n def __init__(self):\n super().__init__()\n\n def forward(\n self,\n q: torch.Tensor,\n k: torch.Tensor,\n v: torch.Tensor,\n gate: torch.Tensor,\n ) -> torch.Tensor:\n return module.gated_delta_attention_cuda(q, k, v, gate)"
}