> For clean Markdown of any page, append .md to the page URL.
> For a complete documentation index, see https://docs.nvidia.com/nemo/automodel/llms.txt.
> For AI client integration (Claude Code, Cursor, etc.), connect to the MCP server at https://docs.nvidia.com/nemo/automodel/_mcp/server.

# nemo_automodel.components.models.deepseek_v41.engram

DeepSeek-V4.1 compressed-token Engram hashes and residual memory lookup.

The tokenizer normalization, per-layer hash seeds, and signed square-root gate
follow DeepSeek's released `inference/engram.py` and `inference/model.py`.
V4.1 omits the short convolution used by earlier Engram architectures.

## Module Contents

### Classes

| Name                                                                                                 | Description                                                               |
| ---------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- |
| [`DeepseekV41Engram`](#nemo_automodel-components-models-deepseek_v41-engram-DeepseekV41Engram)       | Read a row-owner-sharded Engram table and update all HC residual streams. |
| [`DeepseekV41NgramHash`](#nemo_automodel-components-models-deepseek_v41-engram-DeepseekV41NgramHash) | Build all Engram layers' hashes from complete, uncached input sequences.  |

### Functions

| Name                                                                                                   | Description                                                                |
| ------------------------------------------------------------------------------------------------------ | -------------------------------------------------------------------------- |
| [`_compressed_token_map`](#nemo_automodel-components-models-deepseek_v41-engram-_compressed_token_map) | Normalize decoded vocabulary entries using the released token-ID contract. |
| [`_next_prime`](#nemo_automodel-components-models-deepseek_v41-engram-_next_prime)                     | Find the first unused prime above `start` without an optional dependency.  |

### API

```python
class nemo_automodel.components.models.deepseek_v41.engram.DeepseekV41Engram(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig,
    process_group: torch.distributed.ProcessGroup | None = None
)
```

**Bases:** `Module`

Read a row-owner-sharded Engram table and update all HC residual streams.

**Parameters:**

**`config`** `DeepseekV41TextConfig`

Text configuration containing logical table sizes and HC width.

---

**`layer_idx`** `int`

Zero-based decoder layer ID, present in engram\_layer\_ids.

---

**`backend`** `BackendConfig`

Linear backend for the fused key/value projection.

---

**`process_group`** `dist.ProcessGroup | None` — default: None

Runtime row-owner group. None creates the complete table
and is appropriate only when the configuration fits on one device.

---

**`embed`**

---

**`eps`** `= config.rms_norm_eps`

---

**`hash_heads`**

---

**`hc_mult`** `= config.hc_mult`

---

**`hidden_size`** `= config.hidden_size`

---

**`initializer_range`** `= config.initializer_range`

---

**`k_weight`**

---

**`layer_hash_index`** `= tuple(config.engram_layer_ids).index(layer_idx)`

---

**`num_embeddings`**

---

**`q_weight`**

---

**`wkv`**

---

```python
nemo_automodel.components.models.deepseek_v41.engram.DeepseekV41Engram.forward(
    hidden_states: torch.Tensor,
    hash_ids: torch.Tensor,
    token_mask: torch.Tensor | None = None
) -> torch.Tensor
```

Inject the normalized, signed-square-root-gated memory residual.

**Parameters:**

**`hidden_states`** `torch.Tensor`

Tensor of shape \[batch, sequence, hc\_mult, hidden].

---

**`hash_ids`** `torch.Tensor`

Integer tensor of shape \[batch, sequence, hash\_heads]
containing logical table rows for this Engram layer. Under CP
both inputs contain only this rank's local sequence positions.

---

**`token_mask`** `torch.Tensor | None` — default: None

Optional bool tensor of shape \[batch, sequence], with
False for image/padding positions that must remain unchanged.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hc\_mult, hidden] in the input

```python
nemo_automodel.components.models.deepseek_v41.engram.DeepseekV41Engram.init_weights() -> None
```

Initialize the table, projection, and learned branch normalization weights.

```python
class nemo_automodel.components.models.deepseek_v41.engram.DeepseekV41NgramHash(
    config: nemo_automodel.components.models.deepseek_v41.config.DeepseekV41TextConfig,
    tokenizer: transformers.PreTrainedTokenizerFast
)
```

**Bases:** `Module`

Build all Engram layers' hashes from complete, uncached input sequences.

**Parameters:**

**`config`** `DeepseekV41TextConfig`

Text configuration containing the released Engram dimensions.

---

**`tokenizer`** `PreTrainedTokenizerFast`

Fast tokenizer whose normalized vocabulary must have exactly
`config.engram_compressed_vocab_size` entries.

---

**`_multiplier_values`** `= tuple(multipliers)`

---

**`_offset_values`** `= tuple(offsets)`

---

**`_prime_values`** `= tuple(primes)`

---

**`layer_ids`** `= tuple(config.engram_layer_ids)`

---

**`max_ngram_size`** `= config.engram_max_ngram_size`

---

**`n_heads`** `= config.engram_n_heads`

---

**`pad_id`** `= self._token_map_values[config.engram_pad_token_id]`

---

```python
nemo_automodel.components.models.deepseek_v41.engram.DeepseekV41NgramHash.forward(
    input_ids: torch.Tensor,
    token_mask: torch.Tensor | None = None
) -> torch.Tensor
```

Hash unpacked sequences without crossing image or padding boundaries.

**Parameters:**

**`input_ids`** `torch.Tensor`

Integer tensor of shape \[batch, sequence] containing raw
tokenizer IDs for complete, unpacked sequences.

---

**`token_mask`** `torch.Tensor | None` — default: None

Optional boolean tensor of shape \[batch, sequence].
False marks image or padding tokens and blocks all lookback
through those positions. The caller also masks their residual gate.

---

**Returns:** `torch.Tensor`

Integer tensor of shape \[batch, sequence, engram\_layers, hash\_heads],

```python
nemo_automodel.components.models.deepseek_v41.engram.DeepseekV41NgramHash.init_weights() -> None
```

Restore derived integer buffers after meta-device materialization.

```python
nemo_automodel.components.models.deepseek_v41.engram._compressed_token_map(
    tokenizer: transformers.PreTrainedTokenizerFast
) -> tuple[tuple[int, ...], int]
```

Normalize decoded vocabulary entries using the released token-ID contract.

```python
nemo_automodel.components.models.deepseek_v41.engram._next_prime(
    start: int,
    seen: set[int]
) -> int
```

Find the first unused prime above `start` without an optional dependency.