# How Needle 2 Implements a Bounded Memory Architecture with Sliding Window and KV Sinks

> Needle 2 implements bounded memory with a sliding window and KV sinks. Discover how it maintains strict KV-cache limits and preserves tokens for efficient inference.

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

---

**Needle 2 enforces strict KV-cache memory limits through a budget-aware sliding window mechanism that optionally preserves critical tokens via KV-sinks, ensuring inference stays within approximately 11.5 MiB regardless of sequence length.**

The `cactus-compute/needle` repository implements a transformer architecture designed for resource-constrained environments. By combining automatic KV-budget calculation with configurable sliding windows and sink masks, the system guarantees bounded memory usage during generation while allowing users to protect specific tokens (like prompt prefixes) from eviction.

## KV-Budget Calculation and Window Sizing

The foundation of Needle 2's bounded memory architecture lives in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), where the system derives a hard upper bound for the KV-cache size. The `kv_budget_window` function computes how many tokens can fit within the **KV-budget** of approximately 11.5 MiB.

```python
KV_BUDGET_BYTES = 11 * 1024 * 1024 + 512 * 1024   # ≈ 11.5 MiB

KV_GROUP = 32
KV_WINDOW_MIN = 160

def kv_budget_window(config):
    head_dim = (config.attn_dim or config.d_model) // config.num_heads
    kv = config.num_kv_heads * head_dim
    d, L = config.d_model, config.num_layers
    sites = len(tuple(getattr(config, "engram_layers", (2, 15))))
    per_pos = (L * (2 * kv + 2 * (kv // KV_GROUP) * 4)
               + sites * (d + (d // KV_GROUP) * 4))
    window = (KV_BUDGET_BYTES // per_pos) // KV_GROUP * KV_GROUP
    return max(KV_WINDOW_MIN, min(window, config.max_seq_len))

```

This calculation accounts for the model dimensions (`d_model`, `num_layers`, `num_kv_heads`) and **engram layers** to determine bytes per position. The `effective_kv_window` function then selects the smaller value between this budget-derived limit and any user-specified `config.kv_window`, ensuring hardware constraints always take precedence over user preferences.

## Sliding-Window Attention Mask

During the forward pass, Needle 2 applies a **sliding-window mask** to restrict attention to recent tokens only. In [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py) lines 45-50, the `hidden_cells` function constructs a boolean tensor that limits each position to attend only to the `window` most recent tokens.

```python
if window:
    pos = jnp.arange(tokens.shape[1])
    recent = ((pos[:, None] - pos[None, :]) < window)[None, None, :, :]
    keep = recent if sink is None else (recent | sink[:, None, None, :])
    mask = mask & keep

```

The `recent` mask creates a triangular-plus-band structure where positions can only see back `window` steps. When the sequence exceeds this window, older KV entries are effectively masked out and can be dropped from cache, maintaining the bounded memory guarantee. This mechanism operates inside `MultiHeadAttention` through the mask passed to the attention computation.

## KV-Sinks for Token Persistence

To prevent critical tokens from being evicted by the sliding window, Needle 2 implements **KV-sinks**—boolean masks that pin specific positions permanently in the KV-cache. These sinks merge with the sliding-window mask via a logical OR operation.

In `make_causal_packing_mask` (lines 27-30), the system accepts a `prefix` argument that identifies sink tokens:

```python
def make_causal_packing_mask(seg_ids, prefix=None, window=0):
    ...
    if window:
        pos = jnp.arange(T)
        recent = (pos[:, None] - pos[None, :]) < window
        sink = (jnp.zeros_like(seg_ids, dtype=jnp.bool_) if prefix is None
                else prefix > 0)
        mask = mask & (recent[None, :, :] | sink[:, None, :])
    return mask[:, None, :, :]

```

When a `sink` mask is provided to `hidden_cells`, the attention mask becomes `recent | sink`, ensuring that pinned tokens remain visible to all future positions regardless of how far back they reside. This allows the model to maintain access to prompt prefixes or system instructions while still enforcing the bounded window for other tokens.

## Practical Implementation Example

Configuring the bounded memory architecture requires setting the KV-window and optionally defining sink masks. The following example demonstrates automatic budget calculation and explicit sink configuration:

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

# Configure model with automatic KV-budget detection

cfg = TransformerConfig(
    d_model=768,
    num_heads=12,
    num_kv_heads=6,
    num_layers=27,
    kv_window=0,  # 0 lets the budget decide automatically

)

# Query the effective window size

max_window = effective_kv_window(cfg)  # e.g., 2048 tokens

# Create a KV-sink mask protecting the first 10 tokens

seq_len = 32
prefix_mask = jnp.arange(seq_len) < 10  # shape (seq_len,)

# Run inference with 1024-token sliding window and prefix sinks

tokens = jnp.array([[...]])  # shape (batch, seq_len)

logits = model(
    tokens,
    window=1024,        # Sliding-window size

    sink=prefix_mask    # KV-sink protecting prefix

)

```

Higher-level APIs like `encode_contrastive` and `forward_confidence` accept the same `window` and `sink` arguments, delegating to `hidden_cells` for mask composition. For deployment, [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) serializes the `kv_window` value into the binary header for the C++ inference engine, while [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) utilizes these constraints during generation rollouts.

## Summary

- **KV-budget calculation** in `kv_budget_window` derives the maximum safe window size from hardware constraints (≈11.5 MiB) and model architecture.
- **Sliding-window masks** in `hidden_cells` restrict attention to recent tokens, automatically discarding older KV entries when sequences exceed the budget.
- **KV-sinks** allow boolean masks to pin critical tokens (like prefixes) into the cache permanently via logical OR with the window mask.
- **Bounded memory guarantees** are maintained regardless of sequence length, with user overrides always clamped by the hardware-derived budget.

## Frequently Asked Questions

### What happens if I set a kv_window larger than the budget allows?

Needle 2 enforces the hardware budget as the upper bound. The `effective_kv_window` function takes the minimum of your requested window and the budget-calculated maximum, ensuring the KV-cache never exceeds approximately 11.5 MiB.

### Can I use KV-sinks without a sliding window?

Yes. If you provide a `sink` mask but set `window=0`, the system maintains full causal attention while still preserving the sink tokens. However, without an active window, the bounded memory guarantee depends entirely on the sequence length not exceeding the budget.

### How do KV-sinks affect the memory budget calculation?

KV-sinks consume memory within the fixed budget. The `kv_budget_window` calculation assumes worst-case storage for all positions, so enabling sinks reduces the effective tokens available for the sliding window portion of the cache.

### Where is the KV-cache physically limited in the codebase?

The physical limitation occurs in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) through the mask generation logic in `hidden_cells` and `make_causal_packing_mask`. These functions determine which KV entries the attention mechanism can access, effectively controlling which tensors remain in memory during the forward pass.