Implement Grouped-Query Attention (GQA)

Problem Statement

Grouped-Query Attention (GQA) interpolates between Multi-Head Attention (MHA, where queries, keys, and values have the same number of heads) and Multi-Query Attention (MQA, where keys and values share a single head across all query heads). In GQA, query heads are divided into num_heads_kv groups, and each group of query heads shares a single key head and value head.

Implement the forward pass of GQA using PyTorch (or standard tensor operations). Your function should support an optional causal mask.

Function Signature

import torch

def grouped_query_attention(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    is_causal: bool = False,
    scale: float | None = None,
) -> torch.Tensor:
    """
    Computes Grouped-Query Attention.
    
    Args:
        q: Query tensor of shape (batch_size, seq_len_q, num_heads_q, head_dim)
        k: Key tensor of shape (batch_size, seq_len_kv, num_heads_kv, head_dim)
        v: Value tensor of shape (batch_size, seq_len_kv, num_heads_kv, head_dim)
        is_causal: If True, applies an autoregressive causal mask where position i 
                   can only attend to positions j <= i.
        scale: Scaling factor applied before softmax. If None, defaults to 1.0 / sqrt(head_dim).
        
    Returns:
        torch.Tensor of shape (batch_size, seq_len_q, num_heads_q, head_dim)
    """
    pass

Requirements & Constraints

• num_heads_q must be divisible by num_heads_kv. Let group_size = num_heads_q // num_heads_kv.
• Head replication: Expand or repeat each KV head group_size times so that keys and values align with the query heads, preferably using zero-copy view/broadcasting or repeat_interleave.
• Attention formula: softmax(Q @ K.T / sqrt(head_dim) + mask) @ V.
• When is_causal=True, attention scores where query_pos < key_pos must be masked with -inf (or a very large negative value) before computing the softmax.
• Assume q, k, v are float tensors on the same device.
• Sequence lengths seq_len_q and seq_len_kv may differ (e.g., during autoregressive decoding with a KV cache where seq_len_q = 1 and seq_len_kv = T). In that case, causal masking aligns queries to the most recent positions: query index i corresponds to key position seq_len_kv - seq_len_q + i.

Example

batch_size = 2
seq_len = 4
num_heads_q = 8
num_heads_kv = 2
head_dim = 64

q = torch.randn(batch_size, seq_len, num_heads_q, head_dim)
k = torch.randn(batch_size, seq_len, num_heads_kv, head_dim)
v = torch.randn(batch_size, seq_len, num_heads_kv, head_dim)

out = grouped_query_attention(q, k, v, is_causal=True)
assert out.shape == (batch_size, seq_len, num_heads_q, head_dim)