# Engram Key-Value Memory in Needle 2: Architecture and Implementation

> Explore Engram key-value memory in Needle 2. Discover this fixed-size hashed n-gram system for bounded memory usage and retained context via attention-style gating. Learn about its architecture and implementation.

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

---

**Engram key-value memory in Needle 2 is a fixed-size hashed n-gram storage system that provides bounded memory usage (~28 MiB) while allowing the model to retain context from recent tokens through attention-style gating.**

The **engram key-value (KV) memory** mechanism augments Needle 2's transformer architecture with a persistent, sliding-window memory sink. This design enables the model to remember tool calls and recent context across long interactions without linearly increasing RAM usage, making it particularly effective for agentic workflows with extensive tool use.

## What Is Engram Key-Value Memory?

Engram key-value memory consists of fixed-size embedding tables that store vector representations of recent token sequences. Rather than caching every past token, Needle 2 computes **hashed n-gram indices** from the incoming token stream, mapping them to specific slots in compact tables.

The system uses configurable **n-gram orders** (`engram_orders`), **attention heads** (`engram_heads`), and **slot counts** (`engram_slots`) to determine memory capacity. For every incoming token, the `_engram_indices` function derives indices based on the token's n-gram context, addressing tables that hold keys (`k`) and values (`v`). A small convolution operation ("tap") further refines the values by mixing neighboring positions, creating a bounded yet contextually rich memory buffer.

According to the source code in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 81-106), this architecture maintains approximately **28 MiB** of memory regardless of conversation length, using a sliding window limited to 256 tokens.

## How Needle 2 Implements Engram KV Memory

### Index Computation and Hashing

The engram memory lifecycle begins with index generation. The `engram_indices` function (and its internal variant `_engram_indices`) transforms token lists into table indices using configurable geometry:

- **N-gram orders** define the span of context considered (e.g., bigrams, trigrams)
- **Heads per order** allow parallel memory banks
- **Slot count** determines hash table size and collision resistance

These indices address fixed-size tables where vector representations are stored.

### Table Lookup and Value Refinement

Inside the `Engram` class defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), the `__call__` method performs the actual memory retrieval:

1. **Table indexing**: The method accesses specific slots using `tables[jnp.arange(self.num_tables), indices]`
2. **Projection**: Lightweight dense layers (`key_proj` and `value_proj`) transform raw embeddings into key and value vectors
3. **Convolutional mixing**: The `taps` parameter applies a shift-convolution (`_shift_right`) that blends neighboring positions, enriching the value representations with local context

This produces a key-value tuple that carries compressed historical information from the recent token stream.

### Integration with Transformer Blocks

Each `Block` in Needle 2's architecture optionally accepts an `engram_kv` tuple. When present, the block computes an **attention-style gating** mechanism (variable `alpha`) between the current hidden state and the retrieved engram values.

As implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 18-27), this integration occurs before standard self-attention and the Hadamard-MLP steps:

```python

# Conceptual flow inside Block.__call__

alpha = compute_attention_gate(hidden_state, engram_values)
hidden_state = hidden_state + alpha * engram_contribution

# Continue with standard self-attention...

```

This gating allows the model to dynamically weight the importance of historical context against current input signals.

## Configuration and Usage Examples

### Default Configuration

By default, Needle 2 initializes with pre-configured engram KV memory suitable for general tool-using agents:

```python
import needle

# Create agent with built-in engram KV memory (enabled by default)

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

```

The default configuration automatically handles n-gram hashing, table management, and memory windowing without manual intervention.

### Custom Memory Parameters

For specialized workloads requiring larger context windows or different architectural patterns, customize the engram geometry through `needle.config.Config`:

```python
from needle import Needle, config

custom_cfg = config.Config(
    engram_layers=(2, 8, 14),     # Transformer layers receiving engram KV

    engram_slots=16384,           # Hash table size per layer

    engram_orders=(2, 3, 4),      # N-gram context lengths (bigram to 4-gram)

    engram_heads=4,               # Parallel attention heads per order

)

agent = Needle(tools=[...], config=custom_cfg)
result = agent.run("What was the last tool you called?")

```

Increasing `engram_slots` reduces hash collisions at the cost of memory footprint, while adjusting `engram_orders` changes the granularity of historical context captured.

### Low-Level Memory Access

For research or debugging, manually fetch engram KV pairs using the internal API:

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

# Get configuration from initialized agent

cfg = agent.config

# Compute indices from token array

indices = engram_indices(
    tokens, 
    cfg.engram_orders, 
    cfg.engram_heads, 
    cfg.engram_slots
)

# Initialize Engram layer

engram = Engram(
    d_model=cfg.d_model,
    num_tables=len(cfg.engram_orders) * cfg.engram_heads,
    slots=cfg.engram_slots,
    sub_dim=cfg.engram_sub_dim,
    num_layers=cfg.num_layers,
    conv_dilation=1,
)

# Retrieve KV pair

k, v = engram(
    indices, 
    ngram_ok=jnp.ones(indices.shape[:-1]), 
    tap_ok=jnp.ones((cfg.engram_conv_taps,)), 
    quant=False
)

# k and v can be injected into specific Block instances as engram_kv=(k, v)

```

This low-level access is defined in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), while [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) contains the helper `_engram_kv` that automates this wiring during standard forward passes.

## Summary

- **Engram key-value memory** provides a fixed-size (~28 MiB) persistent storage mechanism using hashed n-gram representations.
- The system uses **`engram_indices`** to map tokens to table slots and the **`Engram`** class to project stored values through convolutional taps.
- **Transformer blocks** integrate engram memory via attention-style gating (`alpha`) before standard self-attention.
- Configuration through **`needle.config.Config`** allows adjustment of table size (`engram_slots`), n-gram orders, and injection layers.
- The bounded memory design enables long-context tool use without linear RAM growth.

## Frequently Asked Questions

### What is the memory footprint of engram KV memory in Needle 2?

The engram key-value memory maintains a constant footprint of approximately **28 MiB** regardless of conversation length. This bounded usage stems from fixed-size hash tables and a sliding window limited to 256 recent tokens, making it suitable for long-running agent interactions.

### How does engram memory differ from standard transformer KV caching?

Unlike standard KV caches that grow linearly with sequence length, engram memory uses **hashed n-gram indexing** into fixed-size tables. While traditional caches store every token's key and value vectors, engrams compress history through n-gram hashing and convolutional mixing, providing constant memory usage at the cost of potential hash collisions.

### Which transformer layers receive engram KV memory?

By default, engram memory integrates at specific layers defined by `engram_layers` in the configuration (commonly layers 2, 8, and 14 in default setups). Each `Block` checks for the optional `engram_kv` tuple and applies the gating mechanism before proceeding with standard self-attention and MLP computations.

### Can I disable engram memory for specific inference tasks?

Yes. Since the `Block.__call__` method treats `engram_kv` as an optional parameter, you can run inference without engram contributions by omitting the tuple. However, the default `Needle` agent initializes with engram support enabled, so you would need to manually construct model blocks or modify the configuration to exclude engram layers entirely.