# KV Cache Size in Needle 2's Bounded Memory Architecture: Complete Technical Guide

> Explore Needle 2's KV cache size within its bounded memory architecture. Understand its dense JAX tensor allocation and memory consumption in this technical guide.

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

---

**Needle 2 allocates its KV cache as a dense JAX tensor with shape `(num_layers, batch, num_kv_heads, max_len, head_dim)`, consuming exactly `num_layers × batch × num_kv_heads × max_len × head_dim × 4` bytes in float32 precision.**

The bounded memory architecture in Needle 2, an open-source JAX transformer inference engine from the Cactus Compute project, implements a fixed-capacity KV cache that trades dynamic memory growth for predictable GPU utilization. Understanding this cache geometry is essential for provisioning hardware and configuring model deployments.

## How Needle 2 Defines the KV Cache Shape

The cache initialization occurs in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py). The `init_kv_cache()` function (line 34) constructs identical zero tensors for keys and values from model configuration parameters:

```python
def init_kv_cache(config, batch, max_len):
    head_dim = _attn_width(config) // config.num_heads          # = attn_dim / num_heads

    shape = (config.num_layers, batch, config.num_kv_heads,
             max_len, head_dim)                                 # ← KV‑cache shape

    z = jnp.zeros(shape, jnp.float32)                         # keys and values share the same buffer

    return z, z

```

Each dimension maps directly to a `TransformerConfig` field defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py):

- **`config.num_layers`** — transformer depth (e.g., 27 for the example configuration)
- **`batch`** — concurrent sequence count passed at initialization
- **`config.num_kv_heads`** — GQA/MQA head count, typically half of query heads
- **`max_len`** — hard upper bound on sequence length, drawn from `config.max_seq_len`
- **`head_dim`** — per-head feature dimension, computed as `attn_dim // num_heads`

## Computing the Exact Memory Footprint

The total element count follows directly from the tensor shape. For the **KV cache size in bytes**, multiply by 4 for 32-bit floats:

```

elements = num_layers × batch × num_kv_heads × max_len × head_dim
bytes    = elements × 4

```

Since `init_kv_cache()` returns *separate* tensors for keys and values, **total KV cache memory doubles** this single-tensor figure.

### Worked Example

```python
>>> from needle.model.architecture import TransformerConfig
>>> from needle.model.decode import init_kv_cache
>>> cfg = TransformerConfig(num_layers=27, num_heads=12,
...                        num_kv_heads=6, attn_dim=768,
...                        max_seq_len=2048, kv_window=512)
>>> batch = 1
>>> max_len = cfg.max_seq_len
>>> k_cache, v_cache = init_kv_cache(cfg, batch, max_len)
>>> k_cache.shape
(27, 1, 6, 2048, 64)          # 27 layers, 1 batch, 6 KV‑heads,

                               # 2048 tokens, 64‑dim per head

>>> # Total elements ⇒ 27·1·6·2048·64 = 21 164 928

>>> # Memory (float32) ⇒ ~81 MB per cache (keys + values ≈ 162 MB)

```

This 162 MB allocation remains constant regardless of actual input length—**the bounded-memory guarantee**.

## The Role of `kv_window` in Bounded Memory

The `kv_window` parameter in `TransformerConfig` implements **logical truncation without physical reallocation**. When `kv_window > 0`:

1. The underlying tensor retains its full `max_len` dimension
2. The attention mask in `_attn_cached()` ([`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py)) restricts attention to the most recent `kv_window` positions
3. Memory bandwidth improves by loading only the active window during attention computation

This design preserves the **static allocation benefit**—no fragmentation, no growth reallocations—while allowing efficient long-context inference through sliding-window patterns.

## Key Configuration Parameters

| Parameter | Source Location | Impact on Cache Size |
|-----------|-----------------|----------------------|
| `num_layers` | `TransformerConfig` | Linear multiplier |
| `num_kv_heads` | `TransformerConfig` | Linear multiplier via GQA reduction |
| `max_seq_len` | `TransformerConfig` | Bounds the `max_len` dimension |
| `attn_dim` / `num_heads` | `TransformerConfig` | Determines `head_dim` |
| `kv_window` | `TransformerConfig` | Limits active computation, not allocation |

## Summary

- **KV cache shape**: `(num_layers, batch, num_kv_heads, max_len, head_dim)` as defined in `init_kv_cache()`
- **Memory formula**: `2 × num_layers × batch × num_kv_heads × max_len × head_dim × 4` bytes (keys + values in float32)
- **Bounded memory**: Physical allocation fixed at `max_len`; `kv_window` provides logical sliding-window attention without tensor resizing
- **Source files**: [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) (allocation and attention logic), [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (configuration schema)

## Frequently Asked Questions

### How does `kv_window` affect memory usage in Needle 2?

It does not reduce memory usage. The `kv_window` field only constrains which positions the attention mask allows the model to attend to. The underlying JAX tensor allocated by `init_kv_cache()` always spans the full `max_len` dimension. Memory bandwidth decreases because fewer positions are loaded per attention operation, but GPU memory consumption stays fixed.

### Why does Needle 2 use static allocation instead of dynamic growth?

Static allocation eliminates memory fragmentation and allocation overhead during autoregressive generation. The JAX trace remains stable across variable input lengths, enabling XLA compilation optimizations. This trade-off requires users to provision `max_seq_len` aggressively but guarantees predictable latency without pauses for cache expansion.

### Can I reduce KV cache size by using fewer KV heads?

Yes. The `num_kv_heads` parameter applies **grouped-query attention** compression. Reducing from full `num_heads` to a smaller `num_kv_heads` (e.g., 12 → 6 → 1) linearly decreases cache footprint. In `TransformerConfig`, this is independent of query head count, allowing flexible memory/quality trade-offs.

### What precision does Needle 2 use for KV cache storage?

Float32 (`jnp.float32`) as explicitly specified in `init_kv_cache()`. Each element occupies 4 bytes. Quantization support would require modifications to the initialization dtype and attention kernel compatibility in `_attn_cached()`.