> 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.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`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-Qwen3_8_FlashNextEngramTableConfig)         | Declarative shape and initialization settings for the PLE embedding table.                          |
| [`Qwen3_8_FlashNextNGramEmbedding`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-Qwen3_8_FlashNextNGramEmbedding)               | Hash raw token IDs into the packed Qwen3.8-Flash-Next PLE table.                                    |
| [`Qwen3_8_FlashNextOwnerShardedEmbedding`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-Qwen3_8_FlashNextOwnerShardedEmbedding) | Trainable contiguous row-owner embedding with bidirectional All-to-All.                             |
| [`Qwen3_8_FlashNextPLELayer`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-Qwen3_8_FlashNextPLELayer)                           | Contextualize Qwen3.8-Flash-Next n-gram values and return an HC-sized delta.                        |
| [`_FixedCapacityAllToAll`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-_FixedCapacityAllToAll)                                 | Autograd-aware equal-split All-to-All for compact routed values.                                    |
| [`_PaddedAllToAll`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-_PaddedAllToAll)                                               | Equal-split All-to-All of an already padded `[world_size * capacity, ...]` buffer, both directions. |

### Functions

| Name                                                                                                                   | Description                                                                                                     |
| ---------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------- |
| [`_aligned_capacity`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-_aligned_capacity)                   | Padded rows per peer for a measured maximum count: margin for drift, aligned, never zero.                       |
| [`_fixed_capacity_all_to_all`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-_fixed_capacity_all_to_all) | Exchange compact rank segments through an equal-split All-to-All.                                               |
| [`_segment_rows`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-_segment_rows)                           | Row of every compact row inside the `[world_size * capacity, ...]` padded exchange buffer.                      |
| [`_segment_rows_from_counts`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-_segment_rows_from_counts)   | Row of every compact row inside the `[world_size * capacity, ...]` padded buffer, from device-side counts only. |

### Data

[`QWEN3_8_FLASH_NEXT_LAYER_MULTIPLIERS`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-QWEN3_8_FLASH_NEXT_LAYER_MULTIPLIERS)

[`QWEN3_8_FLASH_NEXT_NGRAM_HEAD_OFFSETS`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-QWEN3_8_FLASH_NEXT_NGRAM_HEAD_OFFSETS)

[`QWEN3_8_FLASH_NEXT_NGRAM_HEAD_VOCAB_SIZES`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-QWEN3_8_FLASH_NEXT_NGRAM_HEAD_VOCAB_SIZES)

[`QWEN3_8_FLASH_NEXT_NGRAM_PADDED_ROWS`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-QWEN3_8_FLASH_NEXT_NGRAM_PADDED_ROWS)

[`_CAPACITY_ALIGN`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-_CAPACITY_ALIGN)

[`_CAPACITY_MARGIN`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-_CAPACITY_MARGIN)

[`_UNTOUCHED`](#nemo_automodel-components-models-qwen3_8_flash_next-engram-_UNTOUCHED)

### API

```python
class nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextEngramTableConfig(
    num_embeddings: int,
    embedding_dim: int,
    initializer_range: float = 0.02,
    strict_validation: bool = True,
    capacity_mode: bool = False
)
```

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`** `float` — default: 0.02

Standard deviation for checkpoint-free normal
initialization.

---

**`strict_validation`** `bool` — default: True

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.

---

**`capacity_mode`** `bool` — default: False

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.

---

**`capacity_mode`** `bool = False`

---

**`embedding_dim`** `int`

---

**`initializer_range`** `float = 0.02`

---

**`num_embeddings`** `int`

---

**`strict_validation`** `bool = True`

---

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

```python
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 | None` — default: None

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

---

**`dtype`** `torch.dtype | None` — default: None

Data type of the rank-local weight tensor.

---

**Returns:** `Qwen3_8_FlashNextOwnerShardedEmbedding`

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

```python
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`** `int` — default: 3

Largest n-gram order, including the current token.

---

**`heads_per_ngram`** `int` — default: 8

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

---

**`eos_token_id`** `int` — default: 248044

Raw tokenizer ID that terminates the preceding segment.

---

**`layer_multipliers`** `tuple[int, ...]` — default: QWEN3_8_FLASH_NEXT_LAYER_MULTIPLIERS

Signed int64 multipliers of shape `[ngram_size]`.

---

**`ngram_heads_vocab_sizes`** `tuple[int, ...]` — default: 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, ...]` — default: 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`.

---

```python
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,
    global_sequence_ids: torch.Tensor | None = None
) -> 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.

---

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

Optional replicated packed segment IDs `[batch, global_sequence]`.

---

**Returns:** `torch.Tensor`

Local PLE values of shape \`\`\[batch, local_sequence,

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextNGramEmbedding._hash_input_ids(
    input_ids: torch.Tensor,
    sequence_ids: torch.Tensor | None = None
) -> 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]`.

---

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

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

```python
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,

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextNGramEmbedding._shift_right_after_eos(
    input_ids: torch.Tensor,
    shift: int,
    sequence_ids: torch.Tensor | None = None
) -> torch.Tensor
```

Read an earlier token without crossing an EOS or packed-segment boundary.

**Parameters:**

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

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

---

**`shift`** `int`

Number of preceding positions to read.

---

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

Optional packed segment IDs with the shape of `input_ids`.

---

**Returns:** `torch.Tensor`

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

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

Return concatenated PLE table values for raw token IDs.

**Parameters:**

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

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

---

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

Optional packed segment IDs `[batch, sequence]`.

---

**Returns:** `torch.Tensor`

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

```python
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 | None` — default: None

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

---

**`dtype`** `torch.dtype | None` — default: None

Data type of the local weight tensor.

---

**`_capacity`** `int | None = None`

---

**`capacity_mode`** `= bool(config.capacity_mode)`

---

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

---

**`strict_validation`** `= config.strict_validation`

---

**`vocab_end_index`**

---

**`vocab_start_index`** `= self.owner_rank * self.num_embeddings_per_rank`

---

**`weight`**

---

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding._apply(
    fn: typing.Callable[[torch.Tensor], torch.Tensor],
    recurse: bool = True
) -> nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding
```

Restore ownership metadata after a module device or dtype conversion.

**Parameters:**

**`fn`** `Callable[[torch.Tensor], torch.Tensor]`

Tensor conversion applied by `Module.to` or `Module.to_empty`.

---

**`recurse`** `bool` — default: True

Whether to convert child modules as well.

---

**Returns:** `Qwen3_8_FlashNextOwnerShardedEmbedding`

This module after the normal PyTorch conversion, with the current

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding._capacity_for(
    count_matrix: torch.Tensor
) -> int
```

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

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

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding._lookup_capacity_mode(
    sorted_global_ids: torch.Tensor,
    send_counts: torch.Tensor,
    unsort_indices: torch.Tensor,
    original_shape: torch.Size
) -> torch.Tensor
```

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.

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding._use_capacity_mode() -> bool
```

Whether lookups run without a device-to-host read (capacity mode): a per-table opt-in, never automatic.

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

---

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding._validate_on_host(
    tensor: torch.Tensor
) -> bool
```

Return whether a lookup is validated with host-side, cross-rank-symmetric checks.

**Parameters:**

**`tensor`** `torch.Tensor`

A tensor of the lookup being validated; its device decides.

---

**Returns:** `bool`

True under `strict_validation` or for CPU tensors (no synchronisation

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextOwnerShardedEmbedding._validate_padded_received_ids(
    received_ids: torch.Tensor,
    capacity: int
) -> None
```

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.

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

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

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

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

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.

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

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

```python
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`** `int` — default: 4

Kernel width of the causal depthwise convolution.

---

**`rms_norm_eps`** `float` — default: 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`**

---

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

```python
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,
    sequence_ids: torch.Tensor | 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 | None` — default: 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.

---

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

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

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram.Qwen3_8_FlashNextPLELayer._segmented_conv(
    padded: torch.Tensor,
    cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None,
    sequence_ids: torch.Tensor
) -> torch.Tensor
```

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:**

**`padded`** `torch.Tensor`

Left history followed by the local sequence, `[batch, history + sequence, channels]`.

---

**`cp_context`** `Qwen3_8_FlashNextCPContext | None`

Optional contiguous CP metadata.

---

**`sequence_ids`** `torch.Tensor`

Packed segment IDs `[batch, sequence]`, or `[batch, global_sequence]` under CP.

---

**Returns:** `torch.Tensor`

Pre-activation output `[batch, sequence, channels]`.

```python
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,
    sequence_ids: torch.Tensor | 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 | None` — default: 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.

---

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

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

```python
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`** `float` — default: 0.02

Standard deviation for projection weights.

---

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

**Bases:** `Function`

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

```python
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, None, None]
```

staticmethod

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

**Parameters:**

**`ctx`** `Any`

PyTorch autograd context populated by `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

```python
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,
    input_rows: torch.Tensor | None = None,
    output_rows: torch.Tensor | None = None
) -> 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

```python
class nemo_automodel.components.models.qwen3_8_flash_next.engram._PaddedAllToAll()
```

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

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

staticmethod

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram._PaddedAllToAll.forward(
    ctx: typing.Any,
    padded: torch.Tensor,
    process_group: torch.distributed.ProcessGroup
) -> torch.Tensor
```

staticmethod

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram._aligned_capacity(
    max_count: int
) -> int
```

Padded rows per peer for a measured maximum count: margin for drift, aligned, never zero.

```python
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,
    input_rows: torch.Tensor | None = None,
    output_rows: torch.Tensor | None = None
) -> 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.

---

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

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

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram._segment_rows(
    split_sizes: tuple[int, ...],
    capacity: int,
    device: torch.device,
    counts: torch.Tensor | None = None
) -> torch.Tensor
```

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

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram._segment_rows_from_counts(
    counts: torch.Tensor,
    capacity: int,
    total: int
) -> torch.Tensor
```

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

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

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

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

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

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram._CAPACITY_ALIGN = 64
```

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram._CAPACITY_MARGIN = 1.25
```

```python
nemo_automodel.components.models.qwen3_8_flash_next.engram._UNTOUCHED = -2
```