# How to Constrain Needle 2 KV Memory to a Fixed Size

> Constrain Needle 2 KV memory to a fixed size by setting kv_window in TransformerConfig. Learn how to limit token usage effectively.

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

---

**To constrain Needle 2's KV memory, set the `kv_window` parameter in `TransformerConfig` to your desired token limit, which the engine will enforce as the minimum between your specified value and the automatic budget calculation.**

Needle 2 implements a sliding-window KV cache to prevent unbounded memory growth during long inference sessions. You can constrain Needle 2 KV memory to a specific size using either automatic budget calculations or manual configuration overrides. This guide explains the dual mechanism—budget-driven defaults and user-specified limits—with specific implementation details from the cactus-compute/needle source code.

## Understanding KV Memory Constraints

Needle 2 determines the effective KV cache size through two mechanisms that work in tandem:

- **Budget-driven window**: Automatically calculated from the model's KV budget (~11 MiB) and architecture dimensions via `kv_budget_window` in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)
- **User-specified override**: Manual cap via the `kv_window` field in the model configuration

The system resolves the final window size using `effective_kv_window` (lines 614-616 in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)), which returns `min(budget_window, kv_window)`. This ensures you never exceed the underlying memory budget while allowing you to enforce stricter limits.

## Enforcing Limits During Decoding

The decoder applies KV constraints when constructing attention masks. In [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py), the `DecodeCfg` object (lines 26-31) receives the resolved `kv_window` value from the configuration. The causal mask construction logic (lines 81-108) then trims attention to only the most recent tokens within the configured window, effectively dropping older KV entries from the cache.

## Methods to Configure KV Window Size

You can constrain KV memory through three methods depending on your deployment scenario.

### Via TransformerConfig Constructor

The most direct approach sets `kv_window` when creating the model configuration:

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

cfg = TransformerConfig(
    d_model=512,
    num_heads=8,
    num_kv_heads=4,
    num_layers=12,
    kv_window=256,  # Cap cache at 256 tokens

)

```

### Via Command-Line Interface

The Needle CLI accepts a `--kv-window` argument that forwards the value to the underlying configuration:

```bash
needle playground --weights my_needle.cact --kv-window 512

```

### Runtime Override

You can adjust the limit dynamically after instantiating the agent:

```python
import needle

agent = needle.Needle(weights="my_needle.cact", config=cfg)
agent.cfg.kv_window = 128  # Reduce window during execution

```

## Key Source Files and Implementation Details

Understanding these files helps debug and customize memory constraints:

- **[`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)** (lines 603-616): Implements `kv_budget_window` (lines 603-612) for automatic budget calculations and `effective_kv_window` (lines 614-616) for resolving the final limit
- **[`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py)** (lines 26-31, 81-108): Enforces the window via `DecodeCfg` initialization and causal mask construction during token generation
- **[`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py)**: Persists the `kv_window` value in the header of exported `.cact` model files
- **[`needle/cli.py`](https://github.com/cactus-compute/needle/blob/main/needle/cli.py)**: Parses command-line arguments including `--kv-window` for user-friendly configuration

## Summary

- Needle 2 uses a sliding-window KV cache with dual constraints: automatic budget calculation (~11 MiB default) and manual `kv_window` limits
- The effective window equals `min(budget_window, kv_window)`, ensuring safe memory usage while respecting user preferences
- Configure constraints via `TransformerConfig(kv_window=N)`, CLI `--kv-window N`, or runtime modification of `agent.cfg.kv_window`
- The architecture enforces limits during mask construction in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) by trimming attention to recent tokens only

## Frequently Asked Questions

### What is the default KV memory budget in Needle 2?

The default KV budget is approximately 11 MiB according to the `kv_budget_window` implementation in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). This value automatically scales based on model dimensions (d_model, num_heads, num_kv_heads) to prevent out-of-memory errors on standard hardware configurations.

### Can I change the KV window after the model has loaded?

Yes, you can modify `agent.cfg.kv_window` at any time after creating the Needle instance. The new limit takes effect immediately for subsequent decode operations, though existing cached KV entries beyond the new window are discarded during the next attention mask construction phase.

### How does Needle handle sequences longer than the configured kv_window?

When input sequences exceed the `kv_window` limit, the causal mask construction logic in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) (lines 81-108) automatically trims attention to only the most recent tokens. The model implements true sliding-window attention, dropping the oldest KV entries and maintaining only the latest entries up to the configured limit.

### Where is the KV window size stored in exported models?

The `kv_window` value persists in the export header when saving models to `.cact` format via [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py). When loading a saved model, the engine reads this value from the file header and applies it to the `DecodeCfg` configuration during initialization, ensuring consistent memory constraints across sessions.