# How Needle 2 Optimizes the KV-Cache for Bounded Memory During Forward Pass

> Discover how Needle 2 optimizes KV-cache for bounded memory during the forward pass. Learn about its hardware-aware budget and sliding window for efficient attention.

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

---

**Needle 2 enforces a hard memory ceiling on the KV-cache by combining a hardware-aware budget calculation with a configurable sliding window, ensuring the cache never exceeds 11 MiB while allowing flexible attention spans via `effective_kv_window` and runtime masking.**

Efficient inference in large language models requires strict memory management, particularly for the key-value cache that grows with sequence length. In the `cactus-compute/needle` repository (Needle 2), the forward pass implements a **bounded KV-cache strategy** that calculates a token budget from hardware constraints and enforces it through dynamic windowing and attention masking.

## Computing the Memory Budget Window

The foundation of Needle 2’s memory optimization lies in calculating the maximum safe cache size before any tokens are processed. This prevents out-of-memory errors during long generation runs.

### The `kv_budget_window` Calculation

In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `kv_budget_window` function derives the maximum storable tokens from a fixed **11 MiB KV-budget** plus safety margin. It accounts for the model’s head dimensions, layer count, quantization group sizes, and special memory sites (engrams).

```python

# needle/model/architecture.py (~line 114)

def kv_budget_window(config):
    head_dim = (getattr(config, "attn_dim", 0) 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))

```

The function computes `per_pos`—the bytes consumed per token position across all layers—then determines the largest multiple of `KV_GROUP` that fits within `KV_BUDGET_BYTES`. The result is clamped to never drop below `KV_WINDOW_MIN` or exceed `config.max_seq_len`.

## Enforcing the Effective Window

Once the budget is calculated, Needle 2 reconciles hardware limits with user preferences to determine the actual attention window used during inference.

### Budget vs. Configuration in `effective_kv_window`

The `effective_kv_window` function in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (around line 163) applies a user-specified override while respecting the memory ceiling:

```python

# needle/model/architecture.py (~line 163)

def effective_kv_window(config):
    budget = kv_budget_window(config)
    return min(budget, config.kv_window) if config.kv_window else budget

```

If `config.kv_window` is unset, the budget-derived window is used directly. If specified, the smaller of the two values is selected, guaranteeing that user configurations cannot accidentally exceed the hardware budget.

## Cache Initialization Strategy

To prevent dynamic allocation overhead and memory fragmentation, Needle 2 allocates the entire cache upfront at the maximum possible length, then views it dynamically.

### Fixed-Size Allocation with `init_kv_cache`

Located in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) (lines 34-38), `init_kv_cache` creates zero-initialized JAX arrays sized to the worst-case scenario:

```python

# needle/model/decode.py (lines 34-38)

def init_kv_cache(config, batch, max_len):
    head_dim = _attn_width(config) // config.num_heads
    shape = (config.num_layers, batch, config.num_kv_heads, max_len, head_dim)
    z = jnp.zeros(shape, jnp.float32)
    return z, z

```

This returns two tensors (`k` and `v`) with shape `[num_layers, batch, num_kv_heads, max_len, head_dim]`. By allocating once with `max_len` derived from `effective_kv_window`, the system reserves exactly the bounded memory block needed for the entire generation run.

## Runtime Forward Pass Optimization

During token generation, the forward pass updates the pre-allocated cache and masks attention to respect the sliding window constraint.

### Dynamic Cache Updates

Each layer in `_forward_cached` writes new keys and values into the pre-allocated buffers using `jax.lax.dynamic_update_slice`. This in-place update avoids memory copies and maintains constant memory usage regardless of how many tokens have been generated previously.

### Sliding Window Attention Masking

The attention mechanism enforces the KV-window boundary at runtime. In [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) (around lines 80-87), the causal mask is intersected with a window constraint:

```python

# needle/model/decode.py – part of attention logic (~lines 80-87)

if cfg.kv_window:
    rows = jnp.arange(S)
    causal = causal & ((rows[:, None] - rows[None, :]) < cfg.kv_window)
    if sink is not None:
        causal = causal | (jnp.tril(jnp.ones((S, S), bool)) & sink[:, :S][:, None, :])

```

The mask ensures positions only attend to the most recent `kv_window` tokens. The optional **sink mask** preserves attention to early document prefixes (e.g., for retrieval-augmented generation) while still bounding the total cache size.

## Practical Implementation Examples

Configure the bounded KV-cache during high-level generation calls:

```python

# Example 1 – Explicit 256-token sliding window

logits, k_cache, v_cache = generate_cached(
    config,               # TransformerConfig with kv_window=256

    params,
    tokenizer,
    prompt="Explain quantum entanglement.",
    max_new_tokens=128,
    kv_window=256)        # Forces 256-token window regardless of budget

```

```python

# Example 2 – Budget-inferred window for batched inference

texts, ids, logps = batched_generate(
    config,               # No kv_window set → uses effective_kv_window

    params,
    tokenizer,
    prompts=["Summarize the article.", "Translate to French."],
    max_new_tokens=64,
    kv_window=0)          # 0 triggers budget calculation

```

## Summary

- **Budget-driven allocation**: `kv_budget_window` in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py) calculates the maximum safe cache size from an 11 MiB hardware limit, accounting for layer count, head dimensions, and quantization.
- **Configurable enforcement**: `effective_kv_window` selects the stricter of the hardware budget or user-specified `kv_window` parameter.
- **Fixed-size initialization**: `init_kv_cache` pre-allocates the entire cache using JAX to prevent runtime memory growth.
- **Runtime windowing**: The forward pass uses `dynamic_update_slice` for cache writes and applies a sliding window mask during attention to ensure strict memory bounds.
- **Sink token support**: Optional sink masks allow early tokens to remain visible for context-heavy tasks without increasing cache storage requirements.

## Frequently Asked Questions

### What is the maximum memory limit for the KV-cache in Needle 2?

Needle 2 enforces a default **11 MiB KV-budget** (plus a small safety margin) calculated in `kv_budget_window`. This limit accounts for all layers, KV heads, and quantization overhead to prevent out-of-memory errors on accelerator devices.

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

When the generated sequence exceeds `effective_kv_window`, the model employs **sliding window attention**. The attention mask in `_attn_cached` restricts each position to attend only to the most recent `kv_window` tokens, effectively treating earlier tokens as if they were evicted from the cache while maintaining constant memory usage.

### Can I override the automatic budget calculation?

Yes. Setting `config.kv_window` to a specific integer overrides the default, but `effective_kv_window` ensures your setting cannot exceed the hardware budget. If you specify a window larger than the budget allows, the function automatically reduces it to the safe maximum.

### What is the purpose of the sink mask in the attention mechanism?

The **sink mask** allows specific early tokens (such as document prefixes or system prompts) to remain attendable even when they fall outside the sliding window. When `cfg.kv_window` is active and a sink is provided, the mask logic adds these sink positions back into the attention matrix via boolean OR operations, preserving critical context without expanding the physical cache storage.