# What Is Engram Key-Value Memory in Needle 2 and How Does It Work?

> Discover Engram Key-Value Memory in Needle 2. Learn how this efficient n-gram system provides persistent tool-call recall using only 28 MiB of memory, regardless of conversation length.

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

---

**Engram key-value memory in Needle 2 is a fixed-size, hashed n-gram memory system that stores compressed representations of recent tokens, enabling persistent tool-call recall with constant 28 MiB memory usage regardless of conversation length.**

Needle 2, an open-source tool-calling language model by cactus-compute, introduces **engram key-value (KV) memory** to solve a critical problem: how to remember tool interactions across long conversations without unbounded memory growth. Unlike standard KV caches that grow linearly with sequence length, the engram KV memory uses hash-based tables with configurable n-gram indexing to provide bounded, persistent storage for critical context.

## How Engram KV Memory Works

The engram KV memory operates as a **content-addressable memory layer** that sits alongside Needle's standard attention mechanism. It transforms recent token history into fixed-size vector tables through a multi-step hashing process.

### Core Components

The architecture centers on three key elements defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py):

- **`embedding`** – Fixed-size hash tables that store vector slots
- **`engram_indices`** – Function that converts token n-grams to table indices
- **`key_proj` and `value_proj`** – Lightweight dense layers that transform table lookups into attention-compatible key and value vectors

For each incoming token, the system computes indices based on n-gram context, retrieves corresponding vectors from the tables, and applies a **tap convolution** (`taps` + `_shift_right`) that blends neighboring positions for spatial coherence.

### Memory Bounds and Configuration

The engram memory remains strictly bounded through three hyperparameters:

| Parameter | Default | Purpose |
|-----------|---------|---------|
| `engram_slots` | 8192 | Hash table size per head |
| `engram_orders` | (2, 3, 4) | N-gram lengths used for indexing |
| `engram_heads` | 2 | Parallel tables per n-gram order |

This configuration yields approximately **28 MiB total memory**, constant for any conversation length up to the 256-token sliding window.

## Implementation: From Indices to KV Tensors

### Step 1: N-Gram Hashing with `_engram_indices`

The indexing process transforms raw tokens into memory addresses. The `_engram_indices` function computes table positions based on:

```python

# From needle/model/architecture.py - engram_indices logic

# (Simplified conceptual view; actual implementation uses JAX)

orders = (2, 3, 4)      # bigrams, trigrams, 4-grams

heads = 4               # parallel hash functions per order

slots = 8192            # table size (must be power of 2)

```

Each n-gram order produces `heads` independent indices, creating multiple "views" of the same context to reduce collision impact.

### Step 2: Table Lookup and Projection in `Engram.__call__`

The `Engram` class executes the memory retrieval:

```python

# Conceptual flow from needle/model/architecture.py lines 81-106

tables[jnp.arange(self.num_tables), indices]  # Gather from hash tables

k = key_proj(lookup)    # Project to key dimension

v = value_proj(lookup)  # Project to value dimension

v = tap_convolution(v)  # Blend neighboring positions

```

The **tap convolution** applies a small learned filter across spatial positions, allowing the memory to capture local structure within the retrieved values.

### Step 3: Block-Level Integration

Each `Block` optionally receives `engram_kv` as a tuple `(k, v)`. Inside `Block.__call__` (lines 18-27 in [`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py)), the integration follows:

```python

# From Block.__call__ - attention-style gating with engram memory

alpha = compute_gate(x, k)          # Attention scores

x = x + alpha * v                   # Gated residual addition

# ... followed by standard self-attention and Hadamard-MLP

```

This placement—**before** regular self-attention—allows the engram memory to seed the hidden state with retrieved context, which standard attention then refines.

## Using Engram KV Memory in Practice

### Default Configuration (Enabled Automatically)

```python
import needle

# Engram KV memory is active by default for tool-capable agents

agent = needle.Needle(tools=[weather_tool, search_tool])
response = agent.run("Summarize the last three tool calls.")
print(response["results"])

```

The default settings optimize for general tool-calling scenarios with balanced memory usage and recall accuracy.

### Custom Configuration for Specific Workloads

```python
from needle import Needle, config

# Tuned for high-frequency tool use with longer n-gram dependencies

custom_cfg = config.Config(
    engram_layers=(2, 8, 14),    # Inject at layers 2, 8, and 14

    engram_slots=16384,          # Double table size for lower collision rate

    engram_orders=(2, 3, 4),     # Keep standard n-gram range

    engram_heads=4,              # More parallel views per n-gram

)

agent = Needle(tools=[...], config=custom_cfg)

```

Increasing `engram_slots` reduces hash collisions at the cost of memory. Increasing `engram_heads` provides more robust retrieval through redundancy. The `engram_layers` tuple controls **where** in the network the memory is injected—earlier layers affect representation learning, later layers influence output generation directly.

### Low-Level Manual Access

For research or debugging, extract raw KV tensors:

```python
import jax.numpy as jnp
from needle.model.architecture import engram_indices, Engram

# Compute indices from token sequence

indices = engram_indices(
    tokens,
    orders=(2, 3, 4),
    heads=4,
    slots=8192
)

# Initialize engram layer

engram = Engram(
    d_model=512,
    num_tables=12,           # 3 orders × 4 heads

    slots=8192,
    sub_dim=64,              # ENGRAM_SUB_DIM

    num_layers=16,
    conv_dilation=1,
)

# Retrieve KV pair with optional quantization bypass

k, v = engram(
    indices,
    ngram_ok=jnp.ones(indices.shape[:-1]),
    tap_ok=jnp.ones((3,)),   # ENGRAM_CONV_TAPS

    quant=False
)

# Use in custom forward pass: engram_kv=(k, v)

```

The `ngram_ok` and `tap_ok` masks enable selective disabling of memory components for ablation studies or partial updates.

## Why Engram KV Memory Matters for Tool-Calling Agents

Traditional language models face a **memory-recall dilemma**: standard KV caches grow with conversation length, eventually exhausting GPU memory, while retrieval-augmented generation requires external databases and latency-heavy search. The engram KV memory offers a third path:

- **Persistence without growth**: Tool call signatures and results remain accessible in fixed-size tables
- **Locality-aware retrieval**: N-gram hashing naturally clusters semantically related token sequences
- **Inference-time efficiency**: No database round-trips; memory access is a single gather operation

According to the Needle 2 source code, this design specifically targets "tool-related tokens" that must survive across long interactions where standard context windows would otherwise overwrite or compress away critical information.

## Source Code Reference Map

| File | Key Components | Lines |
|------|---------------|-------|
| [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) | `Engram` class, `engram_indices`, `Block.__call__` integration | 18-27, 81-106 |
| [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) | `_engram_kv` helper for pipeline wiring | Various |
| [`tests/conftest.py`](https://github.com/cactus-compute/needle/blob/main/tests/conftest.py) | Example hyperparameter configurations | Various |
| [`README.md`](https://github.com/cactus-compute/needle/blob/main/README.md) | High-level architecture description | 13 |

## Summary

- **Engram KV memory** is Needle 2's fixed-size, hashed n-gram storage system for persistent tool-call context
- **Memory remains constant at ~28 MiB** regardless of conversation length through bounded hash tables
- **Three-stage pipeline**: n-gram indexing (`engram_indices`), table lookup with projection (`Engram.__call__`), and attention-style gating (`Block.__call__`)
- **Configurable via** `engram_layers`, `engram_slots`, `engram_orders`, and `engram_heads` in `needle.config.Config`
- **Designed for tool-calling agents** that must recall prior tool interactions without unbounded memory growth

## Frequently Asked Questions

### What makes engram KV memory different from standard KV cache?

Standard KV caches store every token's key and value vectors in a linear buffer, growing with sequence length. Engram KV memory compresses token history into fixed-size hash tables using n-gram indexing, keeping memory constant. The trade-off is approximate retrieval—some information may collide or be overwritten—but the design prioritizes retaining tool-relevant patterns over exact token reconstruction.

### How does the n-gram hashing handle collisions?

The implementation uses multiple strategies: (1) **parallel heads** (`engram_heads`) generate independent hashes for the same n-gram, storing redundant copies; (2) **multiple orders** (`engram_orders`) ensure that longer contexts have multiple access paths; and (3) the **tap convolution** blends spatial neighbors, providing some robustness to individual slot corruption. For critical applications, increasing `engram_slots` directly reduces collision probability.

### Can engram KV memory be disabled or used without tools?

Yes—set `engram_layers=()` in the config to disable entirely, or pass `engram_kv=None` to individual `Block` instances. Without tools, the memory still functions but provides less benefit since standard attention handles short-range dependencies adequately. The engram memory shows strongest advantages in **long-horizon tool interactions** where explicit recall of prior API calls matters.

### What is the performance overhead of engram KV memory?

The overhead is minimal for inference: table lookup is a single `gather` operation, projection layers are small (rank-reduced), and the tap convolution uses a fixed 3-tap filter. The primary cost is the **memory footprint** (~28 MiB by default), which is pre-allocated rather than growing. For training, the hashing is non-differentiable, so gradients flow only through the projection layers and tap convolution—engram indices remain fixed or are updated via separate mechanisms.