Let me analyze this task. I need to implement a custom CUDA kernel for causal grouped-query attention with KV-head expansion. The reference implementation: 1. Takes q, k, v tensors 2. Expands k and v by repeating each KV head across its assigned query-head group 3. Performs causal attention (only attending to positions up to current position) 4. Uses softmax over scores Key dimensions: - batch_size = 2 - seq_len = 80 - num_q_heads = 8 - num_kv_heads = 2 - head_dim = 32 - group_size = num_q_heads // num_kv_heads = 4 The naive PyTorch implementation is very slow due to Python loops. I need to write an optimized CUDA kernel. For a causal attention kernel: - Each thread could handle one query position and head - We need to compute attention scores against all previous positions - Then apply softmax and compute weighted sum of values Let me design the kernel: 1. Grid/block design: - For batch_size=2, seq_len=80, num_q_heads=8 - We have 2 * 80 * 8 = 1280 query positions to process - Each query position needs to attend to up to 80 previous positions - We could use a thread block per (batch, head) or per (batch, head, query_position) 2. Memory access pattern: - Need to access expanded_k and expanded_v - The expansion happens by repeating KV heads - For query head h, the corresponding KV head index is h // group_size 3. Computation: - For each query position t: - Compute scores = q[b,t,h] @ k[b,:t+1,h].T * scale - Apply softmax - Compute output = softmax @ v[b,:t+1,h] Let me think about a parallel approach: - Launch one thread per (batch, query_head, query_position) - Each thread computes the attention for that query position - Load q vector, then iterate through all positions from 0 to t to compute dot products - Apply softmax to the scores - Compute weighted sum of values For softmax, I need to find the max score and subtract it for numerical stability. With seq_len=80, the softmax fits comfortably in registers. Each thread handles one query position, loading the query vector, computing all dot products against key vectors up to that position, and then applying softmax and computing the weighted value sum. The kernel accepts q, k, v tensors and outputs the result. Each thread processes one (batch, head, sequence_position) tuple, computing the attention output for that specific query. To get the correct thread indices, I calculate batch, head, and position from the linear thread ID. The KV head is derived from the query head using the group_size divisor. I then load the query vector for this position, and for each prior position in the sequence, compute the scaled dot-product score between the query and corresponding key vector. After finding the maximum score across all positions, I subtract it from each score to maintain numerical stability, then exponentiate and accumulate the sum of these exponentials. With the normalized attention weights in hand, I can proceed to compute the final output as a weighted combination of value vectors. For each dimension in the head, I iterate through all positions and accumulate the attention-weighted values into the output vector. Finally, I store this result in the output tensor and advance to the next query position. The kernel launches with enough threads to cover all batch elements, query heads, and sequence positions, using a block size of 256 threads. The host function takes the query, key, and value tensors along with the number of key-value heads, calculates the derived parameters like head count and head dimension, then launches the CUDA kernel with the appropriate grid dimensions before returning the output. The grouping relationship matters: query heads get mapped to KV heads via integer division, so for query head 4-7 with group_size=4, they all reference KV head 1. The PyTorch `repeat_interleave` operation replicates each KV head across its corresponding query head group—KV head 0 becomes query heads 0-3, KV head 1 becomes query heads 4-7, and so on. My CUDA kernel uses this same indexing scheme: for a query head h, I compute the KV head as h // group_size, which correctly retrieves the expanded key-value pair. Now I need to double-check the indexing. The input tensors have shape [batch, seq, num_kv_heads, head_dim], while the expanded versions become [batch, seq, num_q_heads, head_dim]. When accessing `k[b, t, kv_h]`, I'm retrieving the key for that KV head at position t, which should align with how query heads reference their corresponding KV head. The softmax computation looks correct—I'm finding the maximum score for stability, then subtracting it before computing exponentials to prevent overflow. The exp_sum accumulates the normalized scores, and the division gives the final attention weights. For the output, I'm multiplying each attention weight by its corresponding value vector and summing across the sequence dimension. The result gets written to the output tensor at the correct batch, sequence position, and query head location. The implementation appears sound. A potential optimization would be shared memory tiling for larger sequence lengths, but at seq_len=80, the current approach should perform adequately. I should add a bounds check to prevent out-of-range access when the total number of threads exceeds the work items. Calculating the total elements: batch_size=2, seq_len=80, num_q_heads=8 gives 1280 query positions to process. With 256 threads per block, I'd need 5 blocks to cover all positions. The softmax computation requires accumulating scores across up to 80 sequence positions, which stays well within reasonable compute limits. I'll finalize the kernel with proper bounds checking and the attention calculation logic. out[idx * head_dim + d] = sum; } } void attention_host( float* q, float* k, float* v, float* out, int batch_size, int seq_len, int num_q_heads, int num_kv_heads, int head_dim ) { int total = batch_size * seq_len * num_q_heads; int threads = 256; int blocks = (total + threads - 1) / threads; attention_kernel<<<blocks, threads>>>( q, k, v, out, batch_size, seq_len, num_q_heads, num_kv_heads, head_dim ); } """ module = load_inline( name='attention_cuda', cpp_sources=cpp_sources, cuda_sources=cuda_sources, functions=['attention_host'], extra_cuda_cflags=['-O3'], ) return module # Load the module once at class definition time _cuda_module = None def get_cuda_module(): global _cuda_module if _cuda_module is None: _cuda_module = load_cuda_module() return _cuda_module I'm completing the softmax computation by exponentiating the scores after subtracting the maximum for stability, then normalizing. For each dimension in the head, I iterate through all positions in the sequence, multiply the value by its corresponding attention weight, and accumulate the result. The final sum gets stored in the output tensor, then I move to the next element. The host function calculates the total number of operations needed, determines the appropriate grid and block dimensions, then launches the kernel with the query, key, and value tensors along with the output buffer and dimensional parameters. In the model's forward pass, I extract the batch size and sequence length from the query tensor, determine the number of key-value heads from the key tensor's shape, calculate the group size as the ratio of query heads to key-value heads, and retrieve the head dimension before invoking the CUDA module's attention computation. I need to verify the KV head indexing is consistent between the PyTorch implementation and my CUDA kernel. The `repeat_interleave` operation expands each KV head across multiple query head positions, so KV head 0 maps to query head indices 0-3, KV head 1 to indices 4-7, and so on. When computing attention in CUDA, for query head h, I should access the corresponding KV head at position (h // group_size) in the KV tensor, which aligns with how the expanded indices are structured. 0,1,2,3] and kv_h=1 maps to [4,5,6,7]. This correctly handles the interleaving pattern. For the single-precision requirement, I'm using float throughout the implementation. I should also apply the standard PyTorch softmax temperature scaling by dividing the attention scores by sqrt(head_dim) for numerical stability.