# What Is the Purpose of Engram Key-Value Memory in Needle 2?

> Discover the purpose of engram key-value memory in Needle 2. This feature caches token representations, enabling retrieval of long-range dependencies beyond the attention window.

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

---

**Engram key-value memory in Needle 2 provides persistent per-layer associative storage that caches token representations as key-value pairs, allowing the model to retrieve long-range dependencies beyond the fixed attention window.**

Needle 2 extends the standard transformer architecture with a differentiable external memory mechanism called Engram. Implemented as lightweight embedding tables attached to specific transformer layers, this key-value memory system enables the model to attend to information from arbitrarily long contexts without increasing the attention window size.

## How Engram Geometry Structures the Memory

The memory layout is determined by helper functions that map model configuration to concrete tensor dimensions. The `engram_geometry` function computes the number of convolution orders, heads, and sub-dimension for the memory slots based on the model configuration. This geometry calculation is defined in [[`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) at lines 120–122](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L120-L122).

To map token positions to specific memory locations, the `engram_indices` function converts sequence positions into slot indices using the precomputed geometry. This indexing logic appears at [line 146](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L146) of the same file.

## The Engram Class: Storage, Projections, and Convolution Taps

Each Engram instance is a self-contained module created for layers specified in `config.engram_layers` (defaulting to layers 2 and 15). The class definition begins at [line 181](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L181) and initializes:

- **Embedding tables**: Dense matrices that store the raw key-value memory contents.
- **Projection layers**: `key_proj` and `value_proj` matrices that map between the model’s hidden dimension and the Engram sub-dimension.
- **Convolution taps**: Optional temporal convolutions specified by `config.engram_orders` that enrich stored values with local context.

These components are initialized in the `Engram.__init__` method around [lines 181–185](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L181-L185).

## Integration into the Forward Pass

During inference, the model accepts an `engram_kv` boolean parameter. When set to `True`, the forward pass queries the Engram modules before processing each layer. The retrieval mechanism, located at [lines 509–517](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L509-L517) of [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py), performs the following steps:

1. Computes indices for the current tokens using `engram_indices`.
2. Looks up the stored key-value pairs from the embedding tables.
3. Projects the retrieved values back to the model dimension.
4. Injects the retrieved context into the token representations via residual connections.

This process allows the model to "remember" information from previous tokens that fall outside the local attention window.

## Configuring Engram Memory

You enable and size the Engram memory through the `Config` object. The following example creates an architecture with Engram layers at positions 2 and 15, allocating 8192 memory slots per layer:

```python
from needle.model import Architecture, Config

cfg = Config(
    d_model=1024,
    engram_layers=(2, 15),   # Layers receiving Engram modules

    engram_slots=8192,       # Total memory slots per layer

    engram_orders=(2, 3),    # Convolution orders for temporal taps

    engram_heads=0,          # Auto-derived from d_model when 0

)

model = Architecture(cfg)   # Engram modules instantiated automatically

```

To activate memory retrieval during generation, pass `engram_kv=True` to the model call:

```python
tokens = tokenizer.encode("The quick brown fox")
output, _ = model(tokens, engram_kv=True)   # Retrieves from Engram cache

```

You can inspect the memory tables directly for debugging:

```python
for i, eng in enumerate(model.engrams):
    layer_idx = cfg.engram_layers[i]
    print(f"Engram layer {layer_idx} – table shape: {eng.tables.shape}")

```

## Key Source Files

The Engram implementation spans several files in the `needle/model` directory:

- **[`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)**: Contains the `Engram` class, `engram_geometry`, `engram_indices`, and the forward-pass integration logic.
- **[`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py)**: Handles parameter naming conversions when quantizing Engram embedding tables and projection matrices.
- **[`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py)**: Serializes Engram tables, projection weights, and convolution taps for model checkpointing.
- **[`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py)**: Utilizes Engram geometry utilities during batched inference to compute KV look-ups efficiently.

## Summary

- Engram key-value memory provides **persistent per-layer storage** that operates alongside the transformer attention mechanism.
- The memory geometry is calculated by `engram_geometry` and indexed via `engram_indices` in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py).
- Each `Engram` instance manages **embedding tables**, **projection layers**, and optional **convolution taps** to store and retrieve context.
- The system is activated via the `engram_kv` parameter during the forward pass, enabling retrieval of long-range dependencies at layers specified in `config.engram_layers`.
- Configuration controls memory capacity through `engram_slots`, `engram_orders`, and `engram_heads`.

## Frequently Asked Questions

### How does Engram memory differ from standard transformer key-value caching?

Standard KV caching stores past key and value tensors for the attention mechanism to avoid recomputation, but it remains limited by the attention window size. Engram memory is a **learned associative store** that compresses historical context into dedicated slots, allowing the model to access information from arbitrarily long sequences without increasing the attention window or computational complexity of the attention layer itself.

### What are the performance implications of enabling `engram_kv=True`?

Engram lookups involve embedding table reads and lightweight projections, adding minimal latency compared to the attention computation. The retrieval is O(1) per token relative to the number of slots, and because the memory is sparse and layer-specific (controlled by `engram_layers`), the overhead scales with the number of Engram layers rather than sequence length.

### Can I use Engram memory with quantized models?

Yes. The [`needle/model/quantize.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/quantize.py) module specifically handles naming conversions for Engram parameters, ensuring that embedding tables and projection layers are correctly mapped during quantization-aware training or post-training quantization. The Engram tables are treated as standard embeddings for quantization purposes.

### Which layers should I select for Engram memory?

The default configuration uses layers `(2, 15)`, placing memory at both early and late stages of the network. Early layers capture low-level syntactic patterns across long distances, while deeper layers retain high-level semantic information. You can adjust `engram_layers` based on your sequence length requirements and available memory, as each listed layer instantiates its own Engram module with `engram_slots` parameters.