# How Needle 2’s Bounded Memory Architecture Prevents Uncontrolled Memory Growth

> Discover how Needle 2’s bounded memory architecture prevents uncontrolled memory growth by enforcing a fixed KV cache budget and masking attention to stay within limits.

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

---

**Needle 2 caps memory usage by enforcing a fixed KV cache budget of approximately 11.5 MiB, dynamically calculating the maximum sequence window that fits within this limit and masking attention to prevent access beyond the bounded window.**

Needle 2 implements a strict bounded memory architecture to ensure transformer inference remains efficient regardless of input length. In the cactus-compute/needle repository, this design centers on budget-driven KV cache management rather than traditional attention mechanisms that allow memory to grow linearly with sequence length.

## The KV Cache Budget Mechanism

At the core of Needle 2’s memory safety is a **fixed byte budget** that strictly limits how much GPU memory the model can allocate for key-value caches during inference. This approach prevents the unbounded growth typical of standard transformer implementations.

### Calculating Per-Position Memory Cost

For each token position, Needle 2 computes the exact storage required for key and value tensors. In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `kv_budget_window` function (lines 100–110) calculates this cost as:

```python
2 * kv + 2 * (kv // KV_GROUP) * 4

```

This formula accounts for the standard key and value pairs plus additional storage for optional en-gram slots. By quantifying the per-position footprint precisely, the architecture can determine exactly how many tokens fit within the memory constraint.

### Enforcing the Byte Budget Limit

The global memory ceiling is defined by `KV_BUDGET_BYTES = 11 MiB + 512 KiB` (approximately 11.5 MiB). The `effective_kv_window` function in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) computes the maximum sequence length that satisfies this budget based on the model configuration—factoring in dimensions like `d_model`, `num_heads`, and `num_layers`.

The implementation clamps the resulting window between `KV_WINDOW_MIN = 160` and the model’s `max_seq_len`, yielding the **effective KV window** that guarantees memory usage never exceeds the predefined limit regardless of input length.

## Masking Strategies for Bounded Windows

Once the effective window is calculated, Needle 2 enforces it at the attention level through specialized masking functions that hide positions outside the allowed range.

### Causal Mask Constraints

The `make_causal_mask` function (lines 887–891 in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py)) generates the standard triangular attention mask used during autoregressive generation. However, this mask is inherently constrained by the effective KV window calculated earlier, ensuring the model only attends to positions that fit within the memory budget.

### Packing Mask Integration

For batched inference scenarios, `make_causal_packing_mask` (lines 192–207) adds a **window clause** that explicitly masks out token positions beyond the effective KV window. This prevents the attention mechanism from accessing cached values that would exceed the allocated memory buffer, effectively implementing a sliding window attention pattern driven by hardware constraints rather than arbitrary hyperparameters.

## Activation Rematerialization

Beyond KV cache management, Needle 2 further bounds memory through **activation rematerialization** during the forward pass.

### Scan-Based Layer Processing

The `Stack` module in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py) (lines 376–425) implements the transformer layers using `nn.scan`, which processes layers sequentially rather than storing all intermediate activations simultaneously. When `cfg.remat` is set to `True`, the `_ScanBody` class (lines 410–418) leverages JAX’s rematerialization to recompute intermediate values during the backward pass instead of retaining them in memory.

This scanned approach ensures that peak memory usage scales with the number of layers only through the loop carry state, not through storing full layer outputs, providing additional headroom for the KV cache budget.

## Practical Implementation

Developers interact with Needle 2’s bounded memory architecture through the `TransformerConfig` class and generation utilities in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py).

### Determining the Effective KV Window

To inspect the memory-constrained window size for a specific model configuration:

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

cfg = TransformerConfig(d_model=768, num_heads=12, num_layers=27, max_seq_len=2048)
window = effective_kv_window(cfg)      # → e.g. 1024 tokens

print(f"KV window limited to {window} positions")

```

*Source*: `TransformerConfig` definition and `effective_kv_window` logic – see **[architecture.py – config & kv window]**(https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L58-L78).

### Integrating Bounded Masks in Generation

During text generation, the bounded mask is applied automatically through the architecture’s internal mechanisms:

```python
from needle.model.run import generate
from needle.model.architecture import SimpleAttentionNetwork, TransformerConfig, make_causal_mask

cfg = TransformerConfig()
model = SimpleAttentionNetwork(cfg)

# The mask automatically respects the bounded KV window

mask = make_causal_mask(cfg.max_seq_len)   # internally limited by effective_kv_window(cfg)

# Pass the mask to the model (the generate function does this internally)

text = generate(model, params, tokenizer, "Hello", max_new_tokens=100)

```

*Source*: `make_causal_mask` creates the triangular mask; `effective_kv_window` is consulted inside `SimpleAttentionNetwork` – see **[architecture.py – mask creation]**(https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L887-L891).

### Optimizing Memory with Rematerialization

To minimize peak memory usage during training or long-context inference, enable rematerialization in the configuration:

```python
cfg = TransformerConfig(remat=True)   # enable rematerialization

model = SimpleAttentionNetwork(cfg)   # Stack will use nn.scan with remat

```

*Source*: `Stack` builds a scanned block that uses `nn.remat` when `cfg.remat` is true – see **[architecture.py – Scan block]**(https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L410-L418).

## Summary

- **Fixed byte budget**: Needle 2 enforces a strict `KV_BUDGET_BYTES` limit of approximately 11.5 MiB, preventing the KV cache from growing beyond hardware constraints.
- **Dynamic window calculation**: The `effective_kv_window` function computes the maximum safe sequence length based on model dimensions and clamps it to a minimum of 160 tokens.
- **Attention masking**: `make_causal_mask` and `make_causal_packing_mask` hide positions outside the effective window, ensuring the model never attempts to access evicted cache entries.
- **Layer rematerialization**: The `Stack` module’s `nn.scan` implementation with optional `cfg.remat` recomputes activations rather than storing them, reducing peak memory during deep forward passes.
- **Configuration-driven**: All constraints are encapsulated in `TransformerConfig`, allowing developers to inspect and adjust memory bounds programmatically.

## Frequently Asked Questions

### What is the exact KV cache budget in Needle 2?

Needle 2 defines `KV_BUDGET_BYTES` as exactly 11 MiB plus 512 KiB, totaling approximately 11.5 MiB. This fixed allocation is hardcoded in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) and represents the maximum GPU memory the KV cache can consume regardless of input sequence length.

### How does Needle 2 handle sequences longer than the effective KV window?

For sequences exceeding the calculated `effective_kv_window`, Needle 2 employs sliding window attention through `make_causal_packing_mask`, which masks out positions beyond the budget limit. The model only attends to the most recent tokens that fit within the memory constraint, effectively treating earlier positions as inaccessible.

### What is the minimum sequence window supported by Needle 2's architecture?

The architecture enforces `KV_WINDOW_MIN = 160` tokens as the absolute floor for the effective KV window. Even if the budget calculation suggests a smaller window would suffice, the implementation clamps the value to 160 to ensure sufficient context for meaningful attention patterns.

### Does enabling rematerialization impact inference speed?

Yes, setting `cfg.remat=True` trades computation for memory by recomputing intermediate activations during the forward or backward pass rather than storing them. While this reduces peak memory usage—particularly beneficial for deep models with 27+ layers—it introduces additional computational overhead that may increase latency during generation.