Needle KV Budget Calculation and Effective Window Size: Managing the 11 MiB Cache Limit

The effective KV window size in Needle is derived by dividing an 11 MiB KV budget by the per-token memory cost across all transformer layers and engram sites, clamping the result between a hardware minimum and the model's maximum sequence length, with optional user override via config.kv_window.

Needle, an open-source transformer inference engine from Cactus Compute, implements a strict memory budgeting system for its key-value cache to enable efficient on-device inference. Understanding the Needle KV budget calculation effective window size is essential for optimizing model exports and preventing out-of-memory errors during inference. The architecture uses two helper functions in needle/model/architecture.py to determine exactly how many tokens can remain cached given the fixed memory constraint.

Understanding the KV Budget Constants

The KV cache calculation relies on several architectural constants defined in the codebase. KV_BUDGET_BYTES is set to approximately 11 MiB, representing the total memory pool shared across all layers, attention heads, and cache positions. This budget must accommodate both standard KV projections and specialized engram layers that store additional activation states.

The calculation also respects KV_GROUP, a grouping constant used for memory alignment, and KV_WINDOW_MIN, a hardware-dependent minimum cache size. These values ensure that the calculated window remains practical for the target deployment hardware while maximizing the available context length.

Calculating the Budget-Derived Window with kv_budget_window

The kv_budget_window(config) function in needle/model/architecture.py (lines 603-612) computes the absolute upper bound on sequence length that fits within the memory budget. It accounts for the unique storage requirements of each layer, including grouped-query attention heads and engram storage sites.

def kv_budget_window(config):
    head_dim = (getattr(config, "attn_dim", 0) or config.d_model) // config.num_heads
    kv = config.num_kv_heads * head_dim
    d, L = config.d_model, config.num_layers
    sites = len(tuple(getattr(config, "engram_layers", (2, 15))))
    per_pos = (L * (2 * kv + 2 * (kv // KV_GROUP) * 4)
               + sites * (d + (d // KV_GROUP) * 4))
    window = (KV_BUDGET_BYTES // per_pos) // KV_GROUP * KV_GROUP
    return max(KV_WINDOW_MIN, min(window, config.max_seq_len))

The function performs four critical calculations:

1. Head dimension derivation – Calculates head_dim from either attn_dim or d_model divided by num_heads, determining the vector size per attention head.

2. KV size computation – Multiplies num_kv_heads by head_dim to get the total KV vector size, which varies for grouped-query attention architectures.

3. Per-position cost calculation – Computes per_pos by summing the storage cost across all num_layers (the L term) and adding the overhead from engram layers (the sites term), with specific arithmetic accounting for KV grouping and quantization overhead.

4. Window alignment and clamping – Divides the total budget by the per-position cost, aligns the result to KV_GROUP boundaries, and clamps between KV_WINDOW_MIN and config.max_seq_len to ensure hardware compatibility.

Respecting User Overrides with effective_kv_window

While kv_budget_window determines the theoretical maximum, effective_kv_window(config) (lines 614-617) returns the actual cache size that will be used during inference. This function respects user configuration while preventing memory overcommitment.

def effective_kv_window(config):
    budget = kv_budget_window(config)
    return min(budget, config.kv_window) if config.kv_window else budget

The logic is straightforward but critical for deployment flexibility. If config.kv_window is explicitly set, the function returns the smaller of the budget-derived value and the user request. If no override exists, it returns the full budgeted window, ensuring the model utilizes all available KV cache memory.

Integration Points in the Needle Pipeline

The effective KV window calculation propagates through multiple stages of the Needle workflow:

• Model Export – In needle/model/export.py (lines 558-560), effective_kv_window(config) determines the cache size embedded in the exported .cact file, ensuring the runtime allocator reserves the correct memory footprint.

• Finetuning – The finetuning script in needle/model/finetune.py (lines 407-430) calls the function when building checkpoints, aligning the training cache size with inference constraints to prevent deployment mismatches.

• Validation – The test suite in tests/test_build.py (lines 43-53) verifies that exported models report the expected KV window, validating that the budget calculations translate correctly to binary artifacts.

Practical Code Examples

Querying the Effective KV Window

To determine the cache capacity for a specific model configuration:

from needle.model.architecture import TransformerConfig, effective_kv_window
import pickle
import pathlib

# Load a checkpoint to extract configuration

ckpt_path = pathlib.Path("tests/fixtures/tiny_checkpoint.pkl")
with ckpt_path.open("rb") as f:
    ckpt = pickle.load(f)

config = TransformerConfig(**ckpt["config"])

# Calculate the effective window size

kv_window = effective_kv_window(config)
print(f"KV cache can hold {kv_window} tokens within the 11 MiB budget")

Overriding the Window During Export

To force a smaller context window while respecting the hardware budget:

from needle.model.export import write_export
from needle.model.tokenizer import get_tokenizer
from needle.model.architecture import effective_kv_window

# Assume params and config are loaded from a checkpoint

tokenizer = get_tokenizer(config.vocab_size)

# Cap at 256 tokens or the budget limit, whichever is smaller

custom_window = min(256, effective_kv_window(config))

write_export(
    params,
    config,
    out_path="optimized_model.cact",
    bits=4,
    tokenizer=tokenizer,
    kv_window=custom_window,
)

Summary

• KV_BUDGET_BYTES (~11 MiB) is the fixed memory pool shared across all layers and positions in Needle's transformer architecture.
• kv_budget_window calculates the theoretical maximum tokens by dividing the budget by per-position costs, including standard attention KV pairs and engram layer storage.
• effective_kv_window applies user overrides via config.kv_window while preventing allocation beyond the hardware budget.
• The calculation accounts for grouped-query attention through KV_GROUP alignment and respects minimum thresholds via KV_WINDOW_MIN.
• Both functions reside in needle/model/architecture.py and are consumed by export, finetuning, and testing pipelines to ensure consistent memory planning.

Frequently Asked Questions

What happens if I set config.kv_window larger than the budget allows?

Needle's effective_kv_window function will return the smaller of the two values, effectively clamping your request to the budget-derived maximum. According to the implementation in needle/model/architecture.py (lines 614-617), the explicit min(budget, config.kv_window) check prevents memory overcommitment that would cause runtime allocation failures.

How do engram layers affect the KV budget calculation?

Engram layers increase the per-token memory cost significantly. In kv_budget_window, the sites variable counts the configured engram layers (defaulting to indices 2 and 15), and the formula adds sites * (d + (d // KV_GROUP) * 4) to the per-position cost. This means models with more engram sites will have proportionally smaller effective context windows to keep the total cache within the 11 MiB budget.

Can I modify KV_BUDGET_BYTES for my specific hardware?

The constant is hardcoded in needle/model/architecture.py based on target device constraints. While you could fork the repository and adjust the value, doing so requires corresponding changes to the runtime memory allocator in the inference engine. For most use cases, use the kv_window configuration parameter to reduce the cache size rather than increasing the budget, as exceeding physical limits will cause out-of-memory errors during model loading.

Why does the window calculation use integer division with KV_GROUP?

The alignment to KV_GROUP boundaries (achieved via (KV_BUDGET_BYTES // per_pos) // KV_GROUP * KV_GROUP) ensures that the KV cache memory layout matches the grouped-query attention implementation. This optimization reduces memory fragmentation and improves cache locality during attention computation, particularly important for quantized models where grouped heads share key-value projections.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}