# Grouped-Query Attention (GQA) in Needle 2: Reducing KV Cache Memory Through Head Grouping

> Discover Grouped-Query Attention GQA in Needle 2. Learn how it slashes KV cache memory by reducing key value heads while preserving query attention diversity for efficient GPU use.

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

---

**Grouped-Query Attention (GQA) in Needle 2 reduces GPU memory consumption by allowing the number of key/value heads to be smaller than the number of query heads, cutting the KV cache size proportionally while maintaining diverse query attention patterns.**

Needle 2 implements Grouped-Query Attention (also known as Multi-Query Attention) through the `MultiHeadAttention` module in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). This mechanism enables the number of query heads (`num_heads`) to exceed the number of key/value heads (`num_kv_heads`), dramatically lowering the memory footprint of the KV cache during inference. By configuring these parameters, developers can optimize the trade-off between model capacity and memory efficiency.

## What Is Grouped-Query Attention?

In standard multi-head attention, every query head maintains its own dedicated key and value projections, resulting in `num_heads` separate KV matrices. Grouped-Query Attention diverges from this pattern by allowing multiple query heads to share the same key and value heads.

In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `MultiHeadAttention` class defines this behavior through two distinct configuration fields:

```python

# needle/model/architecture.py – definition of MultiHeadAttention

class MultiHeadAttention(nn.Module):
    num_heads: int
    num_kv_heads: int          # <‑‑ KV heads can be fewer than query heads

    …

```

When `num_kv_heads < num_heads`, the model enters GQA mode. The KV projections are computed once per group and then repeated to match the query head count, reducing the cache size from `num_heads × head_dim` to `num_kv_heads × head_dim` per token.

## Implementation Details in Needle 2

### Projection Dimensions and Head Sharing

During the forward pass, Needle 2 projects queries to `num_heads × head_dim`, while keys and values are projected to `num_kv_heads × head_dim`. When the Flash-attention kernel is disabled, the implementation repeats the KV tensors to align with the query head count:

```python

# needle/model/architecture.py – KV‑head repetition

if self.flash:
    …
else:
    repeats = self.num_heads // self.num_kv_heads
    if repeats > 1:
        k = jnp.repeat(k, repeats, axis=1)   # share KV across groups

        v = jnp.repeat(v, repeats, axis=1)

```

The `repeats` variable calculates the grouping ratio as `num_heads // num_kv_heads`. For example, with 12 query heads and 2 KV heads, each KV head serves 6 query heads, cutting the KV cache by a factor of 6.

### Flash Attention Integration

When `self.flash` is enabled, Needle 2 bypasses the explicit repetition logic and delegates the grouping to the optimized Flash-attention kernel, which handles the head sharing internally without materializing the repeated tensors in memory.

## Memory Reduction Mechanisms

### KV Cache Size Calculation

The memory savings scale directly with the ratio of query heads to KV heads. Because the KV cache stores activations proportional to `num_kv_heads`, reducing this parameter yields linear memory savings:

| Setting | KV Cache Size (Per Token) |
|---------|---------------------------|
| Standard multi-head (`num_kv_heads = num_heads`) | `num_heads × head_dim` |
| Grouped-query (`num_kv_heads < num_heads`) | `num_kv_heads × head_dim` (shared across groups) |

For a configuration with `num_heads=12` and `num_kv_heads=2`, the KV cache size drops from 12 head dimensions per token to just 2—a **6× memory reduction**.

### Dynamic KV Budget Window

Needle 2 computes a **KV-budget window** that limits the maximum sequence length based on available memory. This budget depends directly on the number of KV heads, meaning GQA configurations can support longer context windows within the same memory constraints:

```python

# needle/model/architecture.py – KV‑budget computation

def kv_budget_window(config):
    head_dim = (config.attn_dim or config.d_model) // config.num_heads
    kv = config.num_kv_heads * head_dim          # <-- KV heads affect memory

    …

```

The `kv_budget_window` function uses `config.num_kv_heads` to calculate the memory cost per token, while `effective_kv_window` (also in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)) determines the actual sequence length capacity. Fewer KV heads translate to larger effective windows.

## Practical Configuration Examples

### Standard Multi-Head Attention (No Grouping)

```python

# Example 1 – Standard multi‑head (no grouping)

config = TransformerConfig(num_heads=12, num_kv_heads=12, d_model=768)
model = SimpleAttentionNetwork(config)

# KV cache size ≈ 12 × head_dim per token

```

### Grouped-Query Attention Configuration

```python

# Example 2 – Grouped‑Query Attention (GQA)

config = TransformerConfig(num_heads=12, num_kv_heads=2, d_model=768)
model = SimpleAttentionNetwork(config)

# KV cache size ≈ 2 × head_dim per token → 6× memory reduction

```

### Inspecting Memory Budgets

```python

# Inspect the KV budget for a given config

from needle.model.architecture import kv_budget_window, effective_kv_window
cfg = TransformerConfig(num_heads=12, num_kv_heads=2, d_model=768,
                        max_seq_len=8192, kv_window=0)
print("Budgeted window:", kv_budget_window(cfg))
print("Effective window:", effective_kv_window(cfg))

# The printed window will be larger for the GQA config than for a

# config where num_kv_heads == num_heads.

```

## Summary

- **File Location**: Needle 2 implements GQA in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) through the `MultiHeadAttention` class and supporting budget functions `kv_budget_window` and `effective_kv_window`.
- **Configuration**: Set `num_kv_heads` lower than `num_heads` in `TransformerConfig` to enable grouped attention.
- **Memory Scaling**: KV cache size is proportional to `num_kv_heads`, offering linear memory reduction (e.g., 6× reduction with 12 query heads and 2 KV heads).
- **Implementation**: Non-Flash paths use `jnp.repeat` to broadcast KV heads, while Flash-attention paths handle grouping internally.
- **Context Length**: Lower KV head counts increase the `kv_budget_window`, allowing longer sequences within fixed GPU memory budgets.

## Frequently Asked Questions

### How does GQA differ from standard multi-head attention in Needle 2?

Standard multi-head attention in Needle 2 sets `num_kv_heads` equal to `num_heads`, allocating separate key and value matrices for every query head. GQA configures `num_kv_heads` to be smaller, causing multiple query heads to share the same KV projections. This sharing is implemented through the `repeats` logic in `MultiHeadAttention` when Flash attention is disabled, or through the kernel when enabled.

### What is the exact memory reduction factor when using GQA?

The memory reduction factor equals the ratio of `num_heads` to `num_kv_heads`. According to the Needle 2 source code in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the KV cache size scales with `num_kv_heads * head_dim`. Therefore, a configuration with 12 query heads and 2 KV heads achieves a 6× reduction in KV cache memory per token compared to standard attention.

### Does enabling GQA affect the model's ability to capture diverse attention patterns?

While GQA reduces the number of distinct key and value representations, Needle 2 maintains `num_heads` separate query projections. This allows the model to attend to the shared KV representations from multiple perspectives (query heads), preserving representational diversity in the attention mechanism while significantly reducing memory overhead.

### How do I determine the optimal `num_kv_heads` value for my Needle 2 model?

The optimal value depends on your GPU memory constraints and required context length. Use the `kv_budget_window` function in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) to calculate the maximum sequence length supported by different `num_kv_heads` configurations. Common practice involves setting `num_kv_heads` to a small integer (such as 2, 4, or 8) that divides evenly into `num_heads`, balancing memory savings against the slight reduction in KV capacity.