# How Engram Key-Value Memory Works in Needle 2 for Contextual Memory

> Discover how Needle 2's Engram key-value memory injects long-range context into transformers efficiently. Learn about its fast, deterministic approach to contextual memory without increased attention cost.

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

---

**Needle 2's Engram module augments transformers with a fast, deterministic key-value memory that injects long-range context without increasing attention cost.**

The Engram key-value memory system in [cactus-compute/needle](https://github.com/cactus-compute/needle) gives large language models persistent, retrievable memory across extremely long sequences. Unlike standard KV caches that grow linearly with sequence length, Engram uses fixed-size hash tables with O(1) lookups—enabling contextual memory that survives far beyond the typical attention window.

## Engram as a Deterministic KV Store

### Table Structure and Hash Indexing

At each **Engram site** (configurable transformer layers), the system maintains embedding tables with shape `(num_tables, slots, sub_dim)`. These tables serve as the raw key-value storage.

Token sequences are hashed into **slot indices** using a deterministic hash function defined by `_ENGRAM_SEED` and `_ENGRAM_PRIME`. The `engram_indices` function in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) computes these indices per n-gram order and per head:

```python

# architecture.py – Engram table lookup

tables = self.param("embedding", default_init(),
                    (self.num_tables, self.slots, self.sub_dim))
fetched = tables[jnp.arange(self.num_tables), indices]      # ← KV retrieval

```

The deterministic hashing guarantees identical token sequences always map to identical slots—critical for reliable memory retrieval.

### Lookup and Masking

Retrieved vectors are optionally filtered by an `ngram_ok` flag that masks invalid positions. This yields a dense tensor of key vectors ready for projection into the transformer's feature space.

## Projection to Transformer-Compatible KV Pairs

Raw table lookups require transformation before injection into attention layers. The Engram applies two learned projections:

- **Key projection**: `nn.Dense(d_model)` → `k`
- **Value projection**: `nn.Dense(d_model)` → `v`

Both projections support optional quantization (`_aq`) for memory-constrained deployments. This step bridges the gap between the hashed table representation and the transformer's expected feature dimensions.

## Temporal Convolution with Taps

Engram incorporates **convolution taps**—small 1-D kernels that blend nearby value vectors temporally. This creates a soft attention window over recent positions:

```python

# architecture.py – convolution over values

v = sum(taps[j] * _shift_right(v, j * self.conv_dilation) * tap_ok[j][..., None]
        for j in range(ENGRAM_CONV_TAPS))

```

The `ENGRAM_CONV_TAPS` parameter controls kernel size, while `conv_dilation` adjusts temporal reach. This lightweight mechanism provides temporal continuity without full self-attention over the memory.

## Integration into Transformer Blocks

Engram KV pairs merge into the main computation via gated addition. In `Block.__call__` (around line 318 of [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py)):

```python

# architecture.py – Block __call__

if engram_kv is not None:
    ek, ev = engram_kv
    alpha = nn.sigmoid(jnp.einsum("btd,sbtd->sbt", _rms_unit(x), _rms_unit(ek))
                       / math.sqrt(self.d_model))
    x = x + jnp.einsum("s,sbt,sbtd->btd", site_flags.astype(jnp.float32),
                       alpha, ev.astype(jnp.float32)).astype(x.dtype)

```

The gating mechanism works as follows:

1. **Similarity computation**: RMS-normalized dot product between hidden states and Engram keys
2. **Sigmoid gating**: Produces per-site attention weights `alpha`
3. **Weighted aggregation**: Engram values scaled by `alpha` and summed across sites
4. **Residual addition**: Result added to main stream

This design lets the model **selectively attend to long-range context** stored in Engram tables while maintaining standard self-attention for local patterns.

## Engram Configuration and Activation

The `Config` class in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py) controls Engram geometry:

| Field | Purpose | Typical Value |
|-------|---------|---------------|
| `engram_layers` | Layer indices hosting Engram sites | `(2, 15)` |
| `engram_slots` | Hash table size per table | `8192` |
| `engram_orders` | N-gram lengths for hashing | `(2, 3)` |
| `engram_heads` | Heads per order (0 = auto) | `0` |
| `engram_conv_taps` | Temporal convolution kernel size | `4` |
| `engram_conv_dilation` | Convolution dilation | `1` |

When `engram_layers` is non-empty, the model constructor (around line 491) instantiates `Engram` objects for each specified layer.

## Practical Usage Examples

### Enabling Engram in a Needle Model

```python
from needle import Needle, Config

cfg = Config(
    d_model=1024,
    num_layers=24,
    # Enable Engram on layers 2 and 15

    engram_layers=(2, 15),
    engram_slots=8192,
    engram_orders=(2, 3),
    engram_heads=0,               # auto-derive from d_model

    engram_conv_taps=4,
    engram_conv_dilation=1,
)

model = Needle(cfg)

```

The configuration creates two Engram sites with 8K slots each, using bigram and trigram hashing with 4-tap convolution.

### Leveraging Contextual Memory Across Turns

```python
prompt = "The capital of France is Paris. "

# Fact stored in Engram memory at layer 2

output_1 = model.generate(prompt, max_new_tokens=20)

# Extended conversation—Engram retrieves despite distance

follow_up = "What is the population of the capital?"
output_2 = model.generate(f"{prompt}{output_1}{follow_up}", max_new_tokens=20)

```

The Engram key-value memory preserves the *Paris* fact independently of intervening tokens, enabling coherent multi-turn dialogue without quadratic attention costs.

### Debugging Engram State

```python

# Inspect raw table contents after forward pass

for i, eng in enumerate(model.engrams):
    tables = eng.variables["params"]["embedding"]   # (num_tables, slots, sub_dim)

    print(f"Engram site {i} – tables shape:", tables.shape)

```

This reveals the actual stored vectors and can help diagnose retrieval failures or memory saturation.

## Key Implementation Files

- **[`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)** – Core `Engram` class, `engram_indices` hashing, and `Block` integration
- **[`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py)** – `_engram_kv` helper for batched KV retrieval during generation
- **[`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py)** – Serialization logic for Engram parameters
- **[`tests/conftest.py`](https://github.com/cactus-compute/needle/blob/main/tests/conftest.py)** – Minimal Engram configuration for testing

## Summary

- **Engram key-value memory** uses fixed-size hash tables with deterministic n-gram indexing for O(1) contextual lookups
- **Two-stage retrieval**: table lookup → projection → convolutional blending → gated integration
- **Configuration via `Config`** fields determines site placement, table capacity, and temporal smoothing
- **Deterministic hashing** ensures repeatable memory access across generation steps
- **Constant memory footprint** regardless of sequence length, unlike standard KV caches

## Frequently Asked Questions

### What makes Engram memory different from a standard KV cache?

Standard KV caches store every token's key and value vectors, growing linearly with sequence length. Engram uses **fixed-size hash tables** where tokens map to slots via deterministic hashing—giving O(1) lookup cost and bounded memory. The trade-off is potential hash collisions versus guaranteed retrieval of exact position content.

### How does the n-gram hashing work in practice?

The `engram_indices` function hashes token sequences of length N (controlled by `engram_orders`) into slot positions. For order 2, token pairs hash together; for order 3, triplets. Multiple orders run in parallel, and `engram_heads` controls the per-order parallelism. The hash combines `_ENGRAM_SEED` and `_ENGRAM_PRIME` for reproducibility across runs.

### Can I use Engram with existing Needle checkpoints?

Engram parameters are standard JAX arrays stored under `"embedding"` and projection layers. The [`needle/model/export.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/export.py) module handles serialization. Loading a checkpoint without Engram config into a model with Engram enabled requires re-initialization of Engram weights—fine-tuning recommended to learn meaningful table contents.

### When should I enable multiple Engram layers?

Place Engram sites **after early feature extraction** (layer 2+) and **before final decoding layers** (layer 15+ in 24-layer models). Early sites capture lexical patterns; later sites capture semantic abstractions. The `engram_layers` tuple accepts arbitrary indices, though each site adds parameters proportional to `slots × sub_dim × num_tables`.