# How GQA Attention Enables Efficient Inference in Needle 2: A Technical Deep Dive

> Discover how GQA Attention powers efficient inference in Needle 2. Learn how it reduces memory bandwidth and cache pressure for large transformer models on edge devices.

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

---

**Grouped-Query-Attention (GQA) reduces memory bandwidth and cache pressure during inference by sharing key and value tensors across multiple query heads, enabling Needle 2 to run large transformer models within a 28 MB memory footprint on edge devices.**

Needle 2 is an open-source transformer engine designed for ultra-low-resource environments, targeting deployment on phones, microcontrollers, and edge GPUs. At the heart of its efficiency lies **GQA Attention**, a memory-optimization technique that dramatically reduces the key-value (KV) cache size during each forward pass. According to the cactus-compute/needle source code, this architectural choice allows the 14 MB engine to achieve a 5×-70× size advantage while maintaining comparable model quality through the **Simple Attention Network** architecture.

## What Is GQA Attention?

In standard **Multi-Head Attention (MHA)**, every attention head independently computes its own query (**Q**), key (**K**), and value (**V**) vectors, resulting in tensors shaped `(B, T, num_heads, head_dim)`. This approach requires materializing full KV tensors for every head, significantly increasing memory traffic during autoregressive inference.

**GQA Attention** modifies this by sharing KV tensors across groups of heads. Only a designated subset of heads—called **KV-heads**—actually produce K and V matrices, while the remaining heads reuse these shared tensors. The ratio between total heads and KV-heads is controlled by the `num_kv_heads` configuration field, directly reducing the memory bandwidth required for each attention step.

## GQA Attention Implementation in Needle 2

The implementation resides in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), where the `TransformerConfig` dataclass and `MultiHeadAttention` module coordinate to enable grouped query patterns.

### TransformerConfig and KV-Head Configuration

The configuration interface defines GQA behavior through two key parameters at lines 58-66:

```python
@dataclass
class TransformerConfig:
    d_model: int
    num_heads: int
    num_kv_heads: int  # Enables GQA when < num_heads

    num_layers: int
    flash: bool = True
    dtype: str = "bfloat16"

```

When `num_kv_heads` is less than `num_heads`, the system activates GQA mode. For example, setting `num_heads=12` and `num_kv_heads=6` creates a 2:1 grouping ratio where every KV tensor is shared between two query heads.

### MultiHeadAttention Module and Tensor Repeats

The `MultiHeadAttention` class implements the actual tensor sharing logic. During initialization (lines 20-30), the module calculates KV dimensions based on `num_kv_heads` rather than the full head count:

```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  # KV size depends on num_kv_heads

```

During the forward pass at lines 57-62, the module replicates KV tensors to match the query head count:

```python
if self.flash:
    # Fast dot-product attention path

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

        v = jnp.repeat(v, repeats, axis=1)  # GQA: share V across groups

```

This `jnp.repeat` operation is computationally cheap compared to generating unique KV tensors, allowing the model to maintain full attention expressiveness while minimizing memory materialization.

## Performance Impact of GQA Attention on Inference

GQA Attention transforms the computational profile of the transformer stack by reducing memory-bound operations that typically bottleneck edge devices.

**Memory Traffic Reduction**

Standard MHA generates KV tensors proportional to `num_heads × head_dim`, while GQA produces only `num_kv_heads × head_dim` values. This reduction by a factor of `num_heads / num_kv_heads` directly decreases RAM usage and the bandwidth required for each attention step.

**Cache Pressure and Latency**

Smaller KV buffers improve cache locality and reduce cache misses during the autoregressive generation loop. The Needle 2 engine leverages this to maintain a 256-token sliding window while keeping the total memory footprint near 28 MB, as documented in the repository README.

**Compute Efficiency**

While the core matrix multiplication `Q·Kᵀ` still executes per head, it operates on the smaller, repeated KV tensors. This structure allows the engine to batch operations effectively, particularly when combined with optimized kernels.

## Integrating GQA Attention with Flash Attention and KV Budgeting

Needle 2 combines GQA with complementary optimizations to maximize throughput on constrained hardware.

**Flash Attention Integration**

When `flash=True` in the configuration, the engine applies flash attention kernels to the GQA-reduced KV tensors. This pairing minimizes both memory movement and compute time, as the flash kernel processes the already-compact KV representation without materializing full attention matrices.

**KV Budget Calculation**

The configuration also supports dynamic cache management through `kv_budget_window` (lines 104-111). This calculation uses `num_kv_heads` to determine the maximum number of tokens that can be retained in the KV cache given memory constraints:

```python

# KV budget uses num_kv_heads to calculate cache capacity

kv_budget = calculate_kv_budget(
    num_kv_heads=self.num_kv_heads,
    head_dim=head_dim,
    max_memory_mb=28
)

```

This integration ensures that GQA Attention works synergistically with memory management policies to prevent out-of-memory errors during long-context inference.

## Practical Configuration Examples

### Configuring a Model with GQA Attention

Instantiate a 12-head model using 6 KV-heads to enable 2:1 GQA grouping:

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

cfg = TransformerConfig(
    d_model=768,
    num_heads=12,
    num_kv_heads=6,      # Enables GQA Attention

    num_layers=27,
    flash=True,
    dtype="bfloat16"
)

model = SimpleAttentionNetwork(cfg)

```

### Running Quantized Inference

Activate fake-quantization during inference to simulate the production engine's 14 MB footprint:

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

tokens = jnp.array([[1, 42, 7, 0, 0]])  # shape (batch, seq_len)

logits = model(tokens, quant=True)       # quant=True activates fake-quant

print(logits.shape)  # → (1, seq_len, vocab_size)

```

### Inspecting KV Tensor Shapes

Debug the GQA configuration by verifying tensor dimensions:

```python
def inspect_kv(cfg):
    head_dim = (cfg.attn_dim or cfg.d_model) // cfg.num_heads
    kv_dim = cfg.num_kv_heads * head_dim
    print(f"Q-head dim: {head_dim}")
    print(f"K/V dim (per KV-head): {head_dim}")
    print(f"Total KV dim: {kv_dim}")

inspect_kv(cfg)

# Output:

# Q-head dim: 64

# K/V dim (per KV-head): 64

# Total KV dim: 384   # 6 KV-heads * 64

```

## Summary

- **GQA Attention** reduces KV cache memory by sharing key and value tensors across multiple query heads, controlled by the `num_kv_heads` parameter in `TransformerConfig`.
- The implementation in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) uses `jnp.repeat` to broadcast KV tensors to query heads, minimizing memory bandwidth during inference.
- Combining GQA with flash attention and KV budgeting allows Needle 2 to operate within a 28 MB memory footprint on edge devices.
- Configuration requires setting `num_kv_heads` lower than `num_heads`, typically achieving 2:1 or 4:1 compression ratios without quality degradation.

## Frequently Asked Questions

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

Standard multi-head attention generates independent K and V tensors for every attention head, resulting in `num_heads × head_dim` parameters per token. GQA Attention generates K and V only for a subset of heads (`num_kv_heads`) and shares these among query heads, reducing memory by the ratio `num_heads / num_kv_heads` while maintaining query diversity.

### How do you configure GQA Attention in Needle 2?

Set `num_kv_heads` in the `TransformerConfig` dataclass to a value less than `num_heads`. For example, `num_heads=12` with `num_kv_heads=6` enables 2:1 grouping. The `MultiHeadAttention` module automatically detects this configuration and applies the repeat logic in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) during the forward pass.

### Does GQA Attention reduce model accuracy?

According to the Needle 2 architecture documentation, GQA Attention retains comparable model quality despite the memory reduction. The technique preserves full query diversity while sharing only the KV representation, allowing the engine to achieve its documented 5×-70× size advantage without significant accuracy loss.

### How does GQA Attention interact with the KV cache budget?

The KV cache budget calculation (`kv_budget_window`) at lines 104-111 uses `num_kv_heads` to determine how many tokens can be cached given memory constraints. Because GQA reduces the per-token KV size, the system can cache more tokens within the same 28 MB limit, supporting longer context windows on memory-constrained devices.