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

> Explore Grouped Query Attention GQA in Needle 2. Learn how GQA enhances performance by sharing KV projections across query groups, reducing memory bandwidth without sacrificing attention quality. Includes implementation and cod...

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: deep-dive
- Published: 2026-09-01

---

**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`](https://github.com/cactus-compute/needle/blob/main/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)](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). The module exposes two distinct count parameters:

```python
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](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L9-L15))*

### Dimensional Separation for Queries vs. KV

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

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

```

*(source: [architecture.py L20-L23](https://github.com/cactus-compute/needle/blob/main/needle/model/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:

```python
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](https://github.com/cactus-compute/needle/blob/main/needle/model/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`:

```python
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](https://github.com/cactus-compute/needle/blob/main/needle/model/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`:

```python
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:

```python
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](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L78-L86). Each block maintains the GQA configuration throughout the stack:

```python
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`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L9-L15)
- **Dimensional separation** reduces KV projection size by using `kv_dim = num_kv_heads * head_dim`
- **KV repetition** via [`jnp.repeat`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L57-L61) 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)](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`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L57-L61) 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`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) 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`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) applies regardless of head grouping.