nemo_automodel.components.models.qwen3_8_flash_next.engram
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
Functions
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
Declarative shape and initialization settings for the PLE embedding table.
Parameters:
Globally padded number of table rows. It must be divisible by the owner process-group size.
Number of values stored in each row.
Standard deviation for checkpoint-free normal initialization.
Validate every lookup on the host with cross-rank
agreement (the default and the original behaviour): each lookup
costs about twenty device-to-host synchronisations and seven small
collectives, which stall the issuing thread on every Engram layer,
forward and recompute. False keeps the same invariants as
device-side asserts (torch._assert_async) on CUDA inputs and
reads the route metadata (the count matrix) back once per lookup;
a violated invariant then surfaces as a CUDA device-side assertion
instead of a symmetric Python exception. DeepSeek-V4.1 opts out
through text_config.engram_strict_validation (default False).
CPU inputs are always validated on the host. The validation path is
chosen per table (device and this flag), never per request count, so
an empty requester takes the same collectives as its peers.
Opt-in for bounded workloads. The padded exchange is sized once, by the first lookup (its per-peer count maximum x 1.25, aligned), read back to the host that one time and fixed from then on, so no later lookup synchronises; a lookup whose per-peer count exceeds that calibration aborts with a device-side assertion. Keep it off (the default, the exact-count path) for ordinary inputs — growing sequence lengths, local-batch padding, distribution shifts — which it would reject; fixed-shape benchmark batches qualify.
Build a local or row-owner-sharded embedding table.
Parameters:
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 on which the rank-local weight of shape
[local_rows, embedding_dim] is allocated.
Data type of the rank-local weight tensor.
Returns: Qwen3_8_FlashNextOwnerShardedEmbedding
An embedding whose input has shape [...] and whose output has
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:
Lookup module accepting global row IDs of shape
[batch, sequence, ngram_heads] and returning values of shape
[batch, sequence, ngram_heads, head_dim].
Largest n-gram order, including the current token.
Number of hash heads for every order from two through
ngram_size.
Raw tokenizer ID that terminates the preceding segment.
Signed int64 multipliers of shape [ngram_size].
Prime modulus for each hash head, shape
[(ngram_size - 1) * heads_per_ngram].
Global packed-table row offset for each hash head,
with the same shape as ngram_heads_vocab_sizes.
Hash a complete raw sequence, then look up only one local slice.
Parameters:
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.
Inclusive global position of the requested shard.
Exclusive global position of the requested shard.
Optional replicated packed segment IDs [batch, global_sequence].
Returns: torch.Tensor
Local PLE values of shape “[batch, local_sequence,
Compute packed global table rows for every n-gram head.
Parameters:
Raw integer tokenizer IDs of shape [batch, sequence].
Optional packed segment IDs of shape [batch, sequence];
n-gram history stays inside a segment.
Returns: torch.Tensor
Global table IDs of shape [batch, sequence, ngram_heads]. Heads
Look up precomputed packed n-gram table rows.
Parameters:
Global table row IDs of shape [batch, sequence, ngram_heads].
Returns: torch.Tensor
Concatenated head values of shape “[batch, sequence,
Read an earlier token without crossing an EOS or packed-segment boundary.
Parameters:
Raw tokenizer IDs of shape [batch, sequence].
Number of preceding positions to read.
Optional packed segment IDs with the shape of input_ids.
Returns: torch.Tensor
Raw IDs of shape [batch, sequence]. Positions lacking valid
Return concatenated PLE table values for raw token IDs.
Parameters:
Raw integer tokenizer IDs of shape [batch, sequence].
Optional packed segment IDs [batch, sequence].
Returns: torch.Tensor
Tensor of shape [batch, sequence, ngram_heads * head_dim]. The
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:
Global table shape and initialization settings.
Runtime owner group. None stores the complete table
locally and performs no collectives.
Device for the local weight of shape
[num_embeddings / owner_world_size, embedding_dim].
Data type of the local weight tensor.
Restore ownership metadata after a module device or dtype conversion.
Parameters:
Tensor conversion applied by Module.to or Module.to_empty.
Whether to convert child modules as well.
Returns: Qwen3_8_FlashNextOwnerShardedEmbedding
This module after the normal PyTorch conversion, with the current
Padded rows per peer: calibrated by the first lookup, fixed from then on.
The first lookup reads the matrix maximum synchronously (one calibration per table, like HybridEP’s
capacity mode). Every later lookup only asserts on the device that its own maximum fits. Nothing about the
padded layout depends on what ran in between, so a checkpointed forward and its recomputation size the
exchange identically; a per-peer count beyond the calibration aborts the step instead of corrupting the
exchange — that is the bounded-workload contract of capacity_mode.
Send global row IDs to their contiguous row owners.
Parameters:
Tensor of shape [request_rows] grouped by
destination owner rank.
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
Owner-sharded lookup with no device-to-host read: padded exchange sized by the calibrated capacity.
Senders pack their owner-sorted IDs with a device-side row map (only their own counts and row total), owners look rows up in the padded layout (padding masked to zero), and the value return is the symmetric equal-split All-to-All, unpacked with the same send-row map. The 64 x 64 count AllGather stays; its result is consumed on the device.
Whether lookups run without a device-to-host read (capacity mode): a per-table opt-in, never automatic.
Validate IDs symmetrically before any variable-sized collective.
Parameters:
Integer tensor of shape [...] containing global,
packed-table row IDs.
Return whether a lookup is validated with host-side, cross-rank-symmetric checks.
Parameters:
A tensor of the lookup being validated; its device decides.
Returns: bool
True under strict_validation or for CPU tensors (no synchronisation
Validate owner membership of a padded [world_size * capacity] receive buffer.
Rows equal to -1 are the senders’ padding and legitimate; _UNTOUCHED rows mean the exchange never wrote
the slot and count as bad. On CUDA this is one device-side assert; on the host path the non-padding rows are
handed to the compact diagnostics in source-rank order.
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:
Global row IDs received from request ranks, with
shape [owned_requests].
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.
Symmetrically verify compact destination segments before routing.
Parameters:
Tensor of shape [request_rows] containing
global IDs grouped by contiguous owner rank.
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.
Look up arbitrary global rows while keeping weights on row owners.
Parameters:
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
Stamp the model-owned contract on the current weight.
DTensor conversion can discard custom attributes even when PyTorch
preserves the Parameter identity. Ordinary module conversions restore
the marker through _apply. Custom parameter-only materialization
bypasses that method, so model initialization stamps it afterwards.
The single-rank reference table needs no contract.
Represent the already-local owner shard as one global DTensor.
The local storage is already the final contiguous row shard, so this
method uses 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:
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
Initialize the rank-local weight with a finite normal distribution.
Bases: Module
Contextualize Qwen3.8-Flash-Next n-gram values and return an HC-sized delta.
Parameters:
Raw-token n-gram embedding whose output has shape
[batch, sequence, ple_embed_dim].
Width of one HyperConnection branch.
Number of persistent HyperConnection branches.
Concatenated n-gram embedding width.
Backend configuration for the key and value projections.
Explicit parameter dtype resolved from the model configuration.
Kernel width of the causal depthwise convolution.
Variance epsilon for branch-local Gemma RMS norms.
Apply a flattened grouped norm while retaining explicit HC branches.
Parameters:
Grouped normalization module with a learned weight of shape
[hc_count * hidden_size].
Tensor of shape
[batch, sequence, hc_count, hidden_size].
Returns: torch.Tensor
Branch-normalized tensor of shape
Apply the PLE causal depthwise convolution with left zero history.
Parameters:
Branch-flattened tensor of shape
[batch, sequence, hc_count * hidden_size].
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.
Optional packed segment IDs of shape [batch, local_sequence] without CP, or replicated [batch, global_sequence] under CP. Each segment then starts with zero
history, as at the start of a row.
Returns: torch.Tensor
Tensor of shape [batch, sequence, hc_count * hidden_size].
Run the convolution over every packed segment as if it started a fresh row.
Each run of one segment in padded gets history zeros ahead of it and
its own F.conv1d, so every tap reads its own segment or zeros.
Parameters:
Left history followed by the local sequence, [batch, history + sequence, channels].
Optional contiguous CP metadata.
Packed segment IDs [batch, sequence], or [batch, global_sequence] under CP.
Returns: torch.Tensor
Pre-activation output [batch, sequence, channels].
Compute the PLE delta injected before the layer’s attention HC read.
Parameters:
True HyperConnection state of shape
[batch, sequence, hc_count * hidden_size].
Raw integer tokenizer IDs of shape [batch, sequence].
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.
Optional packed segment IDs [batch, sequence], or
[batch, global_sequence] under CP: one ID per packed sample, then one
each for the pack’s tail padding and the CP padding, if present. The hash
and the convolution then never read across segments; None keeps the
unpacked path.
Returns: torch.Tensor
PLE delta of shape
Initialize PLE projections and the zero-start causal convolution.
Parameters:
Standard deviation for projection weights.
Bases: Function
Autograd-aware equal-split All-to-All for compact routed values.
Route output gradients back to the ranks that supplied the rows.
Parameters:
PyTorch autograd context populated by forward.
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
Exchange compact rows through fixed-capacity peer segments.
Parameters:
PyTorch autograd context.
Tensor of shape [input_rows, ...]. Axis 0 contains
contiguous per-destination segments described by
input_split_sizes; arbitrary trailing dimensions are kept.
Number of rows sent to every destination rank.
Number of rows received from every source rank.
Global maximum peer-segment row count.
Process group whose rank order defines both split tuples.
Returns: torch.Tensor
Tensor of shape [output_rows, ...], where
Bases: Function
Equal-split All-to-All of an already padded [world_size * capacity, ...] buffer, both directions.
The capacity-mode lookup keeps the owner side in the padded layout, so the value exchange and its backward are the same symmetric collective on identically shaped buffers; nothing is read on the host.
Padded rows per peer for a measured maximum count: margin for drift, aligned, never zero.
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:
Compact tensor of shape [sum(input_split_sizes), ...]
whose axis-0 segments are ordered by destination rank.
Number of rows sent to every destination rank.
Number of rows received from every source rank.
Globally agreed maximum of every source/destination count.
Process group whose rank order defines both count tuples.
Value used for padded and initially untouched output rows.
Optional precomputed _segment_rows(input_split_sizes, capacity, …)
so the packing issues no row-map kernels of its own; output_rows likewise
for the unpacking. Built here from the host tuples when absent.
Returns: torch.Tensor
A compact tensor of shape [sum(output_split_sizes), ...] whose
Row of every compact row inside the [world_size * capacity, ...] padded exchange buffer.
Compact rows are ordered by peer rank; peer r’s segment starts at row r * capacity of the padded buffer.
Built on the device from the caller’s count tensor (the AllGather row / column: no host round trip) so the
packing is a single index_copy_ and the unpacking a single index_select instead of one copy_ launch
per rank — 2 x world_size cudaMemcpyAsync per exchange, forward and backward (1280 per step at EP64 on the
16-node profile). The id exchange builds the send-side and receive-side maps once per lookup; the value
exchange reuses them forward (swapped) and backward (swapped back).
Row of every compact row inside the [world_size * capacity, ...] padded buffer, from device-side counts only.
total (= counts.sum()) is the one host integer; it is the caller’s own row count, so no device-to-host read
is needed. Used by the capacity-mode lookup for its send side (pack and unpack).