# How Needle 2 Uses a Sliding-Window KV Cache to Keep Memory Usage Low

> Discover how Needle 2's sliding window KV cache maintains a constant memory footprint by storing only recent tokens, capping inference memory for efficient AI.

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

---

**Needle 2 caps inference memory by retaining only the most recent *N* tokens in a fixed-size KV cache, turning linear memory growth into a constant footprint.**

The `cactus-compute/needle` repository implements a constrained-context inference engine where long sequences no longer threaten GPU memory. Needle 2 achieves this through a **sliding-window KV cache** that discards stale key-value pairs instead of accumulating them. The result is a predictable, flat memory profile that scales with window size rather than sequence length.

## How the Sliding-Window KV Cache Works in Needle 2

Needle 2’s memory savings come from three coordinated design choices: a configuration flag that enables the window, a runtime buffer that enforces the limit, and a model header that persists the setting across exports.

### Configuring the Window in TransformerConfig

In [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `TransformerConfig` dataclass exposes the fields that control the window:

- `sliding_window: int = 0` — When set to a positive integer, it activates sliding-window attention and tells the runtime how many recent tokens to keep.
- `kv_window: int = 0` — Records the training-time window width for validation during export.

Setting `sliding_window` to a value such as `1024` tells the engine to ignore all tokens beyond the last 1024 positions when computing attention.

### Runtime Allocation in KVCache

At inference time, the `KVCache` class in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) allocates static tensors sized to the window limit rather than the full `max_seq_len`. The `empty` classmethod creates buffers shaped by the configured window. As new tokens are generated, the oldest entries are overwritten in place. Because the attention pattern is restricted to the last *N* positions, the cache never reallocates or copies memory, keeping the footprint constant.

### Memory Footprint Comparison

According to the documentation in [`doc/finetuning.md`](https://github.com/cactus-compute/needle/blob/main/doc/finetuning.md), a standard KV cache grows proportionally to sequence length *L* and the number of KV heads *H*, consuming *L × H* entries. Needle 2’s sliding-window implementation bounds this to *N × H*, where *N* is the `sliding_window` size. For long documents, this replaces unbounded linear growth with a small, fixed buffer.

The API documentation in [`doc/apis.md`](https://github.com/cactus-compute/needle/blob/main/doc/apis.md) reinforces this by noting that old tokens are discarded after the window length, directly reducing memory usage for long sequences.

## Exporting and Loading a Windowed Model

When a checkpoint is serialized, [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) stores the training-time sliding-window width in the binary header field `kv_window`. The inference runtime reads this header to confirm that the requested `sliding_window` matches the exported model’s expectations before allocating the cache. This prevents shape mismatches between the saved weights and the live buffers.

## Practical Configuration Example

You can enable the sliding window by constructing a `TransformerConfig` with a non-zero `sliding_window` value:

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

cfg = TransformerConfig(
    vocab_size=32000,
    d_model=1024,
    num_heads=16,
    num_kv_heads=16,
    num_layers=24,
    max_seq_len=8192,
    sliding_window=1024,   # Retain only the last 1024 tokens per head

)

```

With this configuration, the `KVCache` in [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) will allocate tensors for 1024 positions per head. Even if the generated sequence stretches to 8192 tokens, the cache will continue to overwrite the oldest 1024 slots instead of expanding.

## Summary

- **Fixed-size buffer:** Needle 2 allocates the KV cache once based on `sliding_window`, avoiding dynamic expansion during generation.
- **Circular overwrite:** Old tokens are replaced in place, so memory stays constant regardless of how long the sequence becomes.
- **Config-driven toggle:** The `TransformerConfig.sliding_window` field enables the behavior, while `kv_window` persists the setting in exported checkpoints.
- **Predictable footprint:** Memory usage drops from *L × H* to *N × H*, making long-document inference feasible on limited hardware.

## Frequently Asked Questions

### What does the `sliding_window` parameter do in Needle 2?

The `sliding_window` parameter in `TransformerConfig` sets the maximum number of recent tokens kept in the KV cache for each attention head. When it is greater than zero, Needle 2 switches to a circular buffer that overwrites older entries, preventing the cache from growing with sequence length.

### How is the sliding-window size preserved across model exports?

According to the export logic in [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py), the training-time window width is stored in the checkpoint header field `kv_window`. The runtime reads this value to ensure the inference buffer matches the exported model’s expected window size.

### Why does a sliding window reduce memory instead of a full KV cache?

A full KV cache stores key-value pairs for every token in the sequence, scaling as *L × H*. Needle 2’s sliding-window cache bounds storage to `sliding_window × num_kv_heads`, keeping memory usage flat even when generating thousands of tokens.

### Can I change the window size at inference time?

You can specify a new `sliding_window` value when loading a model, but it must not exceed the `kv_window` stored in the exported header. Exceeding the exported window would require recomputing keys and values for the missing history, which the current runtime does not support.