# What Is Grouped Query Attention (GQA) in Needle 2? Implementation and Configuration Guide

> Discover Grouped Query Attention GQA in Needle 2. This memory efficient mechanism repeats KV tensors to reduce bandwidth while maintaining quality. Learn implementation and configuration.

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

---

**Grouped Query Attention (GQA) in Needle 2 is a memory-efficient attention mechanism implemented in the MultiHeadAttention module that allows query heads to outnumber key/value heads by repeating KV tensors across groups, reducing memory bandwidth while maintaining model quality through independent num_heads and num_kv_heads configuration parameters.**

Grouped Query Attention (GQA) addresses the memory bottleneck in transformer inference by reducing the number of key/value heads relative to query heads. In Needle 2, this optimization is built directly into the `MultiHeadAttention` class located in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), requiring no separate code paths or external dependencies to enable grouped query patterns.

## Core Mechanics of GQA in Needle 2

The implementation centers on flexible head dimension calculations and strategic tensor repetition that allows multiple query heads to share the same key/value representations.

### Independent Head Configuration

Unlike standard multi-head attention where queries, keys, and values share the same head count, Needle 2's `MultiHeadAttention` accepts `num_heads` (queries) and `num_kv_heads` (keys and values) as distinct arguments. This separation occurs at lines 209-210 in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) within the module definition.

When `num_heads` exceeds `num_kv_heads`, the system automatically triggers the grouping behavior. For example, configuring 12 query heads and 4 KV heads creates three query groups per KV head, reducing the KV cache memory requirements by 66%.

### Dimension Calculations and Reshaping

The module calculates dimensions dynamically based on the query head count. At lines 221-223, the implementation derives:

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

```

This ensures that while queries use the full attention dimension divided by the query head count, the KV dimension scales with the reduced KV head count. The reshaping operations at lines 230-232 transform the tensors into head-specific axes:

- Queries reshape to `(B, num_heads, T, head_dim)`
- Keys and values reshape to `(B, num_kv_heads, T, head_dim)`

### The Grouping Mechanism (KV Repetition)

The core GQA logic resides at lines 258-261, where the implementation determines the repetition factor and broadcasts KV heads to match the query head count:

```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)

```

This `jnp.repeat` operation aligns the KV tensor shapes with the query tensor shape, enabling standard attention computation while maintaining the memory savings of fewer KV parameters. Each group of `repeats` query heads shares the same key and value representations.

### Flash Attention Compatibility

Needle 2 applies the same grouping logic regardless of the attention implementation path. At lines 246-254, the code calls `jax.nn.dot_product_attention` (flash attention) when available, but the KV repetition step occurs before this call, ensuring consistent behavior across both the optimized flash path and the explicit matmul fallback at lines 267-271.

## Configuring GQA in Needle 2 Models

You enable GQA by instantiating a `TransformerConfig` with mismatched head counts and passing it to the `Stack` model builder.

### Standard Multi-Head Attention (Baseline)

For comparison, standard attention uses equal head counts:

```python
from needle.model import Stack, TransformerConfig

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

    num_kv_heads=12,  # KV heads (same as queries)

    num_layers=6,
)

model = Stack(cfg)

```

### Grouped Query Attention Configuration

To activate GQA, reduce `num_kv_heads` relative to `num_heads`:

```python
from needle.model import Stack, TransformerConfig
import jax.numpy as jnp

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

    num_kv_heads=4,  # 4 KV heads → each shared by 3 query heads

    num_layers=6,
)

model = Stack(cfg)

# Initialize and run forward pass

x = jnp.ones((1, 128, cfg.d_model))
variables = model.init(jax.random.PRNGKey(0), x)
output = model.apply(variables, x)
print(output.shape)  # (1, 128, 768)

```

### Direct Module Access for Debugging

For experimentation or debugging, instantiate `MultiHeadAttention` directly:

```python
from needle.model import MultiHeadAttention
import jax.numpy as jnp

attn = MultiHeadAttention(
    num_heads=12,
    num_kv_heads=4,
    d_model=768,
    num_layers=1,
    flash=False,  # Force explicit matmul path

)

x = jnp.ones((1, 10, 768))
params = attn.init(jax.random.PRNGKey(0), x)
out = attn.apply(params, x)

```

## Summary

- **Grouped Query Attention (GQA) in Needle 2** reduces memory bandwidth by allowing fewer KV heads than query heads, implemented through the `MultiHeadAttention` module in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).
- **Configuration requires only two parameters**: Set `num_heads` (queries) higher than `num_kv_heads` (keys/values) in `TransformerConfig` to enable automatic grouping.
- **Implementation uses JAX repetition**: The system calculates `repeats = num_heads // num_kv_heads` and uses `jnp.repeat` to broadcast KV tensors across query groups at lines 258-261.
- **Universal compatibility**: The grouping logic works identically with both `jax.nn.dot_product_attention` (flash attention) and the explicit matmul fallback path.
- **No code branching needed**: The same `MultiHeadAttention` class handles standard MHA and GQA based solely on configuration parameters.

## Frequently Asked Questions

### What is the difference between GQA and standard multi-head attention in Needle 2?

Standard multi-head attention uses identical counts for query, key, and value heads, requiring separate projections for each head. In contrast, GQA in Needle 2 allows `num_heads` to be a multiple of `num_kv_heads`, where multiple query heads share the same key/value projections after the KV tensors are repeated via `jnp.repeat`. This reduces the memory footprint of the KV cache during inference while maintaining the expressiveness of multiple query heads.

### How do I configure grouped query attention in my Needle 2 model?

Configure GQA by passing mismatched head counts to `TransformerConfig`. Set `num_heads` to your desired query head count and `num_kv_heads` to a divisor of that number (e.g., 12 queries and 4 KV heads). The `Stack` model and underlying `MultiHeadAttention` module automatically detect the difference and apply the grouping mechanism without additional configuration.

### Does GQA work with Flash Attention in Needle 2?

Yes, GQA is fully compatible with Flash Attention in Needle 2. The KV repetition using `jnp.repeat` occurs before the attention computation, ensuring that both the fast path using `jax.nn.dot_product_attention` (lines 246-254) and the explicit matmul implementation (lines 267-271) receive properly shaped tensors. You can toggle between implementations via the `flash` parameter without modifying your GQA configuration.

### Where is the GQA logic implemented in the Needle 2 source code?

The core GQA implementation resides in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) within the `MultiHeadAttention` class, specifically between lines 221 and 271. The dimension calculations occur at lines 221-223, tensor reshaping at lines 230-232, and the critical grouping logic using `jnp.repeat` at lines 258-261.