> This page is for version 0.7.0 · 26.10.
> For other versions, use one of these documentation indexes:
> - Nightly (default): https://docs.nvidia.com/nemo/automodel/nightly/llms.txt
> - Latest: https://docs.nvidia.com/nemo/automodel/latest/llms.txt
> - 0.7.0 · 26.10: https://docs.nvidia.com/nemo/automodel/v0.7/llms.txt
> - 0.5.0 · 26.06: https://docs.nvidia.com/nemo/automodel/v0.5/llms.txt
> - 0.4.0 · 26.04: https://docs.nvidia.com/nemo/automodel/v0.4/llms.txt

> 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.          |
| [`_engram_gate_core`](#nemo_automodel-components-models-deepseek_v41-engram-_engram_gate_core)         | Inject the normalized, signed-square-root-gated memory value into the fp32 streams. |
| [`_next_prime`](#nemo_automodel-components-models-deepseek_v41-engram-_next_prime)                     | Find the first unused prime above `start` without an optional dependency.           |
| [`compile_engram_gate`](#nemo_automodel-components-models-deepseek_v41-engram-compile_engram_gate)     | Wrap the Engram gate core with `torch.compile` (`BackendConfig.compile_engram`).    |

### Data

[`_ENGRAM_GATE_COMPILED`](#nemo_automodel-components-models-deepseek_v41-engram-_ENGRAM_GATE_COMPILED)

[`_engram_gate`](#nemo_automodel-components-models-deepseek_v41-engram-_engram_gate)

### 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._validates_on_device(
    hash_ids: torch.Tensor
) -> bool
```

Whether hash IDs are checked with device-side asserts (CUDA, strict validation off).

The choice depends on the device and the table's flag only, never on this rank's request count: an empty
requester must take the same path, and the same collectives, as its peers.

**Parameters:**

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

Integer tensor of shape \[batch, sequence, hash_heads]; only its device matters here.

---

**Returns:** `bool`

True for the device-side path, False for the host-side, cross-rank-symmetric checks.

```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]`

---

**`strict_validation`**

---

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

Hash sequences without crossing document, image, or padding boundaries.

**Parameters:**

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

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

---

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

Optional integer document IDs \[batch, sequence]; zero marks padding.

---

**`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._engram_gate_core(
    hidden_states: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    q_weight: torch.Tensor,
    k_weight: torch.Tensor,
    eps: float,
    token_mask: torch.Tensor | None
) -> torch.Tensor
```

Inject the normalized, signed-square-root-gated memory value into the fp32 streams.

**Parameters:**

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

Tensor of shape \[batch, sequence, hc_mult, hidden] in the model dtype.

---

**`key`** `torch.Tensor`

Tensor of shape \[batch, sequence, hc_mult \* hidden], the key half of the `wkv` projection.

---

**`value`** `torch.Tensor`

Tensor of shape \[batch, sequence, hidden], the value half of the `wkv` projection.

---

**`q_weight`** `torch.Tensor`

Learned branch normalization weight of shape \[hc_mult, hidden].

---

**`k_weight`** `torch.Tensor`

Learned branch normalization weight of shape \[hc_mult, hidden].

---

**`eps`** `float`

RMS normalization epsilon.

---

**`token_mask`** `torch.Tensor | None`

Optional bool tensor of shape \[batch, sequence]; masked tokens keep their streams.

---

**Returns:** `torch.Tensor`

Tensor of shape \[batch, sequence, hc_mult, hidden] in the dtype of `hidden_states`.

```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.

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

Wrap the Engram gate core with `torch.compile` (`BackendConfig.compile_engram`).

Runs once per process; static shapes (`dynamic=False`): one specialization per Engram call shape and mask
presence. No `fullgraph`: past dynamo's recompile limit a shape runs eager instead of raising. Inductor fuses
the three reductions over the hidden dimension (stream and key RMS, stream-key dot) into one kernel that reads
the fp32 streams once, the gated residual into another, and the backward through AOTAutograd; the eager chain
is about a dozen elementwise / reduction kernels per Engram layer plus the `pow` backward's full-size clone.
Numerics are allclose to eager, not bitwise (fused reductions change the summation order).

```python
nemo_automodel.components.models.deepseek_v41.engram._ENGRAM_GATE_COMPILED = False
```

```python
nemo_automodel.components.models.deepseek_v41.engram._engram_gate = _engram_gate_core
```