# How Needle 2 Manages the KV Window and Memory Budget

> Discover how Needle 2 manages its KV window and memory budget. Learn about the ~12 MiB cache limit and deterministic group alignment for guaranteed token minimums.

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

---

**Needle 2 imposes a hard ~12 MiB memory ceiling on its key-value cache and derives the maximum sequence length (KV window) through deterministic group alignment, guaranteeing a minimum of 160 tokens regardless of model size.**

Needle 2, developed by [cactus-compute/needle](https://github.com/cactus-compute/needle), implements a rigorous **KV window and memory budget** management system that prioritizes predictable memory usage over dynamic allocation. Rather than allowing the KV cache to grow unbounded with sequence length, the framework pre-calculates the maximum window that fits within a fixed RAM reservation, adjusting for model-specific dimensions and quantization settings.

## Fixed Memory Ceiling: KV_BUDGET_BYTES

At the core of Needle 2's memory management lies the `KV_BUDGET_BYTES` constant defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). This value is hard-coded to exactly 12,064,768 bytes (calculated as `11 * 1024 * 1024 + 512 * 1024`), representing approximately 12 MiB of RAM reserved exclusively for the KV cache [source line 598].

This fixed budget acts as an absolute invariant. Regardless of model complexity or quantization precision, the cache allocator will not exceed this threshold, ensuring that Needle 2 operates within strict memory constraints suitable for edge deployment.

## Group Alignment and Minimum Windows

To optimize memory access patterns and simplify quantization logic, Needle 2 enforces specific alignment constraints and minimum guarantees.

### KV_GROUP Alignment

The cache organizes entries in blocks of **32**, defined by the `KV_GROUP = 32` constant [source line 600]. All KV cache allocations must be multiples of this group size. This alignment simplifies SIMD operations and ensures that quantized representations maintain consistent memory boundaries.

### KV_WINDOW_MIN Floor

Even on extremely small models where the per-position cost is minimal, Needle 2 guarantees a functional sequence length. The `KV_WINDOW_MIN = 160` constant [source line 599] establishes a hard lower bound of 160 tokens for the effective window, preventing degenerate cases where aggressive quantization might otherwise yield impractically small contexts.

## Calculating the Effective KV Window

The `effective_kv_window(config)` function in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) implements the deterministic calculation that maps model parameters to the actual sequence window. This computation follows a three-step process:

1. **Estimate per-position cost** based on hidden size, head count, and group dimensions
2. **Apply group alignment** using integer arithmetic to determine how many complete 32-entry blocks fit within the budget  
3. **Clamp to constraints** ensuring the result respects both the 160-token minimum and the model's absolute sequence limit

The core logic appears at [source lines 610-611]:

```python
window = (KV_BUDGET_BYTES // per_pos) // KV_GROUP * KV_GROUP
return max(KV_WINDOW_MIN, min(window, config.max_seq_len))

```

This formulation guarantees that:
- The result is always a multiple of 32 (`KV_GROUP`)
- The value never exceeds `config.max_seq_len`
- The value never falls below `KV_WINDOW_MIN` (160 tokens)
- The total memory consumption remains strictly below `KV_BUDGET_BYTES`

## Quantization Impact on Memory Budget

While the **KV window and memory budget** calculation maintains a fixed RAM ceiling, quantization settings influence how many tokens fit within that space. The `KV_BITS` configuration (default 0, indicating full precision) and the `set_quant_bits()` function in [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py) [source lines 52-65] allow users to reduce the per-position memory footprint.

Lowering `KV_BITS` decreases the `per_pos` value in the window calculation, effectively increasing the number of tokens that fit within the ~12 MiB budget. However, the system always respects the `KV_GROUP` alignment and `KV_WINDOW_MIN` constraints regardless of quantization level.

## Source Code Architecture

The implementation spans four critical files within the cactus-compute/needle repository:

- **[`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)**: Defines `KV_BUDGET_BYTES`, `KV_GROUP`, `KV_WINDOW_MIN`, and the `effective_kv_window()` calculation logic
- **[`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py)**: Manages `KV_BITS` configuration and group size adjustments via `set_quant_bits()`
- **[`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py)**: Consumes the KV window when packing model parameters for deployment
- **[`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py)**: Utilizes the calculated window during inference to manage cache eviction and access patterns

Together, these components ensure that the KV cache remains within its allocated budget throughout the model lifecycle.

## Practical Implementation Examples

The following examples demonstrate how to inspect and modify **KV window and memory budget** behavior in Needle 2.

### Retrieving the Current KV Window

To calculate the effective window for a specific model configuration:

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

# cfg is a ModelConfig instance loaded from a checkpoint

kv_window = effective_kv_window(cfg)
print(f"Effective KV window = {kv_window} tokens")

```

### Adjusting Quantization to Expand the Window

Reducing KV precision increases the token capacity within the fixed budget:

```python
from needle.model.quantize import set_quant_bits, _KV_GROUP
from needle.model.architecture import effective_kv_window

# Enable 8-bit KV cache entries

set_quant_bits(act_bits=8, kv_bits=8, kv_group=_KV_GROUP)

# Re-calculate with reduced per-position cost

new_window = effective_kv_window(cfg)
print(f"KV window after quantization = {new_window} tokens")

```

### Customizing the Memory Budget

For deployments requiring stricter constraints, override the budget constant (illustrative only):

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

# Reduce budget to 6 MiB

arch.KV_BUDGET_BYTES = 6 * 1024 * 1024

kv_window = arch.effective_kv_window(cfg)
print(f"KV window with 6 MiB budget = {kv_window}")

```

## Summary

- **Fixed Budget**: Needle 2 reserves exactly ~12 MiB (`KV_BUDGET_BYTES`) for the KV cache, regardless of model size or configuration
- **Deterministic Window**: The `effective_kv_window()` function calculates sequence capacity through integer arithmetic that respects group alignment and minimum thresholds
- **32-Entry Alignment**: All cache allocations align to `KV_GROUP = 32` entries to optimize memory access and quantization operations
- **160-Token Minimum**: `KV_WINDOW_MIN` guarantees functional context lengths even on heavily quantized small models
- **Quantization Flexibility**: Adjusting `KV_BITS` via `set_quant_bits()` modifies per-position memory costs, indirectly expanding or contracting the effective window within the fixed budget

## Frequently Asked Questions

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

The default budget is hard-coded to 12,064,768 bytes (approximately 12 MiB) via the `KV_BUDGET_BYTES` constant in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) [source line 598]. This value represents the absolute maximum RAM that Needle 2 will allocate for key-value caching across all model configurations.

### How does quantization affect the KV window size?

Quantization reduces the per-position memory cost (`per_pos`) used in the window calculation. When you call `set_quant_bits()` with lower `kv_bits` values in [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py), the effective number of tokens that fit within the fixed ~12 MiB budget increases proportionally, though the result always remains aligned to 32-entry groups and respects the 160-token minimum.

### Why does Needle 2 align KV cache entries to groups of 32?

The `KV_GROUP = 32` alignment [source line 600] optimizes memory access patterns for SIMD operations and simplifies the implementation of quantized cache formats. This constraint ensures that the cache size is always a multiple of 32 entries, reducing fragmentation and improving computational efficiency during attention operations.

### Where is the KV window calculation implemented in the source code?

The primary calculation logic resides in the `effective_kv_window()` function within [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) [source lines 610-611]. This function consumes model configuration parameters and the global constants `KV_BUDGET_BYTES`, `KV_GROUP`, and `KV_WINDOW_MIN` to derive the final sequence window used by [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) during inference.