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

Data

_ENGRAM_GATE_COMPILED

_engram_gate

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

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]

strict_validation

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

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

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

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.

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

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

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