nemo_automodel.components.models.glm5_next.layers

Native GLM-5.3 text layers.

The released model alternates Kimi Delta Attention (KDA) and GLM-style KPool-compressed Dynamic Sparse Attention (DSA), with manifold-constrained Hyper-Connections (mHC) around both sublayers. FLA owns the production KDA kernel; small pure-Torch fallbacks keep CPU construction and unit tests useful.

Module Contents

Classes

Name	Description
`Glm5NextDecoderLayer`	One mHC decoder block with KDA/DSA and dense/MoE feed-forward.
`Glm5NextHyperConnection`	Manifold-constrained mixer for `hc_mult` residual streams.
`Glm5NextHyperConnectionFp32Params`	Own mHC parameters that must remain fp32 under FSDP mixed precision.
`Glm5NextKDAFp32Params`	Own recurrent-decay parameters that must remain fp32 under FSDP.
`Glm5NextKPoolIndexer`	KPool-compressed DSA indexer for training without a KV cache.
`Glm5NextLinearAttention`	Kimi Delta Attention with released GLM-5.3 checkpoint parameter names.
`Glm5NextRMSNorm`	RMSNorm with fp32 variance accumulation.
`Glm5NextSparseAttention`	NoPE MLA whose visibility is selected by the GLM KPool indexer.
`Glm5NextUnweightedRMSNorm`	Parameter-free fp32 RMS normalization used by mHC.
`_TorchRMSNormGated`	CPU fallback for FLA’s gated per-head RMSNorm.
`_TorchShortConvolution`	Depthwise causal Conv1d matching FLA `ShortConvolution` state keys.

Functions

Name	Description
`_rms_norm_gated`	-
`_short_conv`	-
`_torch_recurrent_kda`	Differentiable reference KDA for CPU/small tensors.

Data

_FLA_MSG

__all__

API

class nemo_automodel.components.models.glm5_next.layers.Glm5NextDecoderLayer(
    config: nemo_automodel.components.models.glm5_next.config.Glm5NextTextConfig,
    layer_idx: int,
    moe_config: nemo_automodel.components.moe.config.MoEConfig,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

One mHC decoder block with KDA/DSA and dense/MoE feed-forward.

attn_hc

= Glm5NextHyperConnection(config)

block_type

= config.layer_types[layer_idx]

ffn_hc

= Glm5NextHyperConnection(config)

input_layernorm

is_linear_attn

= self.block_type == 'linear_attention'

is_moe_layer

= config.mlp_layer_types[layer_idx] == 'sparse'

mlp

post_attention_layernorm

self_attn

nemo_automodel.components.models.glm5_next.layers.Glm5NextDecoderLayer.forward(
    hidden_streams: torch.Tensor,
    packed_context: nemo_automodel.components.models.glm5_next.cp.Glm5NextPackedContext,
    padding_mask: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> torch.Tensor

Transform residual streams [batch, local_sequence, hc_mult, hidden].

nemo_automodel.components.models.glm5_next.layers.Glm5NextDecoderLayer.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None

Initialize all decoder children.

nemo_automodel.components.models.glm5_next.layers.Glm5NextDecoderLayer.update_moe_gate_bias() -> None

Update the correction bias for this layer’s learned MoE router.

class nemo_automodel.components.models.glm5_next.layers.Glm5NextHyperConnection(
    config: nemo_automodel.components.models.glm5_next.config.Glm5NextTextConfig
)

Bases: Module

Manifold-constrained mixer for hc_mult residual streams.

hidden_streams is [batch, sequence, hc_mult, hidden]. The returned tensors are post [batch, sequence, hc_mult], comb [batch, sequence, hc_mult, hc_mult] and collapsed [batch, sequence, hidden].

_fp32_params

= Glm5NextHyperConnectionFp32Params(mix)

base

Parameter

Expose the checkpoint’s flat mHC base parameter.

hc_eps

= config.hc_eps

hc_mult

= config.hc_mult

hc_sinkhorn_iters

= config.hc_sinkhorn_iters

input_norm

= Glm5NextUnweightedRMSNorm(config.rms_norm_eps)

scale

Parameter

Expose the checkpoint’s flat mHC scale parameter.

nemo_automodel.components.models.glm5_next.layers.Glm5NextHyperConnection.forward(
    hidden_streams: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Build Sinkhorn mixing weights and collapse streams for one sublayer.

nemo_automodel.components.models.glm5_next.layers.Glm5NextHyperConnection.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None

Initialize mHC parameters on buffer_device.

class nemo_automodel.components.models.glm5_next.layers.Glm5NextHyperConnectionFp32Params(
    mix_size: int
)

Bases: Module

Own mHC parameters that must remain fp32 under FSDP mixed precision.

base

scale

nemo_automodel.components.models.glm5_next.layers.Glm5NextHyperConnectionFp32Params.forward(
    pre_w: torch.Tensor,
    post_w: torch.Tensor,
    comb_w: torch.Tensor,
    hc: int,
    eps: float,
    sinkhorn_iters: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]

Build FP32 mHC weights while this holder’s FSDP unit is unsharded.

class nemo_automodel.components.models.glm5_next.layers.Glm5NextKDAFp32Params(
    num_heads: int,
    projection_size: int
)

Bases: Module

Own recurrent-decay parameters that must remain fp32 under FSDP.

A_log

dt_bias

nemo_automodel.components.models.glm5_next.layers.Glm5NextKDAFp32Params.forward(
    gate: torch.Tensor,
    head_dim: int,
    lower_bound: float | None
) -> torch.Tensor

Return log-decay gates [batch, sequence, heads, head_dim].

class nemo_automodel.components.models.glm5_next.layers.Glm5NextKPoolIndexer(
    config: nemo_automodel.components.models.glm5_next.config.Glm5NextTextConfig,
    layer_idx: int,
    dtype: torch.dtype
)

Bases: Module

KPool-compressed DSA indexer for training without a KV cache.

always_select_tail

= config.index_kpool_always_select_tail

head_dim

= config.index_head_dim

index_kpool

= config.index_kpool

index_kpool_compress_ape

index_kpool_compress_gate

index_topk

= config.index_topk

k_norm

n_heads

= config.index_n_heads

softmax_scale

= self.head_dim ** -0.5

weights_proj

wq_b

nemo_automodel.components.models.glm5_next.layers.Glm5NextKPoolIndexer.forward(
    full_hidden: torch.Tensor,
    query_hidden: torch.Tensor,
    query_resid: torch.Tensor,
    query_positions: torch.Tensor
) -> torch.Tensor

Prepare one document and select indices for a query chunk.

full_hidden is [1, keys, hidden]; query tensors are [1, queries, ...] and positions are document-local [queries]. Returns [1, queries, index_topk + kpool - 1] int32 indices.

nemo_automodel.components.models.glm5_next.layers.Glm5NextKPoolIndexer.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None

Initialize indexer projections and KPool parameters.

nemo_automodel.components.models.glm5_next.layers.Glm5NextKPoolIndexer.prepare_pools(
    full_hidden: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]

Compress one document’s keys into KPool candidates.

Parameters:

full_hidden

torch.Tensor

Hidden states with shape [1, keys, hidden] for one unpadded document.

Returns: torch.Tensor

Pool keys with shape [complete_pools, index_head_dim] and raw

nemo_automodel.components.models.glm5_next.layers.Glm5NextKPoolIndexer.select(
    query_hidden: torch.Tensor,
    query_resid: torch.Tensor,
    query_positions: torch.Tensor,
    pool_keys: torch.Tensor,
    pool_indices: torch.Tensor,
    key_length: int
) -> torch.Tensor

Select raw key indices for one query chunk.

Parameters:

query_hidden

torch.Tensor

Hidden states with shape [1, queries, hidden].

query_resid

torch.Tensor

Low-rank query states with shape [1, queries, q_lora_rank].

query_positions

torch.Tensor

Document-local positions with shape [queries].

pool_keys

torch.Tensor

Prepared KPool keys with shape [complete_pools, index_head_dim].

pool_indices

torch.Tensor

Prepared raw token indices with shape [complete_pools, index_kpool].

key_length

int

Number of tokens in the unpadded document.

Returns: torch.Tensor

Int32 raw indices with shape

class nemo_automodel.components.models.glm5_next.layers.Glm5NextLinearAttention(
    config: nemo_automodel.components.models.glm5_next.config.Glm5NextTextConfig,
    layer_idx: int
)

Bases: Module

Kimi Delta Attention with released GLM-5.3 checkpoint parameter names.

A_log

Parameter

Expose the checkpoint’s flat A_log parameter name.

_fp32_params

b_proj

conv_size

= config.linear_conv_kernel_dim

dt_bias

Parameter

Expose the checkpoint’s flat dt_bias parameter name.

f_a_proj

f_b_proj

g_a_proj

g_b_proj

head_dim

= config.linear_head_dim

hidden_size

= config.hidden_size

k_conv1d

k_proj

num_heads

= config.linear_num_heads

o_norm

o_proj

projection_size

= self.head_dim * self.num_heads

q_conv1d

q_proj

v_conv1d

v_proj

nemo_automodel.components.models.glm5_next.layers.Glm5NextLinearAttention._core(
    hidden_states: torch.Tensor,
    cu_seqlens: torch.Tensor | None = None,
    cp_context: typing.Any = None
) -> torch.Tensor

Project and execute KDA for one packed row or a regular batch.

nemo_automodel.components.models.glm5_next.layers.Glm5NextLinearAttention.forward(
    hidden_states: torch.Tensor,
    packed_context: nemo_automodel.components.models.glm5_next.cp.Glm5NextPackedContext | None = None,
    padding_mask: torch.Tensor | None = None,
    _: typing.Any = {}
) -> torch.Tensor

Run KDA over [batch, local_sequence, hidden] without crossing documents.

nemo_automodel.components.models.glm5_next.layers.Glm5NextLinearAttention.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None

Initialize KDA while preserving fp32 recurrent parameters.

nemo_automodel.components.models.glm5_next.layers.Glm5NextLinearAttention.setup_cp_attention(
    cp_mesh
) -> None

Attach the one-dimensional contiguous CP mesh.

class nemo_automodel.components.models.glm5_next.layers.Glm5NextRMSNorm(
    hidden_size: int,
    eps: float,
    dtype: torch.dtype
)

Bases: Module

RMSNorm with fp32 variance accumulation.

Input and output have shape [batch, sequence, hidden]; the leading axes may be replaced by any token layout as long as hidden is last.

weight

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

nemo_automodel.components.models.glm5_next.layers.Glm5NextRMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Normalize [..., hidden] and preserve the input dtype.

nemo_automodel.components.models.glm5_next.layers.Glm5NextRMSNorm.reset_parameters() -> None

class nemo_automodel.components.models.glm5_next.layers.Glm5NextSparseAttention(
    config: nemo_automodel.components.models.glm5_next.config.Glm5NextTextConfig,
    layer_idx: int,
    backend: nemo_automodel.components.models.common.BackendConfig
)

Bases: Module

NoPE MLA whose visibility is selected by the GLM KPool indexer.

indexer

= Glm5NextKPoolIndexer(config, layer_idx, dtype)

kv_a_layernorm

kv_a_proj_with_mqa

kv_b_proj

kv_lora_rank

= config.kv_lora_rank

num_heads

= config.num_attention_heads

o_proj

q_a_layernorm

q_a_proj

q_b_proj

q_lora_rank

= config.q_lora_rank

qk_head_dim

= config.qk_nope_head_dim

query_chunk_size

= 32

scaling

= self.qk_head_dim ** -0.5

v_head_dim

= config.v_head_dim

nemo_automodel.components.models.glm5_next.layers.Glm5NextSparseAttention._forward_document(
    full_hidden: torch.Tensor,
    query_start: int,
    query_end: int
) -> torch.Tensor

Execute sparse attention for a local query interval of one full document.

Parameters:

full_hidden

torch.Tensor

Unpadded document states with shape [1, key_tokens, hidden].

query_start

int

Inclusive document-local index of the first local query.

query_end

int

Exclusive document-local index of the final local query.

Returns: torch.Tensor

Projected attention output with shape

nemo_automodel.components.models.glm5_next.layers.Glm5NextSparseAttention._forward_document_cudnn(
    full_hidden: torch.Tensor,
    latent: torch.Tensor,
    pool_keys: torch.Tensor,
    pool_indices: torch.Tensor,
    query_start: int,
    query_end: int
) -> torch.Tensor

Run absorbed latent attention through FlashMLA/cuDNN for one document.

Parameters:

full_hidden

torch.Tensor

Unpadded document states with shape [1, key_tokens, hidden].

latent

torch.Tensor

Normalized shared latent K/V with shape [1, key_tokens, 512].

pool_keys

torch.Tensor

KPool-compressed index keys with shape [complete_pools, index_head_dim].

pool_indices

torch.Tensor

Document-local token indices with shape [complete_pools, index_kpool].

query_start

int

Inclusive document-local index of the first local query.

query_end

int

Exclusive document-local index of the final local query.

Returns: torch.Tensor

Projected attention output with shape

nemo_automodel.components.models.glm5_next.layers.Glm5NextSparseAttention.forward(
    hidden_states: torch.Tensor,
    packed_context: nemo_automodel.components.models.glm5_next.cp.Glm5NextPackedContext,
    padding_mask: torch.Tensor | None = None,
    _: typing.Any = {}
) -> torch.Tensor

Run document-isolated sparse attention on [batch, local_sequence, hidden].

nemo_automodel.components.models.glm5_next.layers.Glm5NextSparseAttention.init_weights(
    buffer_device: torch.device,
    init_std: float
) -> None

Initialize MLA and indexer parameters.

nemo_automodel.components.models.glm5_next.layers.Glm5NextSparseAttention.setup_cp_attention(
    cp_mesh
) -> None

Attach the CP mesh used for differentiable full-sequence gathering.

class nemo_automodel.components.models.glm5_next.layers.Glm5NextUnweightedRMSNorm(
    eps: float
)

Bases: Module

Parameter-free fp32 RMS normalization used by mHC.

nemo_automodel.components.models.glm5_next.layers.Glm5NextUnweightedRMSNorm.forward(
    hidden_states: torch.Tensor
) -> torch.Tensor

Normalize [..., hc_streams * hidden] without a learned weight.

class nemo_automodel.components.models.glm5_next.layers._TorchRMSNormGated()

Bases: Glm5NextRMSNorm

CPU fallback for FLA’s gated per-head RMSNorm.

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

Normalize [batch, sequence, heads, head_dim] then sigmoid-gate it.

class nemo_automodel.components.models.glm5_next.layers._TorchShortConvolution(
    hidden_size: int,
    kernel_size: int,
    dtype: torch.dtype
)

Bases: Module

Depthwise causal Conv1d matching FLA ShortConvolution state keys.

weight

nemo_automodel.components.models.glm5_next.layers._TorchShortConvolution.forward(
    x: torch.Tensor,
    cu_seqlens: torch.Tensor | None = None,
    kwargs: typing.Any = {}
) -> tuple[torch.Tensor, None]

Convolve [batch, sequence, channels] and reset at packed boundaries.

nemo_automodel.components.models.glm5_next.layers._TorchShortConvolution.reset_parameters() -> None

nemo_automodel.components.models.glm5_next.layers._rms_norm_gated(
    hidden_size: int,
    eps: float,
    dtype: torch.dtype
) -> torch.nn.Module

nemo_automodel.components.models.glm5_next.layers._short_conv(
    hidden_size: int,
    kernel_size: int,
    dtype: torch.dtype
) -> torch.nn.Module

nemo_automodel.components.models.glm5_next.layers._torch_recurrent_kda(
    q: torch.Tensor,
    k: torch.Tensor,
    v: torch.Tensor,
    g: torch.Tensor,
    beta: torch.Tensor,
    cu_seqlens: torch.Tensor | None
) -> torch.Tensor

Differentiable reference KDA for CPU/small tensors.

All q/k/v/g tensors are [batch, sequence, heads, head_dim] and beta is [batch, sequence, heads]. Packed boundaries reset the recurrent state.

nemo_automodel.components.models.glm5_next.layers._FLA_MSG = 'GLM-5.3 KDA requires the flash-linear-attention/fla extra for GPU training.'

nemo_automodel.components.models.glm5_next.layers.__all__ = ['Glm5NextDecoderLayer', 'Glm5NextHyperConnection', 'Glm5NextKPoolIndexer', 'Glm...