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`	Read a row-owner-sharded Engram table and update all HC residual streams.
`DeepseekV41NgramHash`	Build all Engram layers’ hashes from complete, uncached input sequences.

Functions

Name	Description
`_compressed_token_map`	Normalize decoded vocabulary entries using the released token-ID contract.
`_next_prime`	Find the first unused prime above `start` without an optional dependency.

API

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 | NoneDefaults to 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

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 | NoneDefaults to 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

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

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

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]

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 | NoneDefaults to 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],

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

Restore derived integer buffers after meta-device materialization.

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.

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.