# Needle 2 Sliding Window Attention: How the 256‑Token KV Cache Limit Works

> Understand Needle 2's 256-token sliding window attention. Learn how constant memory usage is maintained regardless of sequence length by masking self-attention to recent positions.

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

---

**Needle 2 caps the key/value (KV) cache by masking self-attention so that each token only attends to the most recent positions within a user-defined `kv_window`, keeping memory usage constant regardless of total sequence length.**

**Needle 2**, developed in the `cactus-compute/needle` repository, implements a configurable **sliding-window attention** mechanism that bounds the **key/value (KV) cache** to a **256‑token** (or smaller) context limit during inference. While the default budget-based setting retains the most recent 64 tokens (2⁶), the **`TransformerConfig.kv_window`** parameter lets users raise the boundary to 256 tokens or any other size, trading memory usage for long-range dependency coverage.

## How Needle 2 Configures the KV Sliding Window

### `kv_window` and `effective_kv_window`

The window size is controlled by **`TransformerConfig.kv_window`**. When the user leaves this field unset, the function **`effective_kv_window(config)`** defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) computes a budget-based default. This fallback typically yields a 64-token (2⁶) window, but you can override it with any positive integer—such as 128 or 256—depending on your hardware budget and task requirements.

## Mask Creation in the Needle 2 Architecture

### Local attention masking in `hidden_cells` and `_encode_contrastive`

Inside [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the model builds a **recent mask** that keeps only positions whose relative distance is strictly less than the window size. The logic appears in both `hidden_cells` and `_encode_contrastive`:

```python
recent = ((pos[:, None] - pos[None, :]) < window)[None, None, :, :]
mask = mask & recent

```

The expression `pos[:, None] - pos[None, :]` creates a distance matrix, and the comparison `< window` zeros out any attention link older than the window. This matrix is then fused with the base causal mask via a bitwise `&` so that a token cannot attend to future positions or to past positions beyond the sliding boundary.

### Decoder masking in [`decode.py`](https://github.com/cactus-compute/needle/blob/main/decode.py)

During autoregressive generation, the same cutoff is enforced in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py):

```python
if cfg.kv_window:
    causal = causal & ((rows[:, None] - rows[None, :]) < cfg.kv_window)

```

Here, `rows` represents the current sequence positions. If `cfg.kv_window` is set to 256, the decoder restricts attention to the most recent 256 tokens, preventing the KV cache from growing indefinitely as the sequence lengthens.

### Packing masks for retrieval-augmented generation

When multiple segments are packed together—such as for retrieval-augmented generation—the function **`make_causal_packing_mask`** in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) applies the sliding-window rule in the same way. It can also combine the window with a **sink mask** that preserves prefix tokens, ensuring document boundaries remain respected while the cache stays bounded.

## Practical Code Examples

The following snippets show how to invoke and inspect the sliding window in Needle 2.

Use the default 64-token window implicitly:

```python

# Use the default 64-token window (implicitly)

model = needle.SimpleAttentionNetwork(cfg)
tokens = tokenizer.encode("The quick brown fox jumps over the lazy dog …")
logits = model(tokens)  # KV cache will slide over 64-token chunks

```

Override the window size explicitly:

```python

# Override the window size explicitly

cfg = needle.model.TransformerConfig(kv_window=128)  # 2^7-token window

model = needle.SimpleAttentionNetwork(cfg)
logits = model(tokens)

```

You can pass `kv_window=256` (or any other value) to the same constructor to enlarge the context.

Inspect the mask directly for debugging:

```python

# Inspect the mask directly (useful for debugging)

import jax.numpy as jnp

seq_len = 200
window = 64
pos = jnp.arange(seq_len)
recent_mask = (pos[:, None] - pos[None, :]) < window  # shape (seq_len, seq_len)

print(recent_mask.astype(jnp.int32))

```

## Key Source Files

The sliding-window logic is distributed across four main files in the `cactus-compute/needle` codebase:

- **[`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)** — Contains `effective_kv_window`, `kv_budget_window`, `hidden_cells`, `_encode_contrastive`, and `make_causal_packing_mask`. This is where the core mask arithmetic lives.
- **[`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py)** — Applies the KV-window mask during autoregressive decoding via `if cfg.kv_window`.
- **[`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py)** — Records the `kv_window` metadata field so that exported checkpoints remember the training or inference sliding-window width.
- **[`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py)** — Calls `effective_kv_window` during fine-tuning and export, ensuring the configured window is respected across training stages.

## Summary

- Needle 2 bounds the KV cache with a **sliding-window mask** rather than full-history attention.
- The default window size is **64 tokens** (2⁶) computed by `effective_kv_window()` in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), but it can be overridden via `TransformerConfig.kv_window`.
- The mask is built by comparing relative position distances: `((pos[:, None] - pos[None, :]) < window)`.
- Both the encoder paths (`hidden_cells`, `_encode_contrastive`) and the decoder path ([`decode.py`](https://github.com/cactus-compute/needle/blob/main/decode.py)) enforce the same cutoff.
- Packed sequences are supported through `make_causal_packing_mask`, which can optionally layer a sink mask on top of the window.

## Frequently Asked Questions

### What is the default KV cache window size in Needle 2?

By default, Needle 2 uses a budget-based window computed by `effective_kv_window()` in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). This fallback typically produces a **64-token** (2⁶) context. If `TransformerConfig.kv_window` is explicitly provided, that value overrides the budget heuristic entirely.

### How do I configure a 256-token sliding window in Needle 2?

Pass `kv_window=256` to `TransformerConfig` when building the model. The decoder logic in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) will then evaluate `((rows[:, None] - rows[None, :]) < 256)`, allowing each position to attend to the previous 256 tokens while still respecting causality.

### Does the sliding window work with packed sequences in Needle 2?

Yes. The function `make_causal_packing_mask` in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) applies the same distance-based mask across packed segments. It can also combine the sliding window with a sink mask so that prefix tokens remain globally visible even when local attention is enforced.

### How does the sliding window reduce memory during inference?

Without a window, the KV cache grows linearly with sequence length because every new token must store keys and values for all previous positions. Needle 2’s sliding window caps the number of cached entries to the most recent `kv_window` tokens, keeping peak memory roughly constant regardless of how long the generation continues.