nemo_automodel.components.models.qwen3_8_flash_next.layers

Qwen3.8-Flash-Next model-local layers.

The HyperConnection equations in this module follow the Qwen3.8-Flash-Next reference implementation. They are intentionally not shared with DeepSeek-V4: that model uses a different Sinkhorn-based HyperConnection parameterization.

Module Contents

Classes

Name	Description
`Qwen3_8_FlashNextDecoderLayer`	One Qwen3.8-Flash-Next decoder layer with two learned HyperConnections.
`Qwen3_8_FlashNextGatedDeltaNet`	CP-aware GatedDeltaNet with Qwen3.8-Flash-Next’s sigmoid output gate.
`Qwen3_8_FlashNextGroupedRMSNorm`	Gemma-style RMS normalization over fixed-width HyperConnection branches.
`Qwen3_8_FlashNextHyperConnection`	Qwen3.8-Flash-Next gated HyperConnection read/write transform.
`Qwen3_8_FlashNextHyperConnectionResidual`	Residual tensors retained between a HyperConnection read and write.
`Qwen3_8_FlashNextQSAAttention`	Qwen3.8-Flash-Next gated attention with compressed-block QSA routing.
`Qwen3_8_FlashNextRMSNormGated`	Qwen3.8-Flash-Next GDN output normalization with its checkpoint-selected gate.

Functions

Name	Description
`_grouped_rms_norm_fp32`	FP32 grouped RMSNorm as one fusable elementwise chain.
`_qsa_route_extents`	Bound the QSA route width and name the physical K/V extent for one forward.
`_rms_norm_gated_fp32`	FP32 RMSNorm + output gate as one fusable elementwise chain.

Data

_grouped_rms_norm_fp32_compiled

_rms_norm_gated_fp32_compiled

API

class nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextDecoderLayer(
    layer_idx: int,
    config: object,
    moe_config: object,
    backend: nemo_automodel.components.models.common.BackendConfig,
    ple: torch.nn.Module | None = None
)

Bases: Module

One Qwen3.8-Flash-Next decoder layer with two learned HyperConnections.

Tensor layout

The first layer accepts token embeddings [batch, sequence, hidden] and expands them to flattened HC state [batch, sequence, hc_count * hidden]. Every layer returns that flattened HC layout.

Parameters:

layer_idx

int

Zero-based decoder index.

config

object

Qwen3.8-Flash-Next text configuration.

moe_config

object

Native MoE configuration.

backend

BackendConfig

Attention, linear, and expert backend configuration.

ple

nn.Module | NoneDefaults to None

Optional Engram-derived PLE module. The checkpoint installs it only on decoder index 1.

_nemo_disable_activation_checkpointing

= ple is not None

attn_hyper_connection

= Qwen3_8_FlashNextHyperConnection(**hc_kwargs)

hc_count

= int(getattr(config, 'hc_count'))

hidden_size

= int(getattr(config, 'hidden_size'))

layer_type

= str(layer_types[layer_idx])

linear_attn

= Qwen3_8_FlashNextGatedDeltaNet(config, layer_idx)

mlp

= MoE(moe_config, backend)

mlp_hyper_connection

= Qwen3_8_FlashNextHyperConnection(**hc_kwargs)

self_attn

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextDecoderLayer._expand_initial_streams(
    hidden_states: torch.Tensor
) -> torch.Tensor

Expand a one-stream decoder input into the persistent HC layout.

Parameters:

hidden_states

torch.Tensor

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

Returns: torch.Tensor

Tensor of shape

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextDecoderLayer._ple_sequence_ids(
    input_ids: torch.Tensor,
    cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None,
    cu_seqlens: torch.Tensor | None
) -> torch.Tensor | None

staticmethod

Return packed segment IDs for PLE, or None for an unpacked batch.

Parameters:

input_ids

torch.Tensor

Local raw IDs of shape [batch, sequence].

cp_context

Qwen3_8_FlashNextCPContext | None

Optional contiguous CP metadata. Its global_cu_seqlens bound the full, unsharded packed row.

cu_seqlens

torch.Tensor | None

Non-CP boundaries [segments + 1] of the local packed row, built from the loader seq_lens, if any.

Returns: torch.Tensor | None

Segment IDs [1, sequence], or [1, global_sequence] under CP: one per

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextDecoderLayer.forward(
    hidden_states: torch.Tensor,
    input_ids: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    padding_mask: torch.Tensor | None = None,
    position_ids: torch.Tensor | None = None,
    cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None = None,
    attn_kwargs: object = {}
) -> torch.Tensor

Run PLE, attention/GDN, and top-10 MoE updates.

Parameters:

hidden_states

torch.Tensor

One-stream input [batch, sequence, hidden] on the first layer, otherwise flattened HC streams [batch, sequence, hc_count * hidden].

input_ids

torch.Tensor

Raw tokenizer IDs of shape [batch, sequence] used by the PLE hash path.

freqs_cis

torch.Tensor

Composed rotary values [batch, sequence, rotary_dim] whose final axis stores concatenated cosine and sine values.

attention_mask

torch.Tensor | NoneDefaults to None

Optional token mask of shape [batch, sequence] or backend-specific causal attention mask.

padding_mask

torch.Tensor | NoneDefaults to None

Optional [batch, sequence] mask where True marks padding for MoE dispatch.

position_ids

torch.Tensor | NoneDefaults to None

Optional positions of shape [batch, sequence] or [axes, batch, sequence].

cp_context

Qwen3_8_FlashNextCPContext | NoneDefaults to None

Optional contiguous CP metadata. Tensor-bearing fields contain replicated global raw IDs/padding of shape [batch, global_sequence] and identify this rank’s local interval.

**attn_kwargs

objectDefaults to {}

Attention backend metadata.

Returns: torch.Tensor

Flattened HC streams of shape

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextDecoderLayer.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

Initialize this decoder layer for training from scratch.

Parameters:

buffer_device

torch.device

Device used by attention/MoE initializers.

init_std

floatDefaults to 0.02

Standard deviation for dense projection weights.

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextDecoderLayer.nemo_checkpoint_context_fn(
    context_fn: collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None
) -> collections.abc.Callable[[], tuple[contextlib.AbstractContextManager, contextlib.AbstractContextManager]] | None

Extend the activation-checkpoint contexts of this block with QSA route replay.

The MoE parallelizer calls this model-owned hook when it wraps the block in a checkpoint wrapper. QSA layers record the frozen indexer’s routes during the checkpoint forward and replay them during recompute; GatedDeltaNet layers, or a config that disables reuse, return the incoming contexts unchanged.

Parameters:

context_fn

Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None

The parallelizer’s torch.utils.checkpoint context factory, or None.

Returns: Callable[[], tuple[AbstractContextManager, AbstractContextManager]] | None

The (possibly wrapped) context factory.

class nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextGatedDeltaNet(
    config: object,
    layer_idx: int
)

Bases: CPAwareGatedDeltaNet

CP-aware GatedDeltaNet with Qwen3.8-Flash-Next’s sigmoid output gate.

_packed_global_cu_seqlens

Tensor | None = None

norm

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextGatedDeltaNet._forward_with_cp(
    hidden_states: torch.Tensor,
    position_ids: torch.Tensor | None,
    seq_index: torch.Tensor | None,
    blockdiag_state: nemo_automodel.components.distributed.blockdiag_cp.BlockdiagCpModelState | None = None
) -> torch.Tensor

Run the inherited FLA CP core in explicit contiguous sequence order.

Parameters:

hidden_states

torch.Tensor

Local contiguous states of shape [batch, local_sequence, hidden]. Rank r owns global positions [r * local_sequence, (r + 1) * local_sequence).

position_ids

torch.Tensor | None

Local global positions of shape [batch, local_sequence] or repeated text-RoPE axes of shape [axes, batch, local_sequence].

seq_index

torch.Tensor | None

Optional local positions of shape [local_sequence] or [batch, local_sequence]; ignored because Qwen3.8-Flash-Next’s model-owned layout is always contiguous.

blockdiag_state

BlockdiagCpModelState | NoneDefaults to None

External packed CP metadata. Qwen3.8-Flash-Next builds its own state and rejects an externally supplied one.

Returns: torch.Tensor

Local GDN output of shape [batch, local_sequence, hidden] in

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextGatedDeltaNet.forward(
    hidden_states: torch.Tensor,
    kwargs: object = {}
) -> torch.Tensor

Route packed boundaries around the inherited non-CP varlen path.

Without CP, cu_seqlens flows to the inherited FLA varlen forward unchanged. With an active CP mesh the boundaries describe the global packed row, so they are stashed for _forward_with_cp and the inherited dispatcher sees no cu_seqlens.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [batch, sequence, hidden]. Packed non-CP inputs have one physical row; a padded tail must be its own segment.

**kwargs

objectDefaults to {}

Arguments with the tensor layouts documented by CPAwareGatedDeltaNet.forward. cu_seqlens contains document boundaries of shape [documents + 1], global under CP; attention_mask has shape [batch, sequence], with document-ID change points matching cu_seqlens. IDs and indices are normalized to consecutive documents covering the full row.

Returns: torch.Tensor

Tensor of shape [batch, sequence, hidden].

Raises:

ValueError: If non-CP boundaries do not strictly partition one row or the supplied document mask disagrees with those boundaries.

class nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextGroupedRMSNorm(
    hidden_size: int,
    group_size: int,
    eps: float = 1e-06
)

Bases: Module

Gemma-style RMS normalization over fixed-width HyperConnection branches.

Tensor layout

Input and output have shape [..., hidden_size]. The learned weight has shape [hidden_size] and is applied as 1 + weight.

Parameters:

hidden_size

int

Flattened feature width.

group_size

int

Number of features normalized together. A flattened input of shape [..., hidden_size] is viewed as [..., hidden_size // group_size, group_size].

eps

floatDefaults to 1e-06

Variance epsilon.

weight

= nn.Parameter(torch.zeros(hidden_size))

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextGroupedRMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Normalize each branch independently.

Parameters:

hidden_states

torch.Tensor

Flattened branch states of shape [..., hidden_size].

Returns: torch.Tensor

Normalized states of shape [..., hidden_size].

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextGroupedRMSNorm.reset_parameters() -> None

Reset the additive Gemma-style scale to zero.

class nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextHyperConnection(
    hidden_size: int,
    hc_count: int,
    lowrank_size: int,
    rms_norm_eps: float,
    backend: nemo_automodel.components.models.common.BackendConfig,
    use_combine: bool = True,
    dtype: torch.dtype | str | None = None
)

Bases: Module

Qwen3.8-Flash-Next gated HyperConnection read/write transform.

A read normalizes hc_count streams independently, predicts a feature gate, and averages the gated streams to one block input. A write predicts one injection gate per stream and adds the block output back to every stream. The final decoder mixer uses only the read side.

Tensor layout

Reads [..., hc_count * hidden_size] and returns [..., hidden_size]. Writes a block tensor [..., hidden_size] back to a residual tensor [..., hc_count * hidden_size].

Parameters:

hidden_size

int

Width of one HC stream.

hc_count

int

Number of streams.

lowrank_size

int

Bottleneck width used to predict read gates.

rms_norm_eps

float

Variance epsilon for branch normalization.

backend

BackendConfig

Linear backend configuration.

use_combine

boolDefaults to True

Whether to instantiate the write-side injection weight.

dtype

torch.dtype | str | NoneDefaults to None

Parameter dtype override. If omitted, the backend model dtype is resolved by the caller and should be passed explicitly.

block_inject_weight

flat_hidden_size

= hidden_size * hc_count

hc_norm

input_mix_weight_down

input_mix_weight_up

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextHyperConnection.combine(
    block_output: torch.Tensor,
    residual: nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextHyperConnectionResidual
) -> torch.Tensor

Inject one block output into every HC stream.

Parameters:

block_output

torch.Tensor

Attention or MoE output of shape [..., hidden_size].

residual

Qwen3_8_FlashNextHyperConnectionResidual

Flattened pre-block streams and their normalized values, each shaped [..., hc_count * hidden_size].

Returns: torch.Tensor

Updated flattened streams of shape

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextHyperConnection.init_weights(
    init_std: float = 0.02
) -> None

Initialize HC weights for training from scratch.

Parameters:

init_std

floatDefaults to 0.02

Standard deviation for all HC linear weights.

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextHyperConnection.mix(
    hidden_states: torch.Tensor
) -> tuple[torch.Tensor, nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextHyperConnectionResidual]

Collapse HC streams to the input of an attention or MoE block.

Parameters:

hidden_states

torch.Tensor

Flattened HC streams of shape [..., hc_count * hidden_size].

Returns: torch.Tensor

A pair containing the mixed block input of shape

class nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextHyperConnectionResidual(
    hidden_states: torch.Tensor,
    normalized_states: torch.Tensor
)

Dataclass

Residual tensors retained between a HyperConnection read and write.

hidden_states

Tensor

Flattened HC streams of shape [..., hc_count * hidden_size].

normalized_states

Tensor

Branch-normalized streams with the same shape.

class nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextQSAAttention(
    config: object,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Qwen3NextAttention

Qwen3.8-Flash-Next gated attention with compressed-block QSA routing.

The main query/key/value, output gate, and output projection retain the Qwen3-Next equations. A separate frozen indexer returns logical token IDs, then the model-owned QSA dispatcher evaluates only those IDs. CUDA BF16 training selects FlexAttention or SM90 CuTe through backend.attn; CPU and explicit reference backends use the PyTorch oracle. Main Q/K/V remain differentiable.

_cp_mesh

DeviceMesh | None = None

indexer

= Qwen3_8_FlashNextQSAIndexer(config, backend)

reuse_routes_on_recompute

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextQSAAttention._select_routes(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None,
    cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None,
    packed_cu_seqlens: torch.Tensor | None
) -> nemo_automodel.components.models.qwen3_8_flash_next.qsa.QSARouteSelection

Run the frozen indexer, or replay its recorded routes during checkpoint recompute.

Parameters:

torch.Tensor

Block input of shape [batch, sequence, hidden_size].

freqs_cis

torch.Tensor

Rotary values [batch, sequence, rotary_dim].

attention_mask

torch.Tensor | None

Optional binary right-tail mask [batch, sequence]; ignored for packed rows.

cp_context

Qwen3_8_FlashNextCPContext | None

Optional contiguous CP metadata.

packed_cu_seqlens

torch.Tensor | None

Optional packed-document boundaries [num_docs + 1].

Returns: QSARouteSelection

The route IDs [batch, sequence, attention_width] and, on the CUDA flex path,

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextQSAAttention.forward(
    x: torch.Tensor,
    freqs_cis: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None = None,
    attn_kwargs: object = {}
) -> torch.Tensor

Select compressed blocks and run model-owned sparse GQA.

Parameters:

torch.Tensor

Block input of shape [batch, sequence, hidden_size].

freqs_cis

torch.Tensor

Rotary values [batch, sequence, rotary_dim] whose final axis stores concatenated cosine and sine values.

attention_mask

torch.Tensor | NoneDefaults to None

Optional token mask of shape [batch, sequence] or backend-specific causal attention mask.

cp_context

Qwen3_8_FlashNextCPContext | NoneDefaults to None

Optional contiguous CP metadata. When present, x contains local queries while compressed/main K/V are gathered to global rank order.

**attn_kwargs

objectDefaults to {}

Backend attention metadata. Packed layouts provide global document boundaries through cu_seqlens or cp_context.

Returns: torch.Tensor

Attention output with the same shape as x.

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextQSAAttention.init_weights(
    buffer_device: torch.device,
    init_std: float = 0.02
) -> None

Initialize attention and indexer weights.

Parameters:

buffer_device

torch.device

Device retained for the common attention initializer contract.

init_std

floatDefaults to 0.02

Projection initialization standard deviation.

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextQSAAttention.setup_cp_attention(
    cp_mesh: torch.distributed.device_mesh.DeviceMesh
) -> None

Install the contiguous CP mesh used for QSA K/V exchange.

class nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextRMSNormGated(
    hidden_size: int,
    eps: float = 1e-06,
    activation: str = 'sigmoid',
    dtype: torch.dtype | None = None
)

Bases: Module

Qwen3.8-Flash-Next GDN output normalization with its checkpoint-selected gate.

Transformers’ Qwen3.5 GatedDeltaNet hard-codes a SiLU output gate. Qwen3.8-Flash-Next keeps the same projections and delta-rule core but sets output_gate_type='sigmoid'. Keeping this model-local module avoids changing the shared Qwen3.5 execution contract.

weight

= nn.Parameter(torch.ones(hidden_size, dtype=dtype))

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextRMSNormGated.forward(
    hidden_states: torch.Tensor,
    gate: torch.Tensor
) -> torch.Tensor

Normalize in fp32, then apply the configured gate in fp32.

Parameters:

hidden_states

torch.Tensor

Tensor of shape [..., hidden_size] containing GatedDeltaNet values.

gate

torch.Tensor

Tensor of shape [..., hidden_size] containing the elementwise output-gate logits.

Returns: torch.Tensor

Tensor of shape [..., hidden_size] in the input dtype.

nemo_automodel.components.models.qwen3_8_flash_next.layers.Qwen3_8_FlashNextRMSNormGated.reset_parameters() -> None

Match the multiplicative RMSNorm checkpoint convention.

nemo_automodel.components.models.qwen3_8_flash_next.layers._grouped_rms_norm_fp32(
    hidden_states: torch.Tensor,
    weight: torch.Tensor,
    group_size: int,
    eps: float
) -> torch.Tensor

FP32 grouped RMSNorm as one fusable elementwise chain.

nemo_automodel.components.models.qwen3_8_flash_next.layers._qsa_route_extents(
    local_sequence_length: int,
    packed_cu_seqlens: torch.Tensor | None,
    cp_context: nemo_automodel.components.models.qwen3_8_flash_next.cp.Qwen3_8_FlashNextCPContext | None
) -> tuple[int, int]

Bound the QSA route width and name the physical K/V extent for one forward.

Valid route IDs are contiguous and a causal row can never see more than the longest document (packed) or the whole sequence (dense), so the indexer’s fixed-width output can be trimmed without changing attention. Routes are global K/V coordinates: under CP the keys are all-gathered to the global length, so the mask extent follows the CP context whether or not the row is packed.

Parameters:

local_sequence_length

int

Number of query positions on this rank.

packed_cu_seqlens

torch.Tensor | None

Optional packed-document boundaries [num_docs + 1].

cp_context

Qwen3_8_FlashNextCPContext | None

Optional contiguous CP metadata.

Returns: int

(max_visible_tokens, kv_length): the route-width bound and the number of

nemo_automodel.components.models.qwen3_8_flash_next.layers._rms_norm_gated_fp32(
    hidden_states: torch.Tensor,
    gate: torch.Tensor,
    weight: torch.Tensor,
    eps: float,
    use_sigmoid: bool
) -> torch.Tensor

FP32 RMSNorm + output gate as one fusable elementwise chain.

nemo_automodel.components.models.qwen3_8_flash_next.layers._grouped_rms_norm_fp32_compiled = torch.compile(_grouped_rms_norm_fp32, dynamic=True)

nemo_automodel.components.models.qwen3_8_flash_next.layers._rms_norm_gated_fp32_compiled = torch.compile(_rms_norm_gated_fp32, dynamic=True)