# How Grouped Query Attention (GQA) Works in Needle 2: Implementation and Configuration Guide

> Discover how Grouped Query Attention (GQA) works in Needle 2. Learn how decoupling query heads from key/value heads reduces memory usage while preserving model expressivity.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: how-to-guide
- Published: 2026-08-30

---

**Needle 2 implements Grouped Query Attention by decoupling the number of query heads from key/value heads in the `MultiHeadAttention` module, reducing memory usage while maintaining model expressivity through shared KV projections.**

Grouped Query Attention (GQA) is a memory-efficient variant of multi-head attention that reduces the computational overhead of key and value projections. In the Needle 2 deep learning framework, GQA is natively supported through configurable head counts in the core attention mechanism, allowing developers to specify fewer key/value heads than query heads for optimized inference and training.

## What is Grouped Query Attention?

Grouped Query Attention modifies the standard multi-head attention mechanism by allowing query heads to share key and value heads. Instead of maintaining separate key and value projections for every query head, GQA groups query heads to attend to the same set of keys and values. This reduces the memory footprint of the KV cache during inference—a critical optimization for large language models—while preserving the diversity of query representations.

## Implementation Details in Needle 2

The GQA implementation resides in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) within the `MultiHeadAttention` class. The mechanism operates through three distinct phases: configuration, projection, and grouped computation.

### Separate Head Count Configuration

The `MultiHeadAttention` constructor accepts distinct parameters for query and key/value head counts:

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

    num_kv_heads: int        # number of shared KV heads (GQA)

```

As defined in [architecture.py lines 9-13](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L9-L13), this separation allows `num_heads` to be an integer multiple of `num_kv_heads`, establishing the grouping ratio.

### Dimension Derivation and Projection

The forward pass dynamically calculates dimensions based on the head counts:

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

```

These calculations appear in [architecture.py lines 20-23](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L20-L23). The `kv_dim` represents the total dimension for concatenated key and value matrices, which is smaller than the query dimension when `num_kv_heads < num_heads`.

The projection layers reshape inputs into grouped head configurations:

```python
q = nn.Dense(attn_dim, ...)(x)
k = nn.Dense(kv_dim, ...)(x)
v = nn.Dense(kv_dim, ...)(x)

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)

```

This reshaping operation, found in [architecture.py lines 26-33](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L26-L33), ensures queries maintain full head diversity while keys and values use the compressed representation.

### The Grouped Attention Mechanism

Needle 2 provides two execution paths for GQA depending on whether flash attention is enabled.

**Classic Path (Flash Disabled):**
When using the standard attention implementation, the module explicitly repeats KV tensors 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 broadcast operation, located in [architecture.py lines 57-60](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L57-L60), allows each query head to attend to its corresponding shared KV head.

**Flash Attention Path:**
When `flash=True`, the implementation delegates to JAX's optimized kernels:

```python
out = jax.nn.dot_product_attention(
    q.transpose(0, 2, 1, 3),
    k.transpose(0, 2, 1, 3),
    v.transpose(0, 2, 1, 3),
    mask=mask,
    implementation=impl,
)

```

As shown in [architecture.py lines 45-55](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L45-L55), JAX handles the implicit broadcasting of the smaller KV dimension to the larger query dimension, eliminating the explicit repeat operation and improving memory efficiency.

## Code Examples

### Configuring GQA in a Transformer Model

To enable Grouped Query Attention in a Needle 2 model, specify distinct values for `num_heads` and `num_kv_heads` in the configuration:

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

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

    num_kv_heads=4,        # 4 shared KV heads (3:1 grouping ratio)

    num_layers=24,
    flash=True,            # use flash attention when available

)

model = Stack(config=cfg)

# Dummy input: batch-size 2, seq-len 128

x = jnp.ones((2, 128, cfg.d_model), dtype=jnp.bfloat16)
outputs = model(x)  # shape (2, 128, 768)

```

This configuration reduces the KV cache memory by 75% compared to standard multi-head attention while maintaining 12 query heads for representational capacity.

### Inspecting Attention Shapes

To verify the grouping behavior manually:

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

attn = MultiHeadAttention(
    num_heads=8,
    num_kv_heads=2,   # 4x reduction in KV heads

    d_model=512,
    num_layers=1,
    flash=False,      # force classic path to see explicit repeats

)

x = jnp.ones((1, 10, 512))
out = attn(x)
print(out.shape)  # -> (1, 10, 512)

```

When `flash=False`, internal tensors will show the repeat operation expanding the 2 KV heads to match 8 query heads before the attention computation.

## Performance Implications

**Memory Efficiency:** By reducing `num_kv_heads`, Needle 2 decreases the KV cache size proportionally. For a model with 12 query heads and 4 KV heads, the cache requirements drop to 33% of standard multi-head attention during autoregressive generation.

**Computational Overhead:** The explicit `jnp.repeat` operation in the classic path adds minimal overhead compared to the savings from reduced projection parameters. When using flash attention, the grouping incurs virtually no penalty as broadcasting happens within the optimized kernel.

**Model Quality:** GQA maintains performance comparable to full multi-head attention on many tasks while significantly improving throughput, particularly for long sequence lengths where memory bandwidth is the bottleneck.

## Summary

- **Separate Head Configuration:** Needle 2 supports GQA through distinct `num_heads` and `num_kv_heads` parameters in `MultiHeadAttention`.
- **Dynamic Dimension Calculation:** The framework automatically derives `kv_dim` from the reduced head count, compressing key and value projections.
- **Dual Execution Paths:** The implementation adapts between explicit KV repetition for classic attention and implicit broadcasting for JAX flash attention kernels.
- **Memory Optimization:** Reducing KV heads decreases cache memory usage linearly with the grouping ratio (e.g., 4:1 grouping reduces KV cache by 75%).
- **Source Location:** All GQA logic resides in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), specifically within the `MultiHeadAttention` class constructor and forward methods.

## Frequently Asked Questions

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

Standard multi-head attention uses the same number of heads for queries, keys, and values, requiring separate projections for each head. Grouped Query Attention reduces the number of key and value heads, allowing multiple query heads to share the same key/value vectors. In Needle 2, this is controlled by setting `num_kv_heads` lower than `num_heads` in the `TransformerConfig`.

### How does Needle 2 handle GQA with Flash Attention?

When `flash=True` in the configuration, Needle 2 calls `jax.nn.dot_product_attention` directly without explicit tensor repetition. JAX's implementation automatically broadcasts the smaller KV head dimension to match the query head dimension during the matrix multiplication, achieving the grouping effect with zero overhead. This path is defined in [architecture.py lines 45-55](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L45-L55).

### When should I use fewer KV heads than query heads?

Use grouped query attention when deploying large models where memory bandwidth and KV cache size constrain batch sizes or sequence lengths. Common configurations use grouping ratios of 2:1, 4:1, or 8:1 (e.g., 32 query heads with 8 or 4 KV heads). This is particularly beneficial for inference with long contexts, where the KV cache dominates memory usage.

### Where is the GQA logic implemented in the Needle repository?

The core GQA implementation is located in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) within the `MultiHeadAttention` class. The grouping mechanism is configured in the constructor ([lines 9-13](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L9-L13)) and executed in the forward pass, with explicit repetition logic at [lines 57-60](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L57-L60) and flash attention integration at [lines 45-55](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L45-L55).