nemo_automodel.components.models.qwen3_8_flash_next.engram

Qwen3.8-Flash-Next raw-token Engram N-gram Embedding lookup (ple in the checkpoint config).

Module Contents

Classes

Name	Description
`Qwen3_8_FlashNextEngramTableConfig`	Declarative shape and initialization settings for the PLE embedding table.
`Qwen3_8_FlashNextNGramEmbedding`	Hash raw token IDs into the packed Qwen3.8-Flash-Next PLE table.
`Qwen3_8_FlashNextOwnerShardedEmbedding`	Trainable contiguous row-owner embedding with bidirectional All-to-All.
`Qwen3_8_FlashNextPLELayer`	Contextualize Qwen3.8-Flash-Next n-gram values and return an HC-sized delta.
`_FixedCapacityAllToAll`	Autograd-aware equal-split All-to-All for compact routed values.

Functions

Name	Description
`_fixed_capacity_all_to_all`	Exchange compact rank segments through an equal-split All-to-All.

Data

QWEN3_8_FLASH_NEXT_LAYER_MULTIPLIERS

QWEN3_8_FLASH_NEXT_NGRAM_HEAD_OFFSETS

QWEN3_8_FLASH_NEXT_NGRAM_HEAD_VOCAB_SIZES

QWEN3_8_FLASH_NEXT_NGRAM_PADDED_ROWS

API

class nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextEngramTableConfig(
    num_embeddings: int,
    embedding_dim: int,
    initializer_range: float = 0.02
)

Dataclass

Declarative shape and initialization settings for the PLE embedding table.

Parameters:

num_embeddings

int

Globally padded number of table rows. It must be divisible by the owner process-group size.

embedding_dim

int

Number of values stored in each row.

initializer_range

floatDefaults to 0.02

Standard deviation for checkpoint-free normal initialization.

embedding_dim

int

initializer_range

float = 0.02

num_embeddings

int

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextEngramTableConfig.__post_init__() -> None

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextEngramTableConfig.build(
    process_group: torch.distributed.ProcessGroup | None,
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None
) -> nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding

Build a local or row-owner-sharded embedding table.

Parameters:

process_group

dist.ProcessGroup | None

Runtime owner group. Explicitly pass None only for a single-rank reference table containing all global rows. Omitting this argument is intentionally an error, preventing a full 102.4 GB table from being allocated accidentally.

device

torch.device | str | NoneDefaults to None

Device on which the rank-local weight of shape [local_rows, embedding_dim] is allocated.

dtype

torch.dtype | NoneDefaults to None

Data type of the rank-local weight tensor.

Returns: Qwen3_8_FlashNextOwnerShardedEmbedding

An embedding whose input has shape [...] and whose output has

class nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextNGramEmbedding(
    ngram_embedding: torch.nn.Module,
    ngram_size: int = 3,
    heads_per_ngram: int = 8,
    eos_token_id: int = 248044,
    layer_multipliers: tuple[int, ...] = QWEN3_8_FLASH_NEXT_LAYER_MU...,
    ngram_heads_vocab_sizes: tuple[int, ...] = QWEN3_8_FLASH_NEXT_NGRAM_HE...,
    ngram_heads_offsets: tuple[int, ...] = QWEN3_8_FLASH_NEXT_NGRAM_HE...
)

Bases: Module

Hash raw token IDs into the packed Qwen3.8-Flash-Next PLE table.

The first heads_per_ngram heads hash bigrams, the next group hashes trigrams, and so on. Previous-token context resets after an EOS token. Hashing intentionally uses signed int64 overflow, positive remainders, and the checkpoint-provided global head offsets. It does not canonicalize or compress tokenizer IDs as the original DeepSeek Engram implementation does.

Parameters:

ngram_embedding

nn.Module

Lookup module accepting global row IDs of shape [batch, sequence, ngram_heads] and returning values of shape [batch, sequence, ngram_heads, head_dim].

ngram_size

intDefaults to 3

Largest n-gram order, including the current token.

heads_per_ngram

intDefaults to 8

Number of hash heads for every order from two through ngram_size.

eos_token_id

intDefaults to 248044

Raw tokenizer ID that terminates the preceding segment.

layer_multipliers

tuple[int, ...]Defaults to QWEN3_8_FLASH_NEXT_LAYER_MULTIPLIERS

Signed int64 multipliers of shape [ngram_size].

ngram_heads_vocab_sizes

tuple[int, ...]Defaults to QWEN3_8_FLASH_NEXT_NGRAM_HEAD_VOCAB_SIZES

Prime modulus for each hash head, shape [(ngram_size - 1) * heads_per_ngram].

ngram_heads_offsets

tuple[int, ...]Defaults to QWEN3_8_FLASH_NEXT_NGRAM_HEAD_OFFSETS

Global packed-table row offset for each hash head, with the same shape as ngram_heads_vocab_sizes.

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextNGramEmbedding._forward_global_slice(
    global_input_ids: torch.Tensor,
    sequence_start: int,
    sequence_end: int
) -> torch.Tensor

Hash a complete raw sequence, then look up only one local slice.

Parameters:

global_input_ids

torch.Tensor

Replicated raw IDs of shape [batch, global_sequence]. Hashing the full tensor preserves the two preceding raw tokens and EOS resets at a CP boundary.

sequence_start

int

Inclusive global position of the requested shard.

sequence_end

int

Exclusive global position of the requested shard.

Returns: torch.Tensor

Local PLE values of shape “[batch, local_sequence,

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextNGramEmbedding._hash_input_ids(
    input_ids: torch.Tensor
) -> torch.Tensor

Compute packed global table rows for every n-gram head.

Parameters:

input_ids

torch.Tensor

Raw integer tokenizer IDs of shape [batch, sequence].

Returns: torch.Tensor

Global table IDs of shape [batch, sequence, ngram_heads]. Heads

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextNGramEmbedding._lookup_ngram_ids(
    ngram_ids: torch.Tensor
) -> torch.Tensor

Look up precomputed packed n-gram table rows.

Parameters:

ngram_ids

torch.Tensor

Global table row IDs of shape [batch, sequence, ngram_heads].

Returns: torch.Tensor

Concatenated head values of shape “[batch, sequence,

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextNGramEmbedding._shift_right_after_eos(
    input_ids: torch.Tensor,
    shift: int
) -> torch.Tensor

Read an earlier token without crossing an EOS boundary.

Parameters:

input_ids

torch.Tensor

Raw tokenizer IDs of shape [batch, sequence].

shift

int

Number of preceding positions to read.

Returns: torch.Tensor

Raw IDs of shape [batch, sequence]. Positions lacking valid

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextNGramEmbedding.forward(
    input_ids: torch.Tensor
) -> torch.Tensor

Return concatenated PLE table values for raw token IDs.

Parameters:

input_ids

torch.Tensor

Raw integer tokenizer IDs of shape [batch, sequence].

Returns: torch.Tensor

Tensor of shape [batch, sequence, ngram_heads * head_dim]. The

class nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding(
    config: nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextEngramTableConfig,
    process_group: torch.distributed.ProcessGroup | None,
    device: torch.device | str | None = None,
    dtype: torch.dtype | None = None
)

Bases: Module

Trainable contiguous row-owner embedding with bidirectional All-to-All.

Rank r owns rows [r * local_rows, (r + 1) * local_rows). Each request rank groups global row IDs by owner and sends them in the first All-to-All. Owners perform a local embedding lookup. An autograd-aware second All-to-All returns values to request ranks and reverses direction in backward, so only the owner accumulates and updates a row’s gradient.

Parameters:

config

Qwen3_8_FlashNextEngramTableConfig

Global table shape and initialization settings.

process_group

dist.ProcessGroup | None

Runtime owner group. None stores the complete table locally and performs no collectives.

device

torch.device | str | NoneDefaults to None

Device for the local weight of shape [num_embeddings / owner_world_size, embedding_dim].

dtype

torch.dtype | NoneDefaults to None

Data type of the local weight tensor.

embedding_dim

= config.embedding_dim

global_row_end

= self.vocab_end_index

global_row_start

= self.vocab_start_index

initializer_range

= config.initializer_range

local_row_end

= self.num_embeddings_per_rank

local_row_start

= 0

num_embeddings

= config.num_embeddings

num_embeddings_per_rank

= self.num_embeddings // self.owner_world_size

owner_rank

= 0

owner_world_size

= 1

vocab_end_index

vocab_start_index

= self.owner_rank * self.num_embeddings_per_rank

weight

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding._exchange_ids(
    sorted_global_ids: torch.Tensor,
    send_counts: torch.Tensor
) -> tuple[torch.Tensor, tuple[int, ...], tuple[int, ...], int]

Send global row IDs to their contiguous row owners.

Parameters:

sorted_global_ids

torch.Tensor

Tensor of shape [request_rows] grouped by destination owner rank.

send_counts

torch.Tensor

Tensor of shape [owner_world_size] containing the number of IDs sent to each owner.

Returns: torch.Tensor

A tuple containing owner-local received IDs of shape

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding._validate_global_ids(
    global_ids: torch.Tensor
) -> None

Validate IDs symmetrically before any variable-sized collective.

Parameters:

global_ids

torch.Tensor

Integer tensor of shape [...] containing global, packed-table row IDs.

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding._validate_received_ids(
    received_ids: torch.Tensor,
    output_split_sizes: tuple[int, ...]
) -> None

Collectively validate that routed IDs belong to this row owner.

Every rank first contributes its local bad-ID count to an AllReduce. Consequently, all ranks take the same success or failure branch even when only one owner received a misrouted ID. On failure, a compact fixed-size diagnostic from every owner is gathered before raising, so the exception identifies the failing owner rank(s) without leaving peers to enter the value-return All-to-All.

Parameters:

received_ids

torch.Tensor

Global row IDs received from request ranks, with shape [owned_requests].

output_split_sizes

tuple[int, ...]

Number of received IDs from each source owner group rank, in source-rank order.

Raises:

RuntimeError: If any owner received an ID outside its contiguous global row range.

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding._validate_sorted_send_ids(
    sorted_global_ids: torch.Tensor,
    send_counts: torch.Tensor
) -> None

Symmetrically verify compact destination segments before routing.

Parameters:

sorted_global_ids

torch.Tensor

Tensor of shape [request_rows] containing global IDs grouped by contiguous owner rank.

send_counts

torch.Tensor

Tensor of shape [owner_world_size] containing one request count per destination owner.

Raises:

RuntimeError: If any request rank’s segment contains an ID owned by a different destination rank.

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding.forward(
    global_ids: torch.Tensor
) -> torch.Tensor

Look up arbitrary global rows while keeping weights on row owners.

Parameters:

global_ids

torch.Tensor

Integer tensor of shape [...] containing global row IDs in the packed multi-head table.

Returns: torch.Tensor

Tensor of shape [..., embedding_dim] in the original request

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding.mark_sharding_contract() -> None

Stamp the model-owned contract on the current weight.

Meta materialization and dtype casting can replace the Parameter object, and custom tensor attributes do not survive that replacement, so the top-level model calls this again afterwards. The single-rank reference table (plain Parameter, no process group) needs no contract.

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding.parallelize_weight(
    fsdp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> torch.nn.Parameter

Represent the already-local owner shard as one global DTensor.

The local storage is already the final contiguous row shard, so this method uses :meth:DTensor.from_local rather than redistributing or slicing it again. It runs before FSDP records its ignored parameters; the returned parameter identity must be passed unchanged to every FSDP unit containing the table.

Parameters:

fsdp_mesh

DeviceMesh

One-dimensional FSDP shard/CP mesh. Its rank order must exactly match the PLE owner process group.

Returns: nn.Parameter

The registered global [num_embeddings, embedding_dim] DTensor

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding.reset_parameters() -> None

Initialize the rank-local weight with a finite normal distribution.

class nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextPLELayer(
    ple_embedding: nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextNGramEmbedding,
    hidden_size: int,
    hc_count: int,
    ple_embed_dim: int,
    backend: nemo_automodel.components.models.common.BackendConfig,
    dtype: torch.dtype | str,
    conv_kernel_size: int = 4,
    rms_norm_eps: float = 1e-06
)

Bases: Module

Contextualize Qwen3.8-Flash-Next n-gram values and return an HC-sized delta.

Parameters:

ple_embedding

Qwen3_8_FlashNextNGramEmbedding

Raw-token n-gram embedding whose output has shape [batch, sequence, ple_embed_dim].

hidden_size

int

Width of one HyperConnection branch.

hc_count

int

Number of persistent HyperConnection branches.

ple_embed_dim

int

Concatenated n-gram embedding width.

backend

BackendConfig

Backend configuration for the key and value projections.

dtype

torch.dtype | str

Explicit parameter dtype resolved from the model configuration.

conv_kernel_size

intDefaults to 4

Kernel width of the causal depthwise convolution.

rms_norm_eps

floatDefaults to 1e-06

Variance epsilon for branch-local Gemma RMS norms.

conv1d

hc_hidden_size

= hidden_size * hc_count

key_proj

norm_conv

norm_key

norm_query

short_conv_dilation

= ple_embedding.ngram_size

value_proj

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextPLELayer._apply_branch_norm(
    norm: nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextGroupedRMSNorm,
    hidden_states: torch.Tensor
) -> torch.Tensor

Apply a flattened grouped norm while retaining explicit HC branches.

Parameters:

norm

Qwen3_8_FlashNextGroupedRMSNorm

Grouped normalization module with a learned weight of shape [hc_count * hidden_size].

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hc_count, hidden_size].

Returns: torch.Tensor

Branch-normalized tensor of shape

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextPLELayer._causal_short_conv(
    hidden_states: torch.Tensor,
    cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None = None
) -> torch.Tensor

Apply the PLE causal depthwise convolution with left zero history.

Parameters:

hidden_states

torch.Tensor

Branch-flattened tensor of shape [batch, sequence, hc_count * hidden_size].

cp_context

Qwen3_8_FlashNextCPContext | NoneDefaults to None

Optional contiguous CP metadata. Under CP, sequence is local and the method exchanges only the preceding nine-token boundary required by the released dilation/kernel settings.

Returns: torch.Tensor

Tensor of shape [batch, sequence, hc_count * hidden_size].

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextPLELayer.forward(
    hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None = None
) -> torch.Tensor

Compute the PLE delta injected before the layer’s attention HC read.

Parameters:

hidden_states

torch.Tensor

True HyperConnection state of shape [batch, sequence, hc_count * hidden_size].

input_ids

torch.Tensor

Raw integer tokenizer IDs of shape [batch, sequence].

cp_context

Qwen3_8_FlashNextCPContext | NoneDefaults to None

Optional contiguous CP metadata. Its replicated global_input_ids and global_padding_mask fields have shape [batch, global_sequence]; hidden_states and input_ids remain local [batch, sequence, ...] tensors.

Returns: torch.Tensor

PLE delta of shape

nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextPLELayer.init_weights(
    initializer_range: float = 0.02
) -> None

Initialize PLE projections and the zero-start causal convolution.

Parameters:

initializer_range

floatDefaults to 0.02

Standard deviation for projection weights.

class nemo_automodel.components.models.qwen3_8_flash_next.engram._FixedCapacityAllToAll()

Bases: Function

Autograd-aware equal-split All-to-All for compact routed values.

nemo_automodel.components.models.qwen3_8_flash_next.engram._FixedCapacityAllToAll.backward(
    ctx: typing.Any,
    grad_output: torch.Tensor
) -> tuple[torch.Tensor, None, None, None, None]

staticmethod

Route output gradients back to the ranks that supplied the rows.

Parameters:

ctx

Any

PyTorch autograd context populated by :meth:forward.

grad_output

torch.Tensor

Tensor of shape [output_rows, ...] with the same source-rank segmentation as the forward output.

Returns: torch.Tensor

A gradient tensor of shape [input_rows, ...] followed by four

nemo_automodel.components.models.qwen3_8_flash_next.engram._FixedCapacityAllToAll.forward(
    ctx: typing.Any,
    input_tensor: torch.Tensor,
    input_split_sizes: tuple[int, ...],
    output_split_sizes: tuple[int, ...],
    capacity: int,
    process_group: torch.distributed.ProcessGroup
) -> torch.Tensor

staticmethod

Exchange compact rows through fixed-capacity peer segments.

Parameters:

ctx

Any

PyTorch autograd context.

input_tensor

torch.Tensor

Tensor of shape [input_rows, ...]. Axis 0 contains contiguous per-destination segments described by input_split_sizes; arbitrary trailing dimensions are kept.

input_split_sizes

tuple[int, ...]

Number of rows sent to every destination rank.

output_split_sizes

tuple[int, ...]

Number of rows received from every source rank.

capacity

int

Global maximum peer-segment row count.

process_group

dist.ProcessGroup

Process group whose rank order defines both split tuples.

Returns: torch.Tensor

Tensor of shape [output_rows, ...], where

nemo_automodel.components.models.qwen3_8_flash_next.engram._fixed_capacity_all_to_all(
    input_tensor: torch.Tensor,
    input_split_sizes: tuple[int, ...],
    output_split_sizes: tuple[int, ...],
    capacity: int,
    process_group: torch.distributed.ProcessGroup,
    fill_value: int | float
) -> torch.Tensor

Exchange compact rank segments through an equal-split All-to-All.

Padding every peer segment to one globally agreed capacity removes backend-specific uneven-split behavior and gives forward and backward the same symmetric exchange metadata. The compact, source-ordered result expected by the owner lookup is restored after the collective. Transport provider selection remains an independent runtime concern.

Parameters:

input_tensor

torch.Tensor

Compact tensor of shape [sum(input_split_sizes), ...] whose axis-0 segments are ordered by destination rank.

input_split_sizes

tuple[int, ...]

Number of rows sent to every destination rank.

output_split_sizes

tuple[int, ...]

Number of rows received from every source rank.

capacity

int

Globally agreed maximum of every source/destination count.

process_group

dist.ProcessGroup

Process group whose rank order defines both count tuples.

fill_value

int | float

Value used for padded and initially untouched output rows.

Returns: torch.Tensor

A compact tensor of shape [sum(output_split_sizes), ...] whose

nemo_automodel.components.models.qwen3_8_flash_next.engram.QWEN3_8_FLASH_NEXT_LAYER_MULTIPLIERS = (23703573157769, 20109073645365, 8052911324071)

nemo_automodel.components.models.qwen3_8_flash_next.engram.QWEN3_8_FLASH_NEXT_NGRAM_HEAD_OFFSETS = (0, 20000003, 40000026, 60000059, 80000106, 100000165, 120000228, 140000297, 160...

nemo_automodel.components.models.qwen3_8_flash_next.engram.QWEN3_8_FLASH_NEXT_NGRAM_HEAD_VOCAB_SIZES = (20000003, 20000023, 20000033, 20000047, 20000059, 20000063, 20000069, 20000077,...

nemo_automodel.components.models.qwen3_8_flash_next.engram.QWEN3_8_FLASH_NEXT_NGRAM_PADDED_ROWS = 320001536