Grouped Query Attention › Implementation for Inference

class GQAForGeneration(nn.Module):
    """GQA implementation optimized for autoregressive generation."""

    def __init__(
        self,
        d_model: int,
        num_heads: int,
        num_kv_heads: int,
        max_seq_len: int = 4096,
        dropout: float = 0.0,
    ):
        super().__init__()
        assert d_model % num_heads == 0
        assert num_heads % num_kv_heads == 0

        self.d_model = d_model
        self.num_heads = num_heads
        self.num_kv_heads = num_kv_heads
        self.num_queries_per_kv = num_heads // num_kv_heads
        self.head_dim = d_model // num_heads
        self.max_seq_len = max_seq_len

        # Projections
        self.w_q = nn.Linear(d_model, num_heads * self.head_dim, bias=False)
        self.w_k = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
        self.w_v = nn.Linear(d_model, num_kv_heads * self.head_dim, bias=False)
        self.w_o = nn.Linear(d_model, d_model, bias=False)

        self.dropout = nn.Dropout(dropout)
        self.scale = self.head_dim**-0.5

        # Pre-compute causal mask
        self.register_buffer(
            "causal_mask", torch.tril(torch.ones(max_seq_len, max_seq_len))
        )

    def forward(
        self,
        x: torch.Tensor,
        start_pos: int = 0,
        kv_cache: tuple = None,
    ) -> tuple:
        batch_size, seq_len, _ = x.shape

        # Compute queries
        q = self.w_q(x).view(batch_size, seq_len, self.num_heads, self.head_dim)
        q = q.transpose(1, 2)  # (batch, num_heads, seq, head_dim)

        # Compute new keys and values
        k_new = self.w_k(x).view(
            batch_size, seq_len, self.num_kv_heads, self.head_dim
        )
        v_new = self.w_v(x).view(
            batch_size, seq_len, self.num_kv_heads, self.head_dim
        )

        # Update cache
        if kv_cache is not None:
            k_cache, v_cache = kv_cache
            k = torch.cat([k_cache, k_new], dim=1)
            v = torch.cat([v_cache, v_new], dim=1)
        else:
            k = k_new
            v = v_new

        new_kv_cache = (k, v)
        kv_seq_len = k.size(1)

        # Reshape for attention
        k = k.transpose(1, 2)  # (batch, num_kv_heads, kv_seq, head_dim)
        v = v.transpose(1, 2)

        # Expand KV to match query heads
        k = k.repeat_interleave(self.num_queries_per_kv, dim=1)
        v = v.repeat_interleave(self.num_queries_per_kv, dim=1)

        # Attention
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) * self.scale

        # Apply causal mask
        mask = self.causal_mask[start_pos : start_pos + seq_len, :kv_seq_len]
        mask = mask.unsqueeze(0).unsqueeze(0)  # (1, 1, seq, kv_seq)
        attn_scores = attn_scores.masked_fill(mask == 0, float("-inf"))

        attn_weights = F.softmax(attn_scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # Output
        attn_output = torch.matmul(attn_weights, v)
        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.view(batch_size, seq_len, self.d_model)

        return self.w_o(attn_output), new_kv_cache


def simulate_generation(model: GQAForGeneration, prompt_len: int, gen_len: int):
    """Simulate token-by-token generation and track cache growth."""
    batch_size = 1
    d_model = model.d_model

    # Process prompt
    prompt = torch.randn(batch_size, prompt_len, d_model)
    _, kv_cache = model(prompt, start_pos=0)

    cache_sizes = []
    current_pos = prompt_len

    # Generate tokens one by one
    for _ in range(gen_len):
        # Simulate next token embedding
        next_token = torch.randn(batch_size, 1, d_model)

        # Process with cached KV
        _, kv_cache = model(
            next_token, start_pos=current_pos, kv_cache=kv_cache
        )
        current_pos += 1

        # Track cache size
        k_cache, v_cache = kv_cache
        cache_bytes = (k_cache.numel() + v_cache.numel()) * 4  # float32
        cache_sizes.append(
            {
                "position": current_pos,
                "cache_kb": cache_bytes / 1024,
                "k_shape": tuple(k_cache.shape),
            }
        )

    return cache_sizes