Understanding Grouped Query Attention (GQA) in Needle 2: Implementation and Code Walkthrough

Grouped Query Attention (GQA) in Needle 2 decouples the number of query heads from key/value heads, reducing memory bandwidth by sharing KV projections across query groups while maintaining attention quality.

In transformer architectures, the standard multi-head attention (MHA) uses identical counts for query, key, and value heads. Needle 2, an open-source JAX-based transformer implementation by cactus-compute, introduces GQA as a configurable optimization that allows num_heads (queries) to exceed num_kv_heads (keys/values). This article examines the GQA implementation in needle/model/architecture.py, explaining how the code achieves memory efficiency through dimensional separation and KV repetition.

What Is Grouped Query Attention?

GQA is an attention mechanism where multiple query heads share the same key and value representations. Rather than computing separate KV projections for every query head, GQA partitions queries into groups that attend to common KV sets.

The benefits include:

• Reduced KV cache memory during inference — critical for long-context generation
• Lower memory bandwidth during training and inference
• Preserved query capacity for fine-grained attention patterns

Needle 2 implements GQA directly in its MultiHeadAttention class without requiring separate code paths for GQA versus standard MHA.

Core Implementation in MultiHeadAttention

The GQA architecture resides in [needle/model/architecture.py](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). The module exposes two distinct count parameters:

class MultiHeadAttention(nn.Module):
    num_heads: int          # number of query heads

    num_kv_heads: int       # number of key/value heads (GQA when < num_heads)

    # ...

(source: architecture.py L9-L15)

Dimensional Separation for Queries vs. KV

The forward pass computes separate dimensions for queries and KV projections:

head_dim = attn_dim // self.num_heads
kv_dim   = self.num_kv_heads * head_dim

(source: architecture.py L20-L23)

This calculation ensures that:

• Queries project to the full attn_dim (num_heads * head_dim)
• Keys and values project to the reduced kv_dim (num_kv_heads * head_dim)

Projection and Reshaping Logic

The actual tensor transformations implement the dimensional split:

q = nn.Dense(attn_dim, ...)(x)                # (batch, time, attn_dim)

k = nn.Dense(kv_dim, ...)(x)                  # (batch, time, kv_dim)

v = nn.Dense(kv_dim, ...)(x)                  # (batch, time, kv_dim)

# Reshape to separate heads

q = q.reshape(B, -1, self.num_heads, head_dim).transpose(0, 2, 1, 3)
k = k.reshape(B, -1, self.num_kv_heads, head_dim).transpose(0, 2, 1, 3)
v = v.reshape(B, -1, self.num_kv_heads, head_dim).transpose(0, 2, 1, 3)

(source: architecture.py L26-L33)

At this stage, the shapes are:

• q: [B, num_heads, T, head_dim]
• k, v: [B, num_kv_heads, T, head_dim]

KV Head Repetition for Attention Compatibility

Since standard attention requires compatible dimensions between queries and keys/values, Needle 2 repeats KV heads when num_heads > num_kv_heads:

repeats = self.num_heads // self.num_kv_heads
if repeats > 1:
    k = jnp.repeat(k, repeats, axis=1)
    v = jnp.repeat(v, repeats, axis=1)

(source: architecture.py L57-L61)

This broadcasting operation enables each query head to attend to its corresponding KV representation without materializing separate projections. The repetition occurs along axis 1 (the head dimension), expanding KV from [B, num_kv_heads, T, head_dim] to [B, num_heads, T, head_dim].

Configuring GQA in Needle 2 Models

Needle 2 uses a TransformerConfig dataclass to specify architecture hyperparameters. GQA is activated by setting num_kv_heads lower than num_heads:

from needle.model.architecture import TransformerConfig, SimpleAttentionNetwork
import jax.numpy as jnp

# GQA configuration: 12 query heads, 4 KV heads (3:1 grouping ratio)

cfg = TransformerConfig(
    d_model=768,
    num_heads=12,          # query heads

    num_kv_heads=4,        # KV heads — enables GQA

    num_layers=12,
    flash=True,            # uses Flash Attention kernel if available

)

model = SimpleAttentionNetwork(cfg)

# Forward pass with dummy tokens

tokens = jnp.ones((2, 128), dtype=jnp.int32)  # batch=2, seq_len=128

logits = model(tokens)  # shape: (2, 128, vocab_size)

The grouping ratio (num_heads // num_kv_heads) determines memory savings. Common configurations include:

• 4:1 grouping (e.g., 32 query heads, 8 KV heads) — substantial memory reduction
• 2:1 grouping — moderate reduction with minimal quality impact
• 1:1 grouping — standard MHA, no GQA

Reverting to Standard Multi-Head Attention

To disable GQA and use classic MHA, set the counts equal:

cfg = TransformerConfig(
    num_heads=16,
    num_kv_heads=16,  # same as num_heads → standard MHA

    # ...

)

Integration Within Transformer Blocks

The MultiHeadAttention module integrates into Block layers as shown in architecture.py L78-L86. Each block maintains the GQA configuration throughout the stack:

class Block(nn.Module):
    # ...

    def __call__(self, x, mask=None):
        # Pre-norm attention with GQA

        attn_out = self.attention(self.norm1(x), mask)
        x = x + attn_out
        # FFN...

        return x

Performance Characteristics

According to the Needle 2 source code implementation:

Aspect	GQA Impact
KV cache size	Reduced by factor of num_heads / num_kv_heads
Memory bandwidth	Lower during autoregressive decoding
Compute FLOPs	Unaffected in attention matmul (KV repeated before computation)
Training stability	No special requirements beyond standard MHA

The jnp.repeat operation for KV heads adds negligible overhead compared to the savings from reduced projection dimensions and cache memory.

Summary

• GQA in Needle 2 is implemented via separate num_heads and num_kv_heads parameters in MultiHeadAttention
• Dimensional separation reduces KV projection size by using kv_dim = num_kv_heads * head_dim
• KV repetition via jnp.repeat ensures compatibility with query heads during attention computation
• Configuration happens through TransformerConfig — simply set num_kv_heads < num_heads to enable GQA
• Source files: [needle/model/architecture.py](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) contains the complete implementation

Frequently Asked Questions

How does GQA differ from Multi-Query Attention (MQA)?

MQA uses a single KV head shared across all query heads, while GQA uses multiple KV head groups. In Needle 2, MQA is a special case of GQA where num_kv_heads=1. GQA provides a tunable middle ground between MQA's maximum efficiency and MHA's full expressiveness. The jnp.repeat logic handles any grouping ratio uniformly.

What grouping ratio should I use for my model?

Start with 4:1 or 8:1 (num_heads // num_kv_heads) for large language models. Needle 2's implementation places no architectural constraints on the ratio — any integer divisor works. Empirical results suggest 4:1 grouping achieves most of GQA's memory benefits with negligible quality degradation compared to full MHA. The TransformerConfig makes experimentation straightforward.

Does GQA affect Flash Attention compatibility in Needle 2?

No, GQA works transparently with Needle 2's Flash Attention path. The KV repetition occurs before the attention kernel selection, so both standard and Flash Attention implementations receive properly shaped tensors. The flash=True configuration in TransformerConfig applies regardless of head grouping.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}