# How Needle 2 Manages Its KV Cache Budget: Architecture and Implementation

> Discover how Needle 2 manages its KV cache budget with a deterministic system capping usage at 11.5 MiB by calculating per-position costs and aligning window sizes to 32-token granularity.

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

---

**Needle 2 controls memory consumption through a deterministic KV cache budgeting system that caps usage at approximately 11.5 MiB by calculating per-position memory costs and aligning window sizes to 32-token granularity.**

Needle 2, from the cactus-compute/needle repository, implements a strict KV cache budget to prevent memory overflow during transformer inference and fine-tuning. The system uses fixed constants and mathematical alignment rules defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) to compute the maximum context window that fits within the allocated memory budget. This architecture ensures that every model instantiation, regardless of specific hyperparameters, respects the same deterministic memory constraints.

## Core KV Cache Constants and Budget Parameters

The budgeting system relies on three critical constants defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py):

- **KV_BUDGET_BYTES**: Set to `11 * 1024 * 1024 + 512 * 1024` (11 MiB + 512 KiB), this defines the total memory ceiling for the KV cache including metadata overhead.
- **KV_GROUP**: Fixed at `32`, this alignment granularity ensures that all KV entries are padded to multiples of 32 tokens for hardware efficiency.
- **KV_WINDOW_MIN**: Set to `160`, this establishes the absolute minimum context window regardless of budget calculations.

These constants form the foundation of the `kv_budget_window()` function that calculates architecture-specific limits for every model configuration.

## How the KV Budget Is Calculated

The budget calculation follows a deterministic three-step process that converts memory bytes into a valid token window.

### Per-Position Memory Calculation

First, the system computes the memory required for each token position using the formula in `kv_budget_window()`:

```python
per_pos = (L * (2 * kv + 2 * (kv // KV_GROUP) * 4)
           + sites * (d + (d // KV_GROUP) * 4))

```

Where:

- `L` = number of transformer layers
- `kv` = total KV dimension (`num_kv_heads * head_dim`)
- `sites` = number of engram KV layers (default `(2, 15)`)
- `d` = model dimension (`d_model`)

This accounts for both the raw KV data and the 4-byte bookkeeping overhead required per `KV_GROUP` alignment unit.

### Window Size Derivation and Alignment

With the per-position cost established, the raw window is calculated and aligned to hardware-friendly boundaries:

```python
window = (KV_BUDGET_BYTES // per_pos) // KV_GROUP * KV_GROUP

```

The formula performs integer division to determine how many complete positions fit within `KV_BUDGET_BYTES`, then aligns the result downward to the nearest multiple of `KV_GROUP` (32).

### Enforcing Minimum and Maximum Limits

The computed window undergoes clamping to ensure operational viability:

```python
return max(KV_WINDOW_MIN, min(window, config.max_seq_len))

```

This guarantees the window never falls below `KV_WINDOW_MIN` (160 tokens) nor exceeds the model's architectural maximum sequence length defined in the configuration.

## Effective KV Window Resolution

The `effective_kv_window(config)` function in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) serves as the primary interface for retrieving the final context window. It resolves conflicts between the budget-derived limit and user-specified constraints:

```python
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 set, the function returns the stricter of the two values; otherwise, it returns the budget-derived maximum. This allows users to optionally constrain memory usage further without modifying the underlying constants.

## Integration Points in the Codebase

The computed KV window propagates through multiple subsystems to ensure consistent memory management:

- **Model Construction**: In [`tests/test_build.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_build.py), the window passes directly to the transformer constructor via `kv_window=effective_kv_window(config)`.
- **Fine-tuning Pipeline**: [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) invokes the same helper when initializing models for continued training, ensuring training batches respect inference memory constraints.
- **Export Workflows**: Model export utilities embed the effective window into metadata for runtime validation.

These integration points ensure that every instantiated Needle 2 model respects the same deterministic memory budget throughout its lifecycle.

## Practical Usage Examples

### Inspecting the Computed KV Window

To verify the effective window for a specific configuration:

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

config = TransformerConfig(
    d_model=4096,
    num_heads=32,
    num_kv_heads=8,
    num_layers=32,
    max_seq_len=8192,
)

print("Effective KV window:", effective_kv_window(config))

# Output: e.g., 4096 (fits within the 11 MiB budget)

```

### Constructing a Model with Budget Constraints

When building a model instance, pass the calculated window explicitly:

```python
from needle.model.architecture import TransformerConfig, effective_kv_window
from needle.model.run import NeedleModel

cfg = TransformerConfig(d_model=2048, num_layers=16, ...)
model = NeedleModel(cfg, kv_window=effective_kv_window(cfg))

```

This pattern appears throughout the test suite and production inference code.

### Overriding the Default Budget (Advanced)

For specialized hardware or experimental configurations, monkey-patch `KV_BUDGET_BYTES` before computing the window:

```python
import needle.model.architecture as arch

arch.KV_BUDGET_BYTES = 20 * 1024 * 1024  # 20 MiB instead of ~11.5 MiB

cfg = TransformerConfig(...)
print(effective_kv_window(cfg))  # Larger window permitted

```

This advanced technique bypasses the default 11.5 MiB ceiling without modifying source files.

## Summary

- Needle 2 enforces a **deterministic KV cache budget** of approximately 11.5 MiB through constants in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).
- The **`kv_budget_window()`** function calculates per-position memory costs including alignment overhead, then derives a compliant window size.
- **Alignment to 32-token granularity** via `KV_GROUP` ensures hardware-efficient memory layout.
- The **`effective_kv_window()`** wrapper allows users to optionally impose stricter limits via configuration while respecting the hard budget ceiling.
- Budget constraints propagate consistently across **model construction**, **fine-tuning**, and **export** workflows.

## Frequently Asked Questions

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

The default budget is defined by `KV_BUDGET_BYTES` in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) as exactly 11 MiB plus 512 KiB of overhead (approximately 11.5 MiB total). This value represents the absolute memory ceiling for key-value cache storage regardless of model architecture.

### How does KV_GROUP alignment affect the cache window?

`KV_GROUP`, set to 32, forces the calculated window size to be a multiple of 32 tokens. After computing how many positions fit within `KV_BUDGET_BYTES`, the system applies integer division and multiplication to align the result downward: `(budget // per_pos) // 32 * 32`. This padding ensures memory addresses align with hardware cache lines and SIMD widths.

### Can users override the calculated KV window?

Yes. Users may set the optional `kv_window` parameter in `TransformerConfig` to a specific value. The `effective_kv_window()` function selects the minimum of the user-specified value and the budget-derived maximum, allowing stricter constraints without risking memory overflow. Alternatively, advanced users can modify `KV_BUDGET_BYTES` directly before model initialization.

### Where is the KV cache budget enforced during model initialization?

The enforcement occurs in three critical locations: [`tests/test_build.py`](https://github.com/cactus-compute/needle/blob/main/tests/test_build.py) validates the parameter passing during construction, [`needle/model/finetune.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/finetune.py) applies the window when loading models for training, and [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) consumes the value in the core `NeedleModel` constructor. This ensures consistent memory management across inference, training, and export operations.