# Grouped Query Attention GQA Implementation in Needle: MultiHeadAttention Source Analysis

> Discover the Grouped Query Attention GQA implementation in Needle's MultiHeadAttention. Learn how Needle achieves memory efficiency for resource-constrained deployments.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: internals
- Published: 2026-08-15

---

**Needle implements Grouped Query Attention (GQA) by projecting keys and values to fewer heads than queries in `MultiHeadAttention`, then repeating the KV tensors via `jnp.repeat` to match query head counts, enabling memory-efficient attention for resource-constrained deployments.**

Needle is a JAX-based transformer library optimized for edge devices with severe memory constraints. The **Grouped Query Attention GQA implementation in Needle** resides primarily in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), where the `MultiHeadAttention` class separates query head counts from key/value head counts to reduce KV cache memory bandwidth while maintaining full query parallelism.

## Core GQA Architecture in needle/model/architecture.py

The GQA logic is encapsulated within the `MultiHeadAttention` class (lines 210-278) and configured through the `TransformerConfig` dataclass (lines 4-15). This architecture allows the same codebase to run both standard multi-head attention and grouped-query attention depending on configuration parameters.

### TransformerConfig: Defining Head Groups

GQA activation begins in the configuration layer. The `TransformerConfig` dataclass defines `num_heads` for queries and `num_kv_heads` for keys and values:

- **`num_heads`**: Total number of query heads (default: 8 for `needle2`)
- **`num_kv_heads`**: Number of distinct key/value head groups (default: 2 for `needle2`)

When `num_kv_heads < num_heads`, Needle activates the GQA code path. The `Block` class (lines 26-33) instantiates `MultiHeadAttention` with these values, passing them as `self.num_heads` and `self.num_kv_heads` respectively.

### MultiHeadAttention Class Implementation

The `MultiHeadAttention` class in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) implements GQA through differential projection dimensions and tensor repetition:

**Query Projection** uses the full attention dimension:

```python
q = nn.Dense(attn_dim, ...)(x)  # attn_dim = num_heads * head_dim

```

**Key/Value Projection** uses a reduced dimension:

```python
kv_dim = num_kv_heads * head_dim
k = nn.Dense(kv_dim, ...)(x)
v = nn.Dense(kv_dim, ...)(x)

```

This dimensional reduction is the first optimization: by projecting to `num_kv_heads` instead of `num_heads`, the memory footprint of KV tensors drops by the ratio `num_heads / num_kv_heads`.

### The Grouping Mechanism: From Projection to Computation

After projection, the implementation reshapes and repeats tensors to enable grouped attention:

1. **Reshape to Head Dimensions**: Queries split into `num_heads` heads, while keys and values split into `num_kv_heads` heads:
   ```python
   q = q.reshape(..., self.num_heads, head_dim)
   k = k.reshape(..., self.num_kv_heads, head_dim)
   v = v.reshape(..., self.num_kv_heads, head_dim)
   ```

2. **Repeat KV Heads**: When `num_heads > num_kv_heads`, the code calculates `repeats = num_heads // num_kv_heads` and expands the KV tensors:
   ```python
   if repeats > 1:
       k = jnp.repeat(k, repeats, axis=1)
       v = jnp.repeat(v, repeats, axis=1)
   ```

   This `jnp.repeat` operation aligns each query head with its corresponding KV group, implementing the grouped-query pattern without duplicating underlying parameters.

3. **Standard Attention Computation**: The repeated tensors then flow through standard attention weight calculation:
   ```python
   attn_weights = jnp.matmul(q, k.transpose(...)) / scale
   out = jnp.matmul(attn_weights, v)
   ```

When `self.flash` is enabled, the implementation uses optimized kernels; when disabled, the explicit repetition logic ensures compatibility with standard JAX operations while maintaining the GQA memory savings.

## Configuring and Verifying GQA in Practice

You can explicitly configure the GQA ratio when instantiating Needle models. The default `needle2` model uses 8 query heads and 2 KV heads (4:1 grouping), but you can adjust these values via `TransformerConfig`.

### Creating a Custom GQA Configuration

```python
import needle
from needle.model.architecture import TransformerConfig

# Configure 8 query heads with 2 KV heads (4:1 GQA ratio)

cfg = TransformerConfig(
    vocab_size=32000,
    d_model=768,
    num_layers=12,
    num_heads=8,        # full query heads

    num_kv_heads=2,     # grouped KV heads

    flash=True,         # use optimized kernels if available

)

# Instantiate the model with GQA

model = needle.model.SimpleAttentionNetwork(config=cfg)
logits = model(tokens)  # GQA applied internally

```

### Inspecting KV Tensor Shapes

To verify the grouping behavior, inspect the intermediate tensor dimensions:

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

def debug_kv_shapes(x, num_heads=8, num_kv_heads=2):
    """Verify that KV tensors have fewer heads than Q tensors."""
    attn = MultiHeadAttention(
        num_heads=num_heads,
        num_kv_heads=num_kv_heads,
        d_model=768,
        num_layers=12
    )
    variables = attn.init(jax.random.PRNGKey(0), x, quant=False, mask=None, rope=None)
    q, k, v = attn.apply({"params": variables["params"]}, x, 
                         quant=False, mask=None, rope=None, method=attn._project_qkv)
    
    print(f"Q shape: {q.shape}")  # (batch, seq, 8, head_dim)

    print(f"K shape: {k.shape}")  # (batch, seq, 2, head_dim) before repeat

    print(f"V shape: {v.shape}")  # (batch, seq, 2, head_dim) before repeat

# Example with batch=1, seq=128, d_model=768

debug_kv_shapes(jnp.ones((1, 128, 768)))

```

## Memory Efficiency and Performance Impact

The **Grouped Query Attention GQA implementation in Needle** delivers specific advantages for the library's target use case: a 45-million parameter model running on devices with ≤30MiB RAM.

**Memory Bandwidth Reduction**: By using 2 KV heads instead of 8, the KV cache memory bandwidth drops by 75% during inference. This reduction dominates memory savings in long-context scenarios where key and value tensors grow with sequence length.

**Computational Parallelism**: Query projections remain fully parallel across all 8 heads, preserving the model's representational capacity. Only the KV projections are shared, reducing FLOPs for the attention operation while maintaining throughput for the query pathway.

**Configuration Flexibility**: Setting `num_kv_heads == num_heads` disables GQA and reverts to standard multi-head attention using the identical code path. This compatibility ensures `SimpleAttentionNetwork` works seamlessly across different hardware constraints without architectural changes.

## Summary

- **Grouped Query Attention in Needle** is implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) within the `MultiHeadAttention` class (lines 210-278).
- The mechanism relies on **differential head counts**: full `num_heads` for queries, reduced `num_kv_heads` for keys and values.
- **Tensor repetition** via `jnp.repeat` aligns KV groups with query heads during computation, occurring when flash attention is disabled.
- **Configuration** happens through `TransformerConfig` where `num_kv_heads < num_heads` activates GQA; the default `needle2` model uses an 8:2 ratio.
- **Memory savings scale** with the grouping ratio, making the architecture suitable for edge deployment with strict RAM constraints (≤30MiB).

## Frequently Asked Questions

### What is the difference between num_heads and num_kv_heads in Needle?

`num_heads` controls the total number of query heads in the attention mechanism, while `num_kv_heads` specifies how many distinct head groups exist for keys and values. When `num_kv_heads` is smaller than `num_heads`, Needle implements GQA by sharing KV representations across multiple query heads, reducing memory usage by the ratio `num_heads / num_kv_heads`.

### How does Needle handle the case when flash attention is disabled?

When `self.flash` is set to `False` in `MultiHeadAttention`, the implementation explicitly repeats KV tensors using `jnp.repeat(k, repeats, axis=1)` where `repeats = num_heads // num_kv_heads`. This ensures each query head attends to the correct KV group using standard JAX operations, maintaining the GQA memory savings without requiring specialized kernel support.

### Can I use standard multi-head attention instead of GQA in Needle?

Yes. Set `num_kv_heads` equal to `num_heads` in your `TransformerConfig`. When these values match, the `repeats` variable equals 1, the conditional `jnp.repeat` operation is skipped, and the `MultiHeadAttention` class executes standard multi-head attention with unique KV projections for every query head.

### What is the default GQA configuration for the needle2 model?

The default `needle2` model uses `num_heads=8` and `num_kv_heads=2`, creating a 4-to-1 query-to-KV ratio. This configuration reduces the KV cache memory footprint by 75% compared to standard multi-head attention while preserving the full 8-head capacity for query representations, as defined in the `TransformerConfig` instantiation within [`needle/model/__init__.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/__init__.py).